feat(guard): support routing via x_rivet_* query params (#3238)

Author: NathanFlurry

packages/core/guard/server/src/routing/mod.rs

@@ -35,14 +35,17 @@ pub fn create_routing_function(ctx: StandaloneCtx, shared_state: SharedState) ->
 
 					tracing::debug!("Routing request for hostname: {host}, path: {path}");
 
+					// Parse query parameters
+					let query_params = parse_query_params(path);
+
 					// Check if this is a WebSocket upgrade request
 					let is_websocket = headers
 						.get("upgrade")
 						.and_then(|v| v.to_str().ok())
 						.map(|v| v.eq_ignore_ascii_case("websocket"))
 						.unwrap_or(false);
 
-					// Extract target from WebSocket protocol or HTTP header
+					// Extract target from WebSocket protocol, HTTP header, or query param
 					let target = if is_websocket {
 						// For WebSocket, parse the sec-websocket-protocol header
 						headers
@@ -55,15 +58,21 @@ pub fn create_routing_function(ctx: StandaloneCtx, shared_state: SharedState) ->
 									.map(|p| p.trim())
 									.find_map(|p| p.strip_prefix(WS_PROTOCOL_TARGET))
 							})
+							// Fallback to query parameter if protocol not provided
+							.or_else(|| query_params.get("x_rivet_target").map(|s| s.as_str()))
 					} else {
-						// For HTTP, use the x-rivet-target header
-						headers.get(X_RIVET_TARGET).and_then(|x| x.to_str().ok())
+						// For HTTP, use the x-rivet-target header, fallback to query param
+						headers
+							.get(X_RIVET_TARGET)
+							.and_then(|x| x.to_str().ok())
+							.or_else(|| query_params.get("x_rivet_target").map(|s| s.as_str()))
 					};
 
 					// Read target
 					if let Some(target) = target {
 						if let Some(routing_output) =
-							runner::route_request(&ctx, target, host, path, headers).await?
+							runner::route_request(&ctx, target, host, path, headers, &query_params)
+								.await?
 						{
 							return Ok(routing_output);
 						}
@@ -76,6 +85,7 @@ pub fn create_routing_function(ctx: StandaloneCtx, shared_state: SharedState) ->
 							path,
 							headers,
 							is_websocket,
+							&query_params,
 						)
 						.await?
 						{
@@ -109,3 +119,19 @@ pub fn create_routing_function(ctx: StandaloneCtx, shared_state: SharedState) ->
 		},
 	)
 }
+
+/// Parse query parameters from a path string
+fn parse_query_params(path: &str) -> std::collections::HashMap<String, String> {
+	let mut params = std::collections::HashMap::new();
+
+	if let Some(query_start) = path.find('?') {
+		// Strip fragment if present
+		let query = &path[query_start + 1..].split('#').next().unwrap_or("");
+		// Use url::form_urlencoded to properly decode query parameters
+		for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
+			params.insert(key.into_owned(), value.into_owned());
+		}
+	}
+
+	params
+}

packages/core/guard/server/src/routing/pegboard_gateway.rs

@@ -22,13 +22,14 @@ pub async fn route_request(
 	path: &str,
 	headers: &hyper::HeaderMap,
 	is_websocket: bool,
+	query_params: &std::collections::HashMap<String, String>,
 ) -> Result<Option<RoutingOutput>> {
 	// Check target
 	if target != "actor" {
 		return Ok(None);
 	}
 
-	// Extract actor ID from WebSocket protocol or HTTP header
+	// Extract actor ID from WebSocket protocol, HTTP header, or query param
 	let actor_id_str = if is_websocket {
 		// For WebSocket, parse the sec-websocket-protocol header
 		headers
@@ -41,22 +42,26 @@ pub async fn route_request(
 					.map(|p| p.trim())
 					.find_map(|p| p.strip_prefix(WS_PROTOCOL_ACTOR))
 			})
+			// Fallback to query parameter if protocol not provided
+			.or_else(|| query_params.get("x_rivet_actor").map(|s| s.as_str()))
 			.ok_or_else(|| {
 				crate::errors::MissingHeader {
-					header: "`rivet_actor.*` protocol in sec-websocket-protocol".to_string(),
+					header: "`rivet_actor.*` protocol in sec-websocket-protocol or x_rivet_actor query parameter".to_string(),
 				}
 				.build()
 			})?
 	} else {
-		// For HTTP, use the x-rivet-actor header
+		// For HTTP, use the x-rivet-actor header, fallback to query param
 		headers
 			.get(X_RIVET_ACTOR)
 			.map(|x| x.to_str())
 			.transpose()
 			.context("invalid x-rivet-actor header")?
+			// Fallback to query parameter if header not provided
+			.or_else(|| query_params.get("x_rivet_actor").map(|s| s.as_str()))
 			.ok_or_else(|| {
 				crate::errors::MissingHeader {
-					header: X_RIVET_ACTOR.to_string(),
+					header: format!("{} header or x_rivet_actor query parameter", X_RIVET_ACTOR),
 				}
 				.build()
 			})?

packages/core/guard/server/src/routing/runner.rs

@@ -14,6 +14,7 @@ pub async fn route_request(
 	host: &str,
 	path: &str,
 	headers: &hyper::HeaderMap,
+	query_params: &std::collections::HashMap<String, String>,
 ) -> Result<Option<RoutingOutput>> {
 	if target != "runner" {
 		return Ok(None);
@@ -57,7 +58,7 @@ pub async fn route_request(
 
 	// Check auth (if enabled)
 	if let Some(auth) = &ctx.config().auth {
-		// Extract token
+		// Extract token from protocol, header, or query param
 		let token = if is_websocket {
 			headers
 				.get(SEC_WEBSOCKET_PROTOCOL)
@@ -68,19 +69,26 @@ pub async fn route_request(
 						.map(|p| p.trim())
 						.find_map(|p| p.strip_prefix(WS_PROTOCOL_TOKEN))
 				})
+				// Fallback to query parameter if protocol not provided
+				.or_else(|| query_params.get("x_rivet_token").map(|s| s.as_str()))
 				.ok_or_else(|| {
 					crate::errors::MissingHeader {
-						header: "`rivet_token.*` protocol in sec-websocket-protocol".to_string(),
+						header: "`rivet_token.*` protocol in sec-websocket-protocol or x_rivet_token query parameter".to_string(),
 					}
 					.build()
 				})?
 		} else {
 			headers
 				.get(X_RIVET_TOKEN)
 				.and_then(|x| x.to_str().ok())
+				// Fallback to query parameter if header not provided
+				.or_else(|| query_params.get("x_rivet_token").map(|s| s.as_str()))
 				.ok_or_else(|| {
 					crate::errors::MissingHeader {
-						header: X_RIVET_TOKEN.to_string(),
+						header: format!(
+							"{} header or x_rivet_token query parameter",
+							X_RIVET_TOKEN
+						),
 					}
 					.build()
 				})?