diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index 0ac4bb8..66d9f16 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). @@ -423,15 +430,7 @@ async fn run_write_pump( 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); } } @@ -675,12 +674,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 +763,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() } } @@ -933,15 +944,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 +960,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);