diff --git a/src/adapters/from_wss.rs b/src/adapters/from_wss.rs index 76e3a77..7dcc944 100644 --- a/src/adapters/from_wss.rs +++ b/src/adapters/from_wss.rs @@ -196,11 +196,19 @@ impl WssSession { ); } - let (ws, _response) = tokio_tungstenite::connect_async(request) - .await - .map_err(|e| AdapterError::Transport { - message: format!("WSS connect failed: {e}"), - })?; + let (ws, _response) = tokio_tungstenite::connect_async_with_config( + request, + Some( + tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default() + .max_message_size(Some(crate::websocket::INBOUND_WS_MESSAGE_CAP)) + .max_frame_size(Some(crate::websocket::INBOUND_WS_FRAME_CAP)), + ), + false, + ) + .await + .map_err(|e| AdapterError::Transport { + message: format!("WSS connect failed: {e}"), + })?; let (byte_stream, pumps) = split_tungstenite_to_bytes(ws); let conn = Connection::from_bidi(byte_stream, b"alk/channels".to_vec(), None); diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index ecc0b8b..afb5e07 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -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>>` shape exists because `oneshot::Sender` is +/// not `Clone` and two task bodies need the sender. +type WriteErrorSlot = Arc>>>; + +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::>(READ_SLOTS); let (write_tx, mut write_rx) = futures_mpsc::channel::(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::>(READ_SLOTS); let (write_tx, mut write_rx) = futures_mpsc::channel::(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"); + } } diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs index 82925a5..9fe0a0a 100644 --- a/src/websocket/mod.rs +++ b/src/websocket/mod.rs @@ -13,7 +13,8 @@ pub mod byte_adapter; pub mod upgrade; pub use byte_adapter::{ - split_ws_to_bytes, WsByteStream, WsPumps, WS_MESSAGE_CAP, WS_PROTOCOL_ERROR, + split_ws_to_bytes, WsByteStream, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, + MAX_CHUNK_LEN, PENDING_BUFFER_CAP, WS_MESSAGE_CAP, WS_PROTOCOL_ERROR, }; #[cfg(any(test, feature = "wss"))] diff --git a/src/websocket/upgrade.rs b/src/websocket/upgrade.rs index 2bbdb72..af62a94 100644 --- a/src/websocket/upgrade.rs +++ b/src/websocket/upgrade.rs @@ -21,7 +21,7 @@ use axum::extract::ws::WebSocketUpgrade; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; -use super::byte_adapter::split_ws_to_bytes; +use super::byte_adapter::{split_ws_to_bytes, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP}; /// The channels session for an upgraded socket: adapt → `Connection` /// (identity attached) → `ChannelsAdapter::handle`. `policy` gates @@ -140,9 +140,12 @@ pub async fn ws_upgrade_handler( let policy = policy .map(|axum::Extension(p)| p.0) .unwrap_or_else(|| Arc::new(NoCap)); - ws_upgrade.on_upgrade(move |socket| async move { - run_channels_session(socket, registry, identity, policy).await - }) + ws_upgrade + .max_frame_size(INBOUND_WS_FRAME_CAP) + .max_message_size(INBOUND_WS_MESSAGE_CAP) + .on_upgrade(move |socket| async move { + run_channels_session(socket, registry, identity, policy).await + }) } /// Bearer-auth middleware for the WS upgrade route: resolves the token diff --git a/tests/ws_upgrade_session.rs b/tests/ws_upgrade_session.rs index 339c10c..9fc3652 100644 --- a/tests/ws_upgrade_session.rs +++ b/tests/ws_upgrade_session.rs @@ -11,7 +11,8 @@ use alkcall::registry::registration::{ }; use alkcall::registry::spec::{AccessControl, OperationSpec, OperationType, Visibility}; use alkhttp::websocket::{ - frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient, WS_MESSAGE_CAP, + frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient, INBOUND_WS_MESSAGE_CAP, + WS_MESSAGE_CAP, }; use std::collections::HashMap; use std::sync::Arc; @@ -579,3 +580,29 @@ async fn services_list_over_channel0_is_access_control_filtered() { assert!(!names.contains(&"admin/secret"), "got: {names:?}"); ws.close().await; } + +#[tokio::test] +async fn inbound_message_over_cap_fails_the_connection() { + let addr = spawn_ws_server( + echo_registry(), + provider_with(vec![("tok-1", identity("alice", &[]))]), + ) + .await; + let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1") + .await + .unwrap(); + + // One WS binary message larger than INBOUND_WS_MESSAGE_CAP: the + // server's explicit max_message_size (WS-06) must fail the + // connection instead of accepting it (the byte-adapter boundary is + // the chunk header, so the adapter never dissects messages; the WS + // library cap is the enforcement point). + let oversized = vec![b'x'; INBOUND_WS_MESSAGE_CAP + 1]; + ws.send_binary_piece(&oversized).await; + + let close = ws + .next_close(std::time::Duration::from_secs(5)) + .await + .expect("server must fail the oversized-inbound connection (WS-06)"); + assert!(close.is_some(), "expected a close or stream end, got {close:?}"); +}