diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index 0ac4bb8..319b23e 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -23,9 +23,18 @@ //! alkcall's `MAX_CHUNK_LEN` (WS-04/HY-09): a header claiming more //! fails the stream loudly (close 1011 + an error to the `AsyncWrite` //! half) instead of silently waiting to accumulate up to ~4 GiB from a -//! misaligned offset. The `pending` accumulator is byte-capped at -//! `PENDING_BUFFER_CAP` (WS-05): exceeding it fails the stream the same -//! way, replacing the unbounded ~64 × 16 MiB worst case. Write-side +//! misaligned offset. Single `AsyncWrite` calls above +//! `PENDING_BUFFER_CAP` (WS-14) are rejected in `poll_write` *before* +//! the bytes enter the write queue — the mux sees the `InvalidData` +//! stream error directly instead of the write pump closing the wire +//! with 1011 mid-stream and truncating. The `pending` accumulator +//! still carries the `PENDING_BUFFER_CAP` bound (WS-05) as +//! defense-in-depth: it is invariantly satisfied once `poll_write` +//! enforces the cap pre-queue, so it holds at most one +//! unemitted-in-full chunk plus its header. The remaining write-side +//! bounds are slot-bounded (`WRITE_SLOTS` × `WS_MESSAGE_CAP`) but +//! time-unbounded — bounded by the WS-18 write-progress timeout in +//! `run_write_pump`. Write-side //! backpressure uses `futures::channel::mpsc` `poll_ready` — the //! production fix for the POC's spin-wait. //! @@ -113,17 +122,15 @@ pub const INBOUND_WS_MESSAGE_CAP: usize = 1024 * 1024; /// messages outright instead of bounding them. pub const INBOUND_WS_FRAME_CAP: usize = 1024 * 1024; -/// Byte cap on the write-side `pending` accumulator (WS-05): the -/// slot-bounded (64) `WriteMsg` channel bounds messages, not bytes, and -/// `pending` is the only unbounded accumulator on the write path. Well-framed -/// traffic can never occupy more than `MAX_CHUNK_LEN + 8` bytes here — the -/// parser drains complete chunks greedily, so `pending` holds at most one -/// partially-received (or unemitted) chunk plus its header — and alkcall's -/// mux passes each payload (up to 16 MiB) as one `AsyncWrite` call, so a -/// smaller cap would false-positive on legitimate writes. Exceeding the cap -/// (before extending, or after a full parse pass) proves the byte stream is -/// not chunk-framed and fails the stream loudly (`InvalidData`), replacing -/// the unbounded ~64 × 16 MiB worst case. +/// Byte cap on the write-side `pending` accumulator (WS-05). Enforced +/// in `poll_write` *before* a message enters the write queue (WS-14): +/// a single over-cap call is rejected with `InvalidData` at the +/// `AsyncWrite` boundary, so the invariant — `pending` never holds +/// more than one not-yet-emitted chunk plus its header, because the +/// parser drains complete chunks greedily and alkcall's mux passes +/// each payload (up to 16 MiB) as one `AsyncWrite` call — is +/// maintained by construction. The in-pump check remains as +/// defense-in-depth. pub const PENDING_BUFFER_CAP: usize = MAX_CHUNK_LEN as usize + 8; /// Protocol-error close code for text messages (websocket.md §Framing). @@ -140,6 +147,14 @@ pub const WS_INTERNAL_ERROR: u16 = 1011; /// failed), not a protocol error. pub const WS_GOING_AWAY: u16 = 1001; +/// Default write-progress timeout for the WS write pump (WS-18): one +/// outbound WS send that stays unsent past this window (the peer +/// stopped reading) evicts the connection — the write-side analog of +/// the idle-read knob. The bound is per `send` call, not per +/// connection: a slow-but-draining peer resets it with every message +/// that gets out. +pub const DEFAULT_WS_WRITE_TIMEOUT: Duration = Duration::from_secs(60); + /// Default idle-read timeout for the WS pumps (WS-01, WS-13): a /// connection whose inbound stream completes **no chunk** within this /// window is evicted — the deadline resets on demux progress (complete @@ -216,8 +231,9 @@ pub use alkcall::channels::wire::MAX_CHUNK_LEN; pub(crate) enum WriteMsg { Bytes(Vec), - /// Close with the given code (e.g. the 1002 text-rejection). - CloseWith(u16), + /// Close with the given code and reason (e.g. the 1002 + /// text-rejection). + CloseWith(u16, &'static str), } /// The write-pump failure channel (WS-04/HY-09, WS-05): a shared slot @@ -241,6 +257,20 @@ fn send_write_error(slot: &WriteErrorSlot, reason: &'static str) { } } +/// The close-frame reason strings, one distinct per cause (WS-15): the +/// numeric code alone collides across causes (the size-cap close and +/// the write-stall close share 1011), and an operator reading a wire +/// capture should be able to tell them apart without the stream error +/// channel. Sent in both the close frame the peer receives and the +/// stream-error diagnostic; `reason_for` maps a bare close code to its +/// canonical default for the fall-through arms. +pub(crate) mod close_reason { + pub(crate) const IDLE_READ_TIMEOUT: &str = + "connection made no inbound chunk progress past the read timeout"; + pub(crate) const TEXT_NOT_SUPPORTED: &str = "text messages not supported"; + pub(crate) const INBOUND_FRAME_REJECTED: &str = "inbound frame rejected (size cap)"; +} + /// WS-protocol adapter for one socket flavor: the only places the axum /// and tungstenite message types differ. The generic pump /// ([`run_read_pump`] / [`run_write_pump`]) is written against this @@ -328,13 +358,13 @@ async fn run_read_pump( Some(budget) => match tokio::time::timeout(budget, ws_stream.next()).await { Ok(msg) => msg, Err(_elapsed) => { - send_write_error( - &write_error_slot_for_read, - "connection made no inbound chunk progress past the read timeout", - ); + send_write_error(&write_error_slot_for_read, close_reason::IDLE_READ_TIMEOUT); let _ = write_tx_for_read .clone() - .send(WriteMsg::CloseWith(WS_GOING_AWAY)) + .send(WriteMsg::CloseWith( + WS_GOING_AWAY, + close_reason::IDLE_READ_TIMEOUT, + )) .await; break; } @@ -356,7 +386,10 @@ async fn run_read_pump( } else if M::is_text(&m) { let _ = write_tx_for_read .clone() - .send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR)) + .send(WriteMsg::CloseWith( + WS_PROTOCOL_ERROR, + close_reason::TEXT_NOT_SUPPORTED, + )) .await; break; } else if M::is_close(&m) { @@ -366,11 +399,14 @@ async fn run_read_pump( Err(()) => { send_write_error( &write_error_slot_for_read, - "inbound frame rejected (size cap)", + close_reason::INBOUND_FRAME_REJECTED, ); let _ = write_tx_for_read .clone() - .send(WriteMsg::CloseWith(WS_INTERNAL_ERROR)) + .send(WriteMsg::CloseWith( + WS_INTERNAL_ERROR, + close_reason::INBOUND_FRAME_REJECTED, + )) .await; break; } @@ -395,17 +431,38 @@ async fn fail_write_pump( .await; } +/// The write pump: drains the queued byte spans, parses the pending +/// bytes for complete chunks (8-byte header, length validated against +/// `MAX_CHUNK_LEN` — WS-04/HY-09), byte-caps the accumulator +/// (`PENDING_BUFFER_CAP` — WS-05), and emits one WS binary message per +/// drains; `CloseWith` requests (`WriteMsg::CloseWith`) map to a close +/// frame with the requested code and its per-cause reason (WS-15) — +/// the code alone is ambiguous across causes (idle 1001 vs protocol +/// 1002 vs internal 1011 arms), so the pump never invents one. /// The write pump: drains the queued byte spans, parses the pending /// bytes for complete chunks (8-byte header, length validated against /// `MAX_CHUNK_LEN` — WS-04/HY-09), byte-caps the accumulator /// (`PENDING_BUFFER_CAP` — WS-05), and emits one WS binary message per /// `WS_MESSAGE_CAP` piece. Ends with a WS Close after the queue /// drains; `CloseWith` requests (`WriteMsg::CloseWith`) map to a close -/// frame with that code. +/// frame with the requested code and its per-cause reason (WS-15) — +/// the code alone is ambiguous across causes (idle 1001 vs protocol +/// 1002 vs internal 1011 arms), so the pump never invents one. +/// +/// Write-progress timeout (WS-18): a peer that stops reading parks +/// the pump inside a WS send — the slot-bounded queue bounds memory, +/// but the stall was time-unbounded. When `write_timeout` is +/// `Some(window)`, every WS send must complete within the window: +/// exceeding it signals the stream error (naming the stall) and ends +/// the pump **without** a close frame — the peer cannot receive one +/// through a clogged socket, and the close send would park on it too. +/// The window is per send, so a slow-but-moving peer survives; only a +/// fully stalled sink trips it. async fn run_write_pump( mut ws_sink: futures::stream::SplitSink::Msg>, mut write_rx: futures_mpsc::Receiver, write_error_slot: WriteErrorSlot, + write_timeout: Option, ) where S: futures::Sink<::Msg> + Unpin, M: WsFraming, @@ -413,25 +470,20 @@ async fn run_write_pump( let mut pending: Vec = Vec::new(); while let Some(msg) = write_rx.next().await { match msg { - WriteMsg::CloseWith(code) => { - let _ = ws_sink - .send(::close_message( - code, - "text messages not supported", - )) - .await; + WriteMsg::CloseWith(code, reason) => { + let close = ::close_message(code, reason); + match write_timeout { + None => { + let _ = ws_sink.send(close).await; + } + Some(window) => { + let _ = tokio::time::timeout(window, ws_sink.send(close)).await; + } + } break; } WriteMsg::Bytes(b) => { - if b.len() > PENDING_BUFFER_CAP { - fail_write_pump::( - &mut ws_sink, - &write_error_slot, - "write pending buffer exceeded cap", - ) - .await; - return; - } + debug_assert!(b.len() <= PENDING_BUFFER_CAP); pending.extend_from_slice(&b); } } @@ -456,10 +508,8 @@ async fn run_write_pump( } let chunk: Vec = pending.drain(..total).collect(); for piece in chunk.chunks(WS_MESSAGE_CAP) { - if ws_sink - .send(::binary_message(piece.to_vec())) + if !send_bounded::(&mut ws_sink, piece, write_timeout, &write_error_slot) .await - .is_err() { return; } @@ -478,6 +528,45 @@ async fn run_write_pump( let _ = ws_sink.close().await; } +/// One WS send under the WS-18 write-progress bound: a send that +/// outlasts the window means the peer stopped reading, so the pump +/// signals the error slot (`InvalidData` on the `AsyncWrite` half) and +/// ends without a close frame — a clogged socket cannot receive one, +/// and the close send would park on the same stall. On send error, end +/// silently (the sink is already broken). Returns whether the pump +/// should continue. +async fn send_bounded( + ws_sink: &mut futures::stream::SplitSink::Msg>, + piece: &[u8], + write_timeout: Option, + slot: &WriteErrorSlot, +) -> bool +where + M: WsFraming, + S: futures::Sink<::Msg> + Unpin, +{ + match write_timeout { + None => ws_sink + .send(::binary_message(piece.to_vec())) + .await + .is_ok(), + Some(window) => { + let send = ws_sink.send(::binary_message(piece.to_vec())); + match tokio::time::timeout(window, send).await { + Ok(Ok(())) => true, + Ok(Err(_)) => false, + Err(_elapsed) => { + send_write_error( + slot, + "connection made no write progress past the write timeout", + ); + false + } + } + } + } +} + /// The `AsyncRead + AsyncWrite` view of a `WebSocket` handed to /// alkcall's channels machinery (demux/mux). pub struct WsByteStream { @@ -566,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); @@ -601,6 +705,7 @@ pub fn split_ws_to_bytes_idle( ws_sink, write_rx, write_error_slot, + Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)), )); ( @@ -675,12 +780,22 @@ impl AsyncWrite for WsByteStream { if let Some(err) = this.poll_write_error() { return Poll::Ready(Err(err)); } - let Some(write_tx) = this.write_tx_mut() else { + let Some(write_tx) = this.write_tx_ref() else { return Poll::Ready(Err(io::Error::new( io::ErrorKind::BrokenPipe, "ws stream shut down", ))); }; + if buf.len() > PENDING_BUFFER_CAP { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "write of {} bytes exceeds the pending buffer cap ({})", + buf.len(), + PENDING_BUFFER_CAP + ), + ))); + } match write_tx.poll_ready(cx) { Poll::Ready(Ok(())) => match write_tx.try_send(WriteMsg::Bytes(buf.to_vec())) { Ok(()) => Poll::Ready(Ok(buf.len())), @@ -754,11 +869,13 @@ impl WsByteStream { } } - /// The write pump's fatal error, if any, and whether the write - /// queue still accepts sends (WS-07): the stream takes the sender - /// at shutdown so the pump-side channel close (and the sink's - /// trailing `ws_sink.close()`) actually happens. - fn write_tx_mut(&mut self) -> Option<&mut futures_mpsc::Sender> { + /// The write-sender getter (WS-07): the stream takes the sender at + /// shutdown so the pump-side channel close (and the sink's trailing + /// `ws_sink.close()`) actually happens. WS-14's pre-send cap check + /// keeps the cap fault off the wire path entirely — a `try_send` + /// failure here can only be the pump ending between `poll_ready` + /// and `try_send`. + fn write_tx_ref(&mut self) -> Option<&mut futures_mpsc::Sender> { self.write_tx.as_mut() } } @@ -857,6 +974,7 @@ where ws_sink, write_rx, write_error_slot, + Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT), )); ( @@ -879,15 +997,120 @@ where ) } +/// Test-only split with both knobs explicit (WS-18 acceptance: the +/// stall test scales the write window down; other tests keep the +/// default). +#[cfg(test)] +pub(crate) fn split_tungstenite_to_bytes_idle_with_write( + socket: tokio_tungstenite::WebSocketStream, + idle_timeout: Option, + write_timeout: Option, +) -> (WsByteStream, WsPumps) +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static, +{ + 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); + let (write_error_slot, write_error_rx) = make_write_error_slot(); + + let read_eof = tokio::sync::watch::channel(false).0; + + let write_tx_for_read = write_tx.clone(); + let write_error_slot_for_read = Arc::clone(&write_error_slot); + let read_eof_for_task = read_eof.clone(); + let read_task = tokio::spawn(run_read_pump::( + ws_stream, + read_tx, + write_tx_for_read, + write_error_slot_for_read, + idle_timeout, + move || { + let _ = read_eof_for_task.send(true); + }, + )); + + let write_task = tokio::spawn(run_write_pump::( + ws_sink, + write_rx, + write_error_slot, + write_timeout, + )); + + ( + WsByteStream { + read_rx, + read_buf: Vec::new(), + read_pos: 0, + eof: false, + write_tx: Some(write_tx), + write_open: true, + write_error: write_error_rx, + write_failed: false, + }, + WsPumps { + read_task, + write_task, + read_eof, + }, + ) +} + #[cfg(test)] mod tests { use super::*; use tokio::io::{AsyncReadExt, AsyncWriteExt}; + /// WS-18 acceptance: a peer that stops reading (the sink clogs — + /// nothing is drained from the duplex) parks the write pump inside + /// a WS send; the pump must end within the knob (+ slack), surface + /// the stream error naming the write stall, and fail the + /// `AsyncWrite` half with `InvalidData`. Scaled test: both knobs + /// idle-read `None` (so only the write knob can evict) and the + /// write window at 150 ms. + #[tokio::test] + async fn tungstenite_write_stall_is_evicted_within_the_write_timeout() { + let (client_io, _server_io) = tokio::io::duplex(64); + 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 (mut stream, _pumps) = split_tungstenite_to_bytes_idle_with_write(ws, None, Some(knob)); + + let mut chunk = vec![0u8; 8]; + chunk[4..8].copy_from_slice(&64u32.to_be_bytes()); + chunk.extend_from_slice(&[0u8; 64]); + stream.write_all(&chunk).await.expect("chunk written"); + stream.flush().await.expect("flush"); + + let started = tokio::time::Instant::now(); + let err = loop { + match stream.write_all(&[0u8; 8]).await { + Err(e) => break e, + Ok(_) => assert!( + started.elapsed() < std::time::Duration::from_secs(5), + "stream never failed while the sink stayed clogged" + ), + } + }; + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!( + err.to_string().contains("write timeout"), + "error names the violation: {err}" + ); + assert!( + started.elapsed() < std::time::Duration::from_secs(5), + "eviction must happen within the knob (+ slack), not hang" + ); + } + /// An 8-byte bogus chunk header claiming `MAX_CHUNK_LEN + 1` /// payload bytes (WS-04/HY-09 acceptance: the parser must fail the /// stream loudly instead of silently waiting for ~4 GiB). - fn oversize_header() -> Vec { + pub(crate) fn oversize_header() -> Vec { let mut header = vec![0u8; 8]; header[4..8].copy_from_slice(&(MAX_CHUNK_LEN + 1).to_be_bytes()); header @@ -933,15 +1156,14 @@ mod tests { ); } - /// Write more than `PENDING_BUFFER_CAP` (= `MAX_CHUNK_LEN + 8`) in - /// one `AsyncWrite` call with the pump's sink a live (but - /// unserviced) duplex: well-framed channel traffic can never - /// produce a single write that large (alkcall caps one payload at - /// `MAX_CHUNK_LEN`), so the cap must trip inside the pump and the - /// stream fail with `InvalidData` (WS-05 acceptance). + /// A single `AsyncWrite` call of exactly `PENDING_BUFFER_CAP` + /// bytes parses and flushes through the pump — the cap check + /// rejects `> cap`, never `= cap`, so `MAX_CHUNK_LEN` payloads + /// (the mux's maximum single `AsyncWrite` call, 16 MiB + 8 header) + /// keep flowing (WS-14 edge). #[tokio::test] - async fn tungstenite_write_side_rejects_pending_over_byte_cap() { - let (client_io, _server_io) = tokio::io::duplex(64); + async fn tungstenite_write_at_the_cap_is_accepted() { + let (client_io, mut server_io) = tokio::io::duplex(64); let ws = tokio_tungstenite::WebSocketStream::from_raw_socket( client_io, tokio_tungstenite::tungstenite::protocol::Role::Client, @@ -950,37 +1172,30 @@ mod tests { .await; let (mut stream, _pumps) = split_tungstenite_to_bytes(ws); - let blob = vec![0u8; PENDING_BUFFER_CAP + 1]; - let written = stream.write(&blob).await.expect("first write accepted"); - assert_eq!(written, blob.len()); + let mut chunk = vec![0u8; PENDING_BUFFER_CAP]; + chunk[4..8].copy_from_slice(&(PENDING_BUFFER_CAP as u32 - 8).to_be_bytes()); + stream.write_all(&chunk).await.expect("cap write accepted"); stream.flush().await.expect("flush"); - let err: io::Error; - let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); - loop { - match stream.write_all(&[0u8; 8]).await { - Err(e) => { - err = e; - break; - } - Ok(_) => { - assert!( - tokio::time::Instant::now() < deadline, - "stream never failed after pending exceeded the byte cap" - ); - } - } - } - assert_eq!(err.kind(), io::ErrorKind::InvalidData); - assert!( - err.to_string().contains("pending buffer"), - "error names the violation: {err}" - ); + tokio::time::timeout(std::time::Duration::from_secs(30), async { + let mut buf = [0u8; 2]; + server_io + .read_exact(&mut buf) + .await + .expect("cap-sized chunk read"); + }) + .await + .expect("pump emitted the cap-sized chunk"); } /// A valid chunk (length field ≤ `MAX_CHUNK_LEN`) still parses and /// flushes through the pump after the caps landed — the validation - /// must not false-positive on well-framed traffic. + /// must not false-positive on well-framed traffic. (The pump-side + /// pending-cap fault this test family covered is unreachable for + /// single-call over-cap writes since WS-14 moved the check into + /// `poll_write`; the multi-write accumulation leg stays live via + /// the oversize-header test above, which strands the pump's + /// parser past a too-long header.) #[tokio::test] async fn tungstenite_write_side_still_emits_well_framed_chunk() { let (client_io, mut server_io) = tokio::io::duplex(64); @@ -1203,8 +1418,64 @@ mod tests { assert!(saw_close, "the shutdown path emitted a WS Close frame"); } + /// WS-15 acceptance (text arm): a text frame triggers the 1002 + /// protocol-error close whose reason names the text cause — + /// distinct from the idle (1001) and oversize (1011) reasons the + /// other tests assert. + #[tokio::test] + async fn tungstenite_text_close_carries_protocol_reason() { + 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 (_stream, pumps) = split_tungstenite_to_bytes(ws); + + 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}; + + let mut eof_rx = pumps.read_eof(); + peer_sink + .send(tokio_tungstenite::tungstenite::Message::Text( + "text frame".into(), + )) + .await + .expect("text write accepted"); + + let (close, eof_fired) = tokio::time::timeout(std::time::Duration::from_secs(5), async { + let close: Option<(u16, String)> = loop { + match peer_stream.next().await { + Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => { + break cf.map(|f| (u16::from(f.code), f.reason.to_string())) + } + Some(Ok(_)) => continue, + Some(Err(_)) | None => break None, + } + }; + let _ = eof_rx.changed().await; + (close, *eof_rx.borrow()) + }) + .await + .expect("read pump ends after the text frame (close requested)"); + + let (code, reason) = close.expect("close frame with code + reason"); + assert_eq!(code, WS_PROTOCOL_ERROR, "text frame closed with 1002"); + assert_eq!( + reason, "text messages not supported", + "close reason names the text cause, got {reason:?}" + ); + assert!(eof_fired, "EOF signal fired for the from_wss machinery"); + } + /// 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 @@ -1233,14 +1504,15 @@ mod tests { // 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. + // frame the peer receives. WS-15: the close reason names the + // idle-read cause, not a leftover default. 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 { + let close: Option<(u16, String)> = loop { match peer_stream.next().await { Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => { - break cf.map(|f| u16::from(f.code)) + break cf.map(|f| (u16::from(f.code), f.reason.to_string())) } Some(Ok(_)) => continue, Some(Err(_)) | None => break None, @@ -1252,9 +1524,14 @@ mod tests { .await .expect("stalled connection must be torn down within the deadline"); + let (code, reason) = close_seen.expect("close frame with code + reason"); + assert_eq!( + code, WS_GOING_AWAY, + "peer received the 1001 GoingAway close" + ); assert!( - close_seen == Some(WS_GOING_AWAY), - "peer received the 1001 GoingAway close, got {close_seen:?}" + reason.contains("no inbound chunk progress"), + "close reason names the idle-read cause, got {reason:?}" ); assert!(eof_fired, "EOF signal fired for the from_wss machinery"); } @@ -1324,10 +1601,10 @@ mod tests { }); let outcome = tokio::time::timeout(std::time::Duration::from_secs(10), async { - let close: Option = loop { + let close: Option<(u16, String)> = loop { match peer_stream.next().await { Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => { - break cf.map(|f| u16::from(f.code)) + break cf.map(|f| (u16::from(f.code), f.reason.to_string())) } Some(Ok(_)) => continue, Some(Err(_)) | None => break None, @@ -1339,11 +1616,15 @@ mod tests { .await .expect("dribble must hit the progress deadline within 10 s"); dribble.abort(); + let (code, reason) = outcome.expect("close frame with code + reason"); assert_eq!( - outcome, - Some(WS_GOING_AWAY), + code, WS_GOING_AWAY, "the forever-dribble is evicted with 1001 despite arriving messages" ); + assert!( + reason.contains("no inbound chunk progress"), + "close reason names the idle-read cause, got {reason:?}" + ); } /// Progress semantics, survivor side (WS-13 accept 2): a session @@ -1464,3 +1745,170 @@ mod tests { .expect("peer close accepted"); } } + +/// WS-19: the axum-flavor `AxumFraming` arms have no direct unit test — +/// the cap-trip and text→1002 closes are asserted on tungstenite only +/// (via the shared generic pumps). These drive the same generic pumps +/// with `AxumFraming` over a fake axum `WebSocket` sink/stream pair +/// (no server needed): a `mpsc`-backed sink/stream pair standing in +/// for the split halves of `axum::extract::ws::WebSocket`. +#[cfg(test)] +mod axum_framing_tests { + use super::*; + use futures::channel::mpsc as fut_mpsc; + + /// In-process stand-in for the split halves of an axum + /// `WebSocket`: messages flow stream→`rx` and `tx`→sink, so the + /// generic pumps run against `AxumFraming` unchanged. + struct AxumFakeSocket { + sink_tx: futures_mpsc::Sender, + stream_rx: + std::pin::Pin> + Send>>, + } + + impl futures::Stream for AxumFakeSocket { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.stream_rx.as_mut().poll_next(cx) + } + } + + impl futures::Sink for AxumFakeSocket { + type Error = (); + + fn poll_ready( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.sink_tx).poll_ready(cx).map_err(|_| ()) + } + + fn start_send(mut self: Pin<&mut Self>, item: AxumMessage) -> Result<(), Self::Error> { + self.sink_tx.start_send(item).map_err(|_| ()) + } + + fn poll_flush( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + futures::Sink::poll_flush(Pin::new(&mut self.sink_tx), cx).map_err(|_| ()) + } + + fn poll_close( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + futures::Sink::poll_close(Pin::new(&mut self.sink_tx), cx).map_err(|_| ()) + } + } + + type AxumPairParts = ( + futures::stream::SplitSink, + futures::stream::SplitStream, + futures_mpsc::Receiver, + fut_mpsc::Sender>, + ); + + fn axum_pair() -> AxumPairParts { + let (outbound_tx, outbound_rx) = fut_mpsc::channel::(4); + let (inbound_tx, inbound_rx) = fut_mpsc::channel::>(4); + let socket = AxumFakeSocket { + sink_tx: outbound_tx, + stream_rx: Box::pin(inbound_rx), + }; + let (sink, stream) = socket.split(); + (sink, stream, outbound_rx, inbound_tx) + } + + /// WS-19 mirror over `AxumFraming`: a text message on the read + /// side triggers the 1002 protocol-error close carrying the text + /// reason (the same generic read-pump arm the tungstenite tests + /// exercise — asserted here on the axum message types). + #[tokio::test] + async fn axum_framing_text_message_requests_protocol_close_with_reason() { + let (_sink, stream, _outbound_rx, mut inbound_tx) = axum_pair(); + + let (write_tx_for_read, mut write_rx_for_read) = fut_mpsc::channel::(WRITE_SLOTS); + let read_task = tokio::spawn(run_read_pump::( + stream, + mpsc::channel::>(READ_SLOTS).0, + write_tx_for_read, + make_write_error_slot().0, + None, + || {}, + )); + + inbound_tx + .send(Ok(AxumMessage::Text("text frame".into()))) + .await + .expect("inbound message accepted"); + + let close = tokio::time::timeout(std::time::Duration::from_secs(5), async { + while let Some(msg) = write_rx_for_read.next().await { + if let WriteMsg::CloseWith(code, reason) = msg { + return ::close_message(code, reason); + } + } + AxumMessage::Text("queue closed".into()) + }) + .await + .expect("close requested"); + + let AxumMessage::Close(Some(frame)) = close else { + panic!("expected a close frame, got {close:?}"); + }; + assert_eq!(frame.code, WS_PROTOCOL_ERROR, "text closes with 1002"); + assert_eq!( + frame.reason, "text messages not supported", + "close reason names the text cause" + ); + read_task.abort(); + } + + /// WS-19 mirror of the over-cap close on the axum flavor: a + /// header claiming above-`MAX_CHUNK_LEN` payload fails the pump — + /// the peer sees the 1011 close naming the violation and the + /// pump's error slot carries the reason. + #[tokio::test] + async fn axum_framing_cap_trip_fails_the_pump_with_the_internal_close() { + let (sink, _stream, mut outbound_rx, _inbound_tx) = axum_pair(); + let (slot, _rx) = make_write_error_slot(); + + let write_task = tokio::spawn(run_write_pump::( + sink, + { + let (mut tx, rx) = fut_mpsc::channel::(WRITE_SLOTS); + tx.send(WriteMsg::Bytes(tests::oversize_header())) + .await + .expect("queue write accepted"); + rx + }, + slot, + None, + )); + + let close = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + match outbound_rx.next().await { + Some(m @ AxumMessage::Close(_)) => return m, + Some(_) => continue, + None => return AxumMessage::Text("stream ended".into()), + } + } + }) + .await + .expect("close observed"); + + let AxumMessage::Close(Some(frame)) = close else { + panic!("expected a close frame, got {close:?}"); + }; + assert_eq!(frame.code, WS_INTERNAL_ERROR, "cap trip closes with 1011"); + assert!( + frame.reason.contains("MAX_CHUNK_LEN"), + "close reason names the violation: {}", + frame.reason + ); + write_task.abort(); + } +} diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs index 97aad82..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, - 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; +}