From 25975ac2a89852cb47e061c63a99866c60642fe0 Mon Sep 17 00:00:00 2001 From: "glm-5.3-flash" Date: Sat, 29 Aug 2026 13:45:54 +0000 Subject: [PATCH] fix(websocket): configurable WS read idle timeout (WS-01) --- src/gateway/routes.rs | 1 + src/server/adapter.rs | 30 ++++++ src/server/state.rs | 4 + src/websocket/byte_adapter.rs | 195 ++++++++++++++++++++++++++++++++-- src/websocket/mod.rs | 5 +- src/websocket/upgrade.rs | 21 +++- 6 files changed, 241 insertions(+), 15 deletions(-) diff --git a/src/gateway/routes.rs b/src/gateway/routes.rs index 2cd9da8..418ccdc 100644 --- a/src/gateway/routes.rs +++ b/src/gateway/routes.rs @@ -861,6 +861,7 @@ mod tests { ws_session_slots: Arc::new(tokio::sync::Semaphore::new( crate::websocket::DEFAULT_WS_MAX_SESSIONS, )), + ws_idle_timeout: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT), }; let auth_state = Arc::clone(&provider); gateway_router() diff --git a/src/server/adapter.rs b/src/server/adapter.rs index 6ee6b1d..43752d2 100644 --- a/src/server/adapter.rs +++ b/src/server/adapter.rs @@ -90,6 +90,7 @@ pub struct HttpAdapter { ws_sessions: Arc, ws_max_sessions: usize, ws_session_slots: Arc, + ws_idle_timeout: Option, } impl HttpAdapter { @@ -117,6 +118,7 @@ impl HttpAdapter { 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 ws_idle_timeout = Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT); let state = RouterState { registry: Arc::clone(®istry), identity_provider: Arc::clone(&identity_provider), @@ -124,6 +126,7 @@ impl HttpAdapter { openapi_doc: openapi_doc.clone(), ws_sessions: Arc::clone(&ws_sessions), ws_session_slots: Arc::clone(&ws_session_slots), + ws_idle_timeout, }; let router = build_router(state, None); Self { @@ -137,6 +140,7 @@ impl HttpAdapter { ws_sessions, ws_max_sessions, ws_session_slots, + ws_idle_timeout, } } @@ -149,6 +153,7 @@ impl HttpAdapter { openapi_doc: self.openapi_doc.clone(), ws_sessions: Arc::clone(&self.ws_sessions), ws_session_slots: Arc::clone(&self.ws_session_slots), + ws_idle_timeout: self.ws_idle_timeout, }; // `extra_routes` is borrowed, not consumed (SRV-05): a builder // call after `with_extra_routes` must keep the custom routes in @@ -166,6 +171,7 @@ impl HttpAdapter { openapi_doc: self.openapi_doc.clone(), ws_sessions: Arc::clone(&self.ws_sessions), ws_session_slots: Arc::clone(&self.ws_session_slots), + ws_idle_timeout: self.ws_idle_timeout, }; self.router = build_router(state, Some(routes.clone())); self.extra_routes = Some(routes); @@ -188,6 +194,29 @@ impl HttpAdapter { openapi_doc: self.openapi_doc.clone(), ws_sessions: Arc::clone(&self.ws_sessions), ws_session_slots: Self::rebuild_session_slots(max_sessions), + ws_idle_timeout: self.ws_idle_timeout, + }; + self.router = build_router(state, self.extra_routes.clone()); + self + } + + /// The WS idle-read timeout (WS-01): the read pump closes the + /// connection with a 1001 (GoingAway) close frame after this long + /// without an inbound WS message — bounding the demux stall a + /// dribbling (or silently-stalled) peer can pin. `None` disables + /// the knob (not recommended: the stall window is then unbounded). + /// + /// Default: [`crate::websocket::DEFAULT_WS_IDLE_TIMEOUT`] (60 s). + pub fn with_ws_idle_timeout(mut self, idle_timeout: Option) -> Self { + self.ws_idle_timeout = idle_timeout; + 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: Arc::clone(&self.ws_session_slots), + ws_idle_timeout: self.ws_idle_timeout, }; self.router = build_router(state, self.extra_routes.clone()); self @@ -995,6 +1024,7 @@ mod tests { ws_session_slots: Arc::new(tokio::sync::Semaphore::new( crate::websocket::DEFAULT_WS_MAX_SESSIONS, )), + ws_idle_timeout: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT), } } diff --git a/src/server/state.rs b/src/server/state.rs index 5b1251d..1acdc9a 100644 --- a/src/server/state.rs +++ b/src/server/state.rs @@ -40,6 +40,8 @@ pub(crate) struct RouterState { /// 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, + /// Idle-read timeout for the WS pumps (WS-01); `None` disables. + pub(crate) ws_idle_timeout: Option, } impl axum::extract::FromRef for crate::websocket::SessionState { @@ -48,6 +50,7 @@ impl axum::extract::FromRef for crate::websocket::SessionState { Arc::clone(&state.registry), Arc::clone(&state.ws_sessions), Arc::clone(&state.ws_session_slots), + state.ws_idle_timeout, ) } } @@ -98,6 +101,7 @@ mod tests { ws_session_slots: Arc::new(tokio::sync::Semaphore::new( crate::websocket::DEFAULT_WS_MAX_SESSIONS, )), + ws_idle_timeout: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT), }; let extracted: DecoyConfig = axum::extract::FromRef::from_ref(&state); assert!(matches!(extracted, DecoyConfig::Redirect { .. })); diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index 803b1ad..1f42a66 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -47,6 +47,7 @@ use std::{ pin::Pin, sync::{Arc, Mutex as StdMutex}, task::{Context, Poll}, + time::Duration, }; use axum::extract::ws::{CloseFrame, Message as AxumMessage, WebSocket}; @@ -106,6 +107,18 @@ pub const WS_PROTOCOL_ERROR: u16 = 1002; /// violation it cannot recover from (WS-04/HY-09, WS-05). pub const WS_INTERNAL_ERROR: u16 = 1011; +/// Idle-read close code (WS-01): a read side that stays silent past the +/// configured idle timeout is closed with 1001 (Going Away) — a normal +/// connection end from the demux's point of view (EOF → channels +/// cleared, pendings failed), not a protocol error. +pub const WS_GOING_AWAY: u16 = 1001; + +/// Default idle-read timeout for the WS pumps (WS-01): a peer that +/// drips bytes (or a stalled partial chunk header) cannot park the +/// single demux loop longer than this without traffic. Zero disables +/// the timeout. Deployment knob (`HttpAdapter::with_ws_idle_timeout`). +pub const DEFAULT_WS_IDLE_TIMEOUT: Duration = Duration::from_secs(60); + /// Maximum chunk payload length — the channels protocol's 16 MiB wire /// bound (ADR-052 §5), re-exported from alkcall. The write-side chunk /// parser rejects a header claiming a longer payload instead of waiting @@ -190,11 +203,22 @@ where /// feature fires the lossless EOF watch signal through it — the /// from_wss drop monitor's input; its `Sender` semantics are preserved /// EXACTLY: single `send` after the read loop terminates). +/// +/// Idle-read timeout (WS-01): when `idle_timeout` is non-zero, the +/// next-message await is wrapped in a `tokio::time::timeout` — a peer +/// that dribbles (or a stalled header) cannot park the single demux +/// loop beyond the knob. Staleness sends the 1001 GoingAway close to +/// the peer (a normal connection end for the demux/from_wss EOF +/// machinery, not a protocol error) and ends the loop, so the +/// adapter's read half sees EOF and the channels teardown runs. +/// Messages that *do* arrive reset the window; a partial chunk header +/// alone does not (the boundary is a complete WS message). async fn run_read_pump( mut ws_stream: S, read_tx: mpsc::Sender>, write_tx_for_read: futures_mpsc::Sender, write_error_slot_for_read: WriteErrorSlot, + idle_timeout: Option, on_end: F, ) where S: futures::Stream + Unpin, @@ -202,7 +226,27 @@ async fn run_read_pump( S::Item: IntoWsResult, F: FnOnce(), { - while let Some(msg) = ws_stream.next().await { + loop { + let msg = match idle_timeout { + None => ws_stream.next().await, + Some(timeout) => match tokio::time::timeout(timeout, ws_stream.next()).await { + Ok(msg) => msg, + Err(_elapsed) => { + send_write_error( + &write_error_slot_for_read, + "connection idle past the read timeout", + ); + let _ = write_tx_for_read + .clone() + .send(WriteMsg::CloseWith(WS_GOING_AWAY)) + .await; + break; + } + }, + }; + let Some(msg) = msg else { + break; + }; match msg.into_ws_result() { Ok(m) => { if let Some(b) = M::binary(&m) { @@ -409,10 +453,23 @@ impl WsFraming for AxumFraming { } } +/// Split a `WebSocket` into the byte stream + the pump tasks with the +/// default idle-read timeout +/// ([`DEFAULT_WS_IDLE_TIMEOUT`](crate::websocket::DEFAULT_WS_IDLE_TIMEOUT)). +/// See [`split_ws_to_bytes_idle`] for the configurable form. +pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { + split_ws_to_bytes_idle(socket, Some(DEFAULT_WS_IDLE_TIMEOUT)) +} + /// Split a `WebSocket` into the byte stream + the pump tasks. The /// adapter is the single seam between axum's WS and alkcall's /// byte-oriented channels machinery; shared with `from_wss`. -pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { +/// `idle_timeout` (WS-01): `None` = no idle bound, `Some(d)` closes +/// the read with 1001 after `d` without an inbound WS message. +pub fn split_ws_to_bytes_idle( + socket: WebSocket, + idle_timeout: Option, +) -> (WsByteStream, WsPumps) { let (ws_sink, ws_stream) = socket.split(); let (read_tx, read_rx) = mpsc::channel::>(READ_SLOTS); let (write_tx, write_rx) = futures_mpsc::channel::(WRITE_SLOTS); @@ -430,6 +487,7 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { read_tx, write_tx_for_read, write_error_slot_for_read, + idle_timeout, move || { #[cfg(any(test, feature = "wss"))] { @@ -642,17 +700,34 @@ impl WsFraming for TungsteniteFraming { } } -/// Client-side split for a tokio-tungstenite `WebSocketStream` -/// (`from_wss`, ADR-070): the same WS↔byte-stream seam as -/// [`split_ws_to_bytes`], over the tungstenite socket instead of axum's -/// server-side `WebSocket`. Message semantics are identical: binary -/// messages carry the byte stream, text is a protocol error (close -/// 1002), close → read EOF. Both paths run the same generic pumps -/// (WS-11); the socket flavor differs only in the [`WsFraming`] impl. +/// Client-side split for a tokio-tungstenite `WebSocketStream` with +/// the default idle-read timeout +/// ([`DEFAULT_WS_IDLE_TIMEOUT`](crate::websocket::DEFAULT_WS_IDLE_TIMEOUT)). +/// See [`split_tungstenite_to_bytes_idle`] for the configurable form. #[cfg(any(test, feature = "wss"))] pub fn split_tungstenite_to_bytes( socket: tokio_tungstenite::WebSocketStream, ) -> (WsByteStream, WsPumps) +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static, +{ + split_tungstenite_to_bytes_idle(socket, Some(DEFAULT_WS_IDLE_TIMEOUT)) +} + +/// Client-side split for a tokio-tungstenite `WebSocketStream` +/// (`from_wss`, ADR-070): the same WS↔byte-stream seam as +/// [`split_ws_to_bytes_idle`], over the tungstenite socket instead of +/// axum's server-side `WebSocket`. Message semantics are identical: +/// binary messages carry the byte stream, text is a protocol error +/// (close 1002), close → read EOF. Both paths run the same generic +/// pumps (WS-11); the socket flavor differs only in the [`WsFraming`] +/// impl. `idle_timeout` (WS-01): `None` = no idle bound, `Some(d)` +/// closes the read with 1001 after `d` without an inbound WS message. +#[cfg(any(test, feature = "wss"))] +pub fn split_tungstenite_to_bytes_idle( + socket: tokio_tungstenite::WebSocketStream, + idle_timeout: Option, +) -> (WsByteStream, WsPumps) where S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static, { @@ -671,6 +746,7 @@ where read_tx, write_tx_for_read, write_error_slot_for_read, + idle_timeout, move || { let _ = read_eof_for_task.send(true); }, @@ -1025,4 +1101,105 @@ mod tests { .expect("peer observes the close reply after our shutdown"); assert!(saw_close, "the shutdown path emitted a WS Close frame"); } + + /// WS-01 acceptance: a dribbling peer cannot park the read pump + /// past the knob. The idle-read timeout fires with no inbound + /// message inside the window: the pump sends the 1001 GoingAway + /// close to the peer, ends, and fires the EOF watch signal — the + /// from_wss monitor/sweep input — so the demux/channel teardown + /// machinery treats it as a normal connection end. + #[tokio::test] + async fn idle_read_timeout_closes_a_stalled_connection_with_goingaway() { + let (client_io, server_io) = tokio::io::duplex(1 << 16); + let ws = tokio_tungstenite::WebSocketStream::from_raw_socket( + client_io, + tokio_tungstenite::tungstenite::protocol::Role::Client, + None, + ) + .await; + let knob = std::time::Duration::from_millis(150); + let (_stream, pumps) = split_tungstenite_to_bytes_idle(ws, Some(knob)); + + let peer = tokio_tungstenite::WebSocketStream::from_raw_socket( + server_io, + tokio_tungstenite::tungstenite::protocol::Role::Server, + None, + ) + .await; + let (_peer_sink, mut peer_stream) = peer.split(); + use futures::StreamExt; + + // The peer sends nothing (the stall). The read pump must end + // within the knob (+ slack), not hang: observe both the EOF + // signal the from_wss machinery consumes and the 1001 close + // frame the peer receives. + let mut eof_rx = pumps.read_eof(); + let (close_seen, eof_fired) = + tokio::time::timeout(std::time::Duration::from_secs(5), async { + let close: Option = loop { + match peer_stream.next().await { + Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => { + break cf.map(|f| u16::from(f.code)) + } + Some(Ok(_)) => continue, + Some(Err(_)) | None => break None, + } + }; + let _ = eof_rx.changed().await; + (close, *eof_rx.borrow()) + }) + .await + .expect("stalled connection must be torn down within the deadline"); + + assert!( + close_seen == Some(WS_GOING_AWAY), + "peer received the 1001 GoingAway close, got {close_seen:?}" + ); + assert!(eof_fired, "EOF signal fired for the from_wss machinery"); + } + + /// The idle knob resets on traffic: a peer dribbling one message + /// per window (each inside the knob) keeps the connection open — + /// only the *stall* (no message within one full window) closes. + #[tokio::test] + async fn idle_read_timeout_resets_on_traffic() { + let (client_io, server_io) = tokio::io::duplex(1 << 16); + let ws = tokio_tungstenite::WebSocketStream::from_raw_socket( + client_io, + tokio_tungstenite::tungstenite::protocol::Role::Client, + None, + ) + .await; + let knob = std::time::Duration::from_millis(200); + let (mut stream, _pumps) = split_tungstenite_to_bytes_idle(ws, Some(knob)); + + let peer = tokio_tungstenite::WebSocketStream::from_raw_socket( + server_io, + tokio_tungstenite::tungstenite::protocol::Role::Server, + None, + ) + .await; + let (mut peer_sink, mut peer_stream) = peer.split(); + use futures::{SinkExt, StreamExt}; + + // Dribble: a message every 100 ms (half the knob) for 8 rounds — + // 1.6 s total, far past a single 200 ms window. The connection + // must stay open; every message surfaces on the adapter. + let mut buf = [0u8; 1]; + for round in 0u8..8 { + tokio::time::sleep(knob / 2).await; + peer_sink + .send(tokio_tungstenite::tungstenite::Message::Binary( + vec![round].into(), + )) + .await + .expect("dribble write accepted"); + tokio::time::timeout(knob, stream.read_exact(&mut buf)) + .await + .expect("connection alive across the dribble") + .expect("read"); + assert_eq!(buf[0], round); + } + let _ = peer_stream.next().await; + } } diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs index 384df9c..637ffac 100644 --- a/src/websocket/mod.rs +++ b/src/websocket/mod.rs @@ -13,8 +13,9 @@ pub mod byte_adapter; pub mod upgrade; pub use byte_adapter::{ - split_ws_to_bytes, WsByteStream, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, - MAX_CHUNK_LEN, PENDING_BUFFER_CAP, WS_MESSAGE_CAP, WS_PROTOCOL_ERROR, + split_ws_to_bytes, split_ws_to_bytes_idle, WsByteStream, WsPumps, DEFAULT_WS_IDLE_TIMEOUT, + INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, MAX_CHUNK_LEN, PENDING_BUFFER_CAP, WS_GOING_AWAY, + WS_MESSAGE_CAP, WS_PROTOCOL_ERROR, }; #[cfg(any(test, feature = "wss"))] diff --git a/src/websocket/upgrade.rs b/src/websocket/upgrade.rs index 73dc845..b339dc1 100644 --- a/src/websocket/upgrade.rs +++ b/src/websocket/upgrade.rs @@ -25,7 +25,7 @@ use axum::response::{IntoResponse, Response}; use parking_lot::Mutex; use super::byte_adapter::{ - split_ws_to_bytes, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, + split_ws_to_bytes_idle, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, }; /// Registry of live WS session pump handles (WS-08). The upgrade @@ -66,6 +66,8 @@ pub struct SessionState { /// Session cap (WS-09): acquired post-auth, pre-upgrade; a caller /// over the cap is rejected with 503. session_slots: Arc, + /// Idle-read timeout (WS-01): `None` disables the knob. + idle_timeout: Option, } impl SessionState { @@ -78,6 +80,7 @@ impl SessionState { registry: Arc::clone(registry), sessions: Arc::new(WsSessions::new()), session_slots: Arc::new(tokio::sync::Semaphore::new(DEFAULT_WS_MAX_SESSIONS)), + idle_timeout: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT), } } @@ -89,11 +92,13 @@ impl SessionState { registry: Arc, sessions: Arc, session_slots: Arc, + idle_timeout: Option, ) -> Self { Self { registry, sessions, session_slots, + idle_timeout, } } @@ -104,6 +109,10 @@ impl SessionState { pub(crate) fn sessions(&self) -> &Arc { &self.sessions } + + pub(crate) fn idle_timeout(&self) -> Option { + self.idle_timeout + } } impl axum::extract::FromRef for Arc { @@ -157,15 +166,18 @@ impl WsSessions { /// (identity attached) → `ChannelsAdapter::handle`. `policy` gates /// 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`]). +/// the WS-08 eviction lever ([`WsSessions::abort`]). `idle_timeout` +/// bounds the read stall (WS-01): `None` disables the knob, `Some(d)` +/// closes the read with 1001 after `d` without an inbound WS message. pub async fn run_channels_session( socket: axum::extract::ws::WebSocket, registry: Arc, identity: Identity, policy: Arc, sessions: Option, + idle_timeout: Option, ) { - let (byte_stream, pumps) = split_ws_to_bytes(socket); + let (byte_stream, pumps) = split_ws_to_bytes_idle(socket, idle_timeout); let pumps = Arc::new(pumps); let _guard = sessions.map(|sessions| { @@ -316,12 +328,13 @@ pub async fn ws_upgrade_handler( .map(|axum::Extension(p)| p.0) .unwrap_or_else(|| Arc::new(NoCap)); let registry = Arc::clone(state.registry()); + let idle_timeout = state.idle_timeout(); ws_upgrade .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 + run_channels_session(socket, registry, identity, policy, sessions, idle_timeout).await }) }