fix(websocket): configurable WS session cap (WS-09)

src/websocket/upgrade.rs

@@ -47,34 +47,54 @@ pub struct WsSessions {
     sessions: Arc<Mutex<HashMap<u64, Arc<WsPumps>>>>,
 }
 
+/// 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<OperationRegistry>`
-/// (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<OperationRegistry>` (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<OperationRegistry>,
     sessions: Arc<WsSessions>,
+    /// Session cap (WS-09): acquired post-auth, pre-upgrade; a caller
+    /// over the cap is rejected with 503.
+    session_slots: Arc<tokio::sync::Semaphore>,
 }
 
 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<OperationRegistry>) -> 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<OperationRegistry>, sessions: Arc<WsSessions>) -> 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<OperationRegistry>,
+        sessions: Arc<WsSessions>,
+        session_slots: Arc<tokio::sync::Semaphore>,
+    ) -> Self {
+        Self {
+            registry,
+            sessions,
+            session_slots,
+        }
     }
 
     pub(crate) fn registry(&self) -> &Arc<OperationRegistry> {
@@ -266,6 +286,13 @@ pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);
 /// 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::Extension<WsSessions>>,
     axum::extract::State(state): axum::extract::State<SessionState>,
@@ -273,6 +300,13 @@ pub async fn ws_upgrade_handler(
     policy: Option<axum::Extension<ChannelsPolicy>>,
     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
         })
 }