fix(websocket): configurable WS read idle timeout (WS-01)

src/websocket/byte_adapter.rs

@@ -47,6 +47,7 @@ use std::{
     pin::Pin,
     sync::{Arc, Mutex as StdMutex},
     task::{Context, Poll},
+    time::Duration,
 };
 
 use axum::extract::ws::{CloseFrame, Message as AxumMessage, WebSocket};
@@ -106,6 +107,18 @@ pub const WS_PROTOCOL_ERROR: u16 = 1002;
 /// violation it cannot recover from (WS-04/HY-09, WS-05).
 pub const WS_INTERNAL_ERROR: u16 = 1011;
 
+/// Idle-read close code (WS-01): a read side that stays silent past the
+/// configured idle timeout 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 idle-read timeout for the WS pumps (WS-01): a peer that
+/// drips bytes (or a stalled partial chunk header) cannot park the
+/// single demux loop longer than this without traffic. Zero disables
+/// the timeout. Deployment knob (`HttpAdapter::with_ws_idle_timeout`).
+pub const DEFAULT_WS_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
+
 /// 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
@@ -190,11 +203,22 @@ where
 /// 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): when `idle_timeout` is non-zero, the
+/// next-message await is wrapped in a `tokio::time::timeout` — a peer
+/// that dribbles (or a stalled header) cannot park the single demux
+/// loop beyond the knob. Staleness 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.
+/// Messages that *do* arrive reset the window; a partial chunk header
+/// alone does not (the boundary is a complete WS message).
 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,
@@ -202,7 +226,27 @@ async fn run_read_pump<M, S, F>(
     S::Item: IntoWsResult<M>,
     F: FnOnce(),
 {
-    while let Some(msg) = ws_stream.next().await {
+    loop {
+        let msg = match idle_timeout {
+            None => ws_stream.next().await,
+            Some(timeout) => match tokio::time::timeout(timeout, ws_stream.next()).await {
+                Ok(msg) => msg,
+                Err(_elapsed) => {
+                    send_write_error(
+                        &write_error_slot_for_read,
+                        "connection idle past the read timeout",
+                    );
+                    let _ = write_tx_for_read
+                        .clone()
+                        .send(WriteMsg::CloseWith(WS_GOING_AWAY))
+                        .await;
+                    break;
+                }
+            },
+        };
+        let Some(msg) = msg else {
+            break;
+        };
         match msg.into_ws_result() {
             Ok(m) => {
                 if let Some(b) = M::binary(&m) {
@@ -409,10 +453,23 @@ impl WsFraming for AxumFraming {
     }
 }
 
+/// Split a `WebSocket` into the byte stream + the pump tasks with the
+/// default idle-read timeout
+/// ([`DEFAULT_WS_IDLE_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`.
-pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
+/// `idle_timeout` (WS-01): `None` = no idle bound, `Some(d)` closes
+/// the read with 1001 after `d` without an inbound WS message.
+pub fn split_ws_to_bytes_idle(
+    socket: WebSocket,
+    idle_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);
@@ -430,6 +487,7 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
         read_tx,
         write_tx_for_read,
         write_error_slot_for_read,
+        idle_timeout,
         move || {
             #[cfg(any(test, feature = "wss"))]
             {
@@ -642,17 +700,34 @@ impl WsFraming for TungsteniteFraming {
     }
 }
 
-/// Client-side split for a tokio-tungstenite `WebSocketStream`
-/// (`from_wss`, ADR-070): the same WS↔byte-stream seam as
-/// [`split_ws_to_bytes`], over the tungstenite socket instead of axum's
-/// server-side `WebSocket`. Message semantics are identical: binary
-/// messages carry the byte stream, text is a protocol error (close
-/// 1002), close → read EOF. Both paths run the same generic pumps
-/// (WS-11); the socket flavor differs only in the [`WsFraming`] impl.
+/// 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,
 {
@@ -671,6 +746,7 @@ where
         read_tx,
         write_tx_for_read,
         write_error_slot_for_read,
+        idle_timeout,
         move || {
             let _ = read_eof_for_task.send(true);
         },
@@ -1025,4 +1101,105 @@ mod tests {
         .expect("peer observes the close reply after our shutdown");
         assert!(saw_close, "the shutdown path emitted a WS Close frame");
     }
+
+    /// 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
+    /// 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.
+        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 {
+                    match peer_stream.next().await {
+                        Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => {
+                            break cf.map(|f| u16::from(f.code))
+                        }
+                        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");
+
+        assert!(
+            close_seen == Some(WS_GOING_AWAY),
+            "peer received the 1001 GoingAway close, got {close_seen:?}"
+        );
+        assert!(eof_fired, "EOF signal fired for the from_wss machinery");
+    }
+
+    /// The idle knob resets on traffic: a peer dribbling one message
+    /// per window (each inside the knob) keeps the connection open —
+    /// only the *stall* (no message within one full window) closes.
+    #[tokio::test]
+    async fn idle_read_timeout_resets_on_traffic() {
+        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 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};
+
+        // Dribble: a message every 100 ms (half the knob) for 8 rounds —
+        // 1.6 s total, far past a single 200 ms window. The connection
+        // must stay open; every message surfaces on the adapter.
+        let mut buf = [0u8; 1];
+        for round in 0u8..8 {
+            tokio::time::sleep(knob / 2).await;
+            peer_sink
+                .send(tokio_tungstenite::tungstenite::Message::Binary(
+                    vec![round].into(),
+                ))
+                .await
+                .expect("dribble write accepted");
+            tokio::time::timeout(knob, stream.read_exact(&mut buf))
+                .await
+                .expect("connection alive across the dribble")
+                .expect("read");
+            assert_eq!(buf[0], round);
+        }
+        let _ = peer_stream.next().await;
+    }
 }