glm-5.3-flash
bd87c4613b
- AxumFraming arms were exercised only indirectly via tungstenite (shared generic pumps); drive the axum message types directly with an in-process fake WebSocket (futures mpsc-backed Sink+Stream stand-in for the split halves) - text test: read pump maps a text message to the WriteMsg close carrying 1002 + the text reason - cap-trip test: an above-MAX_CHUNK_LEN header through the write pump closes with 1011 naming the violation Verification: scripts/verify.sh OK (345 passed), test-support suite ok, clippy -D warnings clean, fmt clean
1915 lines
74 KiB
Rust
1915 lines
74 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. Single `AsyncWrite` calls above
|
||
//! `PENDING_BUFFER_CAP` (WS-14) are rejected in `poll_write` *before*
|
||
//! the bytes enter the write queue — the mux sees the `InvalidData`
|
||
//! stream error directly instead of the write pump closing the wire
|
||
//! with 1011 mid-stream and truncating. The `pending` accumulator
|
||
//! still carries the `PENDING_BUFFER_CAP` bound (WS-05) as
|
||
//! defense-in-depth: it is invariantly satisfied once `poll_write`
|
||
//! enforces the cap pre-queue, so it holds at most one
|
||
//! unemitted-in-full chunk plus its header. The remaining write-side
|
||
//! bounds are slot-bounded (`WRITE_SLOTS` × `WS_MESSAGE_CAP`) but
|
||
//! time-unbounded — bounded by the WS-18 write-progress timeout in
|
||
//! `run_write_pump`. 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).
|
||
//!
|
||
//! Idle-read timeout (WS-01, WS-13 semantics): the knob
|
||
//! (`DEFAULT_WS_IDLE_TIMEOUT`, deployment-adjustable via
|
||
//! `HttpAdapter::with_ws_idle_timeout`, disable via `None`) evicts a
|
||
//! connection whose inbound stream produces **no completed chunk for
|
||
//! the whole window** — the deadline resets on demux progress (bytes
|
||
//! forwarded into `read_tx` that complete 8-byte-header-framed chunks),
|
||
//! never on WS message arrival, so a forever-dribble inside a declared
|
||
//! chunk hits the deadline even though messages keep arriving, while a
|
||
//! peer delivering complete chunks — however slowly per-message —
|
||
//! re-arms the window with each one.
|
||
//!
|
||
//! Legitimate silence (WS-13 decision, recorded — option (b)): there
|
||
//! is deliberately **no WS ping/pong keepalive**. A keepalive can only
|
||
//! rescue app-silence by re-arming the deadline, which would reopen the
|
||
//! dribble hole it exists to seal (pong = traffic from the attacker's
|
||
//! point of view); instead, 60 s of *no chunk progress* is an
|
||
//! intentional eviction line even for a silent subscription — a
|
||
//! long-lived quiet subscription that must survive past the window
|
||
//! (with server-side pushes; see the keep-alive discussion in
|
||
//! `websocket.md`) is exactly the deployment that dials
|
||
//! `with_ws_idle_timeout(None)` and leans on the other bounds
|
||
//! (`WsSessions::abort` eviction, the write-side caps). Read eviction
|
||
//! closes with 1001 (Going Away) — a normal connection end from the
|
||
//! demux's point of view (EOF → channels cleared, pendings failed),
|
||
//! not a protocol error.
|
||
//!
|
||
//! 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},
|
||
time::Duration,
|
||
};
|
||
|
||
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).
|
||
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). Enforced
|
||
/// in `poll_write` *before* a message enters the write queue (WS-14):
|
||
/// a single over-cap call is rejected with `InvalidData` at the
|
||
/// `AsyncWrite` boundary, so the invariant — `pending` never holds
|
||
/// more than one not-yet-emitted chunk plus its header, because the
|
||
/// parser drains complete chunks greedily and alkcall's mux passes
|
||
/// each payload (up to 16 MiB) as one `AsyncWrite` call — is
|
||
/// maintained by construction. The in-pump check remains as
|
||
/// defense-in-depth.
|
||
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;
|
||
|
||
/// Idle-read close code (WS-01/WS-13): a read side that produces no
|
||
/// demux progress (no completed inbound chunk) past the configured
|
||
/// window is closed with 1001 (Going Away) — a normal connection end
|
||
/// from the demux's point of view (EOF → channels cleared, pendings
|
||
/// failed), not a protocol error.
|
||
pub const WS_GOING_AWAY: u16 = 1001;
|
||
|
||
/// Default write-progress timeout for the WS write pump (WS-18): one
|
||
/// outbound WS send that stays unsent past this window (the peer
|
||
/// stopped reading) evicts the connection — the write-side analog of
|
||
/// the idle-read knob. The bound is per `send` call, not per
|
||
/// connection: a slow-but-draining peer resets it with every message
|
||
/// that gets out.
|
||
pub const DEFAULT_WS_WRITE_TIMEOUT: Duration = Duration::from_secs(60);
|
||
|
||
/// Default idle-read timeout for the WS pumps (WS-01, WS-13): a
|
||
/// connection whose inbound stream completes **no chunk** within this
|
||
/// window is evicted — the deadline resets on demux progress (complete
|
||
/// chunks forwarded into `read_tx`), not on WS message arrival, so a
|
||
/// forever-dribble inside a declared chunk still hits it.
|
||
///
|
||
/// This is an intentional no-progress eviction line, *not* a
|
||
/// transport-idle bound: there is no WS ping/pong keepalive, and
|
||
/// app-silence that outlasts the window (a quiet subscription) is
|
||
/// evicted with 1001 by design — see the module doc's "Legitimate
|
||
/// silence" decision. A deployment running long-lived silent
|
||
/// subscriptions disables the knob with
|
||
/// `HttpAdapter::with_ws_idle_timeout(None)` (`None`, not zero — zero
|
||
/// is not a meaningful window).
|
||
pub const DEFAULT_WS_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
|
||
|
||
/// Observes the inbound byte stream for chunk framing (WS-13): counts
|
||
/// complete chunks whose bytes have been forwarded into `read_tx` —
|
||
/// the progress signal the idle-read deadline resets on. The parse
|
||
/// walk mirrors the demux loop's byte-for-byte (8-byte header →
|
||
/// payload skip; an over-`MAX_CHUNK_LEN` length is skipped like the
|
||
/// demux's TooLarge arm and parsing continues), with O(1) state: no
|
||
/// byte copying, and the forwarded byte stream is unmodified.
|
||
struct InboundChunkProgress {
|
||
header: [u8; 8],
|
||
header_fill: u8,
|
||
payload_remaining: u64,
|
||
}
|
||
|
||
impl InboundChunkProgress {
|
||
fn new() -> Self {
|
||
Self {
|
||
header: [0u8; 8],
|
||
header_fill: 0,
|
||
payload_remaining: 0,
|
||
}
|
||
}
|
||
|
||
fn observe(&mut self, bytes: &[u8]) -> u32 {
|
||
let mut completed = 0u32;
|
||
for &byte in bytes {
|
||
if self.payload_remaining > 0 {
|
||
self.payload_remaining -= 1;
|
||
if self.payload_remaining == 0 {
|
||
completed += 1;
|
||
}
|
||
continue;
|
||
}
|
||
self.header[self.header_fill as usize] = byte;
|
||
self.header_fill += 1;
|
||
if self.header_fill == 8 {
|
||
self.header_fill = 0;
|
||
let len = u32::from_be_bytes([
|
||
self.header[4],
|
||
self.header[5],
|
||
self.header[6],
|
||
self.header[7],
|
||
]);
|
||
self.payload_remaining = len as u64;
|
||
if len == 0 {
|
||
completed += 1;
|
||
}
|
||
}
|
||
}
|
||
completed
|
||
}
|
||
}
|
||
|
||
/// 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 and reason (e.g. the 1002
|
||
/// text-rejection).
|
||
CloseWith(u16, &'static str),
|
||
}
|
||
|
||
/// 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);
|
||
}
|
||
}
|
||
|
||
/// The close-frame reason strings, one distinct per cause (WS-15): the
|
||
/// numeric code alone collides across causes (the size-cap close and
|
||
/// the write-stall close share 1011), and an operator reading a wire
|
||
/// capture should be able to tell them apart without the stream error
|
||
/// channel. Sent in both the close frame the peer receives and the
|
||
/// stream-error diagnostic; `reason_for` maps a bare close code to its
|
||
/// canonical default for the fall-through arms.
|
||
pub(crate) mod close_reason {
|
||
pub(crate) const IDLE_READ_TIMEOUT: &str =
|
||
"connection made no inbound chunk progress past the read timeout";
|
||
pub(crate) const TEXT_NOT_SUPPORTED: &str = "text messages not supported";
|
||
pub(crate) const INBOUND_FRAME_REJECTED: &str = "inbound frame rejected (size cap)";
|
||
}
|
||
|
||
/// 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).
|
||
///
|
||
/// Idle-read timeout (WS-01, WS-13 semantics): when `idle_timeout` is
|
||
/// non-zero, each next-message await is bounded by the *remaining*
|
||
/// budget — the window minus the time since the last **demux
|
||
/// progress** event (bytes actually forwarded into `read_tx` that
|
||
/// complete inbound chunks, i.e. frames the demux will route). WS
|
||
/// message arrival resets nothing: a peer dribbling bytes into a
|
||
/// declared chunk forever hits the deadline even though messages keep
|
||
/// arriving; a peer delivering complete chunks (even slowly, one per
|
||
/// message) refreshes the window with each chunk. Budget exhaustion
|
||
/// sends the 1001 GoingAway close to the peer (a normal connection
|
||
/// end for the demux/from_wss EOF machinery, not a protocol error)
|
||
/// and ends the loop, so the adapter's read half sees EOF and the
|
||
/// channels teardown runs.
|
||
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,
|
||
idle_timeout: Option<Duration>,
|
||
on_end: F,
|
||
) where
|
||
S: futures::Stream + Unpin,
|
||
M: WsFraming,
|
||
S::Item: IntoWsResult<M>,
|
||
F: FnOnce(),
|
||
{
|
||
let mut progress = InboundChunkProgress::new();
|
||
let mut last_progress_at = tokio::time::Instant::now();
|
||
loop {
|
||
let budget = idle_timeout.map(|window| window.saturating_sub(last_progress_at.elapsed()));
|
||
let msg = match budget {
|
||
None => ws_stream.next().await,
|
||
Some(budget) => match tokio::time::timeout(budget, ws_stream.next()).await {
|
||
Ok(msg) => msg,
|
||
Err(_elapsed) => {
|
||
send_write_error(&write_error_slot_for_read, close_reason::IDLE_READ_TIMEOUT);
|
||
let _ = write_tx_for_read
|
||
.clone()
|
||
.send(WriteMsg::CloseWith(
|
||
WS_GOING_AWAY,
|
||
close_reason::IDLE_READ_TIMEOUT,
|
||
))
|
||
.await;
|
||
break;
|
||
}
|
||
},
|
||
};
|
||
let Some(msg) = msg else {
|
||
break;
|
||
};
|
||
match msg.into_ws_result() {
|
||
Ok(m) => {
|
||
if let Some(b) = M::binary(&m) {
|
||
let completed = progress.observe(b.as_ref());
|
||
if read_tx.send(b.as_ref().to_vec()).await.is_err() {
|
||
break;
|
||
}
|
||
if completed > 0 {
|
||
last_progress_at = tokio::time::Instant::now();
|
||
}
|
||
} else if M::is_text(&m) {
|
||
let _ = write_tx_for_read
|
||
.clone()
|
||
.send(WriteMsg::CloseWith(
|
||
WS_PROTOCOL_ERROR,
|
||
close_reason::TEXT_NOT_SUPPORTED,
|
||
))
|
||
.await;
|
||
break;
|
||
} else if M::is_close(&m) {
|
||
break;
|
||
}
|
||
}
|
||
Err(()) => {
|
||
send_write_error(
|
||
&write_error_slot_for_read,
|
||
close_reason::INBOUND_FRAME_REJECTED,
|
||
);
|
||
let _ = write_tx_for_read
|
||
.clone()
|
||
.send(WriteMsg::CloseWith(
|
||
WS_INTERNAL_ERROR,
|
||
close_reason::INBOUND_FRAME_REJECTED,
|
||
))
|
||
.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
|
||
/// drains; `CloseWith` requests (`WriteMsg::CloseWith`) map to a close
|
||
/// frame with the requested code and its per-cause reason (WS-15) —
|
||
/// the code alone is ambiguous across causes (idle 1001 vs protocol
|
||
/// 1002 vs internal 1011 arms), so the pump never invents one.
|
||
/// 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 the requested code and its per-cause reason (WS-15) —
|
||
/// the code alone is ambiguous across causes (idle 1001 vs protocol
|
||
/// 1002 vs internal 1011 arms), so the pump never invents one.
|
||
///
|
||
/// Write-progress timeout (WS-18): a peer that stops reading parks
|
||
/// the pump inside a WS send — the slot-bounded queue bounds memory,
|
||
/// but the stall was time-unbounded. When `write_timeout` is
|
||
/// `Some(window)`, every WS send must complete within the window:
|
||
/// exceeding it signals the stream error (naming the stall) and ends
|
||
/// the pump **without** a close frame — the peer cannot receive one
|
||
/// through a clogged socket, and the close send would park on it too.
|
||
/// The window is per send, so a slow-but-moving peer survives; only a
|
||
/// fully stalled sink trips it.
|
||
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,
|
||
write_timeout: Option<Duration>,
|
||
) 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, reason) => {
|
||
let close = <M as WsFraming>::close_message(code, reason);
|
||
match write_timeout {
|
||
None => {
|
||
let _ = ws_sink.send(close).await;
|
||
}
|
||
Some(window) => {
|
||
let _ = tokio::time::timeout(window, ws_sink.send(close)).await;
|
||
}
|
||
}
|
||
break;
|
||
}
|
||
WriteMsg::Bytes(b) => {
|
||
debug_assert!(b.len() <= PENDING_BUFFER_CAP);
|
||
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 !send_bounded::<M, _>(&mut ws_sink, piece, write_timeout, &write_error_slot)
|
||
.await
|
||
{
|
||
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;
|
||
}
|
||
|
||
/// One WS send under the WS-18 write-progress bound: a send that
|
||
/// outlasts the window means the peer stopped reading, so the pump
|
||
/// signals the error slot (`InvalidData` on the `AsyncWrite` half) and
|
||
/// ends without a close frame — a clogged socket cannot receive one,
|
||
/// and the close send would park on the same stall. On send error, end
|
||
/// silently (the sink is already broken). Returns whether the pump
|
||
/// should continue.
|
||
async fn send_bounded<M, S>(
|
||
ws_sink: &mut futures::stream::SplitSink<S, <M as WsFraming>::Msg>,
|
||
piece: &[u8],
|
||
write_timeout: Option<Duration>,
|
||
slot: &WriteErrorSlot,
|
||
) -> bool
|
||
where
|
||
M: WsFraming,
|
||
S: futures::Sink<<M as WsFraming>::Msg> + Unpin,
|
||
{
|
||
match write_timeout {
|
||
None => ws_sink
|
||
.send(<M as WsFraming>::binary_message(piece.to_vec()))
|
||
.await
|
||
.is_ok(),
|
||
Some(window) => {
|
||
let send = ws_sink.send(<M as WsFraming>::binary_message(piece.to_vec()));
|
||
match tokio::time::timeout(window, send).await {
|
||
Ok(Ok(())) => true,
|
||
Ok(Err(_)) => false,
|
||
Err(_elapsed) => {
|
||
send_write_error(
|
||
slot,
|
||
"connection made no write progress past the write timeout",
|
||
);
|
||
false
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 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 {
|
||
/// Abort both pump tasks (forced session teardown; the remote sees
|
||
/// an abrupt close, not a graceful one).
|
||
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 = axum::body::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 with the
|
||
/// default idle-read timeout (`crate::websocket::DEFAULT_WS_IDLE_TIMEOUT`).
|
||
/// See [`split_ws_to_bytes_idle`] for the configurable form.
|
||
pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
|
||
split_ws_to_bytes_idle(socket, Some(DEFAULT_WS_IDLE_TIMEOUT))
|
||
}
|
||
|
||
/// 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`.
|
||
/// `idle_timeout` (WS-01): `None` = no idle bound, `Some(d)` closes
|
||
/// the read with 1001 after `d` without inbound chunk progress. The
|
||
/// write pump runs with [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`]
|
||
/// — see [`split_ws_to_bytes_idle_with_write`] for the WS-18
|
||
/// configurable form.
|
||
pub fn split_ws_to_bytes_idle(
|
||
socket: WebSocket,
|
||
idle_timeout: Option<Duration>,
|
||
) -> (WsByteStream, WsPumps) {
|
||
split_ws_to_bytes_idle_with_write(socket, idle_timeout, None)
|
||
}
|
||
|
||
/// [`split_ws_to_bytes_idle`] with an explicit WS-18 write-progress
|
||
/// window: `None` = the crate default
|
||
/// ([`DEFAULT_WS_WRITE_TIMEOUT`](crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)),
|
||
/// `Some(d)` a deployment-set window.
|
||
pub fn split_ws_to_bytes_idle_with_write(
|
||
socket: WebSocket,
|
||
idle_timeout: Option<Duration>,
|
||
write_timeout: Option<Duration>,
|
||
) -> (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,
|
||
idle_timeout,
|
||
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,
|
||
Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)),
|
||
));
|
||
|
||
(
|
||
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_ref() else {
|
||
return Poll::Ready(Err(io::Error::new(
|
||
io::ErrorKind::BrokenPipe,
|
||
"ws stream shut down",
|
||
)));
|
||
};
|
||
if buf.len() > PENDING_BUFFER_CAP {
|
||
return Poll::Ready(Err(io::Error::new(
|
||
io::ErrorKind::InvalidData,
|
||
format!(
|
||
"write of {} bytes exceeds the pending buffer cap ({})",
|
||
buf.len(),
|
||
PENDING_BUFFER_CAP
|
||
),
|
||
)));
|
||
}
|
||
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-sender getter (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. WS-14's pre-send cap check
|
||
/// keeps the cap fault off the wire path entirely — a `try_send`
|
||
/// failure here can only be the pump ending between `poll_ready`
|
||
/// and `try_send`.
|
||
fn write_tx_ref(&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 = axum::body::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` with
|
||
/// the default idle-read timeout
|
||
/// ([`DEFAULT_WS_IDLE_TIMEOUT`](crate::websocket::DEFAULT_WS_IDLE_TIMEOUT)).
|
||
/// See [`split_tungstenite_to_bytes_idle`] for the configurable form.
|
||
#[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,
|
||
{
|
||
split_tungstenite_to_bytes_idle(socket, Some(DEFAULT_WS_IDLE_TIMEOUT))
|
||
}
|
||
|
||
/// Client-side split for a tokio-tungstenite `WebSocketStream`
|
||
/// (`from_wss`, ADR-070): the same WS↔byte-stream seam as
|
||
/// [`split_ws_to_bytes_idle`], 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. `idle_timeout` (WS-01): `None` = no idle bound, `Some(d)`
|
||
/// closes the read with 1001 after `d` without an inbound WS message.
|
||
#[cfg(any(test, feature = "wss"))]
|
||
pub fn split_tungstenite_to_bytes_idle<S>(
|
||
socket: tokio_tungstenite::WebSocketStream<S>,
|
||
idle_timeout: Option<Duration>,
|
||
) -> (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,
|
||
idle_timeout,
|
||
move || {
|
||
let _ = read_eof_for_task.send(true);
|
||
},
|
||
));
|
||
|
||
let write_task = tokio::spawn(run_write_pump::<TungsteniteFraming, _>(
|
||
ws_sink,
|
||
write_rx,
|
||
write_error_slot,
|
||
Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT),
|
||
));
|
||
|
||
(
|
||
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,
|
||
},
|
||
)
|
||
}
|
||
|
||
/// Test-only split with both knobs explicit (WS-18 acceptance: the
|
||
/// stall test scales the write window down; other tests keep the
|
||
/// default).
|
||
#[cfg(test)]
|
||
pub(crate) fn split_tungstenite_to_bytes_idle_with_write<S>(
|
||
socket: tokio_tungstenite::WebSocketStream<S>,
|
||
idle_timeout: Option<Duration>,
|
||
write_timeout: Option<Duration>,
|
||
) -> (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,
|
||
idle_timeout,
|
||
move || {
|
||
let _ = read_eof_for_task.send(true);
|
||
},
|
||
));
|
||
|
||
let write_task = tokio::spawn(run_write_pump::<TungsteniteFraming, _>(
|
||
ws_sink,
|
||
write_rx,
|
||
write_error_slot,
|
||
write_timeout,
|
||
));
|
||
|
||
(
|
||
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,
|
||
read_eof,
|
||
},
|
||
)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
|
||
/// WS-18 acceptance: a peer that stops reading (the sink clogs —
|
||
/// nothing is drained from the duplex) parks the write pump inside
|
||
/// a WS send; the pump must end within the knob (+ slack), surface
|
||
/// the stream error naming the write stall, and fail the
|
||
/// `AsyncWrite` half with `InvalidData`. Scaled test: both knobs
|
||
/// idle-read `None` (so only the write knob can evict) and the
|
||
/// write window at 150 ms.
|
||
#[tokio::test]
|
||
async fn tungstenite_write_stall_is_evicted_within_the_write_timeout() {
|
||
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 knob = std::time::Duration::from_millis(150);
|
||
let (mut stream, _pumps) = split_tungstenite_to_bytes_idle_with_write(ws, None, Some(knob));
|
||
|
||
let mut chunk = vec![0u8; 8];
|
||
chunk[4..8].copy_from_slice(&64u32.to_be_bytes());
|
||
chunk.extend_from_slice(&[0u8; 64]);
|
||
stream.write_all(&chunk).await.expect("chunk written");
|
||
stream.flush().await.expect("flush");
|
||
|
||
let started = tokio::time::Instant::now();
|
||
let err = loop {
|
||
match stream.write_all(&[0u8; 8]).await {
|
||
Err(e) => break e,
|
||
Ok(_) => assert!(
|
||
started.elapsed() < std::time::Duration::from_secs(5),
|
||
"stream never failed while the sink stayed clogged"
|
||
),
|
||
}
|
||
};
|
||
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
|
||
assert!(
|
||
err.to_string().contains("write timeout"),
|
||
"error names the violation: {err}"
|
||
);
|
||
assert!(
|
||
started.elapsed() < std::time::Duration::from_secs(5),
|
||
"eviction must happen within the knob (+ slack), not hang"
|
||
);
|
||
}
|
||
|
||
/// 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).
|
||
pub(crate) 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}"
|
||
);
|
||
}
|
||
|
||
/// A single `AsyncWrite` call of exactly `PENDING_BUFFER_CAP`
|
||
/// bytes parses and flushes through the pump — the cap check
|
||
/// rejects `> cap`, never `= cap`, so `MAX_CHUNK_LEN` payloads
|
||
/// (the mux's maximum single `AsyncWrite` call, 16 MiB + 8 header)
|
||
/// keep flowing (WS-14 edge).
|
||
#[tokio::test]
|
||
async fn tungstenite_write_at_the_cap_is_accepted() {
|
||
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; PENDING_BUFFER_CAP];
|
||
chunk[4..8].copy_from_slice(&(PENDING_BUFFER_CAP as u32 - 8).to_be_bytes());
|
||
stream.write_all(&chunk).await.expect("cap write accepted");
|
||
stream.flush().await.expect("flush");
|
||
|
||
tokio::time::timeout(std::time::Duration::from_secs(30), async {
|
||
let mut buf = [0u8; 2];
|
||
server_io
|
||
.read_exact(&mut buf)
|
||
.await
|
||
.expect("cap-sized chunk read");
|
||
})
|
||
.await
|
||
.expect("pump emitted the cap-sized chunk");
|
||
}
|
||
|
||
/// 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. (The pump-side
|
||
/// pending-cap fault this test family covered is unreachable for
|
||
/// single-call over-cap writes since WS-14 moved the check into
|
||
/// `poll_write`; the multi-write accumulation leg stays live via
|
||
/// the oversize-header test above, which strands the pump's
|
||
/// parser past a too-long header.)
|
||
#[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");
|
||
}
|
||
|
||
/// WS-15 acceptance (text arm): a text frame triggers the 1002
|
||
/// protocol-error close whose reason names the text cause —
|
||
/// distinct from the idle (1001) and oversize (1011) reasons the
|
||
/// other tests assert.
|
||
#[tokio::test]
|
||
async fn tungstenite_text_close_carries_protocol_reason() {
|
||
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 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};
|
||
|
||
let mut eof_rx = pumps.read_eof();
|
||
peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Text(
|
||
"text frame".into(),
|
||
))
|
||
.await
|
||
.expect("text write accepted");
|
||
|
||
let (close, eof_fired) = tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||
let close: Option<(u16, String)> = loop {
|
||
match peer_stream.next().await {
|
||
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => {
|
||
break cf.map(|f| (u16::from(f.code), f.reason.to_string()))
|
||
}
|
||
Some(Ok(_)) => continue,
|
||
Some(Err(_)) | None => break None,
|
||
}
|
||
};
|
||
let _ = eof_rx.changed().await;
|
||
(close, *eof_rx.borrow())
|
||
})
|
||
.await
|
||
.expect("read pump ends after the text frame (close requested)");
|
||
|
||
let (code, reason) = close.expect("close frame with code + reason");
|
||
assert_eq!(code, WS_PROTOCOL_ERROR, "text frame closed with 1002");
|
||
assert_eq!(
|
||
reason, "text messages not supported",
|
||
"close reason names the text cause, got {reason:?}"
|
||
);
|
||
assert!(eof_fired, "EOF signal fired for the from_wss machinery");
|
||
}
|
||
|
||
/// WS-01 acceptance: a dribbling peer cannot park the read pump
|
||
/// message inside the window: the pump sends the 1001 GoingAway
|
||
/// close to the peer, ends, and fires the EOF watch signal — the
|
||
/// from_wss monitor/sweep input — so the demux/channel teardown
|
||
/// machinery treats it as a normal connection end.
|
||
#[tokio::test]
|
||
async fn idle_read_timeout_closes_a_stalled_connection_with_goingaway() {
|
||
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 knob = std::time::Duration::from_millis(150);
|
||
let (_stream, pumps) = split_tungstenite_to_bytes_idle(ws, Some(knob));
|
||
|
||
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
|
||
server_io,
|
||
tokio_tungstenite::tungstenite::protocol::Role::Server,
|
||
None,
|
||
)
|
||
.await;
|
||
let (_peer_sink, mut peer_stream) = peer.split();
|
||
use futures::StreamExt;
|
||
|
||
// The peer sends nothing (the stall). The read pump must end
|
||
// within the knob (+ slack), not hang: observe both the EOF
|
||
// signal the from_wss machinery consumes and the 1001 close
|
||
// frame the peer receives. WS-15: the close reason names the
|
||
// idle-read cause, not a leftover default.
|
||
let mut eof_rx = pumps.read_eof();
|
||
let (close_seen, eof_fired) =
|
||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||
let close: Option<(u16, String)> = loop {
|
||
match peer_stream.next().await {
|
||
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => {
|
||
break cf.map(|f| (u16::from(f.code), f.reason.to_string()))
|
||
}
|
||
Some(Ok(_)) => continue,
|
||
Some(Err(_)) | None => break None,
|
||
}
|
||
};
|
||
let _ = eof_rx.changed().await;
|
||
(close, *eof_rx.borrow())
|
||
})
|
||
.await
|
||
.expect("stalled connection must be torn down within the deadline");
|
||
|
||
let (code, reason) = close_seen.expect("close frame with code + reason");
|
||
assert_eq!(
|
||
code, WS_GOING_AWAY,
|
||
"peer received the 1001 GoingAway close"
|
||
);
|
||
assert!(
|
||
reason.contains("no inbound chunk progress"),
|
||
"close reason names the idle-read cause, got {reason:?}"
|
||
);
|
||
assert!(eof_fired, "EOF signal fired for the from_wss machinery");
|
||
}
|
||
|
||
/// The forever-dribble stall is bounded by *progress*, not
|
||
/// message arrival (WS-02 acceptance / WS-13 semantics): a peer
|
||
/// declaring one chunk (`[0: u32 BE][64: u32 BE]`) and dribbling
|
||
/// its 64 payload bytes one byte per message, half a window apart,
|
||
/// never completes the chunk — the read pump must end with the
|
||
/// 1001 GoingAway close within the knob even though messages keep
|
||
/// arriving, so the demux channel machinery sees EOF.
|
||
#[tokio::test]
|
||
async fn idle_read_timeout_bounds_a_forever_dribble_with_no_chunk_progress() {
|
||
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 knob = std::time::Duration::from_millis(200);
|
||
let (_stream, pumps) = split_tungstenite_to_bytes_idle(ws, Some(knob));
|
||
|
||
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};
|
||
|
||
let mut eof_rx = pumps.read_eof();
|
||
let dribble = tokio::spawn(async move {
|
||
let mut header = vec![0u8; 8];
|
||
header[4..8].copy_from_slice(&64u32.to_be_bytes());
|
||
peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Binary(
|
||
header.into(),
|
||
))
|
||
.await
|
||
.expect("header write accepted");
|
||
for byte in 0u8..64 {
|
||
tokio::time::sleep(knob / 4).await;
|
||
if peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Binary(
|
||
vec![byte].into(),
|
||
))
|
||
.await
|
||
.is_err()
|
||
{
|
||
return;
|
||
}
|
||
}
|
||
loop {
|
||
tokio::time::sleep(knob).await;
|
||
if peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Binary(
|
||
vec![0u8].into(),
|
||
))
|
||
.await
|
||
.is_err()
|
||
{
|
||
return;
|
||
}
|
||
}
|
||
});
|
||
|
||
let outcome = tokio::time::timeout(std::time::Duration::from_secs(10), async {
|
||
let close: Option<(u16, String)> = loop {
|
||
match peer_stream.next().await {
|
||
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => {
|
||
break cf.map(|f| (u16::from(f.code), f.reason.to_string()))
|
||
}
|
||
Some(Ok(_)) => continue,
|
||
Some(Err(_)) | None => break None,
|
||
}
|
||
};
|
||
let _ = eof_rx.changed().await;
|
||
close
|
||
})
|
||
.await
|
||
.expect("dribble must hit the progress deadline within 10 s");
|
||
dribble.abort();
|
||
let (code, reason) = outcome.expect("close frame with code + reason");
|
||
assert_eq!(
|
||
code, WS_GOING_AWAY,
|
||
"the forever-dribble is evicted with 1001 despite arriving messages"
|
||
);
|
||
assert!(
|
||
reason.contains("no inbound chunk progress"),
|
||
"close reason names the idle-read cause, got {reason:?}"
|
||
);
|
||
}
|
||
|
||
/// Progress semantics, survivor side (WS-13 accept 2): a session
|
||
/// whose inbound traffic keeps *making progress* (complete chunks
|
||
/// arriving, each within the window) is NOT disconnected even
|
||
/// though it never goes quiet in a way the WS-01 shape would
|
||
/// reward — every complete chunk re-arms the deadline from zero,
|
||
/// which is the "messages flowing that make progress" survival
|
||
/// the knob's spirit prescribes.
|
||
#[tokio::test]
|
||
async fn idle_read_timeout_survives_slow_productive_chunks() {
|
||
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 knob = std::time::Duration::from_millis(200);
|
||
let (mut stream, pumps) = split_tungstenite_to_bytes_idle(ws, Some(knob));
|
||
let mut evict_rx = pumps.read_eof();
|
||
|
||
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
|
||
server_io,
|
||
tokio_tungstenite::tungstenite::protocol::Role::Server,
|
||
None,
|
||
)
|
||
.await;
|
||
let (mut peer_sink, _peer_stream) = peer.split();
|
||
use futures::{SinkExt, StreamExt};
|
||
|
||
for round in 0u8..6u8 {
|
||
tokio::time::sleep(knob / 2).await;
|
||
let mut chunk = vec![0u8; 12];
|
||
chunk[0..4].copy_from_slice(&0u32.to_be_bytes());
|
||
chunk[4..8].copy_from_slice(&4u32.to_be_bytes());
|
||
chunk[8..12].copy_from_slice(&[round; 4]);
|
||
peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Binary(
|
||
chunk.clone().into(),
|
||
))
|
||
.await
|
||
.expect("chunk write accepted");
|
||
let mut buf = vec![0u8; 12];
|
||
tokio::time::timeout(knob, stream.read_exact(&mut buf))
|
||
.await
|
||
.expect("connection alive between productive chunks")
|
||
.expect("read");
|
||
assert_eq!(buf, chunk, "chunk bytes surfaced in order");
|
||
}
|
||
let evicted_early =
|
||
tokio::time::timeout(std::time::Duration::from_millis(0), evict_rx.changed()).await;
|
||
assert!(
|
||
evicted_early.is_err(),
|
||
"no eviction while chunks keep landing inside the window"
|
||
);
|
||
let mut header = vec![0u8; 8];
|
||
header[4..8].copy_from_slice(&0u32.to_be_bytes());
|
||
peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Binary(
|
||
header.into(),
|
||
))
|
||
.await
|
||
.expect("EOF chunk write accepted");
|
||
tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||
loop {
|
||
match stream.read(&mut [0u8; 1]).await {
|
||
Ok(0) => return,
|
||
Ok(_) => continue,
|
||
Err(e) => panic!("read failed: {e}"),
|
||
}
|
||
}
|
||
})
|
||
.await
|
||
.expect("demux EOF observed after the final chunk");
|
||
}
|
||
|
||
/// `idle_timeout = None` disables the deadline entirely — the
|
||
/// "no bound" arm of the WS-01 knob surface: no timer runs, so
|
||
/// messages arriving without chunk progress never evict.
|
||
#[tokio::test]
|
||
async fn idle_read_timeout_disabled_when_none() {
|
||
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_idle(ws, None);
|
||
|
||
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
|
||
server_io,
|
||
tokio_tungstenite::tungstenite::protocol::Role::Server,
|
||
None,
|
||
)
|
||
.await;
|
||
let (mut peer_sink, _peer_stream) = peer.split();
|
||
use futures::SinkExt;
|
||
|
||
let eof_rx = pumps.read_eof();
|
||
for _ in 0..4 {
|
||
tokio::time::sleep(std::time::Duration::from_millis(150)).await;
|
||
peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Binary(
|
||
vec![0u8].into(),
|
||
))
|
||
.await
|
||
.expect("dribble write accepted");
|
||
}
|
||
assert!(
|
||
!*eof_rx.borrow(),
|
||
"no idle timer ran with the knob disabled: arriving messages evict nothing"
|
||
);
|
||
peer_sink
|
||
.send(tokio_tungstenite::tungstenite::Message::Close(None))
|
||
.await
|
||
.expect("peer close accepted");
|
||
}
|
||
}
|
||
|
||
/// WS-19: the axum-flavor `AxumFraming` arms have no direct unit test —
|
||
/// the cap-trip and text→1002 closes are asserted on tungstenite only
|
||
/// (via the shared generic pumps). These drive the same generic pumps
|
||
/// with `AxumFraming` over a fake axum `WebSocket` sink/stream pair
|
||
/// (no server needed): a `mpsc`-backed sink/stream pair standing in
|
||
/// for the split halves of `axum::extract::ws::WebSocket`.
|
||
#[cfg(test)]
|
||
mod axum_framing_tests {
|
||
use super::*;
|
||
use futures::channel::mpsc as fut_mpsc;
|
||
|
||
/// In-process stand-in for the split halves of an axum
|
||
/// `WebSocket`: messages flow stream→`rx` and `tx`→sink, so the
|
||
/// generic pumps run against `AxumFraming` unchanged.
|
||
struct AxumFakeSocket {
|
||
sink_tx: futures_mpsc::Sender<AxumMessage>,
|
||
stream_rx:
|
||
std::pin::Pin<Box<dyn futures::Stream<Item = Result<AxumMessage, AxumMessage>> + Send>>,
|
||
}
|
||
|
||
impl futures::Stream for AxumFakeSocket {
|
||
type Item = Result<AxumMessage, AxumMessage>;
|
||
|
||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||
self.stream_rx.as_mut().poll_next(cx)
|
||
}
|
||
}
|
||
|
||
impl futures::Sink<AxumMessage> for AxumFakeSocket {
|
||
type Error = ();
|
||
|
||
fn poll_ready(
|
||
mut self: Pin<&mut Self>,
|
||
cx: &mut Context<'_>,
|
||
) -> Poll<Result<(), Self::Error>> {
|
||
Pin::new(&mut self.sink_tx).poll_ready(cx).map_err(|_| ())
|
||
}
|
||
|
||
fn start_send(mut self: Pin<&mut Self>, item: AxumMessage) -> Result<(), Self::Error> {
|
||
self.sink_tx.start_send(item).map_err(|_| ())
|
||
}
|
||
|
||
fn poll_flush(
|
||
mut self: Pin<&mut Self>,
|
||
cx: &mut Context<'_>,
|
||
) -> Poll<Result<(), Self::Error>> {
|
||
futures::Sink::poll_flush(Pin::new(&mut self.sink_tx), cx).map_err(|_| ())
|
||
}
|
||
|
||
fn poll_close(
|
||
mut self: Pin<&mut Self>,
|
||
cx: &mut Context<'_>,
|
||
) -> Poll<Result<(), Self::Error>> {
|
||
futures::Sink::poll_close(Pin::new(&mut self.sink_tx), cx).map_err(|_| ())
|
||
}
|
||
}
|
||
|
||
type AxumPairParts = (
|
||
futures::stream::SplitSink<AxumFakeSocket, AxumMessage>,
|
||
futures::stream::SplitStream<AxumFakeSocket>,
|
||
futures_mpsc::Receiver<AxumMessage>,
|
||
fut_mpsc::Sender<Result<AxumMessage, AxumMessage>>,
|
||
);
|
||
|
||
fn axum_pair() -> AxumPairParts {
|
||
let (outbound_tx, outbound_rx) = fut_mpsc::channel::<AxumMessage>(4);
|
||
let (inbound_tx, inbound_rx) = fut_mpsc::channel::<Result<AxumMessage, AxumMessage>>(4);
|
||
let socket = AxumFakeSocket {
|
||
sink_tx: outbound_tx,
|
||
stream_rx: Box::pin(inbound_rx),
|
||
};
|
||
let (sink, stream) = socket.split();
|
||
(sink, stream, outbound_rx, inbound_tx)
|
||
}
|
||
|
||
/// WS-19 mirror over `AxumFraming`: a text message on the read
|
||
/// side triggers the 1002 protocol-error close carrying the text
|
||
/// reason (the same generic read-pump arm the tungstenite tests
|
||
/// exercise — asserted here on the axum message types).
|
||
#[tokio::test]
|
||
async fn axum_framing_text_message_requests_protocol_close_with_reason() {
|
||
let (_sink, stream, _outbound_rx, mut inbound_tx) = axum_pair();
|
||
|
||
let (write_tx_for_read, mut write_rx_for_read) = fut_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
|
||
let read_task = tokio::spawn(run_read_pump::<AxumFraming, _, _>(
|
||
stream,
|
||
mpsc::channel::<Vec<u8>>(READ_SLOTS).0,
|
||
write_tx_for_read,
|
||
make_write_error_slot().0,
|
||
None,
|
||
|| {},
|
||
));
|
||
|
||
inbound_tx
|
||
.send(Ok(AxumMessage::Text("text frame".into())))
|
||
.await
|
||
.expect("inbound message accepted");
|
||
|
||
let close = tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||
while let Some(msg) = write_rx_for_read.next().await {
|
||
if let WriteMsg::CloseWith(code, reason) = msg {
|
||
return <AxumFraming as WsFraming>::close_message(code, reason);
|
||
}
|
||
}
|
||
AxumMessage::Text("queue closed".into())
|
||
})
|
||
.await
|
||
.expect("close requested");
|
||
|
||
let AxumMessage::Close(Some(frame)) = close else {
|
||
panic!("expected a close frame, got {close:?}");
|
||
};
|
||
assert_eq!(frame.code, WS_PROTOCOL_ERROR, "text closes with 1002");
|
||
assert_eq!(
|
||
frame.reason, "text messages not supported",
|
||
"close reason names the text cause"
|
||
);
|
||
read_task.abort();
|
||
}
|
||
|
||
/// WS-19 mirror of the over-cap close on the axum flavor: a
|
||
/// header claiming above-`MAX_CHUNK_LEN` payload fails the pump —
|
||
/// the peer sees the 1011 close naming the violation and the
|
||
/// pump's error slot carries the reason.
|
||
#[tokio::test]
|
||
async fn axum_framing_cap_trip_fails_the_pump_with_the_internal_close() {
|
||
let (sink, _stream, mut outbound_rx, _inbound_tx) = axum_pair();
|
||
let (slot, _rx) = make_write_error_slot();
|
||
|
||
let write_task = tokio::spawn(run_write_pump::<AxumFraming, _>(
|
||
sink,
|
||
{
|
||
let (mut tx, rx) = fut_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
|
||
tx.send(WriteMsg::Bytes(tests::oversize_header()))
|
||
.await
|
||
.expect("queue write accepted");
|
||
rx
|
||
},
|
||
slot,
|
||
None,
|
||
));
|
||
|
||
let close = tokio::time::timeout(std::time::Duration::from_secs(5), async {
|
||
loop {
|
||
match outbound_rx.next().await {
|
||
Some(m @ AxumMessage::Close(_)) => return m,
|
||
Some(_) => continue,
|
||
None => return AxumMessage::Text("stream ended".into()),
|
||
}
|
||
}
|
||
})
|
||
.await
|
||
.expect("close observed");
|
||
|
||
let AxumMessage::Close(Some(frame)) = close else {
|
||
panic!("expected a close frame, got {close:?}");
|
||
};
|
||
assert_eq!(frame.code, WS_INTERNAL_ERROR, "cap trip closes with 1011");
|
||
assert!(
|
||
frame.reason.contains("MAX_CHUNK_LEN"),
|
||
"close reason names the violation: {}",
|
||
frame.reason
|
||
);
|
||
write_task.abort();
|
||
}
|
||
}
|