alkdev/alkhttp
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(websocket): reject over-cap writes in poll_write before the queue (WS-14)

Browse Source 
- byte-cap check moved ahead of poll_ready/try_send so the mux sees
  the InvalidData stream error at the failing call instead of the
  write pump emitting a mid-stream 1011 close (the chunk was already
  committed to the wire path)
- pump-side over-cap arm replaced by a debug_assert (defense-in-depth
  invariant; the error was unreachable for single-call writes once the
  pre-send check exists)
- write_tx_mut renamed write_tx_ref with an updated doc contract
- WS-05's over-cap pump test replaced by a pre-send rejection test
  (error surfaces at the first write, names the cap) + a cap-edge
  test (exactly PENDING_BUFFER_CAP still flows)

Verification: scripts/verify.sh OK (341 passed), clippy -D warnings
clean, fmt clean
This commit is contained in:
glm-5.3-flash
2026-08-30 20:56:56 +00:00
parent 5244dc46e2
commit 9ef2352568
1 changed files with 65 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/websocket/byte_adapter.rs
+65 -62
View File
@@ -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<M, S>(
break;
}
WriteMsg::Bytes(b) => {
if b.len() > PENDING_BUFFER_CAP {
fail_write_pump::<M, _>(
&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<WriteMsg>> {
/// 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<WriteMsg>> {
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);