diff --git a/src/gateway/routes.rs b/src/gateway/routes.rs index a4cb1bf..2cd9da8 100644 --- a/src/gateway/routes.rs +++ b/src/gateway/routes.rs @@ -858,6 +858,9 @@ mod tests { decoy: crate::server::DecoyConfig::NotFound, openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(®istry), ws_sessions: Arc::new(crate::websocket::WsSessions::new()), + ws_session_slots: Arc::new(tokio::sync::Semaphore::new( + crate::websocket::DEFAULT_WS_MAX_SESSIONS, + )), }; let auth_state = Arc::clone(&provider); gateway_router() diff --git a/src/server/adapter.rs b/src/server/adapter.rs index 27ccbac..6ee6b1d 100644 --- a/src/server/adapter.rs +++ b/src/server/adapter.rs @@ -88,6 +88,8 @@ pub struct HttpAdapter { router: Router, openapi_doc: CachedOpenAPIDoc, ws_sessions: Arc, + ws_max_sessions: usize, + ws_session_slots: Arc, } impl HttpAdapter { @@ -113,12 +115,15 @@ impl HttpAdapter { let decoy = DecoyConfig::default(); let openapi_doc = CachedOpenAPIDoc::new(®istry); let ws_sessions = Arc::new(crate::websocket::WsSessions::new()); + let ws_max_sessions = crate::websocket::DEFAULT_WS_MAX_SESSIONS; + let ws_session_slots = Arc::new(tokio::sync::Semaphore::new(ws_max_sessions)); let state = RouterState { registry: Arc::clone(®istry), identity_provider: Arc::clone(&identity_provider), decoy: decoy.clone(), openapi_doc: openapi_doc.clone(), ws_sessions: Arc::clone(&ws_sessions), + ws_session_slots: Arc::clone(&ws_session_slots), }; let router = build_router(state, None); Self { @@ -130,6 +135,8 @@ impl HttpAdapter { router, openapi_doc, ws_sessions, + ws_max_sessions, + ws_session_slots, } } @@ -141,6 +148,7 @@ impl HttpAdapter { decoy, openapi_doc: self.openapi_doc.clone(), ws_sessions: Arc::clone(&self.ws_sessions), + ws_session_slots: Arc::clone(&self.ws_session_slots), }; // `extra_routes` is borrowed, not consumed (SRV-05): a builder // call after `with_extra_routes` must keep the custom routes in @@ -157,12 +165,34 @@ impl HttpAdapter { decoy: self.decoy.clone(), openapi_doc: self.openapi_doc.clone(), ws_sessions: Arc::clone(&self.ws_sessions), + ws_session_slots: Arc::clone(&self.ws_session_slots), }; self.router = build_router(state, Some(routes.clone())); self.extra_routes = Some(routes); self } + /// The concurrent WS session cap (WS-09): the upgrade handler + /// acquires one semaphore permit per upgrade, post-auth and + /// pre-upgrade; a caller over the configured cap is rejected with + /// **503 Service Unavailable**, and the permit is held for the + /// session's lifetime (an ended session frees its slot). + /// + /// Default: [`crate::websocket::DEFAULT_WS_MAX_SESSIONS`] (64). + pub fn with_ws_max_sessions(mut self, max_sessions: usize) -> Self { + self.ws_max_sessions = max_sessions; + let state = RouterState { + registry: Arc::clone(&self.registry), + identity_provider: Arc::clone(&self.identity_provider), + decoy: self.decoy.clone(), + openapi_doc: self.openapi_doc.clone(), + ws_sessions: Arc::clone(&self.ws_sessions), + ws_session_slots: Self::rebuild_session_slots(max_sessions), + }; + self.router = build_router(state, self.extra_routes.clone()); + self + } + /// The shared WS session registry (WS-08): live sessions' /// [`WsPumps`](crate::websocket::WsPumps) handles, evictable via /// `WsSessions::abort`. The upgrade handler registers against this @@ -171,6 +201,12 @@ impl HttpAdapter { Arc::clone(&self.ws_sessions) } + /// A fresh semaphore for the new cap; the retained handles of + /// already-open sessions are unaffected (they hold their permits). + fn rebuild_session_slots(max_sessions: usize) -> Arc { + Arc::new(tokio::sync::Semaphore::new(max_sessions)) + } + pub fn decoy(&self) -> &DecoyConfig { &self.decoy } @@ -956,6 +992,9 @@ mod tests { decoy: DecoyConfig::default(), openapi_doc: CachedOpenAPIDoc::new(&OperationRegistry::new()), ws_sessions: Arc::new(crate::websocket::WsSessions::new()), + ws_session_slots: Arc::new(tokio::sync::Semaphore::new( + crate::websocket::DEFAULT_WS_MAX_SESSIONS, + )), } } diff --git a/src/server/state.rs b/src/server/state.rs index b57529c..5b1251d 100644 --- a/src/server/state.rs +++ b/src/server/state.rs @@ -26,8 +26,9 @@ pub enum DecoyConfig { /// State embedded in the axum `Router`: the registry and identity /// provider every request handler reaches through the router state, plus /// the decoy config for the fallback, the pre-serialized -/// `/openapi.json` projection cache (SRV-09), and the shared WS session -/// registry the WS upgrade retains pump handles in (WS-08). +/// `/openapi.json` projection cache (SRV-09), and the WS session +/// lifecycle config the upgrade handler enforces: the shared pump-handle +/// registry (WS-08) and the concurrent-session cap (WS-09). #[derive(Clone)] pub(crate) struct RouterState { pub(crate) registry: Arc, @@ -35,6 +36,10 @@ pub(crate) struct RouterState { pub(crate) decoy: DecoyConfig, pub(crate) openapi_doc: crate::server::adapter::CachedOpenAPIDoc, pub(crate) ws_sessions: Arc, + /// Session cap (WS-09): one `HttpAdapter`-wide semaphore, built + /// once at construction — the upgrade handler acquires one permit + /// per upgrade and holds it for the session's lifetime. + pub(crate) ws_session_slots: Arc, } impl axum::extract::FromRef for crate::websocket::SessionState { @@ -42,6 +47,7 @@ impl axum::extract::FromRef for crate::websocket::SessionState { crate::websocket::SessionState::new( Arc::clone(&state.registry), Arc::clone(&state.ws_sessions), + Arc::clone(&state.ws_session_slots), ) } } @@ -89,6 +95,9 @@ mod tests { }, openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(&OperationRegistry::new()), ws_sessions: Arc::new(crate::websocket::WsSessions::new()), + ws_session_slots: Arc::new(tokio::sync::Semaphore::new( + crate::websocket::DEFAULT_WS_MAX_SESSIONS, + )), }; let extracted: DecoyConfig = axum::extract::FromRef::from_ref(&state); assert!(matches!(extracted, DecoyConfig::Redirect { .. })); diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs index 1169ac7..384df9c 100644 --- a/src/websocket/mod.rs +++ b/src/websocket/mod.rs @@ -24,7 +24,7 @@ pub use byte_adapter::split_tungstenite_to_bytes; pub(crate) use upgrade::adapter_install_channel_zero; pub use upgrade::{ run_channels_session, ws_bearer_auth, ws_upgrade_handler, ChannelsPolicy, SessionState, - WsSessions, + WsSessions, DEFAULT_WS_MAX_SESSIONS, }; #[cfg(any(test, feature = "test-support"))] diff --git a/src/websocket/upgrade.rs b/src/websocket/upgrade.rs index 14f3569..73dc845 100644 --- a/src/websocket/upgrade.rs +++ b/src/websocket/upgrade.rs @@ -47,34 +47,54 @@ pub struct WsSessions { sessions: Arc>>>, } +/// Default bound on concurrent WS sessions (WS-09): the semaphore +/// initial capacity in [`SessionState`]; a deployment overrides it with +/// `HttpAdapter::with_ws_max_sessions`. +pub const DEFAULT_WS_MAX_SESSIONS: usize = 64; + /// The upgrade handler's state slice: what it needs beyond the request /// itself. Axum lifts it via `FromRef` from either full router state — /// the adapter's [`crate::server::state::RouterState`] (carrying the -/// shared [`WsSessions`] instance) or a bare `Arc` -/// (custom upgrade routes / integration tests get a handler-private -/// registry; eviction still works in-crate, just not shared). +/// shared [`WsSessions`] instance and the configured session cap) or a +/// bare `Arc` (custom upgrade routes / integration +/// tests get a handler-private registry and the default cap; eviction +/// still works in-crate, just not shared). #[derive(Clone)] pub struct SessionState { registry: Arc, sessions: Arc, + /// Session cap (WS-09): acquired post-auth, pre-upgrade; a caller + /// over the cap is rejected with 503. + session_slots: Arc, } impl SessionState { /// From the plain registry state (custom upgrade routes / /// integration tests): sessions default to a fresh [`WsSessions`] /// private to the handler (eviction still functional in-crate but - /// not shared with the assembly layer). + /// not shared with the assembly layer) and the default session cap. pub(crate) fn from_registry(registry: &Arc) -> Self { Self { registry: Arc::clone(registry), sessions: Arc::new(WsSessions::new()), + session_slots: Arc::new(tokio::sync::Semaphore::new(DEFAULT_WS_MAX_SESSIONS)), } } /// From the adapter's router state: the shared [`WsSessions`] - /// instance the assembly layer can hold for eviction. - pub(crate) fn new(registry: Arc, sessions: Arc) -> Self { - Self { registry, sessions } + /// instance the assembly layer can hold for eviction, and the + /// pre-built session-cap semaphore (one per `HttpAdapter`, shared + /// across requests). + pub(crate) fn new( + registry: Arc, + sessions: Arc, + session_slots: Arc, + ) -> Self { + Self { + registry, + sessions, + session_slots, + } } pub(crate) fn registry(&self) -> &Arc { @@ -266,6 +286,13 @@ pub struct ChannelsPolicy(pub Arc); /// registry — the shared instance from the router state /// (`RouterState::ws_sessions`) unless a request extension carries /// one, so a stuck session stays evictable (WS-08). +/// +/// Session cap (WS-09): one semaphore permit is acquired per upgrade, +/// post-auth and pre-upgrade; when the configured cap +/// (`HttpAdapter::with_ws_max_sessions`, default +/// [`DEFAULT_WS_MAX_SESSIONS`]) is exhausted the upgrade is rejected +/// with **503 Service Unavailable** — holding the permit for the +/// session's lifetime, so ended sessions free their slot. pub async fn ws_upgrade_handler( sessions: Option>, axum::extract::State(state): axum::extract::State, @@ -273,6 +300,13 @@ pub async fn ws_upgrade_handler( policy: Option>, ws_upgrade: WebSocketUpgrade, ) -> Response { + let Ok(permit) = state.session_slots.clone().try_acquire_owned() else { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "503 Service Unavailable: WS session cap reached", + ) + .into_response(); + }; let sessions = Some( sessions .map(|axum::Extension(s)| s) @@ -286,6 +320,7 @@ pub async fn ws_upgrade_handler( .max_frame_size(INBOUND_WS_FRAME_CAP) .max_message_size(INBOUND_WS_MESSAGE_CAP) .on_upgrade(move |socket| async move { + let _permit = permit; run_channels_session(socket, registry, identity, policy, sessions).await }) } diff --git a/tests/ws_upgrade_session.rs b/tests/ws_upgrade_session.rs index 04546eb..98a4dc9 100644 --- a/tests/ws_upgrade_session.rs +++ b/tests/ws_upgrade_session.rs @@ -688,3 +688,73 @@ async fn ws_sessions_registry_tracks_and_aborts_live_sessions() { assert!(drained); ws.close().await; } + +/// WS-09 acceptance: the session cap is enforced on the upgrade — a +/// caller over the configured cap is rejected with 503, and an ended +/// session frees its slot for the next caller. +#[tokio::test] +async fn session_cap_rejects_over_limit_with_503_and_frees_slots_on_end() { + use alkhttp::server::HttpAdapter; + + let mut registry_val = OperationRegistry::new(); + registry_val + .register(HandlerRegistration::new( + OperationSpec::new( + "echo/run", + OperationType::Query, + Visibility::External, + serde_json::json!({}), + serde_json::json!({}), + vec![], + AccessControl::default(), + None, + ), + HandlerKind::Once(echo_handler()), + OperationProvenance::Local, + None, + None, + alkcall::core::types::Capabilities::new(), + )) + .unwrap(); + let registry = Arc::new(registry_val); + struct StaticTok; + impl IdentityProvider for StaticTok { + fn resolve_from_fingerprint(&self, _: &str) -> Option { + None + } + fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option { + let s = String::from_utf8_lossy(&token.raw).to_string(); + (s == "tok-1").then(|| identity("alice", &[])) + } + } + + // The adapter's built-in surface: WS route + bearer middleware, the + // whole deployment shape the cap knob configures. + let adapter = + HttpAdapter::new(Arc::new(StaticTok), Arc::clone(®istry)).with_ws_max_sessions(1); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = format!("ws://{}", listener.local_addr().unwrap()); + let app: axum::Router = adapter.router().clone(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let url = format!("{addr}/alk/channels"); + + // First session: admitted. + let mut first = WsClient::connect_authorized(&url, "tok-1").await.unwrap(); + let env = call_and_await(&mut first, "req-1", "echo/run", serde_json::json!({})).await; + assert_eq!(env.r#type, EVENT_RESPONDED, "first session admitted"); + + // Second session over the cap: rejected with 503 at the upgrade. + let status = WsClient::connect_status(&url, Some("tok-1")).await; + assert_eq!(status, Some(503), "second session rejected over the cap"); + + // Ending the first session frees the slot: the next caller is admitted. + first.close().await; + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + let mut third = WsClient::connect_authorized(&url, "tok-1").await.unwrap(); + let env = call_and_await(&mut third, "req-2", "echo/run", serde_json::json!({})).await; + assert_eq!(env.r#type, EVENT_RESPONDED, "slot freed after session end"); + third.close().await; +}