From 1b2fb5401acdf98315f6fb6f53f4b4b412192840 Mon Sep 17 00:00:00 2001 From: black-binary Date: Fri, 4 Sep 2026 00:47:31 +0800 Subject: [PATCH 1/3] test: expose missing stream half-close semantics --- src/lib.rs | 111 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 111 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index 198d651..1eed3b6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -425,6 +425,117 @@ mod tests { stream2.write_all(&data).await.unwrap_err(); } + #[tokio::test] + async fn test_half_close_shutdown_keeps_read_open() { + let (client, mut server) = get_duplex_mux_pair().await; + let (mut client_reader, mut client_writer) = tokio::io::split(client); + + client_writer.write_all(b"request").await.unwrap(); + client_writer.shutdown().await.unwrap(); + assert!( + client_writer.write_all(b"late request").await.is_err(), + "writes must fail after the local write half is shut down" + ); + + let mut request = Vec::new(); + tokio::time::timeout(Duration::from_secs(1), server.read_to_end(&mut request)) + .await + .expect("server did not observe EOF from the client write half") + .unwrap(); + assert_eq!(request, b"request"); + + server.write_all(b"response").await.unwrap(); + server.shutdown().await.unwrap(); + + let mut response = Vec::new(); + tokio::time::timeout( + Duration::from_secs(1), + client_reader.read_to_end(&mut response), + ) + .await + .expect("client read half closed before the server response") + .unwrap(); + assert_eq!(response, b"response"); + } + + #[tokio::test] + async fn test_half_close_local_fin_waits_for_remote_read_eof() { + let (mut local, mut peer) = get_duplex_mux_pair().await; + + local.shutdown().await.unwrap(); + + let read_is_pending = poll_fn(|cx| { + let mut byte = [0u8; 1]; + let mut read_buf = ReadBuf::new(&mut byte); + Poll::Ready( + Pin::new(&mut local) + .poll_read(cx, &mut read_buf) + .is_pending(), + ) + }) + .await; + assert!( + read_is_pending, + "closing the local write half must not create local read EOF" + ); + + peer.write_all(b"x").await.unwrap(); + peer.flush().await.unwrap(); + let mut byte = [0u8; 1]; + local.read_exact(&mut byte).await.unwrap(); + assert_eq!(byte, *b"x"); + + peer.shutdown().await.unwrap(); + assert_eq!(local.read(&mut byte).await.unwrap(), 0); + } + + #[tokio::test] + async fn test_half_close_remote_fin_only_closes_read_direction() { + let (mut local, mut peer) = get_duplex_mux_pair().await; + + peer.write_all(b"done").await.unwrap(); + peer.shutdown().await.unwrap(); + + let mut received = Vec::new(); + local.read_to_end(&mut received).await.unwrap(); + assert_eq!(received, b"done"); + + local.write_all(b"reply after EOF").await.unwrap(); + local.shutdown().await.unwrap(); + + let mut reply = Vec::new(); + peer.read_to_end(&mut reply).await.unwrap(); + assert_eq!(reply, b"reply after EOF"); + } + + #[tokio::test] + async fn test_half_close_stream_is_closed_only_after_both_fins() { + let (a, b) = tokio::io::duplex(4096); + let (connector_a, _acceptor_a, worker_a) = MuxBuilder::client().with_connection(a).build(); + let (connector_b, mut acceptor_b, worker_b) = + MuxBuilder::server().with_connection(b).build(); + tokio::spawn(worker_a); + tokio::spawn(worker_b); + + let mut local = connector_a.connect().unwrap(); + let mut peer = acceptor_b.accept().await.unwrap(); + local.shutdown().await.unwrap(); + + let mut byte = [0u8; 1]; + assert_eq!(peer.read(&mut byte).await.unwrap(), 0); + assert!(!local.is_closed()); + assert!(!peer.is_closed()); + assert_eq!(connector_a.get_num_streams(), 1); + assert_eq!(connector_b.get_num_streams(), 1); + + peer.shutdown().await.unwrap(); + assert_eq!(local.read(&mut byte).await.unwrap(), 0); + assert!(local.is_closed()); + assert!(peer.is_closed()); + assert_eq!(connector_a.get_num_streams(), 0); + assert_eq!(connector_b.get_num_streams(), 0); + } + #[tokio::test] async fn test_timeout() { let (a, b) = get_tcp_pair().await; From e552716074ca4890dbec9f01e99c6200ff3adba6 Mon Sep 17 00:00:00 2001 From: black-binary Date: Fri, 4 Sep 2026 01:11:59 +0800 Subject: [PATCH 2/3] feat: support stream half-close --- README.md | 11 ++- src/lib.rs | 226 +++++++++++++++++++++++++++++++++++++---------------- src/mux.rs | 176 +++++++++++++++++++++++++++++------------ 3 files changed, 297 insertions(+), 116 deletions(-) diff --git a/README.md b/README.md index da8bb6d..924429c 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,15 @@ The worker exits when: sink is closed. It also works without the worker being polled — useful in test setups or sync teardown paths. +`MuxStream` supports TCP-style half-close. Calling +`AsyncWriteExt::shutdown()` sends FIN and closes only the local write +direction; reads remain open until the peer sends its FIN. Likewise, a +peer FIN produces read EOF without preventing a response from being +written. Dropping a stream drains writes already accepted by +`poll_write`, sends FIN if needed, and releases the local stream handle. +The wire format is unchanged; a peer must also interpret FIN as directional +EOF to send data after receiving it. + ## Configuration ```rust,ignore @@ -107,7 +116,7 @@ VERSION(1B) | CMD(1B) | LENGTH(2B LE) | STREAMID(4B LE) | DATA(LENGTH) VERSION: 1 CMD: SYN(0) open stream (LENGTH must be 0) - FIN(1) close stream (LENGTH must be 0) + FIN(1) close sender's write direction (LENGTH must be 0) PSH(2) payload NOP(3) keep-alive (LENGTH must be 0; STREAMID is 0) ``` diff --git a/src/lib.rs b/src/lib.rs index 1eed3b6..a3f2e64 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -46,6 +46,14 @@ //! and it is cancellation-safe: dropping the future mid-flight hands //! control back to the worker without wedging it. //! +//! [`MuxStream`] supports TCP-style half-close. Calling +//! [`AsyncWriteExt::shutdown`](tokio::io::AsyncWriteExt::shutdown) sends FIN +//! and closes only the local write direction; reads remain open until the +//! peer sends its FIN. A peer FIN likewise produces read EOF without +//! preventing a response from being written. This does not change the wire +//! format; peers must also interpret FIN as directional EOF to write after +//! receiving it. +//! //! # Configuration //! //! See [`MuxBuilder`] for the available knobs: `with_keep_alive_interval`, @@ -396,19 +404,20 @@ mod tests { stream2.write_all(&data).await.unwrap(); stream2.shutdown().await.unwrap(); - tokio::time::sleep(Duration::from_secs(1)).await; - - stream1.write_all(&[0, 1, 2, 3]).await.unwrap_err(); - stream1.flush().await.unwrap_err(); let mut buf = vec![0; 4]; stream1.read_exact(&mut buf).await.unwrap(); assert_eq!(buf, data); assert_eq!(stream1.read(&mut buf).await.unwrap(), 0); + stream1.write_all(&[4, 5, 6, 7]).await.unwrap(); + stream1.shutdown().await.unwrap(); + stream2.read_exact(&mut buf).await.unwrap(); + assert_eq!(buf, [4, 5, 6, 7]); + assert_eq!(stream2.read(&mut buf).await.unwrap(), 0); drop(acceptor_a); let mut stream = connector_b.connect().unwrap(); assert_eq!(stream.read(&mut buf).await.unwrap(), 0); - stream.flush().await.unwrap_err(); + stream.flush().await.unwrap(); stream.shutdown().await.unwrap(); let mut stream1 = connector_a.connect().unwrap(); @@ -422,7 +431,7 @@ mod tests { stream2.read_exact(&mut buf).await.unwrap(); assert!(buf == data); stream2.read_exact(&mut buf).await.unwrap_err(); - stream2.write_all(&data).await.unwrap_err(); + stream2.shutdown().await.unwrap(); } #[tokio::test] @@ -536,6 +545,96 @@ mod tests { assert_eq!(connector_b.get_num_streams(), 0); } + #[tokio::test] + async fn test_half_close_fin_before_accept_allows_response() { + let (a, b) = tokio::io::duplex(4096); + let (connector_a, _acceptor_a, worker_a) = MuxBuilder::client().with_connection(a).build(); + let (_connector_b, mut acceptor_b, worker_b) = + MuxBuilder::server().with_connection(b).build(); + tokio::spawn(worker_a); + tokio::spawn(worker_b); + + let mut client = connector_a.connect().unwrap(); + client.write_all(b"request before accept").await.unwrap(); + client.shutdown().await.unwrap(); + + let mut server = acceptor_b.accept().await.unwrap(); + let mut request = Vec::new(); + server.read_to_end(&mut request).await.unwrap(); + assert_eq!(request, b"request before accept"); + + server.write_all(b"response after accept").await.unwrap(); + server.shutdown().await.unwrap(); + let mut response = Vec::new(); + client.read_to_end(&mut response).await.unwrap(); + assert_eq!(response, b"response after accept"); + } + + #[tokio::test] + async fn test_half_close_repeated_shutdown_sends_one_fin() { + let (a, mut peer) = tokio::io::duplex(1024); + let (connector, acceptor, worker) = MuxBuilder::client().with_connection(a).build(); + tokio::spawn(worker); + + let mut stream = connector.connect().unwrap(); + let stream_id = stream.get_stream_id(); + stream.shutdown().await.unwrap(); + stream.shutdown().await.unwrap(); + drop(stream); + drop(connector); + drop(acceptor); + + let mut wire = Vec::new(); + tokio::time::timeout(Duration::from_secs(1), peer.read_to_end(&mut wire)) + .await + .expect("session did not close after dropping every public handle") + .unwrap(); + assert_eq!( + wire.len(), + 16, + "repeated shutdown or Drop emitted extra data" + ); + assert_eq!(wire[1], 0, "first frame was not SYN"); + assert_eq!( + u32::from_le_bytes(wire[4..8].try_into().unwrap()), + stream_id + ); + assert_eq!(wire[9], 1, "second frame was not FIN"); + assert_eq!( + u32::from_le_bytes(wire[12..16].try_into().unwrap()), + stream_id + ); + } + + #[tokio::test] + async fn test_half_close_simultaneous_fin_preserves_crossed_data() { + let (mut left, mut right) = get_duplex_mux_pair().await; + + left.write_all(b"from left").await.unwrap(); + right.write_all(b"from right").await.unwrap(); + let (left_shutdown, right_shutdown) = tokio::join!(left.shutdown(), right.shutdown()); + left_shutdown.unwrap(); + right_shutdown.unwrap(); + + let mut received_left = Vec::new(); + let mut received_right = Vec::new(); + let (left_read, right_read) = tokio::time::timeout(Duration::from_secs(1), async { + tokio::join!( + left.read_to_end(&mut received_left), + right.read_to_end(&mut received_right) + ) + }) + .await + .expect("crossed FINs did not produce EOF"); + left_read.unwrap(); + right_read.unwrap(); + + assert_eq!(received_left, b"from right"); + assert_eq!(received_right, b"from left"); + assert!(left.is_closed()); + assert!(right.is_closed()); + } + #[tokio::test] async fn test_timeout() { let (a, b) = get_tcp_pair().await; @@ -552,7 +651,7 @@ mod tests { }); let stream1 = connector_a.connect().unwrap(); - let stream2 = acceptor_b.accept().await.unwrap(); + let mut stream2 = acceptor_b.accept().await.unwrap(); tokio::time::sleep(Duration::from_secs(1)).await; assert!(!stream1.is_closed()); assert!(!stream2.is_closed()); @@ -560,6 +659,10 @@ mod tests { tokio::time::sleep(Duration::from_secs(5)).await; assert!(stream1.is_closed()); + assert!(!stream2.is_closed()); + let mut byte = [0u8; 1]; + assert_eq!(stream2.read(&mut byte).await.unwrap(), 0); + stream2.shutdown().await.unwrap(); assert!(stream2.is_closed()); } @@ -609,20 +712,27 @@ mod tests { } #[tokio::test] - async fn test_connection_drop() { + async fn test_stream_drop_delivers_read_eof() { let (a, b) = get_tcp_pair().await; let (connector_a, _, worker_a) = MuxBuilder::client().with_connection(a).build(); let (_, mut acceptor_b, worker_b) = MuxBuilder::server().with_connection(b).build(); tokio::spawn(worker_a); tokio::spawn(worker_b); - let mut _stream1 = connector_a.connect().unwrap(); + let stream1 = connector_a.connect().unwrap(); let mut stream2 = acceptor_b.accept().await.unwrap(); - drop(_stream1); - tokio::time::sleep(Duration::from_secs(1)).await; + drop(stream1); - assert!(stream2.write_all(b"1234").await.is_err()); + let mut byte = [0u8; 1]; + assert_eq!(stream2.read(&mut byte).await.unwrap(), 0); + assert!( + !stream2.is_closed(), + "peer drop closes only this stream's read direction" + ); + + stream2.shutdown().await.unwrap(); + assert!(stream2.is_closed()); } #[tokio::test] @@ -748,43 +858,38 @@ mod tests { assert_eq!(buf, data); } - // BUG: After local shutdown we returned EOF, but a subsequent peer Push - // (peer hadn't seen our FIN yet) would still populate rx_queue and a later - // poll_read would surface that data, breaking AsyncRead EOF monotonicity. + // Once a peer FIN has produced EOF, a protocol-invalid late PSH must not + // make a subsequent read return data or implicitly close our write half. #[tokio::test(flavor = "multi_thread")] - async fn test_eof_monotonic_after_local_shutdown() { - let (a, b) = get_tcp_pair().await; - let (connector_a, _, worker_a) = MuxBuilder::client().with_connection(a).build(); - let (_, mut acceptor_b, worker_b) = MuxBuilder::server().with_connection(b).build(); - tokio::spawn(worker_a); - tokio::spawn(worker_b); - - let mut stream1 = connector_a.connect().unwrap(); - let mut stream2 = acceptor_b.accept().await.unwrap(); + async fn test_eof_monotonic_after_remote_fin() { + let (a, mut peer) = tokio::io::duplex(4096); + let (_connector, mut acceptor, worker) = MuxBuilder::client().with_connection(a).build(); + tokio::spawn(worker); - // Local shutdown -> our handle.closed=true, FIN enqueued globally. - stream1.shutdown().await.unwrap(); + let stream_id = 2; + peer.write_all(&raw_frame(0, stream_id, &[])).await.unwrap(); + let mut stream = acceptor.accept().await.unwrap(); + peer.write_all(&raw_frame(1, stream_id, &[])).await.unwrap(); - // First read sees EOF immediately (rx_queue empty + handle.closed). let mut buf = [0u8; 4]; - let n = stream1.read(&mut buf).await.unwrap(); - assert_eq!(n, 0, "first read after shutdown must be EOF"); + assert_eq!(stream.read(&mut buf).await.unwrap(), 0); - // Peer races and writes data before processing our FIN. Inject a - // PSH frame on the wire from b's side - but stream2's writer is - // still open from b's perspective. - let _ = stream2.write_all(b"late").await; - let _ = stream2.flush().await; + peer.write_all(&raw_frame(2, stream_id, b"late")) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; - // Give the dispatcher time to receive the late PSH. - tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!(stream.read(&mut buf).await.unwrap(), 0); - // EOF must remain EOF. - let n = stream1.read(&mut buf).await.unwrap(); - assert_eq!( - n, 0, - "subsequent read must remain EOF, not return late data" - ); + stream.write_all(b"reply").await.unwrap(); + stream.flush().await.unwrap(); + let mut wire = [0u8; 13]; + tokio::time::timeout(Duration::from_secs(1), peer.read_exact(&mut wire)) + .await + .expect("reply was not sent after remote FIN") + .unwrap(); + assert_eq!(wire[1], 2, "late PSH caused an implicit local FIN"); + assert_eq!(&wire[8..], b"reply"); } // BUG: poll_flush returns Err the moment the stream is locally closed, @@ -803,8 +908,8 @@ mod tests { let data = b"hello world"; stream1.write_all(data).await.unwrap(); - // shutdown will mark our handle.closed=true and enqueue FIN. Pending - // PSH frames from the write_all above should still go out. + // shutdown closes the local write half and enqueues FIN. Pending PSH + // frames from the write_all above should still go out first. stream1.shutdown().await.unwrap(); let mut buf = vec![0u8; data.len()]; @@ -1439,13 +1544,12 @@ mod tests { peer.write_all(&raw_frame(1, stream_id, &[])).await.unwrap(); tokio::spawn(worker); - tokio::time::timeout(Duration::from_secs(1), async { - while !stream.is_closed() { - tokio::task::yield_now().await; - } - }) - .await - .expect("remote FIN was not dispatched"); + let mut byte = [0u8; 1]; + let n = tokio::time::timeout(Duration::from_secs(1), stream.read(&mut byte)) + .await + .expect("remote FIN was not dispatched") + .unwrap(); + assert_eq!(n, 0); drop(stream); drop(connector); @@ -1458,7 +1562,7 @@ mod tests { .unwrap(); let payload_frames = payload.len().div_ceil(MAX_PAYLOAD_SIZE); - let expected = 8 + payload_frames * 8 + payload.len(); // SYN + PSHs + let expected = 8 + payload_frames * 8 + payload.len() + 8; // SYN + PSHs + FIN assert_eq!( wire.len(), expected, @@ -1864,10 +1968,8 @@ mod tests { let mut actual = vec![0; response_len]; stream.read_exact(&mut actual).await.unwrap(); assert_eq!(actual, expected, "response mismatch for case {case_id}"); - // Tell the responder it may send FIN. Without this - // application-level handshake, its valid early FIN can - // race our final request flush and make the stress test - // assert a stronger half-close contract than smux has. + // Coordinate completion so both FIN paths are exercised + // deterministically before the tasks finish. stream.write_all(&[0xac]).await.unwrap(); stream.shutdown().await.unwrap(); }); @@ -1909,17 +2011,7 @@ mod tests { let response_len = payload_len / 2 + case_id as usize % 257; let response = stress_payload(!payload_seed, response_len); write_stress_chunks(&mut stream, &response, !case_id).await; - if let Err(error) = stream.flush().await { - assert_eq!( - error.kind(), - std::io::ErrorKind::ConnectionReset, - "unexpected response flush error for case {case_id}: {error}" - ); - // The client only sends FIN after verifying the - // complete response, so this race still proves - // that every accepted byte reached the peer. - return; - } + stream.flush().await.unwrap(); let mut ack = [0u8; 1]; stream.read_exact(&mut ack).await.unwrap(); assert_eq!(ack, [0xac], "invalid completion ack for case {case_id}"); diff --git a/src/mux.rs b/src/mux.rs index 12f1b6c..1e9e9ec 100644 --- a/src/mux.rs +++ b/src/mux.rs @@ -210,7 +210,7 @@ impl MuxConnector { .lock() .handles .values() - .filter(|handle| !handle.closed) + .filter(|handle| !handle.is_fully_closed()) .count() } @@ -367,7 +367,7 @@ impl Future for MuxTimer { .handles .iter() .filter_map(|(id, h)| { - if !h.closed && now.duration_since(h.last_active) >= timeout { + if !h.is_fully_closed() && now.duration_since(h.last_active) >= timeout { Some(*id) } else { None @@ -384,20 +384,22 @@ impl Future for MuxTimer { // An unaccepted stream cannot have local writes, so // its FIN can go directly to the control queue before // the handle is reaped. - state.try_mark_finish(stream_id); - state.send_finish(stream_id); + if state.mark_fully_closed(stream_id) { + state.send_finish(stream_id); + } state.accept_queue.remove(position); state.remove_stream(stream_id); } else { // Keep FIN behind PSHs already accepted for an active // stream. Putting it in the global control queue would // let it overtake the stream's pending data. - state.try_mark_finish(stream_id); - state.enqueue_frame_stream( - stream_id, - MuxFrame::new(MuxCommand::Finish, stream_id, Bytes::new()), - ); - state.notify_should_tx(); + if state.mark_fully_closed(stream_id) { + state.enqueue_frame_stream( + stream_id, + MuxFrame::new(MuxCommand::Finish, stream_id, Bytes::new()), + ); + state.notify_should_tx(); + } state.notify_rx_consumed(); } } @@ -501,7 +503,7 @@ impl Future for MuxDispatcher { state.notify_accept_stream(); } MuxCommand::Finish => { - state.try_mark_finish(frame.header.stream_id); + state.mark_read_closed(frame.header.stream_id); } MuxCommand::Push => { let stream_id = frame.header.stream_id; @@ -554,7 +556,7 @@ impl Future for MuxWorker { /// `AsyncRead + AsyncWrite + Unpin`, so it can be used anywhere a /// `TcpStream` would. Dropping it without `shutdown()` is fine — the /// stream's pending tx queue is moved to the global queue, and a FIN is -/// enqueued when neither the peer nor the session has already closed it. +/// enqueued unless its local write half or the session is already closed. pub struct MuxStream { stream_id: u32, state: Arc>>, @@ -575,7 +577,7 @@ impl Drop for MuxStream { .map(|h| { ( std::mem::take(&mut h.tx_queue), - may_send_finish && !h.closed, + may_send_finish && !h.write_closed, ) }) .unwrap_or_default(); @@ -656,9 +658,9 @@ impl AsyncWrite for MuxStream { } let mut state = self.state.lock(); - if state.is_closed(self.stream_id) { + if state.is_write_closed(self.stream_id) { return Poll::Ready(Err(new_io_err( - StdIo::ErrorKind::ConnectionReset, + StdIo::ErrorKind::BrokenPipe, "stream tx is already closed", ))); } @@ -680,19 +682,16 @@ impl AsyncWrite for MuxStream { fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut state = self.state.lock(); - // Always try to drain the stream's tx_queue first, even if the - // stream has been remotely closed. Bytes that the user already - // accepted via poll_write should reach the wire on a best-effort - // basis; only after the queue is empty do we surface the close as - // an error. + // Always drain bytes accepted before the local write half closed. + // A peer FIN only closes our read half and must not make flush fail. match state.poll_flush_stream_frames(cx, self.stream_id) { Poll::Ready(Ok(())) => {} Poll::Ready(Err(e)) => return Poll::Ready(Err(mux_to_io_err(e))), Poll::Pending => return Poll::Pending, } - if state.is_closed(self.stream_id) { + if state.is_write_closed(self.stream_id) { return Poll::Ready(Err(new_io_err( - StdIo::ErrorKind::ConnectionReset, + StdIo::ErrorKind::BrokenPipe, "stream tx is already closed", ))); } @@ -706,11 +705,11 @@ impl AsyncWrite for MuxStream { .poll_flush_stream_frames(cx, self.stream_id) .map_err(mux_to_io_err))?; - if state.is_closed(self.stream_id) { + if state.is_write_closed(self.stream_id) { return Poll::Ready(Ok(())); } - state.try_mark_finish(self.stream_id); + state.mark_write_closed(self.stream_id); state.enqueue_frame_stream( self.stream_id, MuxFrame::new(MuxCommand::Finish, self.stream_id, Bytes::new()), @@ -721,10 +720,13 @@ impl AsyncWrite for MuxStream { } impl MuxStream { - /// True once the stream has received FIN, the session has hard- - /// closed, or an idle timeout has closed the stream. + /// True once both stream directions have closed, the session has hard- + /// closed, or an idle timeout has fully closed the stream. + /// + /// A single FIN only closes one direction, so a half-closed stream still + /// returns `false` while it remains usable for reading or writing. pub fn is_closed(&self) -> bool { - self.state.lock().is_closed(self.stream_id) + self.state.lock().is_stream_fully_closed(self.stream_id) } /// The 32-bit smux stream id assigned to this stream. @@ -734,7 +736,11 @@ impl MuxStream { } struct StreamHandle { - closed: bool, + /// The peer sent FIN, so no more inbound payload is valid once the + /// receive queue has been drained. + read_closed: bool, + /// We sent (or queued) FIN, so no more writes may be accepted locally. + write_closed: bool, tx_queue: VecDeque, tx_done_waker: Option, @@ -755,7 +761,8 @@ struct StreamHandle { impl StreamHandle { fn new() -> Self { Self { - closed: false, + read_closed: false, + write_closed: false, tx_queue: VecDeque::new(), tx_done_waker: None, unflushed: false, @@ -766,6 +773,11 @@ impl StreamHandle { } } + #[inline] + fn is_fully_closed(&self) -> bool { + self.read_closed && self.write_closed + } + #[inline] fn register_tx_done_waker(&mut self, cx: &Context<'_>) { self.tx_done_waker = Some(cx.waker().clone()); @@ -963,30 +975,67 @@ impl MuxState { fn close_unaccepted_streams(&mut self) { while let Some(stream_id) = self.accept_queue.pop_front() { - self.try_mark_finish(stream_id); - self.send_finish(stream_id); + if self.mark_fully_closed(stream_id) { + self.send_finish(stream_id); + } self.remove_stream(stream_id); } } #[inline] - fn try_mark_finish(&mut self, stream_id: u32) { + fn mark_read_closed(&mut self, stream_id: u32) { + if let Some(h) = self.handles.get_mut(&stream_id) { + h.last_active = Instant::now(); + if !h.read_closed { + h.read_closed = true; + h.notify_rx_ready(); + } + } + } + + /// Mark the local write half closed. Returns true only for the transition + /// that must enqueue FIN. + #[inline] + fn mark_write_closed(&mut self, stream_id: u32) -> bool { + if let Some(h) = self.handles.get_mut(&stream_id) { + if h.write_closed { + return false; + } + h.write_closed = true; + h.notify_tx_done(); + true + } else { + false + } + } + + /// Fully close a stream locally. Returns whether a FIN still needs to be + /// sent for its write half. + #[inline] + fn mark_fully_closed(&mut self, stream_id: u32) -> bool { if let Some(h) = self.handles.get_mut(&stream_id) { - h.closed = true; + let should_send_finish = !h.write_closed; + h.read_closed = true; + h.write_closed = true; h.notify_rx_ready(); h.notify_tx_done(); + should_send_finish + } else { + false } } fn recv_push(&mut self, frame: MuxFrame) -> bool { if let Some(handle) = self.handles.get_mut(&frame.header.stream_id) { - // If we've already locally closed the stream (shutdown / FIN - // received), drop incoming PSH instead of silently appending it - // to rx_queue. Otherwise EOF would not be monotonic: a reader - // who already saw 0 bytes could then surface bytes that arrived - // afterwards. - if handle.closed { - return false; + // FIN from the peer closes only our read direction. Discard PSH + // after that EOF boundary, but continue accepting peer responses + // after our own write half has been shut down. + if handle.read_closed { + // The stream still exists and its local write half may remain + // usable. Silently discard this protocol-invalid late data; + // replying with FIN here would close that write half on the + // wire while the local API continued accepting writes. + return true; } handle.last_active = Instant::now(); // A zero-length PSH carries no data. Queuing it would make a @@ -1028,16 +1077,30 @@ impl MuxState { } } - fn is_closed(&self, stream_id: u32) -> bool { + fn is_stream_fully_closed(&self, stream_id: u32) -> bool { // Invariant: when a MuxStream holds `stream_id`, its handle must // be present. Catch invariant breaks in debug builds; in release // builds, treat a missing handle as closed so we don't panic in // the user's hot path. debug_assert!( self.handles.contains_key(&stream_id), - "is_closed called with unknown stream id {stream_id}" + "is_stream_fully_closed called with unknown stream id {stream_id}" + ); + self.closed + || self + .handles + .get(&stream_id) + .is_none_or(StreamHandle::is_fully_closed) + } + + fn is_write_closed(&self, stream_id: u32) -> bool { + debug_assert!( + self.handles.contains_key(&stream_id), + "is_write_closed called with unknown stream id {stream_id}" ); - self.handles.get(&stream_id).is_none_or(|h| h.closed) + self.closed + || self.shutdown_requested + || self.handles.get(&stream_id).is_none_or(|h| h.write_closed) } fn poll_next_frame(&mut self, cx: &mut Context<'_>) -> Poll> { @@ -1088,7 +1151,7 @@ impl MuxState { Poll::Ready(Ok(Some(f))) } else if self.closed { Poll::Ready(Err(MuxError::ConnectionClosed)) - } else if handle.closed { + } else if handle.read_closed { // EOF Poll::Ready(Ok(None)) } else { @@ -1107,7 +1170,9 @@ impl MuxState { let Some(handle) = self.handles.get_mut(&stream_id) else { return Poll::Ready(Err(MuxError::StreamClosed(stream_id))); }; - if handle.tx_queue.len() >= self.max_tx_queue { + if handle.write_closed { + Poll::Ready(Err(MuxError::StreamClosed(stream_id))) + } else if handle.tx_queue.len() >= self.max_tx_queue { // A stream's tx queue is full handle.register_tx_done_waker(cx); // Notify the worker to transfer data now @@ -1223,7 +1288,8 @@ impl MuxState { // reads/writes wake immediately. Their already-accepted tx queues are // retained and will still be drained by poll_flush_frames. for handle in self.handles.values_mut() { - handle.closed = true; + handle.read_closed = true; + handle.write_closed = true; handle.notify_rx_ready(); handle.notify_tx_done(); } @@ -1257,7 +1323,8 @@ impl MuxState { self.notify_should_tx(); self.notify_close_waiters(); for h in self.handles.values_mut() { - h.closed = true; + h.read_closed = true; + h.write_closed = true; h.tx_queue = VecDeque::new(); h.notify_rx_ready(); h.notify_tx_done(); @@ -1442,7 +1509,7 @@ mod alloc_tests { Bytes::from_static(b"queued") ))); - s.try_mark_finish(stream_id); + s.mark_read_closed(stream_id); assert_eq!( s.get_rx_pending(), @@ -1450,4 +1517,17 @@ mod alloc_tests { "unread payload stopped counting as soon as FIN arrived" ); } + + #[test] + fn remote_fin_refreshes_stream_idle_activity() { + let mut s = fresh_state(StreamIdType::Odd); + let stream_id = 2; + s.process_sync(stream_id, Direction::Rx).unwrap(); + let stale = Instant::now() - Duration::from_secs(60); + s.handles.get_mut(&stream_id).unwrap().last_active = stale; + + s.mark_read_closed(stream_id); + + assert!(s.handles[&stream_id].last_active > stale); + } } From 8f9d956c5ce26e18665a02888b04dda954a7fbd7 Mon Sep 17 00:00:00 2001 From: black-binary Date: Fri, 4 Sep 2026 10:47:07 +0800 Subject: [PATCH 3/3] chore: bump version to 0.4.0 --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index a76ace8..c69505c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "async_smux" -version = "0.3.4" +version = "0.4.0" authors = ["black-binary "] description = "Asynchronous smux multiplexing library" license = "MIT"