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

fix(websocket): byte-based caps for pending buffer + inbound WS sizes (WS-05, WS-06)

Browse Source
This commit is contained in:
glm-5.3-flash
2026-08-29 12:00:04 +00:00
parent 5024d99862
commit 92cc11a74f
5 changed files with 264 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/websocket/byte_adapter.rs
+214 -12
View File
@@ -7,7 +7,10 @@
//!
//! Inbound: a WS read task pushes binary-message bytes into a bounded
//! mpsc (64 slots); the `AsyncRead` half drains it. Backpressure = mpsc
//! capacity (OQ-01a): the read task awaits `send` when full.
//! capacity (OQ-01a): the read task awaits `send` when full. Inbound
//! message/frame sizes are explicitly capped (`INBOUND_WS_MESSAGE_CAP` /
//! `INBOUND_WS_FRAME_CAP`, WS-06) on both the axum upgrade and the
//! tungstenite dial instead of the libraries' 64 MiB defaults.
//!
//! Outbound: the `AsyncWrite` half queues byte spans; a writer task
//! parses the pending bytes for complete chunks (8-byte header → payload
@@ -16,8 +19,15 @@
//! chunk header, not the message; a chunk may span messages). Chunk
//! parsing is required because a logical write above the mux (channel
//! 0's `write_frame` issues prefix+body separately) surfaces as multiple
//! mux payloads. Write-side backpressure uses `futures::channel::mpsc`
//! `poll_ready` — the production fix for the POC's spin-wait.
//! mux payloads. The chunk parser validates the length field against
//! 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
//! backpressure uses `futures::channel::mpsc` `poll_ready` — the
//! production fix for the POC's spin-wait.
//!
//! Text WS messages are rejected with a protocol-level close (code
//! 1002); all frames are binary (websocket.md §Framing).
@@ -33,6 +43,7 @@
use std::{
io,
pin::Pin,
sync::{Arc, Mutex as StdMutex},
task::{Context, Poll},
};
@@ -55,6 +66,37 @@ pub const READ_SLOTS: usize = 64;
const WRITE_SLOTS: usize = 64;
/// Inbound WS message size cap (WS-06): explicitly configured on both
/// the axum upgrade and the tungstenite client instead of the
/// libraries' 64 MiB defaults, so one connection cannot pin ~4 GiB
/// while the demux drains slower than the socket delivers (WS-01's
/// stall shape). Chunks larger than this must arrive split across
/// multiple WS messages — legal, since the chunk header is the only
/// receiver boundary.
pub const INBOUND_WS_MESSAGE_CAP: usize = 1024 * 1024;
/// Inbound WS frame size cap (WS-06). Equal to the message cap by
/// design: the WS↔byte-stream path does not use WS fragmentation (each
/// WS binary message is a single unfragmented frame up to
/// `INBOUND_WS_MESSAGE_CAP`), and tungstenite enforces
/// `max_frame_size` against a single frame's payload *before*
/// continuation reassembly — a smaller cap would reject legal ≥1 MiB
/// 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.
pub const PENDING_BUFFER_CAP: usize = MAX_CHUNK_LEN as usize + 8;
/// Protocol-error close code for text messages (websocket.md §Framing).
pub const WS_PROTOCOL_ERROR: u16 = 1002;
@@ -74,6 +116,27 @@ pub(crate) enum WriteMsg {
CloseWith(u16),
}
/// The write-pump failure channel (WS-04/HY-09, WS-05): a shared slot
/// holding the pump's one-shot error sender. Both pump tasks (write
/// pump and read task, on the inbound-error arm) can surface a fatal
/// reason; the stream side drains it with `try_recv` semantics. The
/// `Arc<Mutex<Option<_>>>` shape exists because `oneshot::Sender` is
/// not `Clone` and two task bodies need the sender.
type WriteErrorSlot = Arc<StdMutex<Option<oneshot::Sender<&'static str>>>>;
fn make_write_error_slot() -> (WriteErrorSlot, oneshot::Receiver<&'static str>) {
let (tx, rx) = oneshot::channel::<&'static str>();
(Arc::new(StdMutex::new(Some(tx))), rx)
}
/// Send the pump's fatal reason; a no-op once the slot is taken
/// (first failure wins — the reason text is diagnostic only).
fn send_write_error(slot: &WriteErrorSlot, reason: &'static str) {
if let Some(tx) = slot.lock().unwrap_or_else(|e| e.into_inner()).take() {
let _ = tx.send(reason);
}
}
/// The `AsyncRead + AsyncWrite` view of a `WebSocket` handed to
/// alkcall's channels machinery (demux/mux).
pub struct WsByteStream {
@@ -122,12 +185,13 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
let (mut ws_sink, mut ws_stream) = socket.split();
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let (write_error_tx, write_error_rx) = oneshot::channel::<&'static str>();
let (write_error_slot, write_error_rx) = make_write_error_slot();
#[cfg(feature = "wss")]
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);
#[cfg(feature = "wss")]
let read_eof_for_task = read_eof.clone();
let read_task = tokio::spawn(async move {
@@ -145,7 +209,15 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
.await;
break;
}
Ok(AxumMessage::Close(_)) | Err(_) => break,
Ok(AxumMessage::Close(_)) => break,
Err(_) => {
send_write_error(&write_error_slot_for_read, "inbound frame rejected (size cap)");
let _ = write_tx_for_read
.clone()
.send(WriteMsg::CloseWith(WS_INTERNAL_ERROR))
.await;
break;
}
Ok(_) => {}
}
}
@@ -168,7 +240,19 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
.await;
break;
}
WriteMsg::Bytes(b) => pending.extend_from_slice(&b),
WriteMsg::Bytes(b) => {
if b.len() > PENDING_BUFFER_CAP {
send_write_error(&write_error_slot, "write pending buffer exceeded cap");
let _ = ws_sink
.send(AxumMessage::Close(Some(CloseFrame {
code: WS_INTERNAL_ERROR,
reason: "write pending buffer exceeded cap".into(),
})))
.await;
return;
}
pending.extend_from_slice(&b);
}
}
loop {
if pending.len() < 8 {
@@ -177,7 +261,7 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
let len_bytes = [pending[4], pending[5], pending[6], pending[7]];
let len = u32::from_be_bytes(len_bytes);
if len > MAX_CHUNK_LEN {
let _ = write_error_tx.send("chunk length exceeds MAX_CHUNK_LEN");
send_write_error(&write_error_slot, "chunk length exceeds MAX_CHUNK_LEN");
let _ = ws_sink
.send(AxumMessage::Close(Some(CloseFrame {
code: WS_INTERNAL_ERROR,
@@ -201,6 +285,16 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
}
}
}
if pending.len() > PENDING_BUFFER_CAP {
send_write_error(&write_error_slot, "write pending buffer exceeded cap");
let _ = ws_sink
.send(AxumMessage::Close(Some(CloseFrame {
code: WS_INTERNAL_ERROR,
reason: "write pending buffer exceeded cap".into(),
})))
.await;
return;
}
}
let _ = ws_sink.close().await;
});
@@ -357,11 +451,12 @@ where
let (mut ws_sink, mut ws_stream) = socket.split();
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let (write_error_tx, write_error_rx) = oneshot::channel::<&'static str>();
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(async move {
while let Some(msg) = ws_stream.next().await {
@@ -378,7 +473,15 @@ where
.await;
break;
}
Ok(tokio_tungstenite::tungstenite::Message::Close(_)) | Err(_) => break,
Ok(tokio_tungstenite::tungstenite::Message::Close(_)) => break,
Err(_) => {
send_write_error(&write_error_slot_for_read, "inbound frame rejected (size cap)");
let _ = write_tx_for_read
.clone()
.send(WriteMsg::CloseWith(WS_INTERNAL_ERROR))
.await;
break;
}
Ok(_) => {}
}
}
@@ -400,7 +503,21 @@ where
.await;
break;
}
WriteMsg::Bytes(b) => pending.extend_from_slice(&b),
WriteMsg::Bytes(b) => {
if b.len() > PENDING_BUFFER_CAP {
send_write_error(&write_error_slot, "write pending buffer exceeded cap");
let _ = ws_sink
.send(tokio_tungstenite::tungstenite::Message::Close(Some(
tokio_tungstenite::tungstenite::protocol::CloseFrame {
code: WS_INTERNAL_ERROR.into(),
reason: "write pending buffer exceeded cap".into(),
},
)))
.await;
return;
}
pending.extend_from_slice(&b);
}
}
loop {
if pending.len() < 8 {
@@ -409,7 +526,7 @@ where
let len_bytes = [pending[4], pending[5], pending[6], pending[7]];
let len = u32::from_be_bytes(len_bytes);
if len > MAX_CHUNK_LEN {
let _ = write_error_tx.send("chunk length exceeds MAX_CHUNK_LEN");
send_write_error(&write_error_slot, "chunk length exceeds MAX_CHUNK_LEN");
let _ = ws_sink
.send(tokio_tungstenite::tungstenite::Message::Close(Some(
tokio_tungstenite::tungstenite::protocol::CloseFrame {
@@ -437,6 +554,18 @@ where
}
}
}
if pending.len() > PENDING_BUFFER_CAP {
send_write_error(&write_error_slot, "write pending buffer exceeded cap");
let _ = ws_sink
.send(tokio_tungstenite::tungstenite::Message::Close(Some(
tokio_tungstenite::tungstenite::protocol::CloseFrame {
code: WS_INTERNAL_ERROR.into(),
reason: "write pending buffer exceeded cap".into(),
},
)))
.await;
return;
}
}
let _ = ws_sink.close().await;
});
@@ -464,7 +593,7 @@ where
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::AsyncWriteExt;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
/// An 8-byte bogus chunk header claiming `MAX_CHUNK_LEN + 1`
/// payload bytes (WS-04/HY-09 acceptance: the parser must fail the
@@ -511,4 +640,77 @@ mod tests {
"error names the violation: {err}"
);
}
/// 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).
#[tokio::test]
async fn tungstenite_write_side_rejects_pending_over_byte_cap() {
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 (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());
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}"
);
}
/// 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.
#[tokio::test]
async fn tungstenite_write_side_still_emits_well_framed_chunk() {
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,
None,
)
.await;
let (mut stream, _pumps) = split_tungstenite_to_bytes(ws);
let mut chunk = vec![0u8; 8];
chunk[4..8].copy_from_slice(&4u32.to_be_bytes());
chunk.extend_from_slice(b"payload");
stream.write_all(&chunk).await.expect("write accepted");
stream.flush().await.expect("flush");
tokio::time::timeout(std::time::Duration::from_secs(5), async {
let mut buf = [0u8; 2];
server_io.read_exact(&mut buf).await.expect("mask scan read");
})
.await
.expect("pump emitted the framed chunk");
}
}