diff --git a/src/websocket/byte_adapter.rs b/src/websocket/byte_adapter.rs index 5bd570b..803b1ad 100644 --- a/src/websocket/byte_adapter.rs +++ b/src/websocket/byte_adapter.rs @@ -341,7 +341,7 @@ pub struct WsByteStream { read_buf: Vec, read_pos: usize, eof: bool, - write_tx: futures_mpsc::Sender, + write_tx: Option>, write_open: bool, write_error: oneshot::Receiver<&'static str>, write_failed: bool, @@ -450,7 +450,7 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) { read_buf: Vec::new(), read_pos: 0, eof: false, - write_tx, + write_tx: Some(write_tx), write_open: true, write_error: write_error_rx, write_failed: false, @@ -516,8 +516,14 @@ impl AsyncWrite for WsByteStream { if let Some(err) = this.poll_write_error() { return Poll::Ready(Err(err)); } - match this.write_tx.poll_ready(cx) { - Poll::Ready(Ok(())) => match this.write_tx.try_send(WriteMsg::Bytes(buf.to_vec())) { + let Some(write_tx) = this.write_tx_mut() else { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "ws stream shut down", + ))); + }; + match write_tx.poll_ready(cx) { + Poll::Ready(Ok(())) => match write_tx.try_send(WriteMsg::Bytes(buf.to_vec())) { Ok(()) => Poll::Ready(Ok(buf.len())), Err(_disconnected_or_full_race) => Poll::Ready(Err(io::Error::new( io::ErrorKind::BrokenPipe, @@ -540,7 +546,17 @@ impl AsyncWrite for WsByteStream { let this = self.get_mut(); if this.write_open { this.write_open = false; - drop(this.write_tx.clone()); + // 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(())) } @@ -578,6 +594,14 @@ impl WsByteStream { Err(oneshot::error::TryRecvError::Empty) => None, } } + + /// The write pump's fatal error, if any, and whether the write + /// queue still accepts sends (WS-07): the stream takes the sender + /// at shutdown so the pump-side channel close (and the sink's + /// trailing `ws_sink.close()`) actually happens. + fn write_tx_mut(&mut self) -> Option<&mut futures_mpsc::Sender> { + self.write_tx.as_mut() + } } /// tokio-tungstenite's message types as a [`WsFraming`] flavor. @@ -664,7 +688,7 @@ where read_buf: Vec::new(), read_pos: 0, eof: false, - write_tx, + write_tx: Some(write_tx), write_open: true, write_error: write_error_rx, write_failed: false, @@ -922,4 +946,83 @@ mod tests { 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"); + } }