chore(guard): move token routing out off inner router for path-based routing

Author: NathanFlurry

engine/packages/guard/src/routing/mod.rs

@@ -16,6 +16,8 @@ pub(crate) const X_RIVET_TOKEN: HeaderName = HeaderName::from_static("x-rivet-to
 pub(crate) const SEC_WEBSOCKET_PROTOCOL: HeaderName =
 	HeaderName::from_static("sec-websocket-protocol");
 pub(crate) const WS_PROTOCOL_TARGET: &str = "rivet_target.";
+pub(crate) const WS_PROTOCOL_ACTOR: &str = "rivet_actor.";
+pub(crate) const WS_PROTOCOL_TOKEN: &str = "rivet_token.";
 
 #[derive(Debug, Clone)]
 pub struct ActorPathInfo {

engine/packages/guard/src/routing/pegboard_gateway.rs

@@ -5,32 +5,28 @@ use gas::prelude::*;
 use hyper::header::HeaderName;
 use rivet_guard_core::proxy_service::{RouteConfig, RouteTarget, RoutingOutput, RoutingTimeout};
 
-use super::SEC_WEBSOCKET_PROTOCOL;
+use super::{SEC_WEBSOCKET_PROTOCOL, WS_PROTOCOL_ACTOR, WS_PROTOCOL_TOKEN, X_RIVET_TOKEN};
 use crate::{errors, shared_state::SharedState};
 
 const ACTOR_READY_TIMEOUT: Duration = Duration::from_secs(10);
 pub const X_RIVET_ACTOR: HeaderName = HeaderName::from_static("x-rivet-actor");
 pub const X_RIVET_AMESPACE: HeaderName = HeaderName::from_static("x-rivet-namespace");
-const WS_PROTOCOL_ACTOR: &str = "rivet_actor.";
-const WS_PROTOCOL_TOKEN: &str = "rivet_token.";
 
 /// Route requests to actor services using path-based routing
 #[tracing::instrument(skip_all)]
 pub async fn route_request_path_based(
 	ctx: &StandaloneCtx,
 	shared_state: &SharedState,
 	actor_id_str: &str,
-	_token: Option<&str>,
+	token: Option<&str>,
 	path: &str,
 	_headers: &hyper::HeaderMap,
 	_is_websocket: bool,
 ) -> Result<Option<RoutingOutput>> {
-	// NOTE: Token validation implemented in EE
-
 	// Parse actor ID
 	let actor_id = Id::parse(actor_id_str).context("invalid actor id in path")?;
 
-	route_request_inner(ctx, shared_state, actor_id, path).await
+	route_request_inner(ctx, shared_state, actor_id, path, token).await
 }
 
 /// Route requests to actor services based on headers
@@ -49,28 +45,39 @@ pub async fn route_request(
 		return Ok(None);
 	}
 
-	// Extract actor ID from WebSocket protocol or HTTP header
-	let actor_id_str = if is_websocket {
+	// Extract actor ID and token from WebSocket protocol or HTTP headers
+	let (actor_id_str, token) = if is_websocket {
 		// For WebSocket, parse the sec-websocket-protocol header
-		headers
+		let protocols_header = headers
 			.get(SEC_WEBSOCKET_PROTOCOL)
 			.and_then(|protocols| protocols.to_str().ok())
-			.and_then(|protocols| {
-				// Parse protocols to find actor.{id}
-				protocols
-					.split(',')
-					.map(|p| p.trim())
-					.find_map(|p| p.strip_prefix(WS_PROTOCOL_ACTOR))
-			})
+			.ok_or_else(|| {
+				crate::errors::MissingHeader {
+					header: "sec-websocket-protocol".to_string(),
+				}
+				.build()
+			})?;
+
+		let protocols: Vec<&str> = protocols_header.split(',').map(|p| p.trim()).collect();
+
+		let actor_id = protocols
+			.iter()
+			.find_map(|p| p.strip_prefix(WS_PROTOCOL_ACTOR))
 			.ok_or_else(|| {
 				crate::errors::MissingHeader {
 					header: "`rivet_actor.*` protocol in sec-websocket-protocol".to_string(),
 				}
 				.build()
-			})?
+			})?;
+
+		let token = protocols
+			.iter()
+			.find_map(|p| p.strip_prefix(WS_PROTOCOL_TOKEN));
+
+		(actor_id, token)
 	} else {
-		// For HTTP, use the x-rivet-actor header
-		headers
+		// For HTTP, use headers
+		let actor_id = headers
 			.get(X_RIVET_ACTOR)
 			.map(|x| x.to_str())
 			.transpose()
@@ -80,21 +87,32 @@ pub async fn route_request(
 					header: X_RIVET_ACTOR.to_string(),
 				}
 				.build()
-			})?
+			})?;
+
+		let token = headers
+			.get(X_RIVET_TOKEN)
+			.map(|x| x.to_str())
+			.transpose()
+			.context("invalid x-rivet-token header")?;
+
+		(actor_id, token)
 	};
 
 	// Find actor to route to
 	let actor_id = Id::parse(actor_id_str).context("invalid x-rivet-actor header")?;
 
-	route_request_inner(ctx, shared_state, actor_id, path).await
+	route_request_inner(ctx, shared_state, actor_id, path, token).await
 }
 
 async fn route_request_inner(
 	ctx: &StandaloneCtx,
 	shared_state: &SharedState,
 	actor_id: Id,
 	path: &str,
+	_token: Option<&str>,
 ) -> Result<Option<RoutingOutput>> {
+	// NOTE: Token validation implemented in EE
+
 	// Route to peer dc where the actor lives
 	if actor_id.label() != ctx.config().dc_label() {
 		tracing::debug!(peer_dc_label=?actor_id.label(), "re-routing actor to peer dc");