1029 lines
38 KiB
Rust
1029 lines
38 KiB
Rust
//! The WS ↔ byte-stream adapter (OQ-01) — the single seam between axum's
|
||
//! message-oriented `WebSocket` and alkcall's byte-oriented channels
|
||
//! machinery (`AsyncRead` + `AsyncWrite`).
|
||
//!
|
||
//! Validated by the ws-byte-adapter POC (`/workspace/ws-byte-adapter-poc/`,
|
||
//! OQ-01 GO); this is the production shape.
|
||
//!
|
||
//! 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. 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
|
||
//! length) and emits one WS binary message per chunk, splitting chunks
|
||
//! over the 1 MiB message cap (legal — the receiver's boundary is the
|
||
//! 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. 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).
|
||
//!
|
||
//! Close mapping: WS close (either side) → read EOF → the demux clears
|
||
//! all channels (REQ-CH-02) and the dispatch loop fails outstanding
|
||
//! pendings. `AsyncWrite::shutdown` closes the WS sink after the queued
|
||
//! bytes drain (the mux's EOF sentinels ride the same queue).
|
||
//!
|
||
//! Shared with the `from_wss` consumer path (ADR-070): one
|
||
//! implementation, both directions ([`WsFraming`] + the generic
|
||
//! [`run_read_pump`] / [`run_write_pump`] instantiated once per socket
|
||
//! flavor — WS-11/COV-03).
|
||
|
||
use std::{
|
||
io,
|
||
pin::Pin,
|
||
sync::{Arc, Mutex as StdMutex},
|
||
task::{Context, Poll},
|
||
};
|
||
|
||
use axum::extract::ws::{CloseFrame, Message as AxumMessage, WebSocket};
|
||
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
|
||
/// is the chunk header, not the message. alkcall's MAX_CHUNK_LEN is
|
||
/// 16 MiB.
|
||
pub const WS_MESSAGE_CAP: usize = 1024 * 1024;
|
||
|
||
/// Inbound buffer: slots × in-flight message bytes. The WS read task
|
||
/// awaits `send` when full — the backpressure mechanism (OQ-01a).
|
||
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;
|
||
|
||
/// 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).
|
||
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<Mutex<Option<_>>>` shape exists because `oneshot::Sender` is
|
||
/// not `Clone` and two task bodies need the sender.
|
||
type WriteErrorSlot = Arc<StdMutex<Option<oneshot::Sender<&'static str>>>>;
|
||
|
||
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);
|
||
}
|
||
}
|
||
|
||
/// 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 {
|
||
read_rx: mpsc::Receiver<Vec<u8>>,
|
||
read_buf: Vec<u8>,
|
||
read_pos: usize,
|
||
eof: bool,
|
||
write_tx: Option<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
|
||
/// semantics); they end on their own when the socket halves close.
|
||
/// `abort()` is available for forced teardown.
|
||
pub struct WsPumps {
|
||
read_task: tokio::task::JoinHandle<()>,
|
||
write_task: tokio::task::JoinHandle<()>,
|
||
#[cfg(any(test, feature = "wss"))]
|
||
read_eof: tokio::sync::watch::Sender<bool>,
|
||
}
|
||
|
||
impl WsPumps {
|
||
pub fn abort(&self) {
|
||
self.read_task.abort();
|
||
self.write_task.abort();
|
||
}
|
||
|
||
/// A lossless receiver for the WS read-EOF signal (socket close from
|
||
/// either side): the watch channel retains the latest value, so an
|
||
/// EOF signaled at any point — including before the receiver is
|
||
/// 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(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 (ws_sink, ws_stream) = socket.split();
|
||
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
|
||
let (write_tx, write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
|
||
let (write_error_slot, write_error_rx) = make_write_error_slot();
|
||
|
||
#[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(any(test, feature = "wss"))]
|
||
let read_eof_for_task = read_eof.clone();
|
||
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);
|
||
}
|
||
},
|
||
));
|
||
|
||
let write_task = tokio::spawn(run_write_pump::<AxumFraming, _>(
|
||
ws_sink,
|
||
write_rx,
|
||
write_error_slot,
|
||
));
|
||
|
||
(
|
||
WsByteStream {
|
||
read_rx,
|
||
read_buf: Vec::new(),
|
||
read_pos: 0,
|
||
eof: false,
|
||
write_tx: Some(write_tx),
|
||
write_open: true,
|
||
write_error: write_error_rx,
|
||
write_failed: false,
|
||
},
|
||
WsPumps {
|
||
read_task,
|
||
write_task,
|
||
#[cfg(any(test, feature = "wss"))]
|
||
read_eof,
|
||
},
|
||
)
|
||
}
|
||
|
||
impl AsyncRead for WsByteStream {
|
||
fn poll_read(
|
||
self: Pin<&mut Self>,
|
||
cx: &mut Context<'_>,
|
||
buf: &mut tokio::io::ReadBuf<'_>,
|
||
) -> Poll<io::Result<()>> {
|
||
let this = self.get_mut();
|
||
loop {
|
||
if this.read_pos < this.read_buf.len() {
|
||
let n = (this.read_buf.len() - this.read_pos).min(buf.remaining());
|
||
let end = this.read_pos + n;
|
||
buf.put_slice(&this.read_buf[this.read_pos..end]);
|
||
this.read_pos = end;
|
||
if this.read_pos == this.read_buf.len() {
|
||
this.read_buf.clear();
|
||
this.read_pos = 0;
|
||
}
|
||
return Poll::Ready(Ok(()));
|
||
}
|
||
if this.eof {
|
||
return Poll::Ready(Ok(()));
|
||
}
|
||
match this.read_rx.poll_recv(cx) {
|
||
Poll::Ready(Some(bytes)) => {
|
||
this.read_buf = bytes;
|
||
this.read_pos = 0;
|
||
}
|
||
Poll::Ready(None) => {
|
||
this.eof = true;
|
||
}
|
||
Poll::Pending => return Poll::Pending,
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
impl AsyncWrite for WsByteStream {
|
||
fn poll_write(
|
||
self: Pin<&mut Self>,
|
||
cx: &mut Context<'_>,
|
||
buf: &[u8],
|
||
) -> Poll<io::Result<usize>> {
|
||
let this = self.get_mut();
|
||
if !this.write_open {
|
||
return Poll::Ready(Err(io::Error::new(
|
||
io::ErrorKind::BrokenPipe,
|
||
"ws stream shut down",
|
||
)));
|
||
}
|
||
if let Some(err) = this.poll_write_error() {
|
||
return Poll::Ready(Err(err));
|
||
}
|
||
let Some(write_tx) = this.write_tx_mut() else {
|
||
return Poll::Ready(Err(io::Error::new(
|
||
io::ErrorKind::BrokenPipe,
|
||
"ws stream shut down",
|
||
)));
|
||
};
|
||
match write_tx.poll_ready(cx) {
|
||
Poll::Ready(Ok(())) => match write_tx.try_send(WriteMsg::Bytes(buf.to_vec())) {
|
||
Ok(()) => Poll::Ready(Ok(buf.len())),
|
||
Err(_disconnected_or_full_race) => Poll::Ready(Err(io::Error::new(
|
||
io::ErrorKind::BrokenPipe,
|
||
"ws writer closed",
|
||
))),
|
||
},
|
||
Poll::Ready(Err(_send_error)) => Poll::Ready(Err(io::Error::new(
|
||
io::ErrorKind::BrokenPipe,
|
||
"ws writer closed",
|
||
))),
|
||
Poll::Pending => Poll::Pending,
|
||
}
|
||
}
|
||
|
||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||
Poll::Ready(Ok(()))
|
||
}
|
||
|
||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||
let this = self.get_mut();
|
||
if this.write_open {
|
||
this.write_open = false;
|
||
// WS-07: drop the *held* sender, not a fresh clone — the
|
||
// clone left the pump-side channel open for the stream's
|
||
// lifetime and the write pump never observed the queue
|
||
// close, so the documented trailing `ws_sink.close()` (the
|
||
// WS Close frame at shutdown) never ran. Dropping the held
|
||
// sender means the channel closes for the pump once the
|
||
// read task's own clone is also gone (its loop has ended —
|
||
// peer close or read error; at a local shutdown the read
|
||
// half ends with the connection), letting the write pump
|
||
// drain the queued bytes and close the sink.
|
||
drop(this.write_tx.take());
|
||
}
|
||
Poll::Ready(Ok(()))
|
||
}
|
||
}
|
||
|
||
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,
|
||
}
|
||
}
|
||
|
||
/// The write pump's fatal error, if any, and whether the write
|
||
/// queue still accepts sends (WS-07): the stream takes the sender
|
||
/// at shutdown so the pump-side channel close (and the sink's
|
||
/// trailing `ws_sink.close()`) actually happens.
|
||
fn write_tx_mut(&mut self) -> Option<&mut futures_mpsc::Sender<WriteMsg>> {
|
||
self.write_tx.as_mut()
|
||
}
|
||
}
|
||
|
||
/// 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. 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>,
|
||
) -> (WsByteStream, WsPumps)
|
||
where
|
||
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
|
||
{
|
||
let (ws_sink, ws_stream) = socket.split();
|
||
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_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;
|
||
|
||
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(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(run_write_pump::<TungsteniteFraming, _>(
|
||
ws_sink,
|
||
write_rx,
|
||
write_error_slot,
|
||
));
|
||
|
||
(
|
||
WsByteStream {
|
||
read_rx,
|
||
read_buf: Vec::new(),
|
||
read_pos: 0,
|
||
eof: false,
|
||
write_tx: Some(write_tx),
|
||
write_open: true,
|
||
write_error: write_error_rx,
|
||
write_failed: false,
|
||
},
|
||
WsPumps {
|
||
read_task,
|
||
write_task,
|
||
#[cfg(any(test, feature = "wss"))]
|
||
read_eof,
|
||
},
|
||
)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
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
|
||
/// 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}"
|
||
);
|
||
}
|
||
|
||
/// 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");
|
||
}
|
||
|
||
/// 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");
|
||
}
|
||
|
||
/// WS-07 acceptance: `AsyncWrite::poll_shutdown` drops the *held*
|
||
/// write sender, so the write pump observes the queue close and
|
||
/// runs its trailing `ws_sink.close()` — the peer receives the WS
|
||
/// Close frame after the queued chunk drains. Both senders must be
|
||
/// gone for the close: the stream takes its own at shutdown and
|
||
/// the read pump's clone drops when its loop ends (the peer's
|
||
/// Close frame below). With the bug (a fresh clone dropped
|
||
/// instead) the pump-side channel never closed and the peer never
|
||
/// saw a Close.
|
||
#[tokio::test]
|
||
async fn tungstenite_shutdown_closes_the_write_sink_toward_the_peer() {
|
||
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 peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
|
||
server_io,
|
||
tokio_tungstenite::tungstenite::protocol::Role::Server,
|
||
None,
|
||
)
|
||
.await;
|
||
let (mut peer_sink, mut peer_stream) = peer.split();
|
||
use futures::{SinkExt, StreamExt};
|
||
|
||
// One well-framed chunk so the pump emits one binary message.
|
||
let mut chunk = vec![0u8; 8];
|
||
chunk[4..8].copy_from_slice(&4u32.to_be_bytes());
|
||
chunk.extend_from_slice(b"pay!");
|
||
stream.write_all(&chunk).await.expect("write accepted");
|
||
stream.flush().await.expect("flush");
|
||
|
||
// The pump emits the chunk promptly; the peer consumes it.
|
||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||
loop {
|
||
match peer_stream.next().await {
|
||
Some(Ok(tokio_tungstenite::tungstenite::Message::Binary(b))) => {
|
||
assert_eq!(&*b, &chunk, "queued bytes drain before the close");
|
||
return;
|
||
}
|
||
Some(Ok(_)) => continue,
|
||
other => panic!("peer stream ended early: {other:?}"),
|
||
}
|
||
}
|
||
})
|
||
.await
|
||
.expect("chunk emitted within deadline");
|
||
|
||
// Local shutdown takes the stream's own sender.
|
||
stream.shutdown().await.expect("shutdown runs");
|
||
|
||
// End the read loop: the peer sends its Close frame (its stream
|
||
// stays readable). The read pump's sender clone drops; the
|
||
// write queue closes; the pump runs its trailing ws_sink.close()
|
||
// and the peer observes the close reply.
|
||
peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Close(None))
|
||
.await
|
||
.expect("peer sends close");
|
||
|
||
let saw_close = tokio::time::timeout(std::time::Duration::from_secs(5), async move {
|
||
loop {
|
||
match peer_stream.next().await {
|
||
None => return false,
|
||
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_))) => return true,
|
||
Some(Ok(_)) => continue,
|
||
Some(Err(_)) => return false,
|
||
}
|
||
}
|
||
})
|
||
.await
|
||
.expect("peer observes the close reply after our shutdown");
|
||
assert!(saw_close, "the shutdown path emitted a WS Close frame");
|
||
}
|
||
}
|