From 5c8bd68aa6f02d0344b67f77c034c3715d2d1776 Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:49:56 -0400 Subject: [PATCH] Retire closed transport stream ids --- src/bin/dosh-client.rs | 10 +- src/bin/dosh-server.rs | 16 ++- src/transport.rs | 315 ++++++++++++++++++++++++++++++++++++----- 3 files changed, 303 insertions(+), 38 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index eb6dc69..0c784a2 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -4986,7 +4986,9 @@ async fn run_terminal( continue; }; last_packet_at = Instant::now(); - stream_pending_opens.remove(&ok.stream_id); + if stream_pending_opens.remove(&ok.stream_id).is_none() { + continue; + } opened_streams.insert(ok.stream_id); stream_send_credit.entry(ok.stream_id).or_insert(STREAM_INITIAL_WINDOW); stream_next_send_offset.entry(ok.stream_id).or_insert(0); @@ -5044,6 +5046,9 @@ async fn run_terminal( }; last_packet_at = Instant::now(); let stream_id = data.stream_id; + if !opened_streams.contains(&stream_id) { + continue; + } let (writes, consumed, received_offset) = accept_stream_data(&mut stream_next_recv_offset, &mut stream_recv_pending, data); if let Some(writer) = stream_writers.get_mut(&stream_id) { @@ -5078,6 +5083,9 @@ async fn run_terminal( continue; }; last_packet_at = Instant::now(); + if !opened_streams.contains(&adjust.stream_id) { + continue; + } ack_stream_data( &mut stream_sent_data, &mut stream_send_credit, diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 9104348..2b4aba7 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -2053,6 +2053,9 @@ async fn handle_stream_data( } client.endpoint = peer; client.last_seen = Instant::now(); + if !client.opened_streams.contains(&data.stream_id) { + return Ok(()); + } let writer = client.stream_writers.get(&data.stream_id).cloned(); let (writes, consumed, received_offset) = accept_stream_data(client, data); (writer, writes, consumed, received_offset) @@ -2062,9 +2065,9 @@ async fn handle_stream_data( let _ = writer.send(bytes).await; } } - // Always return flow-control credit, even when the stream's writer is already - // gone (closed/unknown). The peer debited its send window for these bytes, so - // skipping the adjust would wedge the stream at a permanent credit deficit. + // Always return flow-control credit for live streams, even when the stream's + // writer is already gone. The peer debited its send window for these bytes, + // so skipping the adjust would wedge the stream at a permanent credit deficit. send_stream_window_adjust_to_client( state, socket, @@ -2101,7 +2104,9 @@ async fn handle_stream_open_ok( } client.endpoint = peer; client.last_seen = Instant::now(); - client.stream_pending_opens.remove(&ok.stream_id); + if client.stream_pending_opens.remove(&ok.stream_id).is_none() { + return Ok(()); + } client.opened_streams.insert(ok.stream_id); client .stream_send_credit @@ -2161,6 +2166,9 @@ async fn handle_stream_window_adjust( } client.endpoint = peer; client.last_seen = Instant::now(); + if !client.opened_streams.contains(&adjust.stream_id) { + return Ok(()); + } ack_stream_data(client, adjust.stream_id, adjust.received_offset); } flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, adjust.stream_id) diff --git a/src/transport.rs b/src/transport.rs index 58c455e..866db14 100644 --- a/src/transport.rs +++ b/src/transport.rs @@ -141,6 +141,9 @@ pub enum TransportEvent { Close { stream_id: u64, }, + Ignored { + stream_id: u64, + }, } #[derive(Debug, Clone)] @@ -170,6 +173,7 @@ pub struct StreamMux { next_send_offset: HashMap, next_recv_offset: HashMap, recv_pending: HashMap>>, + retired_streams: HashSet, } #[derive(Debug, Clone, Default)] @@ -225,6 +229,7 @@ impl StreamMux { next_send_offset: HashMap::new(), next_recv_offset: HashMap::new(), recv_pending: HashMap::new(), + retired_streams: HashSet::new(), } } @@ -235,7 +240,10 @@ impl StreamMux { target_port: u16, ) -> Result { let target_host = target_host.into(); - if self.opened_streams.contains(&stream_id) || self.pending_opens.contains_key(&stream_id) { + if self.opened_streams.contains(&stream_id) + || self.pending_opens.contains_key(&stream_id) + || self.retired_streams.contains(&stream_id) + { bail!("stream {stream_id} already exists"); } self.send_credit @@ -262,6 +270,9 @@ impl StreamMux { } pub fn accept_open(&mut self, open: StreamOpen) -> Result { + if self.retired_streams.contains(&open.stream_id) { + bail!("stream {} was already closed", open.stream_id); + } if self.pending_opens.contains_key(&open.stream_id) { bail!( "stream {} collides with a local pending open", @@ -295,19 +306,39 @@ impl StreamMux { ) } - pub fn handle_open_ok(&mut self, ok: StreamOpenOk) -> Result> { - self.pending_opens.remove(&ok.stream_id); + pub fn reject_open( + &mut self, + stream_id: u64, + reason: impl Into, + ) -> Result { + if self.has_stream_state(stream_id) { + self.retire_stream(stream_id); + } + Self::open_reject(stream_id, reason) + } + + pub fn handle_open_ok( + &mut self, + ok: StreamOpenOk, + ) -> Result>> { + if self.pending_opens.remove(&ok.stream_id).is_none() { + return Ok(None); + } self.opened_streams.insert(ok.stream_id); self.send_credit .entry(ok.stream_id) .or_insert(self.config.initial_window); self.next_send_offset.entry(ok.stream_id).or_insert(0); self.next_recv_offset.entry(ok.stream_id).or_insert(0); - self.flush_pending_data(ok.stream_id) + self.flush_pending_data(ok.stream_id).map(Some) } - pub fn handle_open_reject(&mut self, reject: StreamOpenReject) { - self.cleanup_stream(reject.stream_id); + pub fn handle_open_reject(&mut self, reject: StreamOpenReject) -> bool { + if !self.has_stream_state(reject.stream_id) { + return false; + } + self.retire_stream(reject.stream_id); + true } pub fn send_data( @@ -344,8 +375,11 @@ impl StreamMux { Ok(vec![self.send_data_now(stream_id, bytes)?]) } - pub fn handle_data(&mut self, data: StreamData) -> Result { + pub fn handle_data(&mut self, data: StreamData) -> Result> { let stream_id = data.stream_id; + if !self.opened_streams.contains(&stream_id) { + return Ok(None); + } let (chunks, consumed, received_offset) = self.accept_data(data); let window_adjust = encode_packet( PacketKind::StreamWindowAdjust, @@ -355,37 +389,52 @@ impl StreamMux { bytes: consumed.min(u32::MAX as usize) as u32, }, )?; - Ok(IncomingStreamData { + Ok(Some(IncomingStreamData { stream_id, chunks, consumed, received_offset, window_adjust, - }) + })) } pub fn handle_window_adjust( &mut self, adjust: StreamWindowAdjust, - ) -> Result> { + ) -> Result>> { + if !self.opened_streams.contains(&adjust.stream_id) { + return Ok(None); + } let acked_bytes = self.ack_data(adjust.stream_id, adjust.received_offset); self.add_credit(adjust.stream_id, acked_bytes); - self.flush_pending_data(adjust.stream_id) + self.flush_pending_data(adjust.stream_id).map(Some) } pub fn close_stream(&mut self, stream_id: u64) -> Result { - self.cleanup_stream(stream_id); + if !self.has_stream_state(stream_id) { + bail!("stream {stream_id} is not open"); + } + self.retire_stream(stream_id); encode_packet(PacketKind::StreamClose, &StreamClose { stream_id }) } - pub fn handle_close(&mut self, close: StreamClose) { - self.cleanup_stream(close.stream_id); + pub fn handle_close(&mut self, close: StreamClose) -> bool { + if !self.has_stream_state(close.stream_id) { + return false; + } + self.retire_stream(close.stream_id); + true } pub fn handle_packet(&mut self, kind: PacketKind, body: &[u8]) -> Result { match kind { PacketKind::StreamOpen => { let open: StreamOpen = protocol::from_body(body)?; + if self.retired_streams.contains(&open.stream_id) { + return Ok(TransportEvent::Ignored { + stream_id: open.stream_id, + }); + } if self.opened_streams.contains(&open.stream_id) { return Ok(TransportEvent::DuplicateOpen { stream_id: open.stream_id, @@ -397,32 +446,47 @@ impl StreamMux { PacketKind::StreamOpenOk => { let ok: StreamOpenOk = protocol::from_body(body)?; let stream_id = ok.stream_id; - let flushed = self.handle_open_ok(ok)?; - Ok(TransportEvent::OpenOk { stream_id, flushed }) + match self.handle_open_ok(ok)? { + Some(flushed) => Ok(TransportEvent::OpenOk { stream_id, flushed }), + None => Ok(TransportEvent::Ignored { stream_id }), + } } PacketKind::StreamOpenReject => { let reject: StreamOpenReject = protocol::from_body(body)?; - self.handle_open_reject(reject.clone()); - Ok(TransportEvent::OpenReject { - stream_id: reject.stream_id, - reason: reject.reason, - }) + let stream_id = reject.stream_id; + if self.handle_open_reject(reject.clone()) { + Ok(TransportEvent::OpenReject { + stream_id, + reason: reject.reason, + }) + } else { + Ok(TransportEvent::Ignored { stream_id }) + } } PacketKind::StreamData => { let data: StreamData = protocol::from_body(body)?; - self.handle_data(data).map(TransportEvent::Data) + let stream_id = data.stream_id; + match self.handle_data(data)? { + Some(data) => Ok(TransportEvent::Data(data)), + None => Ok(TransportEvent::Ignored { stream_id }), + } } PacketKind::StreamWindowAdjust => { let adjust: StreamWindowAdjust = protocol::from_body(body)?; let stream_id = adjust.stream_id; - let flushed = self.handle_window_adjust(adjust)?; - Ok(TransportEvent::WindowAdjust { stream_id, flushed }) + match self.handle_window_adjust(adjust)? { + Some(flushed) => Ok(TransportEvent::WindowAdjust { stream_id, flushed }), + None => Ok(TransportEvent::Ignored { stream_id }), + } } PacketKind::StreamClose => { let close: StreamClose = protocol::from_body(body)?; let stream_id = close.stream_id; - self.handle_close(close); - Ok(TransportEvent::Close { stream_id }) + if self.handle_close(close) { + Ok(TransportEvent::Close { stream_id }) + } else { + Ok(TransportEvent::Ignored { stream_id }) + } } _ => bail!("packet kind {kind:?} is not a stream transport packet"), } @@ -610,6 +674,22 @@ impl StreamMux { *credit = credit.saturating_add(bytes).min(self.config.initial_window); } + fn has_stream_state(&self, stream_id: u64) -> bool { + self.opened_streams.contains(&stream_id) + || self.pending_opens.contains_key(&stream_id) + || self.send_credit.contains_key(&stream_id) + || self.pending_data.contains_key(&stream_id) + || self.sent_data.contains_key(&stream_id) + || self.next_send_offset.contains_key(&stream_id) + || self.next_recv_offset.contains_key(&stream_id) + || self.recv_pending.contains_key(&stream_id) + } + + fn retire_stream(&mut self, stream_id: u64) { + self.cleanup_stream(stream_id); + self.retired_streams.insert(stream_id); + } + fn cleanup_stream(&mut self, stream_id: u64) { self.opened_streams.remove(&stream_id); self.pending_opens.remove(&stream_id); @@ -727,8 +807,8 @@ impl DoshTransport { } pub async fn reject_stream(&mut self, stream_id: u64, reason: impl Into) -> Result<()> { - self.send_outgoing(StreamMux::open_reject(stream_id, reason)?) - .await + let packet = self.mux.reject_open(stream_id, reason)?; + self.send_outgoing(packet).await } pub async fn send(&mut self, stream_id: u64, bytes: impl Into>) -> Result<()> { @@ -853,7 +933,8 @@ impl DoshTransport { } TransportEvent::Open(_) | TransportEvent::OpenReject { .. } - | TransportEvent::Close { .. } => {} + | TransportEvent::Close { .. } + | TransportEvent::Ignored { .. } => {} } Ok(()) } @@ -976,7 +1057,10 @@ mod tests { assert!(mux.send_data(7, b"hello".to_vec()).unwrap().is_empty()); assert_eq!(mux.pending_bytes(7), 5); - let flushed = mux.handle_open_ok(StreamOpenOk { stream_id: 7 }).unwrap(); + let flushed = mux + .handle_open_ok(StreamOpenOk { stream_id: 7 }) + .unwrap() + .unwrap(); assert_eq!(flushed.len(), 1); assert_eq!(flushed[0].kind, PacketKind::StreamData); let data: StreamData = decode(&flushed[0]); @@ -999,6 +1083,7 @@ mod tests { offset: 5, bytes: b"world".to_vec(), }) + .unwrap() .unwrap(); assert!(first.chunks.is_empty()); assert_eq!(first.received_offset, 0); @@ -1009,6 +1094,7 @@ mod tests { offset: 0, bytes: b"hello".to_vec(), }) + .unwrap() .unwrap(); assert_eq!(second.chunks, vec![b"hello".to_vec(), b"world".to_vec()]); assert_eq!(second.consumed, 10); @@ -1023,7 +1109,9 @@ mod tests { ..TransportConfig::default() }); mux.open_stream(1, "@dosh-test", 0).unwrap(); - mux.handle_open_ok(StreamOpenOk { stream_id: 1 }).unwrap(); + mux.handle_open_ok(StreamOpenOk { stream_id: 1 }) + .unwrap() + .unwrap(); let sent = mux.send_data(1, b"hello".to_vec()).unwrap(); assert_eq!(sent.len(), 1); assert!(mux.send_data(1, b"!".to_vec()).unwrap().is_empty()); @@ -1034,6 +1122,7 @@ mod tests { received_offset: 5, bytes: 5, }) + .unwrap() .unwrap(); assert_eq!(flushed.len(), 1); let data: StreamData = decode(&flushed[0]); @@ -1049,7 +1138,9 @@ mod tests { ..TransportConfig::default() }); mux.open_stream(1, "@dosh-test", 0).unwrap(); - mux.handle_open_ok(StreamOpenOk { stream_id: 1 }).unwrap(); + mux.handle_open_ok(StreamOpenOk { stream_id: 1 }) + .unwrap() + .unwrap(); assert_eq!(mux.send_data(1, b"hello".to_vec()).unwrap().len(), 1); assert_eq!(mux.send_credit[&1], 0); @@ -1059,6 +1150,7 @@ mod tests { received_offset: 5, bytes: 0, }) + .unwrap() .unwrap(); assert!(flushed.is_empty()); assert_eq!(mux.send_credit[&1], 5); @@ -1119,13 +1211,113 @@ mod tests { let mut mux = StreamMux::new(TransportConfig::default()); mux.open_stream(1, "@dosh-test", 0).unwrap(); assert!(mux.send_data(1, b"queued".to_vec()).unwrap().is_empty()); - mux.handle_open_ok(StreamOpenOk { stream_id: 1 }).unwrap(); + mux.handle_open_ok(StreamOpenOk { stream_id: 1 }) + .unwrap() + .unwrap(); mux.close_stream(1).unwrap(); assert!(mux.send_data(1, b"lost".to_vec()).is_err()); assert!(mux.send_data(99, b"lost".to_vec()).is_err()); } + #[test] + fn rejecting_incoming_stream_retires_local_state() { + let mut mux = StreamMux::new(TransportConfig::default()); + mux.accept_open(StreamOpen { + stream_id: 7, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + assert!(mux.is_open(7)); + + let reject = mux.reject_open(7, "nope").unwrap(); + + assert_eq!(reject.kind, PacketKind::StreamOpenReject); + assert!(!mux.is_open(7)); + assert!(mux.retired_streams.contains(&7)); + assert!(mux.send_data(7, b"must-not-send".to_vec()).is_err()); + assert!( + mux.accept_open(StreamOpen { + stream_id: 7, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .is_err() + ); + } + + #[test] + fn stale_packets_after_close_do_not_recreate_stream_state() { + let mut mux = StreamMux::new(TransportConfig::default()); + mux.open_stream(1, "@dosh-test", 0).unwrap(); + mux.handle_open_ok(StreamOpenOk { stream_id: 1 }) + .unwrap() + .unwrap(); + assert_eq!(mux.send_data(1, b"hello".to_vec()).unwrap().len(), 1); + mux.close_stream(1).unwrap(); + + let stale_data = protocol::to_body(&StreamData { + stream_id: 1, + offset: 0, + bytes: b"stale".to_vec(), + }) + .unwrap(); + assert_eq!( + mux.handle_packet(PacketKind::StreamData, &stale_data) + .unwrap(), + TransportEvent::Ignored { stream_id: 1 } + ); + + let stale_open_ok = protocol::to_body(&StreamOpenOk { stream_id: 1 }).unwrap(); + assert_eq!( + mux.handle_packet(PacketKind::StreamOpenOk, &stale_open_ok) + .unwrap(), + TransportEvent::Ignored { stream_id: 1 } + ); + + let stale_adjust = protocol::to_body(&StreamWindowAdjust { + stream_id: 1, + received_offset: 5, + bytes: 5, + }) + .unwrap(); + assert_eq!( + mux.handle_packet(PacketKind::StreamWindowAdjust, &stale_adjust) + .unwrap(), + TransportEvent::Ignored { stream_id: 1 } + ); + + let stale_open = protocol::to_body(&StreamOpen { + stream_id: 1, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + assert_eq!( + mux.handle_packet(PacketKind::StreamOpen, &stale_open) + .unwrap(), + TransportEvent::Ignored { stream_id: 1 } + ); + + let stale_close = protocol::to_body(&StreamClose { stream_id: 1 }).unwrap(); + assert_eq!( + mux.handle_packet(PacketKind::StreamClose, &stale_close) + .unwrap(), + TransportEvent::Ignored { stream_id: 1 } + ); + + assert!(!mux.is_open(1)); + assert!(!mux.pending_opens.contains_key(&1)); + assert!(!mux.send_credit.contains_key(&1)); + assert!(!mux.pending_data.contains_key(&1)); + assert!(!mux.sent_data.contains_key(&1)); + assert!(!mux.next_send_offset.contains_key(&1)); + assert!(!mux.next_recv_offset.contains_key(&1)); + assert!(!mux.recv_pending.contains_key(&1)); + assert!(mux.retired_streams.contains(&1)); + } + #[tokio::test] async fn session_transport_partitions_client_and_server_stream_ids() { let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); @@ -1291,6 +1483,63 @@ mod tests { } } + #[tokio::test] + async fn session_transport_reject_stream_retires_server_state() { + let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let client_addr = client_socket.local_addr().unwrap(); + let server_addr = server_socket.local_addr().unwrap(); + let key = [46u8; 32]; + let conn_id = [10u8; 16]; + let mut client = DoshTransport::new_owned( + client_socket, + SessionTransportConfig { + role: SessionRole::Client, + conn_id, + session_key: key, + peer_addr: server_addr, + initial_send_seq: 1, + initial_ack: 0, + stream: TransportConfig::default(), + }, + ); + let mut server = DoshTransport::new_owned( + server_socket, + SessionTransportConfig { + role: SessionRole::Server, + conn_id, + session_key: key, + peer_addr: client_addr, + initial_send_seq: 1, + initial_ack: 0, + stream: TransportConfig::default(), + }, + ); + + let stream_id = client.open_service("app").await.unwrap(); + match server.recv().await.unwrap() { + SessionEvent::Stream(TransportEvent::Open(open)) => { + assert_eq!(open.stream_id, stream_id); + assert!(server.mux().is_open(stream_id)); + server.reject_stream(stream_id, "disabled").await.unwrap(); + } + other => panic!("unexpected event {other:?}"), + } + + assert!(!server.mux().is_open(stream_id)); + assert!( + server + .send(stream_id, b"must-not-send".to_vec()) + .await + .is_err() + ); + assert!(matches!( + client.recv().await.unwrap(), + SessionEvent::Stream(TransportEvent::OpenReject { stream_id: rejected, reason }) + if rejected == stream_id && reason == "disabled" + )); + } + #[tokio::test] async fn session_transport_roams_to_authenticated_new_peer_addr() { let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();