fix: allow token through ws protocols (#2976)

Author: MasterPtato

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

@@ -6,11 +6,11 @@ use hyper::header::HeaderName;
 use rivet_guard_core::proxy_service::{RouteConfig, RouteTarget, RoutingOutput, RoutingTimeout};
 use universaldb::utils::IsolationLevel::*;
 
+use super::SEC_WEBSOCKET_PROTOCOL;
 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");
-const SEC_WEBSOCKET_PROTOCOL: HeaderName = HeaderName::from_static("sec-websocket-protocol");
 const WS_PROTOCOL_ACTOR: &str = "rivet_actor.";
 
 /// Route requests to actor services based on hostname and path

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

@@ -3,7 +3,8 @@ use gas::prelude::*;
 use rivet_guard_core::proxy_service::RoutingOutput;
 use std::sync::Arc;
 
-use super::X_RIVET_TOKEN;
+use super::{SEC_WEBSOCKET_PROTOCOL, X_RIVET_TOKEN};
+pub(crate) const WS_PROTOCOL_TOKEN: &str = "rivet_target.";
 
 /// Route requests to the API service
 #[tracing::instrument(skip_all)]
@@ -18,11 +19,33 @@ pub async fn route_request(
 		return Ok(None);
 	}
 
+	let is_websocket = headers
+		.get("upgrade")
+		.and_then(|v| v.to_str().ok())
+		.map(|v| v.eq_ignore_ascii_case("websocket"))
+		.unwrap_or(false);
+
 	// Check auth (if enabled)
 	if let Some(auth) = &ctx.config().auth {
 		let token = headers
 			.get(X_RIVET_TOKEN)
 			.and_then(|x| x.to_str().ok())
+			// Fallback to checking websocket protocol if rivet token is not set
+			.or_else(|| {
+				if is_websocket {
+					headers
+						.get(SEC_WEBSOCKET_PROTOCOL)
+						.and_then(|protocols| protocols.to_str().ok())
+						.and_then(|protocols| {
+							protocols
+								.split(',')
+								.map(|p| p.trim())
+								.find_map(|p| p.strip_prefix(WS_PROTOCOL_TOKEN))
+						})
+				} else {
+					None
+				}
+			})
 			.ok_or_else(|| {
 				crate::errors::MissingHeader {
 					header: X_RIVET_TOKEN.to_string(),