Retire closed transport stream ids

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,

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
+    // Always return flow-control credit for live streams, even when the stream's
-    // gone (closed/unknown). The peer debited its send window for these bytes, so
+    // writer is already gone. The peer debited its send window for these bytes,
-    // skipping the adjust would wedge the stream at a permanent credit deficit.
+    // 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)

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<u64, u64>,
     next_recv_offset: HashMap<u64, u64>,
     recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
+    retired_streams: HashSet<u64>,
 }
 
 #[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<OutgoingStreamPacket> {
         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<IncomingStreamOpen> {
+        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<Vec<OutgoingStreamPacket>> {
+    pub fn reject_open(
-        self.pending_opens.remove(&ok.stream_id);
+        &mut self,
+        stream_id: u64,
+        reason: impl Into<String>,
+    ) -> Result<OutgoingStreamPacket> {
+        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<Option<Vec<OutgoingStreamPacket>>> {
+        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) {
+    pub fn handle_open_reject(&mut self, reject: StreamOpenReject) -> bool {
-        self.cleanup_stream(reject.stream_id);
+        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<IncomingStreamData> {
+    pub fn handle_data(&mut self, data: StreamData) -> Result<Option<IncomingStreamData>> {
         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<Vec<OutgoingStreamPacket>> {
+    ) -> Result<Option<Vec<OutgoingStreamPacket>>> {
+        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<OutgoingStreamPacket> {
-        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) {
+    pub fn handle_close(&mut self, close: StreamClose) -> bool {
-        self.cleanup_stream(close.stream_id);
+        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<TransportEvent> {
         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)?;
+                match self.handle_open_ok(ok)? {
-                Ok(TransportEvent::OpenOk { stream_id, flushed })
+                    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());
+                let stream_id = reject.stream_id;
+                if self.handle_open_reject(reject.clone()) {
                     Ok(TransportEvent::OpenReject {
-                    stream_id: reject.stream_id,
+                        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)?;
+                match self.handle_window_adjust(adjust)? {
-                Ok(TransportEvent::WindowAdjust { stream_id, flushed })
+                    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);
+                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<String>) -> Result<()> {
-        self.send_outgoing(StreamMux::open_reject(stream_id, reason)?)
+        let packet = self.mux.reject_open(stream_id, reason)?;
-            .await
+        self.send_outgoing(packet).await
     }
 
     pub async fn send(&mut self, stream_id: u64, bytes: impl Into<Vec<u8>>) -> 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();