refactor(websocket): one generic pump implementation (WS-11, COV-03)

This commit is contained in:
2026-08-29 12:11:50 +00:00
parent 92cc11a74f
commit 1db0ea88e5
+435 -226
View File
@@ -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<u8>) -> Self::Msg;
}
/// The read stream's item for a flavor: `Result<Msg, Error>` (both
/// flavors surface read failures this way; the error type is flavor
/// specific and only ever inspected as "failed").
trait IntoWsResult<M: WsFraming> {
fn into_ws_result(self) -> Result<M::Msg, ()>;
}
impl<M, E> IntoWsResult<M> for Result<M::Msg, E>
where
M: WsFraming,
{
fn into_ws_result(self) -> Result<M::Msg, ()> {
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<M, S, F>(
mut ws_stream: S,
read_tx: mpsc::Sender<Vec<u8>>,
write_tx_for_read: futures_mpsc::Sender<WriteMsg>,
write_error_slot_for_read: WriteErrorSlot,
on_end: F,
) where
S: futures::Stream + Unpin,
M: WsFraming,
S::Item: IntoWsResult<M>,
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<M, S>(
ws_sink: &mut futures::stream::SplitSink<S, <M as WsFraming>::Msg>,
slot: &WriteErrorSlot,
reason: &'static str,
) where
M: WsFraming,
S: futures::Sink<<M as WsFraming>::Msg> + Unpin,
{
send_write_error(slot, reason);
let _ = ws_sink
.send(<M as WsFraming>::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<M, S>(
mut ws_sink: futures::stream::SplitSink<S, <M as WsFraming>::Msg>,
mut write_rx: futures_mpsc::Receiver<WriteMsg>,
write_error_slot: WriteErrorSlot,
) where
S: futures::Sink<<M as WsFraming>::Msg> + Unpin,
M: WsFraming,
{
let mut pending: Vec<u8> = Vec::new();
while let Some(msg) = write_rx.next().await {
match msg {
WriteMsg::CloseWith(code) => {
let _ = ws_sink
.send(<M as WsFraming>::close_message(
code,
"text messages not supported",
))
.await;
break;
}
WriteMsg::Bytes(b) => {
if b.len() > PENDING_BUFFER_CAP {
fail_write_pump::<M, _>(
&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::<M, _>(
&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<u8> = pending.drain(..total).collect();
for piece in chunk.chunks(WS_MESSAGE_CAP) {
if ws_sink
.send(<M as WsFraming>::binary_message(piece.to_vec()))
.await
.is_err()
{
return;
}
}
}
if pending.len() > PENDING_BUFFER_CAP {
fail_write_pump::<M, _>(
&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<bool>,
}
@@ -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<bool> {
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<u8>) -> 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::<Vec<u8>>(READ_SLOTS);
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let (write_tx, write_rx) = futures_mpsc::channel::<WriteMsg>(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::<AxumFraming, _, _>(
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<u8> = 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<u8> = 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::<AxumFraming, _>(
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<u8>) -> 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<S>(
socket: tokio_tungstenite::WebSocketStream<S>,
@@ -448,9 +632,9 @@ pub fn split_tungstenite_to_bytes<S>(
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::<Vec<u8>>(READ_SLOTS);
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let (write_tx, write_rx) = futures_mpsc::channel::<WriteMsg>(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::<TungsteniteFraming, _, _>(
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<u8> = 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<u8> = 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::<TungsteniteFraming, _>(
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");
}
}