Retire closed transport stream ids

2026-07-09 19:49:56 -04:00
parent 68cf411143
commit 5c8bd68aa6
3 changed files with 303 additions and 38 deletions
@@ -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,
@@ -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)
@@ -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>> {
-        self.pending_opens.remove(&ok.stream_id);
+    pub fn reject_open(
+        &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) {
-        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<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) {
-        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<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)?;
-                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<String>) -> 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<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();