From 4747a12c023e2d3d02d61c754f151843ab2a8b3b Mon Sep 17 00:00:00 2001 From: "glm-5.3-flash" Date: Sat, 29 Aug 2026 13:31:33 +0000 Subject: [PATCH] fix(websocket): retain WsPumps handle on the server path (WS-08) --- src/gateway/routes.rs | 1 + src/server/adapter.rs | 15 ++++ src/server/state.rs | 16 +++- src/websocket/mod.rs | 5 +- src/websocket/upgrade.rs | 152 ++++++++++++++++++++++++++++++++++-- tests/ws_upgrade_session.rs | 79 +++++++++++++++++++ 6 files changed, 260 insertions(+), 8 deletions(-) diff --git a/src/gateway/routes.rs b/src/gateway/routes.rs index 7cdc9bb..a4cb1bf 100644 --- a/src/gateway/routes.rs +++ b/src/gateway/routes.rs @@ -857,6 +857,7 @@ mod tests { identity_provider: Arc::clone(&provider), decoy: crate::server::DecoyConfig::NotFound, openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(®istry), + ws_sessions: Arc::new(crate::websocket::WsSessions::new()), }; let auth_state = Arc::clone(&provider); gateway_router() diff --git a/src/server/adapter.rs b/src/server/adapter.rs index e50a093..27ccbac 100644 --- a/src/server/adapter.rs +++ b/src/server/adapter.rs @@ -87,6 +87,7 @@ pub struct HttpAdapter { alpn: &'static [u8], router: Router, openapi_doc: CachedOpenAPIDoc, + ws_sessions: Arc, } impl HttpAdapter { @@ -111,11 +112,13 @@ impl HttpAdapter { ) -> Self { let decoy = DecoyConfig::default(); let openapi_doc = CachedOpenAPIDoc::new(®istry); + let ws_sessions = Arc::new(crate::websocket::WsSessions::new()); 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), }; let router = build_router(state, None); Self { @@ -126,6 +129,7 @@ impl HttpAdapter { alpn, router, openapi_doc, + ws_sessions, } } @@ -136,6 +140,7 @@ impl HttpAdapter { identity_provider: Arc::clone(&self.identity_provider), decoy, openapi_doc: self.openapi_doc.clone(), + ws_sessions: Arc::clone(&self.ws_sessions), }; // `extra_routes` is borrowed, not consumed (SRV-05): a builder // call after `with_extra_routes` must keep the custom routes in @@ -151,12 +156,21 @@ impl HttpAdapter { identity_provider: Arc::clone(&self.identity_provider), decoy: self.decoy.clone(), openapi_doc: self.openapi_doc.clone(), + ws_sessions: Arc::clone(&self.ws_sessions), }; self.router = build_router(state, Some(routes.clone())); self.extra_routes = Some(routes); 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 + /// instance. + pub fn ws_sessions(&self) -> Arc { + Arc::clone(&self.ws_sessions) + } + pub fn decoy(&self) -> &DecoyConfig { &self.decoy } @@ -941,6 +955,7 @@ mod tests { identity_provider: idp, decoy: DecoyConfig::default(), openapi_doc: CachedOpenAPIDoc::new(&OperationRegistry::new()), + ws_sessions: Arc::new(crate::websocket::WsSessions::new()), } } diff --git a/src/server/state.rs b/src/server/state.rs index 954f143..b57529c 100644 --- a/src/server/state.rs +++ b/src/server/state.rs @@ -25,14 +25,25 @@ 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 and the pre-serialized -/// `/openapi.json` projection cache (SRV-09). +/// 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). #[derive(Clone)] pub(crate) struct RouterState { pub(crate) registry: Arc, pub(crate) identity_provider: Arc, pub(crate) decoy: DecoyConfig, pub(crate) openapi_doc: crate::server::adapter::CachedOpenAPIDoc, + pub(crate) ws_sessions: Arc, +} + +impl axum::extract::FromRef for crate::websocket::SessionState { + fn from_ref(state: &RouterState) -> Self { + crate::websocket::SessionState::new( + Arc::clone(&state.registry), + Arc::clone(&state.ws_sessions), + ) + } } impl axum::extract::FromRef for DecoyConfig { @@ -77,6 +88,7 @@ mod tests { to: "https://example.com".to_string(), }, openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(&OperationRegistry::new()), + ws_sessions: Arc::new(crate::websocket::WsSessions::new()), }; 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 9fe0a0a..1169ac7 100644 --- a/src/websocket/mod.rs +++ b/src/websocket/mod.rs @@ -22,7 +22,10 @@ pub use byte_adapter::split_tungstenite_to_bytes; #[cfg(any(test, feature = "wss"))] #[allow(unused_imports)] pub(crate) use upgrade::adapter_install_channel_zero; -pub use upgrade::{run_channels_session, ws_bearer_auth, ws_upgrade_handler, ChannelsPolicy}; +pub use upgrade::{ + run_channels_session, ws_bearer_auth, ws_upgrade_handler, ChannelsPolicy, SessionState, + WsSessions, +}; #[cfg(any(test, feature = "test-support"))] pub use upgrade::test_support::{frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient}; diff --git a/src/websocket/upgrade.rs b/src/websocket/upgrade.rs index af62a94..14f3569 100644 --- a/src/websocket/upgrade.rs +++ b/src/websocket/upgrade.rs @@ -10,6 +10,8 @@ //! `CallConnection` (the identity rides the channels-layer //! `Connection`) and runs the shared `Dispatcher::run_loop_single_stream`. +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use alkcall::channels::adapter::ChannelsAdapter; @@ -20,19 +22,137 @@ use alkcall::registry::registration::OperationRegistry; use axum::extract::ws::WebSocketUpgrade; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; +use parking_lot::Mutex; -use super::byte_adapter::{split_ws_to_bytes, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP}; +use super::byte_adapter::{ + split_ws_to_bytes, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, +}; + +/// Registry of live WS session pump handles (WS-08). The upgrade +/// handler retains each session's [`WsPumps`] handle here for the +/// session's lifetime so a stuck session is evictable in-crate: +/// [`WsSessions::abort`] forces the pumps' tasks to end (the socket +/// halves close and the channels session unwinds). The handle for a +/// session is removed when the session task finishes (self-removing +/// guard), so the registry only holds live sessions. Assembly layers +/// share one instance via [`crate::server::state::RouterState`] (the +/// upgrade handler registers against it) or per-route request +/// extensions (an extension clone takes precedence). +/// +/// Un-registered (default) — the upgrade runs fine and simply keeps no +/// eviction lever, matching the pre-WS-08 behavior. +#[derive(Clone, Default)] +pub struct WsSessions { + counter: Arc, + sessions: Arc>>>, +} + +/// 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). +#[derive(Clone)] +pub struct SessionState { + registry: Arc, + sessions: 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). + pub(crate) fn from_registry(registry: &Arc) -> Self { + Self { + registry: Arc::clone(registry), + sessions: Arc::new(WsSessions::new()), + } + } + + /// 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 } + } + + pub(crate) fn registry(&self) -> &Arc { + &self.registry + } + + pub(crate) fn sessions(&self) -> &Arc { + &self.sessions + } +} + +impl axum::extract::FromRef for Arc { + fn from_ref(state: &SessionState) -> Self { + Arc::clone(&state.registry) + } +} + +/// `FromRef` chain: a bare `Arc` router state lifts +/// into the handler's [`SessionState`]; a full [`RouterState`] carries +/// the shared [`WsSessions`] instance and lifts through its own impl. +impl axum::extract::FromRef> for SessionState { + fn from_ref(registry: &Arc) -> Self { + SessionState::from_registry(registry) + } +} + +impl WsSessions { + pub fn new() -> Self { + Self::default() + } + + /// Abort every live session's pump tasks (forced teardown). + pub fn abort(&self) { + for (_, pumps) in self.sessions.lock().drain() { + pumps.abort(); + } + } + + /// Number of live sessions currently tracked. + pub fn len(&self) -> usize { + self.sessions.lock().len() + } + + pub fn is_empty(&self) -> bool { + self.sessions.lock().is_empty() + } + + fn insert(&self, pumps: Arc) -> u64 { + let id = self.counter.fetch_add(1, Ordering::Relaxed); + self.sessions.lock().insert(id, pumps); + id + } + + fn remove(&self, id: u64) { + self.sessions.lock().remove(&id); + } +} /// The channels session for an upgraded socket: adapt → `Connection` /// (identity attached) → `ChannelsAdapter::handle`. `policy` gates -/// data-channel opens (ADR-041). +/// data-channel opens (ADR-041). When `sessions` is `Some`, the +/// session's pump handle is registered for the session's lifetime — +/// the WS-08 eviction lever ([`WsSessions::abort`]). pub async fn run_channels_session( socket: axum::extract::ws::WebSocket, registry: Arc, identity: Identity, policy: Arc, + sessions: Option, ) { - let (byte_stream, _pumps) = split_ws_to_bytes(socket); + let (byte_stream, pumps) = split_ws_to_bytes(socket); + let pumps = Arc::new(pumps); + + let _guard = sessions.map(|sessions| { + let id = sessions.insert(Arc::clone(&pumps)); + SessionGuard { sessions, id } + }); + let conn = Connection::from_bidi(byte_stream, b"alk/channels".to_vec(), None); let _ = conn.set_identity(identity.clone()); @@ -48,6 +168,17 @@ pub async fn run_channels_session( } } +struct SessionGuard { + sessions: WsSessions, + id: u64, +} + +impl Drop for SessionGuard { + fn drop(&mut self) { + self.sessions.remove(self.id); + } +} + /// The `install_channel_zero` hook: split channel 0's `BiStream` into /// the shared writer + reader, construct the single-stream /// `CallConnection`, and run the shared `Dispatcher`'s single-stream @@ -131,20 +262,31 @@ pub struct ChannelsPolicy(pub Arc); /// /// The channel lifecycle policy comes from the /// [`ChannelsPolicy`] request extension when present, else `NoCap`. +/// The session pump handles are retained in the [`WsSessions`] +/// 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). pub async fn ws_upgrade_handler( - axum::extract::State(registry): axum::extract::State>, + sessions: Option>, + axum::extract::State(state): axum::extract::State, axum::Extension(identity): axum::Extension, policy: Option>, ws_upgrade: WebSocketUpgrade, ) -> Response { + let sessions = Some( + sessions + .map(|axum::Extension(s)| s) + .unwrap_or_else(|| WsSessions::clone(state.sessions())), + ); let policy = policy .map(|axum::Extension(p)| p.0) .unwrap_or_else(|| Arc::new(NoCap)); + let registry = Arc::clone(state.registry()); ws_upgrade .max_frame_size(INBOUND_WS_FRAME_CAP) .max_message_size(INBOUND_WS_MESSAGE_CAP) .on_upgrade(move |socket| async move { - run_channels_session(socket, registry, identity, policy).await + run_channels_session(socket, registry, identity, policy, sessions).await }) } diff --git a/tests/ws_upgrade_session.rs b/tests/ws_upgrade_session.rs index 7742eda..04546eb 100644 --- a/tests/ws_upgrade_session.rs +++ b/tests/ws_upgrade_session.rs @@ -609,3 +609,82 @@ async fn inbound_message_over_cap_fails_the_connection() { "expected a close or stream end, got {close:?}" ); } + +/// WS-08 acceptance: the upgrade retains each session's pump handle in +/// the shared `WsSessions` registry; `abort()` evicts a stuck session +/// (the peer's socket closes) and the entry is removed when the +/// session task ends, so the registry only tracks live sessions. +#[tokio::test] +async fn ws_sessions_registry_tracks_and_aborts_live_sessions() { + use alkhttp::websocket::WsSessions; + + let sessions = WsSessions::new(); + let registry = echo_registry(); + let provider = provider_with(vec![("tok-1", identity("alice", &[]))]); + + let app = axum::Router::new() + .route( + "/alk/channels", + axum::routing::get(alkhttp::websocket::ws_upgrade_handler), + ) + .layer(axum::middleware::from_fn_with_state( + Arc::clone(&provider), + alkhttp::websocket::ws_bearer_auth, + )) + .with_state(registry) + .layer(axum::Extension(sessions.clone())); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = format!("ws://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + let env = call_and_await(&mut ws, "req-1", "echo/run", serde_json::json!({})).await; + assert_eq!(env.r#type, EVENT_RESPONDED, "session established"); + + // The upgrade registered the session. + let saw_session = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if !sessions.is_empty() { + return true; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("session registered in the shared registry"); + assert!(saw_session); + assert_eq!(sessions.len(), 1); + + // Forced teardown: abort() closes the peer's socket. + sessions.abort(); + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10); + let close = loop { + assert!( + tokio::time::Instant::now() < deadline, + "abort never reached the peer" + ); + match ws.next_close(std::time::Duration::from_millis(250)).await { + Some(x) => break Some(x), + None => continue, + } + }; + assert!(close.is_some(), "peer observes the teardown: {close:?}"); + + // The self-removing guard drops the entry once the session task ends. + let drained = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if sessions.is_empty() { + return true; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("session entry removed after the session task ends"); + assert!(drained); + ws.close().await; +}