diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index 135238a..7c5be69 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -147,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 @@ -431,10 +439,30 @@ async fn fail_write_pump( /// 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 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, @@ -443,9 +471,15 @@ async fn run_write_pump( while let Some(msg) = write_rx.next().await { match msg { WriteMsg::CloseWith(code, reason) => { - let _ = ws_sink - .send(::close_message(code, reason)) - .await; + 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) => { @@ -474,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; } @@ -496,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 { @@ -619,6 +690,7 @@ pub fn split_ws_to_bytes_idle( ws_sink, write_rx, write_error_slot, + Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT), )); ( @@ -887,6 +959,7 @@ where ws_sink, write_rx, write_error_slot, + Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT), )); ( @@ -909,11 +982,116 @@ 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). diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs index 97aad82..a8705e5 100644 --- a/src/websocket/mod.rs +++ b/src/websocket/mod.rs @@ -14,8 +14,8 @@ 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, + 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"))]