diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index d859bcc..ecc0b8b 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -41,6 +41,7 @@ use futures::channel::mpsc as futures_mpsc; use futures::{SinkExt, StreamExt}; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::sync::mpsc; +use tokio::sync::oneshot; /// Practical per-WS-message cap. Chunks larger than this are split /// across multiple WS messages — legal, since the receiver's boundary @@ -57,6 +58,16 @@ const WRITE_SLOTS: usize = 64; /// Protocol-error close code for text messages (websocket.md §Framing). pub const WS_PROTOCOL_ERROR: u16 = 1002; +/// Internal-error close code used when the write pump hits a protocol +/// violation it cannot recover from (WS-04/HY-09, WS-05). +pub const WS_INTERNAL_ERROR: u16 = 1011; + +/// Maximum chunk payload length — the channels protocol's 16 MiB wire +/// bound (ADR-052 §5), re-exported from alkcall. The write-side chunk +/// parser rejects a header claiming a longer payload instead of waiting +/// to accumulate up to ~4 GiB from a misaligned offset. +pub use alkcall::channels::wire::MAX_CHUNK_LEN; + pub(crate) enum WriteMsg { Bytes(Vec), /// Close with the given code (e.g. the 1002 text-rejection). @@ -72,6 +83,8 @@ pub struct WsByteStream { eof: bool, write_tx: futures_mpsc::Sender, write_open: bool, + write_error: oneshot::Receiver<&'static str>, + write_failed: bool, } /// The WS pump tasks. Dropping this guard detaches them (tokio @@ -109,6 +122,7 @@ 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>(); #[cfg(feature = "wss")] let read_eof = tokio::sync::watch::channel(false).0; @@ -160,9 +174,19 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { if pending.len() < 8 { break; } - let len = - u32::from_be_bytes([pending[4], pending[5], pending[6], pending[7]]) as usize; - let total = 8 + len; + 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"); + let _ = ws_sink + .send(AxumMessage::Close(Some(CloseFrame { + code: WS_INTERNAL_ERROR, + reason: "chunk length exceeds MAX_CHUNK_LEN".into(), + }))) + .await; + return; + } + let total = 8usize.saturating_add(len as usize); if pending.len() < total { break; } @@ -189,6 +213,8 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { eof: false, write_tx, write_open: true, + write_error: write_error_rx, + write_failed: false, }, WsPumps { read_task, @@ -248,6 +274,9 @@ impl AsyncWrite for WsByteStream { "ws stream shut down", ))); } + if let Some(err) = this.poll_write_error() { + return Poll::Ready(Err(err)); + } match this.write_tx.poll_ready(cx) { Poll::Ready(Ok(())) => match this.write_tx.try_send(WriteMsg::Bytes(buf.to_vec())) { Ok(()) => Poll::Ready(Ok(buf.len())), @@ -278,6 +307,40 @@ impl AsyncWrite for WsByteStream { } } +impl WsByteStream { + /// Fail the stream with the write pump's protocol error, if one + /// has arrived on the error channel (WS-04/HY-09, WS-05) — checked + /// before every `poll_write`, so a failed write task surfaces + /// loudly instead of leaving the stream pending/flushed silently. + fn poll_write_error(&mut self) -> Option { + if self.write_failed { + return Some(io::Error::new( + io::ErrorKind::InvalidData, + "ws write pump failed (protocol violation)", + )); + } + match self.write_error.try_recv() { + Ok(reason) => { + self.write_failed = true; + self.write_open = false; + Some(io::Error::new( + io::ErrorKind::InvalidData, + format!("ws write pump failed: {reason}"), + )) + } + Err(oneshot::error::TryRecvError::Closed) => { + self.write_failed = true; + self.write_open = false; + Some(io::Error::new( + io::ErrorKind::BrokenPipe, + "ws write pump ended", + )) + } + Err(oneshot::error::TryRecvError::Empty) => None, + } + } +} + /// Client-side split for a tokio-tungstenite `WebSocketStream` /// (`from_wss`, ADR-070): the same WS↔byte-stream seam as /// [`split_ws_to_bytes`], over the tungstenite socket instead of axum's @@ -294,6 +357,7 @@ 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 read_eof = tokio::sync::watch::channel(false).0; @@ -342,9 +406,21 @@ where if pending.len() < 8 { break; } - let len = - u32::from_be_bytes([pending[4], pending[5], pending[6], pending[7]]) as usize; - let total = 8 + len; + 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"); + let _ = ws_sink + .send(tokio_tungstenite::tungstenite::Message::Close(Some( + tokio_tungstenite::tungstenite::protocol::CloseFrame { + code: WS_INTERNAL_ERROR.into(), + reason: "chunk length exceeds MAX_CHUNK_LEN".into(), + }, + ))) + .await; + return; + } + let total = 8usize.saturating_add(len as usize); if pending.len() < total { break; } @@ -373,6 +449,8 @@ where eof: false, write_tx, write_open: true, + write_error: write_error_rx, + write_failed: false, }, WsPumps { read_task, @@ -382,3 +460,55 @@ where }, ) } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::AsyncWriteExt; + + /// 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). + fn oversize_header() -> Vec { + let mut header = vec![0u8; 8]; + header[4..8].copy_from_slice(&(MAX_CHUNK_LEN + 1).to_be_bytes()); + header + } + + #[tokio::test] + async fn tungstenite_write_side_rejects_oversized_chunk_header() { + 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); + + stream.write_all(&oversize_header()).await.expect("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 the oversized header" + ); + } + } + } + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!( + err.to_string().contains("MAX_CHUNK_LEN"), + "error names the violation: {err}" + ); + } +}