fix(websocket): shutdown drops the held write sender (WS-07)

src/websocket/byte_adapter.rs

@@ -341,7 +341,7 @@ pub struct WsByteStream {
     read_buf: Vec<u8>,
     read_pos: usize,
     eof: bool,
-    write_tx: futures_mpsc::Sender<WriteMsg>,
+    write_tx: Option<futures_mpsc::Sender<WriteMsg>>,
     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<WriteMsg>> {
+        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");
+    }
 }