diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index 7c5be69..91452b8 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -655,10 +655,25 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { /// adapter is the single seam between axum's WS and alkcall's /// byte-oriented channels machinery; shared with `from_wss`. /// `idle_timeout` (WS-01): `None` = no idle bound, `Some(d)` closes -/// the read with 1001 after `d` without an inbound WS message. +/// the read with 1001 after `d` without inbound chunk progress. The +/// write pump runs with [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] +/// — see [`split_ws_to_bytes_idle_with_write`] for the WS-18 +/// configurable form. pub fn split_ws_to_bytes_idle( socket: WebSocket, idle_timeout: Option, +) -> (WsByteStream, WsPumps) { + split_ws_to_bytes_idle_with_write(socket, idle_timeout, None) +} + +/// [`split_ws_to_bytes_idle`] with an explicit WS-18 write-progress +/// window: `None` = the crate default +/// ([`DEFAULT_WS_WRITE_TIMEOUT`](crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)), +/// `Some(d)` a deployment-set window. +pub fn split_ws_to_bytes_idle_with_write( + socket: WebSocket, + idle_timeout: Option, + write_timeout: Option, ) -> (WsByteStream, WsPumps) { let (ws_sink, ws_stream) = socket.split(); let (read_tx, read_rx) = mpsc::channel::>(READ_SLOTS); @@ -690,7 +705,7 @@ pub fn split_ws_to_bytes_idle( ws_sink, write_rx, write_error_slot, - Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT), + Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)), )); ( diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs index a8705e5..30e1044 100644 --- a/src/websocket/mod.rs +++ b/src/websocket/mod.rs @@ -13,9 +13,10 @@ pub mod byte_adapter; pub mod upgrade; pub use byte_adapter::{ - split_ws_to_bytes, split_ws_to_bytes_idle, WsByteStream, WsPumps, DEFAULT_WS_IDLE_TIMEOUT, - DEFAULT_WS_WRITE_TIMEOUT, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, MAX_CHUNK_LEN, - PENDING_BUFFER_CAP, WS_GOING_AWAY, WS_MESSAGE_CAP, WS_PROTOCOL_ERROR, + split_ws_to_bytes, split_ws_to_bytes_idle, split_ws_to_bytes_idle_with_write, WsByteStream, + WsPumps, DEFAULT_WS_IDLE_TIMEOUT, DEFAULT_WS_WRITE_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"))] @@ -25,7 +26,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, DEFAULT_WS_MAX_SESSIONS, + WsSessions, WsTimeouts, DEFAULT_WS_MAX_SESSIONS, }; #[cfg(any(test, feature = "test-support"))] diff --git a/src/websocket/upgrade.rs b/src/websocket/upgrade.rs index 1a3df1e..cf93954 100644 --- a/src/websocket/upgrade.rs +++ b/src/websocket/upgrade.rs @@ -51,7 +51,7 @@ use axum::response::{IntoResponse, Response}; use parking_lot::Mutex; use super::byte_adapter::{ - split_ws_to_bytes_idle, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, + split_ws_to_bytes_idle_with_write, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, }; /// Registry of live WS session pump handles (WS-08). The upgrade @@ -196,7 +196,11 @@ impl WsSessions { /// session's pump handle is registered for the session's lifetime — /// 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. +/// closes the read with 1001 after `d` without inbound chunk +/// progress. `write_timeout` bounds one outbound WS send (WS-18): +/// `None` = the crate default window, `Some(d)` a deployment-set +/// window. +#[allow(clippy::too_many_arguments)] pub async fn run_channels_session( socket: axum::extract::ws::WebSocket, registry: Arc, @@ -204,8 +208,10 @@ pub async fn run_channels_session( policy: Arc, sessions: Option, idle_timeout: Option, + write_timeout: Option, ) { - let (byte_stream, pumps) = split_ws_to_bytes_idle(socket, idle_timeout); + let (byte_stream, pumps) = + split_ws_to_bytes_idle_with_write(socket, idle_timeout, write_timeout); let pumps = Arc::new(pumps); let _guard = sessions.map(|sessions| { @@ -316,6 +322,40 @@ impl alkcall::core::auth::IdentityProvider for NoopProvider { #[derive(Clone)] pub struct ChannelsPolicy(pub Arc); +/// Per-request WS pump timeouts (WS-17): a deployment inserts +/// `WsTimeouts` into the request extensions (a route layer on the WS +/// route, mirroring [`ChannelsPolicy`]) to set the pump knobs per +/// route instead of the router-state defaults. `idle` (WS-01) bounds +/// the read staleness (no completed inbound chunk for the window → +/// 1001 eviction); `write` (WS-18) bounds one outbound WS send (a peer +/// that stops reading is evicted once a single send outlasts it); +/// `None` disables a knob. Without the extension the +/// [`SessionState`] values apply — the built-in surface carries +/// `HttpAdapter::with_ws_idle_timeout` / +/// `DEFAULT_WS_IDLE_TIMEOUT` for the read side and +/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] for the write side, +/// which is also what bare-registry routes (custom upgrade routes on +/// a plain `Arc` state) get: 60 s idle + 60 s +/// write windows, 64-session semaphore, handler-private +/// [`WsSessions`]. +#[derive(Clone, Copy, Debug)] +pub struct WsTimeouts { + /// Idle-read window (WS-01); `None` disables the eviction. + pub idle: Option, + /// Write-progress window (WS-18); `None` selects the crate + /// default (`DEFAULT_WS_WRITE_TIMEOUT`). + pub write: Option, +} + +impl Default for WsTimeouts { + fn default() -> Self { + Self { + idle: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT), + write: None, + } + } +} + /// The upgrade handler. Requires the resolved identity in request /// extensions (stashed by [`ws_bearer_auth`]) — a WS session without /// an identity cannot run `AccessControl::check`. @@ -327,6 +367,16 @@ pub struct ChannelsPolicy(pub Arc); /// (`RouterState::ws_sessions`) unless a request extension carries /// one, so a stuck session stays evictable (WS-08). /// +/// The pump timeout knobs (WS-01 read / WS-18 write) come from the +/// [`WsTimeouts`] request extension when present, else from the +/// router state (`HttpAdapter::with_ws_idle_timeout` for the read; +/// the write side is fixed at +/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] on the built-in +/// surface). Bare-registry routes (`Arc` state) +/// carry no configurable state — they run the documented defaults +/// above; a deployment overriding them inserts the `WsTimeouts` +/// extension on its upgrade route. +/// /// 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 @@ -338,6 +388,7 @@ pub async fn ws_upgrade_handler( axum::extract::State(state): axum::extract::State, axum::Extension(identity): axum::Extension, policy: Option>, + timeouts: Option>, ws_upgrade: WebSocketUpgrade, ) -> Response { let Ok(permit) = state.session_slots.clone().try_acquire_owned() else { @@ -356,13 +407,26 @@ 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(); + let idle_timeout = match timeouts { + Some(axum::Extension(t)) => t.idle, + None => state.idle_timeout(), + }; + let write_timeout = timeouts.and_then(|t| t.write); 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, idle_timeout).await + run_channels_session( + socket, + registry, + identity, + policy, + sessions, + idle_timeout, + Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)), + ) + .await }) } diff --git a/tests/ws_upgrade_session.rs b/tests/ws_upgrade_session.rs index dfd6afa..88d5442 100644 --- a/tests/ws_upgrade_session.rs +++ b/tests/ws_upgrade_session.rs @@ -918,3 +918,54 @@ async fn idle_progress_knob_none_disables_eviction_over_axum_upgrade() { ); ws.close().await; } + +/// WS-17 acceptance: a bare-registry upgrade route takes the pump +/// knobs per request via the `WsTimeouts` extension (mirroring +/// `ChannelsPolicy`) — a client that completes no chunk is evicted +/// with 1001 inside the extension's short idle window, not the 60 s +/// default. +#[tokio::test] +async fn ws_timeouts_extension_sets_the_idle_window_on_a_bare_registry_route() { + let registry = std::sync::Arc::new(OperationRegistry::new()); + 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", &[])) + } + } + let app = axum::Router::new() + .route( + "/alk/channels", + axum::routing::get(alkhttp::websocket::ws_upgrade_handler), + ) + .layer(axum::middleware::from_fn_with_state( + std::sync::Arc::new(StaticTok) as std::sync::Arc, + alkhttp::websocket::ws_bearer_auth, + )) + .layer(axum::Extension(alkhttp::websocket::WsTimeouts { + idle: Some(std::time::Duration::from_millis(150)), + write: None, + })) + .with_state(registry); + 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 close = ws.next_close(std::time::Duration::from_secs(5)).await; + assert_eq!( + close, + Some(Some(alkhttp::websocket::WS_GOING_AWAY)), + "evicted with 1001 inside the extension's idle window" + ); + ws.close().await; +}