fix(websocket): validate chunk length on write side (WS-04, HY-09)

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<u8>),
     /// 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<WriteMsg>,
     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::<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>();
 
     #[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<io::Error> {
+        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::<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 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<u8> {
+        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}"
+        );
+    }
+}