diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index afb5e07..5bd570b 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -38,7 +38,9 @@ //! bytes drain (the mux's EOF sentinels ride the same queue). //! //! Shared with the `from_wss` consumer path (ADR-070): one -//! implementation, both directions. +//! implementation, both directions ([`WsFraming`] + the generic +//! [`run_read_pump`] / [`run_write_pump`] instantiated once per socket +//! flavor — WS-11/COV-03). use std::{ io, @@ -137,6 +139,201 @@ fn send_write_error(slot: &WriteErrorSlot, reason: &'static str) { } } +/// WS-protocol adapter for one socket flavor: the only places the axum +/// and tungstenite message types differ. The generic pump +/// ([`run_read_pump`] / [`run_write_pump`]) is written against this +/// trait so both paths share one implementation (WS-11, "one +/// implementation, both directions"). +trait WsFraming: Sized { + /// The WS library's binary-message byte type. + type Bytes: AsRef<[u8]>; + /// The sink/stream message type. + type Msg; + + /// `Some(bytes)` for a binary message (the byte stream's carrier), + /// `None` for anything else. + fn binary(msg: &Self::Msg) -> Option<&Self::Bytes>; + + /// `true` for a text message (a protocol error: close 1002). + fn is_text(msg: &Self::Msg) -> bool; + + /// `true` for a Close message from the peer. + fn is_close(msg: &Self::Msg) -> bool; + + /// A Close message with the given code and reason. + fn close_message(code: u16, reason: &'static str) -> Self::Msg; + + /// A binary message carrying `bytes` (chunk pieces up to + /// `WS_MESSAGE_CAP`). + fn binary_message(bytes: Vec) -> Self::Msg; +} + +/// The read stream's item for a flavor: `Result` (both +/// flavors surface read failures this way; the error type is flavor +/// specific and only ever inspected as "failed"). +trait IntoWsResult { + fn into_ws_result(self) -> Result; +} + +impl IntoWsResult for Result +where + M: WsFraming, +{ + fn into_ws_result(self) -> Result { + self.map_err(|_| ()) + } +} + +/// The read pump: forwards binary-message bytes into `read_tx`, turns +/// text into a 1002 close request, and treats peer close / read error +/// as the connection end. `on_end` runs after the loop (the `wss` +/// feature fires the lossless EOF watch signal through it — the +/// from_wss drop monitor's input; its `Sender` semantics are preserved +/// EXACTLY: single `send` after the read loop terminates). +async fn run_read_pump( + mut ws_stream: S, + read_tx: mpsc::Sender>, + write_tx_for_read: futures_mpsc::Sender, + write_error_slot_for_read: WriteErrorSlot, + on_end: F, +) where + S: futures::Stream + Unpin, + M: WsFraming, + S::Item: IntoWsResult, + F: FnOnce(), +{ + while let Some(msg) = ws_stream.next().await { + match msg.into_ws_result() { + Ok(m) => { + if let Some(b) = M::binary(&m) { + if read_tx.send(b.as_ref().to_vec()).await.is_err() { + break; + } + } else if M::is_text(&m) { + let _ = write_tx_for_read + .clone() + .send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR)) + .await; + break; + } else if M::is_close(&m) { + 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; + } + } + } + on_end(); +} + +/// Emit the pump's fatal close + error, then end the write task +/// (WS-04/HY-09, WS-05: fail loudly instead of accumulating/hanging). +async fn fail_write_pump( + ws_sink: &mut futures::stream::SplitSink::Msg>, + slot: &WriteErrorSlot, + reason: &'static str, +) where + M: WsFraming, + S: futures::Sink<::Msg> + Unpin, +{ + send_write_error(slot, reason); + let _ = ws_sink + .send(::close_message(WS_INTERNAL_ERROR, reason)) + .await; +} + +/// The write pump: drains the queued byte spans, parses the pending +/// bytes for complete chunks (8-byte header, length validated against +/// `MAX_CHUNK_LEN` — WS-04/HY-09), byte-caps the accumulator +/// (`PENDING_BUFFER_CAP` — WS-05), and emits one WS binary message per +/// `WS_MESSAGE_CAP` piece. Ends with a WS Close after the queue +/// drains; `CloseWith` requests (`WriteMsg::CloseWith`) map to a close +/// frame with that code. +async fn run_write_pump( + mut ws_sink: futures::stream::SplitSink::Msg>, + mut write_rx: futures_mpsc::Receiver, + write_error_slot: WriteErrorSlot, +) where + S: futures::Sink<::Msg> + Unpin, + M: WsFraming, +{ + let mut pending: Vec = Vec::new(); + while let Some(msg) = write_rx.next().await { + match msg { + WriteMsg::CloseWith(code) => { + let _ = ws_sink + .send(::close_message( + code, + "text messages not supported", + )) + .await; + break; + } + WriteMsg::Bytes(b) => { + if b.len() > PENDING_BUFFER_CAP { + fail_write_pump::( + &mut ws_sink, + &write_error_slot, + "write pending buffer exceeded cap", + ) + .await; + return; + } + pending.extend_from_slice(&b); + } + } + loop { + if pending.len() < 8 { + break; + } + let len_bytes = [pending[4], pending[5], pending[6], pending[7]]; + let len = u32::from_be_bytes(len_bytes); + if len > MAX_CHUNK_LEN { + fail_write_pump::( + &mut ws_sink, + &write_error_slot, + "chunk length exceeds MAX_CHUNK_LEN", + ) + .await; + return; + } + let total = 8usize.saturating_add(len as usize); + if pending.len() < total { + break; + } + let chunk: Vec = pending.drain(..total).collect(); + for piece in chunk.chunks(WS_MESSAGE_CAP) { + if ws_sink + .send(::binary_message(piece.to_vec())) + .await + .is_err() + { + return; + } + } + } + if pending.len() > PENDING_BUFFER_CAP { + fail_write_pump::( + &mut ws_sink, + &write_error_slot, + "write pending buffer exceeded cap", + ) + .await; + return; + } + } + let _ = ws_sink.close().await; +} + /// The `AsyncRead + AsyncWrite` view of a `WebSocket` handed to /// alkcall's channels machinery (demux/mux). pub struct WsByteStream { @@ -156,7 +353,7 @@ pub struct WsByteStream { pub struct WsPumps { read_task: tokio::task::JoinHandle<()>, write_task: tokio::task::JoinHandle<()>, - #[cfg(feature = "wss")] + #[cfg(any(test, feature = "wss"))] read_eof: tokio::sync::watch::Sender, } @@ -172,132 +369,80 @@ impl WsPumps { /// taken or the observer starts awaiting — is still observed, and /// may be observed repeatedly. Used by `from_wss`'s /// connection-drop monitor (ADR-070). - #[cfg(feature = "wss")] + #[cfg(any(test, feature = "wss"))] pub(crate) fn read_eof(&self) -> tokio::sync::watch::Receiver { self.read_eof.subscribe() } } +/// axum's message/CloseFrame types as a [`WsFraming`] flavor. +struct AxumFraming; + +impl WsFraming for AxumFraming { + type Bytes = bytes::Bytes; + type Msg = AxumMessage; + + fn binary(msg: &AxumMessage) -> Option<&Self::Bytes> { + match msg { + AxumMessage::Binary(b) => Some(b), + _ => None, + } + } + + fn is_text(msg: &AxumMessage) -> bool { + matches!(msg, AxumMessage::Text(_)) + } + + fn is_close(msg: &AxumMessage) -> bool { + matches!(msg, AxumMessage::Close(_)) + } + + fn close_message(code: u16, reason: &'static str) -> AxumMessage { + AxumMessage::Close(Some(CloseFrame { + code, + reason: reason.into(), + })) + } + + fn binary_message(bytes: Vec) -> AxumMessage { + AxumMessage::Binary(bytes.into()) + } +} + /// Split a `WebSocket` into the byte stream + the pump tasks. The /// adapter is the single seam between axum's WS and alkcall's /// byte-oriented channels machinery; shared with `from_wss`. pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { - let (mut ws_sink, mut ws_stream) = socket.split(); + let (ws_sink, 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_tx, write_rx) = futures_mpsc::channel::(WRITE_SLOTS); let (write_error_slot, write_error_rx) = make_write_error_slot(); - #[cfg(feature = "wss")] + #[cfg(any(test, 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")] + #[cfg(any(test, feature = "wss"))] let read_eof_for_task = read_eof.clone(); - let read_task = tokio::spawn(async move { - while let Some(msg) = ws_stream.next().await { - match msg { - Ok(AxumMessage::Binary(b)) => { - if read_tx.send(b.to_vec()).await.is_err() { - break; - } - } - Ok(AxumMessage::Text(_)) => { - let _ = write_tx_for_read - .clone() - .send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR)) - .await; - 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(_) => {} + let read_task = tokio::spawn(run_read_pump::( + ws_stream, + read_tx, + write_tx_for_read, + write_error_slot_for_read, + move || { + #[cfg(any(test, feature = "wss"))] + { + let _ = read_eof_for_task.send(true); } - } - #[cfg(feature = "wss")] - { - let _ = read_eof_for_task.send(true); - } - }); + }, + )); - let write_task = tokio::spawn(async move { - let mut pending: Vec = Vec::new(); - while let Some(msg) = write_rx.next().await { - match msg { - WriteMsg::CloseWith(code) => { - let _ = ws_sink - .send(AxumMessage::Close(Some(CloseFrame { - code, - reason: "text messages not supported".into(), - }))) - .await; - break; - } - 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 { - break; - } - let len_bytes = [pending[4], pending[5], pending[6], pending[7]]; - let len = u32::from_be_bytes(len_bytes); - if len > 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, - reason: "chunk length exceeds MAX_CHUNK_LEN".into(), - }))) - .await; - return; - } - let total = 8usize.saturating_add(len as usize); - if pending.len() < total { - break; - } - let chunk: Vec = pending.drain(..total).collect(); - for piece in chunk.chunks(WS_MESSAGE_CAP) { - if ws_sink - .send(AxumMessage::Binary(piece.to_vec().into())) - .await - .is_err() - { - return; - } - } - } - 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; - }); + let write_task = tokio::spawn(run_write_pump::( + ws_sink, + write_rx, + write_error_slot, + )); ( WsByteStream { @@ -313,7 +458,7 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { WsPumps { read_task, write_task, - #[cfg(feature = "wss")] + #[cfg(any(test, feature = "wss"))] read_eof, }, ) @@ -435,12 +580,51 @@ impl WsByteStream { } } +/// tokio-tungstenite's message types as a [`WsFraming`] flavor. +#[cfg(any(test, feature = "wss"))] +struct TungsteniteFraming; + +#[cfg(any(test, feature = "wss"))] +impl WsFraming for TungsteniteFraming { + type Bytes = bytes::Bytes; + type Msg = tokio_tungstenite::tungstenite::Message; + + fn binary(msg: &Self::Msg) -> Option<&Self::Bytes> { + match msg { + tokio_tungstenite::tungstenite::Message::Binary(b) => Some(b), + _ => None, + } + } + + fn is_text(msg: &Self::Msg) -> bool { + matches!(msg, tokio_tungstenite::tungstenite::Message::Text(_)) + } + + fn is_close(msg: &Self::Msg) -> bool { + matches!(msg, tokio_tungstenite::tungstenite::Message::Close(_)) + } + + fn close_message(code: u16, reason: &'static str) -> Self::Msg { + tokio_tungstenite::tungstenite::Message::Close(Some( + tokio_tungstenite::tungstenite::protocol::CloseFrame { + code: code.into(), + reason: reason.into(), + }, + )) + } + + fn binary_message(bytes: Vec) -> Self::Msg { + tokio_tungstenite::tungstenite::Message::Binary(bytes.into()) + } +} + /// 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 /// server-side `WebSocket`. Message semantics are identical: binary /// messages carry the byte stream, text is a protocol error (close -/// 1002), close → read EOF. +/// 1002), close → read EOF. Both paths run the same generic pumps +/// (WS-11); the socket flavor differs only in the [`WsFraming`] impl. #[cfg(any(test, feature = "wss"))] pub fn split_tungstenite_to_bytes( socket: tokio_tungstenite::WebSocketStream, @@ -448,9 +632,9 @@ pub fn split_tungstenite_to_bytes( where S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static, { - let (mut ws_sink, mut ws_stream) = socket.split(); + let (ws_sink, 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_tx, write_rx) = futures_mpsc::channel::(WRITE_SLOTS); let (write_error_slot, write_error_rx) = make_write_error_slot(); let read_eof = tokio::sync::watch::channel(false).0; @@ -458,117 +642,21 @@ where 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 { - match msg { - Ok(tokio_tungstenite::tungstenite::Message::Binary(b)) => { - if read_tx.send(b.to_vec()).await.is_err() { - break; - } - } - Ok(tokio_tungstenite::tungstenite::Message::Text(_)) => { - let _ = write_tx_for_read - .clone() - .send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR)) - .await; - 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(_) => {} - } - } - let _ = read_eof_for_task.send(true); - }); + let read_task = tokio::spawn(run_read_pump::( + ws_stream, + read_tx, + write_tx_for_read, + write_error_slot_for_read, + move || { + let _ = read_eof_for_task.send(true); + }, + )); - let write_task = tokio::spawn(async move { - let mut pending: Vec = Vec::new(); - while let Some(msg) = write_rx.next().await { - match msg { - WriteMsg::CloseWith(code) => { - let _ = ws_sink - .send(tokio_tungstenite::tungstenite::Message::Close(Some( - tokio_tungstenite::tungstenite::protocol::CloseFrame { - code: code.into(), - reason: "text messages not supported".into(), - }, - ))) - .await; - break; - } - 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 { - break; - } - let len_bytes = [pending[4], pending[5], pending[6], pending[7]]; - let len = u32::from_be_bytes(len_bytes); - if len > 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 { - 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; - } - let chunk: Vec = pending.drain(..total).collect(); - for piece in chunk.chunks(WS_MESSAGE_CAP) { - if ws_sink - .send(tokio_tungstenite::tungstenite::Message::Binary( - piece.to_vec().into(), - )) - .await - .is_err() - { - return; - } - } - } - 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; - }); + let write_task = tokio::spawn(run_write_pump::( + ws_sink, + write_rx, + write_error_slot, + )); ( WsByteStream { @@ -584,7 +672,7 @@ where WsPumps { read_task, write_task, - #[cfg(feature = "wss")] + #[cfg(any(test, feature = "wss"))] read_eof, }, ) @@ -615,7 +703,10 @@ mod tests { .await; let (mut stream, _pumps) = split_tungstenite_to_bytes(ws); - stream.write_all(&oversize_header()).await.expect("write accepted"); + stream + .write_all(&oversize_header()) + .await + .expect("write accepted"); stream.flush().await.expect("flush"); let err: io::Error; @@ -708,9 +799,127 @@ mod tests { 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"); + server_io + .read_exact(&mut buf) + .await + .expect("mask scan read"); }) .await .expect("pump emitted the framed chunk"); } + + /// Read-pump binary passthrough on the tungstenite path (COV-03): + /// a peer binary byte-message surfaces on the `AsyncRead` half of + /// the adapter over a raw duplex pair (no network), with a server + /// role peer writing the frame. + #[tokio::test] + async fn tungstenite_read_side_forwards_binary_message_bytes() { + let (client_io, server_io) = tokio::io::duplex(1 << 16); + 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 write_side = tokio::spawn(async move { + let peer = tokio_tungstenite::WebSocketStream::from_raw_socket( + server_io, + tokio_tungstenite::tungstenite::protocol::Role::Server, + None, + ) + .await; + let (mut sink, _reader) = peer.split(); + use futures::SinkExt; + let _ = sink + .send(tokio_tungstenite::tungstenite::Message::Binary( + b"from-peer".to_vec().into(), + )) + .await; + }); + write_side.await.expect("peer write task completes"); + + use tokio::io::AsyncReadExt as _; + let mut buf = [0u8; 9]; + tokio::time::timeout( + std::time::Duration::from_secs(5), + stream.read_exact(&mut buf), + ) + .await + .expect("read within deadline") + .expect("read"); + assert_eq!(&buf, b"from-peer"); + } + + /// The lossless read-EOF signal on the tungstenite path (COV-03, + /// the WS-11 constraint): the peer disappearing surfaces on + /// `pumps.read_eof()` as `watch` = true — observable both before + /// and after the fact, the from_wss drop-monitor contract. + #[tokio::test] + async fn tungstenite_read_eof_signal_fires_and_is_retained() { + let (client_io, server_io) = tokio::io::duplex(1 << 16); + let ws = tokio_tungstenite::WebSocketStream::from_raw_socket( + client_io, + tokio_tungstenite::tungstenite::protocol::Role::Client, + None, + ) + .await; + let (_stream, pumps) = split_tungstenite_to_bytes(ws); + let mut eof_rx = pumps.read_eof(); + + drop(server_io); + + let changed = tokio::time::timeout(std::time::Duration::from_secs(5), eof_rx.changed()) + .await + .expect("EOF observed within deadline"); + assert!(changed.is_ok() || *eof_rx.borrow(), "EOF flagged"); + assert!(*eof_rx.borrow(), "EOF signal retained as true"); + + let mut late_rx = pumps.read_eof(); + assert!( + *late_rx.borrow_and_update(), + "a receiver taken after EOF still observes the signal" + ); + } + + /// Text frames over the tungstenite path are a protocol error: the + /// pump bridges text → 1002 close on the socket (COV-03 read-pump + /// coverage; pitcher on a raw duplex pair). + #[tokio::test] + async fn tungstenite_read_side_text_message_requests_protocol_close() { + let (client_io, server_io) = tokio::io::duplex(1 << 16); + let ws = tokio_tungstenite::WebSocketStream::from_raw_socket( + client_io, + tokio_tungstenite::tungstenite::protocol::Role::Client, + None, + ) + .await; + let (_stream, pumps) = split_tungstenite_to_bytes(ws); + + let write_side = tokio::spawn(async move { + let peer = tokio_tungstenite::WebSocketStream::from_raw_socket( + server_io, + tokio_tungstenite::tungstenite::protocol::Role::Server, + None, + ) + .await; + let (mut sink, mut reader) = peer.split(); + use futures::{SinkExt, StreamExt}; + let _ = sink + .send(tokio_tungstenite::tungstenite::Message::Text( + "text frame".into(), + )) + .await; + let _ = reader.next().await; + }); + + let mut eof_rx = pumps.read_eof(); + let changed = tokio::time::timeout(std::time::Duration::from_secs(5), eof_rx.changed()) + .await + .expect("read pump ends after the text frame (close requested)"); + let _ = changed; + + write_side.await.expect("peer task completes"); + } }