alkdev/alkhttp
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'wt/review-002-con18b-ws-polish'

Browse Source
This commit is contained in:
glm-5.3-flash
2026-08-30 22:04:03 +00:00
parent 765dfc0858 bd87c4613b
commit 57ba9100c9
4 changed files with 668 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/websocket/byte_adapter.rs
+543 -95
View File
@@ -23,9 +23,18 @@
//! 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
//! 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.
//!
@@ -113,17 +122,15 @@ pub const INBOUND_WS_MESSAGE_CAP: usize = 1024 * 1024;
/// 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.
/// 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).
@@ -140,6 +147,14 @@ pub const WS_INTERNAL_ERROR: u16 = 1011;
/// 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
@@ -216,8 +231,9 @@ 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),
/// 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
@@ -241,6 +257,20 @@ fn send_write_error(slot: &WriteErrorSlot, reason: &'static str) {
}
}
/// 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
@@ -328,13 +358,13 @@ async fn run_read_pump<M, S, F>(
Some(budget) => match tokio::time::timeout(budget, ws_stream.next()).await {
Ok(msg) => msg,
Err(_elapsed) => {
send_write_error(
&write_error_slot_for_read,
"connection made no inbound chunk progress past the read timeout",
);
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))
.send(WriteMsg::CloseWith(
WS_GOING_AWAY,
close_reason::IDLE_READ_TIMEOUT,
))
.await;
break;
}
@@ -356,7 +386,10 @@ async fn run_read_pump<M, S, F>(
} else if M::is_text(&m) {
let _ = write_tx_for_read
.clone()
.send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR))
.send(WriteMsg::CloseWith(
WS_PROTOCOL_ERROR,
close_reason::TEXT_NOT_SUPPORTED,
))
.await;
break;
} else if M::is_close(&m) {
@@ -366,11 +399,14 @@ async fn run_read_pump<M, S, F>(
Err(()) => {
send_write_error(
&write_error_slot_for_read,
"inbound frame rejected (size cap)",
close_reason::INBOUND_FRAME_REJECTED,
);
let _ = write_tx_for_read
.clone()
.send(WriteMsg::CloseWith(WS_INTERNAL_ERROR))
.send(WriteMsg::CloseWith(
WS_INTERNAL_ERROR,
close_reason::INBOUND_FRAME_REJECTED,
))
.await;
break;
}
@@ -395,17 +431,38 @@ async fn fail_write_pump<M, S>(
.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 that code.
/// 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,
@@ -413,25 +470,20 @@ async fn run_write_pump<M, S>(
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;
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) => {
if b.len() > PENDING_BUFFER_CAP {
fail_write_pump::<M, _>(
&mut ws_sink,
&write_error_slot,
"write pending buffer exceeded cap",
)
.await;
return;
}
debug_assert!(b.len() <= PENDING_BUFFER_CAP);
pending.extend_from_slice(&b);
}
}
@@ -456,10 +508,8 @@ async fn run_write_pump<M, S>(
}
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()))
if !send_bounded::<M, _>(&mut ws_sink, piece, write_timeout, &write_error_slot)
.await
.is_err()
{
return;
}
@@ -478,6 +528,45 @@ async fn run_write_pump<M, S>(
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 {
@@ -566,10 +655,25 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
/// 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 an inbound WS message.
/// 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);
@@ -601,6 +705,7 @@ pub fn split_ws_to_bytes_idle(
ws_sink,
write_rx,
write_error_slot,
Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)),
));
(
@@ -675,12 +780,22 @@ impl AsyncWrite for WsByteStream {
if let Some(err) = this.poll_write_error() {
return Poll::Ready(Err(err));
}
let Some(write_tx) = this.write_tx_mut() else {
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())),
@@ -754,11 +869,13 @@ impl WsByteStream {
}
}
/// 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>> {
/// 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()
}
}
@@ -857,6 +974,7 @@ where
ws_sink,
write_rx,
write_error_slot,
Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT),
));
(
@@ -879,15 +997,120 @@ where
)
}
/// 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).
fn oversize_header() -> Vec<u8> {
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
@@ -933,15 +1156,14 @@ mod tests {
);
}
/// 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).
/// 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_side_rejects_pending_over_byte_cap() {
let (client_io, _server_io) = tokio::io::duplex(64);
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,
@@ -950,37 +1172,30 @@ mod tests {
.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());
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");
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}"
);
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.
/// 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);
@@ -1203,8 +1418,64 @@ mod tests {
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
/// past the knob. The idle-read timeout fires with no inbound
/// 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
@@ -1233,14 +1504,15 @@ mod tests {
// 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.
// 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> = loop {
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))
break cf.map(|f| (u16::from(f.code), f.reason.to_string()))
}
Some(Ok(_)) => continue,
Some(Err(_)) | None => break None,
@@ -1252,9 +1524,14 @@ mod tests {
.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!(
close_seen == Some(WS_GOING_AWAY),
"peer received the 1001 GoingAway close, got {close_seen:?}"
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");
}
@@ -1324,10 +1601,10 @@ mod tests {
});
let outcome = tokio::time::timeout(std::time::Duration::from_secs(10), async {
let close: Option<u16> = loop {
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))
break cf.map(|f| (u16::from(f.code), f.reason.to_string()))
}
Some(Ok(_)) => continue,
Some(Err(_)) | None => break None,
@@ -1339,11 +1616,15 @@ mod tests {
.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!(
outcome,
Some(WS_GOING_AWAY),
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
@@ -1464,3 +1745,170 @@ mod tests {
.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();
}
}
src/websocket/mod.rs
+5 -4
View File
@@ -13,9 +13,10 @@ pub mod byte_adapter;
pub mod upgrade;
pub use byte_adapter::{
split_ws_to_bytes, split_ws_to_bytes_idle, WsByteStream, WsPumps, DEFAULT_WS_IDLE_TIMEOUT,
INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP, MAX_CHUNK_LEN, PENDING_BUFFER_CAP, WS_GOING_AWAY,
WS_MESSAGE_CAP, WS_PROTOCOL_ERROR,
split_ws_to_bytes, split_ws_to_bytes_idle, split_ws_to_bytes_idle_with_write, WsByteStream,
WsPumps, DEFAULT_WS_IDLE_TIMEOUT, DEFAULT_WS_WRITE_TIMEOUT, INBOUND_WS_FRAME_CAP,
INBOUND_WS_MESSAGE_CAP, MAX_CHUNK_LEN, PENDING_BUFFER_CAP, WS_GOING_AWAY, WS_MESSAGE_CAP,
WS_PROTOCOL_ERROR,
};
#[cfg(any(test, feature = "wss"))]
@@ -25,7 +26,7 @@ pub use byte_adapter::split_tungstenite_to_bytes;
pub(crate) use upgrade::adapter_install_channel_zero;
pub use upgrade::{
run_channels_session, ws_bearer_auth, ws_upgrade_handler, ChannelsPolicy, SessionState,
WsSessions, DEFAULT_WS_MAX_SESSIONS,
WsSessions, WsTimeouts, DEFAULT_WS_MAX_SESSIONS,
};
#[cfg(any(test, feature = "test-support"))]
src/websocket/upgrade.rs
+69 -5
View File
@@ -51,7 +51,7 @@ use axum::response::{IntoResponse, Response};
use parking_lot::Mutex;
use super::byte_adapter::{
split_ws_to_bytes_idle, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP,
split_ws_to_bytes_idle_with_write, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP,
};
/// Registry of live WS session pump handles (WS-08). The upgrade
@@ -196,7 +196,11 @@ impl WsSessions {
/// session's pump handle is registered for the session's lifetime —
/// the WS-08 eviction lever ([`WsSessions::abort`]). `idle_timeout`
/// bounds the read stall (WS-01): `None` disables the knob, `Some(d)`
/// closes the read with 1001 after `d` without an inbound WS message.
/// closes the read with 1001 after `d` without inbound chunk
/// progress. `write_timeout` bounds one outbound WS send (WS-18):
/// `None` = the crate default window, `Some(d)` a deployment-set
/// window.
#[allow(clippy::too_many_arguments)]
pub async fn run_channels_session(
socket: axum::extract::ws::WebSocket,
registry: Arc<OperationRegistry>,
@@ -204,8 +208,10 @@ pub async fn run_channels_session(
policy: Arc<dyn ChannelLifecyclePolicy>,
sessions: Option<WsSessions>,
idle_timeout: Option<std::time::Duration>,
write_timeout: Option<std::time::Duration>,
) {
let (byte_stream, pumps) = split_ws_to_bytes_idle(socket, idle_timeout);
let (byte_stream, pumps) =
split_ws_to_bytes_idle_with_write(socket, idle_timeout, write_timeout);
let pumps = Arc::new(pumps);
let _guard = sessions.map(|sessions| {
@@ -316,6 +322,40 @@ impl alkcall::core::auth::IdentityProvider for NoopProvider {
#[derive(Clone)]
pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);
/// Per-request WS pump timeouts (WS-17): a deployment inserts
/// `WsTimeouts` into the request extensions (a route layer on the WS
/// route, mirroring [`ChannelsPolicy`]) to set the pump knobs per
/// route instead of the router-state defaults. `idle` (WS-01) bounds
/// the read staleness (no completed inbound chunk for the window →
/// 1001 eviction); `write` (WS-18) bounds one outbound WS send (a peer
/// that stops reading is evicted once a single send outlasts it);
/// `None` disables a knob. Without the extension the
/// [`SessionState`] values apply — the built-in surface carries
/// `HttpAdapter::with_ws_idle_timeout` /
/// `DEFAULT_WS_IDLE_TIMEOUT` for the read side and
/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] for the write side,
/// which is also what bare-registry routes (custom upgrade routes on
/// a plain `Arc<OperationRegistry>` state) get: 60 s idle + 60 s
/// write windows, 64-session semaphore, handler-private
/// [`WsSessions`].
#[derive(Clone, Copy, Debug)]
pub struct WsTimeouts {
/// Idle-read window (WS-01); `None` disables the eviction.
pub idle: Option<std::time::Duration>,
/// Write-progress window (WS-18); `None` selects the crate
/// default (`DEFAULT_WS_WRITE_TIMEOUT`).
pub write: Option<std::time::Duration>,
}
impl Default for WsTimeouts {
fn default() -> Self {
Self {
idle: Some(crate::websocket::DEFAULT_WS_IDLE_TIMEOUT),
write: None,
}
}
}
/// The upgrade handler. Requires the resolved identity in request
/// extensions (stashed by [`ws_bearer_auth`]) — a WS session without
/// an identity cannot run `AccessControl::check`.
@@ -327,6 +367,16 @@ pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);
/// (`RouterState::ws_sessions`) unless a request extension carries
/// one, so a stuck session stays evictable (WS-08).
///
/// The pump timeout knobs (WS-01 read / WS-18 write) come from the
/// [`WsTimeouts`] request extension when present, else from the
/// router state (`HttpAdapter::with_ws_idle_timeout` for the read;
/// the write side is fixed at
/// [`crate::websocket::DEFAULT_WS_WRITE_TIMEOUT`] on the built-in
/// surface). Bare-registry routes (`Arc<OperationRegistry>` state)
/// carry no configurable state — they run the documented defaults
/// above; a deployment overriding them inserts the `WsTimeouts`
/// extension on its upgrade route.
///
/// Session cap (WS-09): one semaphore permit is acquired per upgrade,
/// post-auth and pre-upgrade; when the configured cap
/// (`HttpAdapter::with_ws_max_sessions`, default
@@ -338,6 +388,7 @@ pub async fn ws_upgrade_handler(
axum::extract::State(state): axum::extract::State<SessionState>,
axum::Extension(identity): axum::Extension<Identity>,
policy: Option<axum::Extension<ChannelsPolicy>>,
timeouts: Option<axum::Extension<WsTimeouts>>,
ws_upgrade: WebSocketUpgrade,
) -> Response {
let Ok(permit) = state.session_slots.clone().try_acquire_owned() else {
@@ -356,13 +407,26 @@ pub async fn ws_upgrade_handler(
.map(|axum::Extension(p)| p.0)
.unwrap_or_else(|| Arc::new(NoCap));
let registry = Arc::clone(state.registry());
let idle_timeout = state.idle_timeout();
let idle_timeout = match timeouts {
Some(axum::Extension(t)) => t.idle,
None => state.idle_timeout(),
};
let write_timeout = timeouts.and_then(|t| t.write);
ws_upgrade
.max_frame_size(INBOUND_WS_FRAME_CAP)
.max_message_size(INBOUND_WS_MESSAGE_CAP)
.on_upgrade(move |socket| async move {
let _permit = permit;
run_channels_session(socket, registry, identity, policy, sessions, idle_timeout).await
run_channels_session(
socket,
registry,
identity,
policy,
sessions,
idle_timeout,
Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)),
)
.await
})
}
tests/ws_upgrade_session.rs
+51
View File
@@ -918,3 +918,54 @@ async fn idle_progress_knob_none_disables_eviction_over_axum_upgrade() {
);
ws.close().await;
}
/// WS-17 acceptance: a bare-registry upgrade route takes the pump
/// knobs per request via the `WsTimeouts` extension (mirroring
/// `ChannelsPolicy`) — a client that completes no chunk is evicted
/// with 1001 inside the extension's short idle window, not the 60 s
/// default.
#[tokio::test]
async fn ws_timeouts_extension_sets_the_idle_window_on_a_bare_registry_route() {
let registry = std::sync::Arc::new(OperationRegistry::new());
struct StaticTok;
impl IdentityProvider for StaticTok {
fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
None
}
fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
let s = String::from_utf8_lossy(&token.raw).to_string();
(s == "tok-1").then(|| identity("alice", &[]))
}
}
let app = axum::Router::new()
.route(
"/alk/channels",
axum::routing::get(alkhttp::websocket::ws_upgrade_handler),
)
.layer(axum::middleware::from_fn_with_state(
std::sync::Arc::new(StaticTok) as std::sync::Arc<dyn IdentityProvider>,
alkhttp::websocket::ws_bearer_auth,
))
.layer(axum::Extension(alkhttp::websocket::WsTimeouts {
idle: Some(std::time::Duration::from_millis(150)),
write: None,
}))
.with_state(registry);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = format!("ws://{}", listener.local_addr().unwrap());
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
.await
.unwrap();
let close = ws.next_close(std::time::Duration::from_secs(5)).await;
assert_eq!(
close,
Some(Some(alkhttp::websocket::WS_GOING_AWAY)),
"evicted with 1001 inside the extension's idle window"
);
ws.close().await;
}