Harden stream cleanup state

src/bin/dosh-client.rs

@@ -5022,13 +5022,17 @@ async fn run_terminal(
                         {
                             let _ = writer.write_all(&[5, 5, 0, 1, 0, 0, 0, 0, 0, 0]).await;
                         }
-                        opened_streams.remove(&reject.stream_id);
-                        stream_send_credit.remove(&reject.stream_id);
-                        stream_pending_data.remove(&reject.stream_id);
-                        stream_next_send_offset.remove(&reject.stream_id);
-                        stream_sent_data.remove(&reject.stream_id);
-                        stream_next_recv_offset.remove(&reject.stream_id);
-                        stream_recv_pending.remove(&reject.stream_id);
+                        cleanup_stream_state(
+                            reject.stream_id,
+                            &mut opened_streams,
+                            &mut stream_send_credit,
+                            &mut stream_pending_data,
+                            &mut stream_pending_opens,
+                            &mut stream_next_send_offset,
+                            &mut stream_sent_data,
+                            &mut stream_next_recv_offset,
+                            &mut stream_recv_pending,
+                        );
                         stream_writers.remove(&reject.stream_id);
                     }
                     PacketKind::StreamData => {
@@ -5101,14 +5105,17 @@ async fn run_terminal(
                         };
                         last_packet_at = Instant::now();
                         pending_socks_replies.remove(&close.stream_id);
-                        stream_pending_opens.remove(&close.stream_id);
-                        opened_streams.remove(&close.stream_id);
-                        stream_send_credit.remove(&close.stream_id);
-                        stream_pending_data.remove(&close.stream_id);
-                        stream_next_send_offset.remove(&close.stream_id);
-                        stream_sent_data.remove(&close.stream_id);
-                        stream_next_recv_offset.remove(&close.stream_id);
-                        stream_recv_pending.remove(&close.stream_id);
+                        cleanup_stream_state(
+                            close.stream_id,
+                            &mut opened_streams,
+                            &mut stream_send_credit,
+                            &mut stream_pending_data,
+                            &mut stream_pending_opens,
+                            &mut stream_next_send_offset,
+                            &mut stream_sent_data,
+                            &mut stream_next_recv_offset,
+                            &mut stream_recv_pending,
+                        );
                         stream_writers.remove(&close.stream_id);
                         #[cfg(unix)]
                         agent_writers.remove(&close.stream_id);
@@ -5160,14 +5167,17 @@ async fn run_terminal(
                     }
                     Some(ForwardEvent::Close { stream_id }) => {
                         pending_socks_replies.remove(&stream_id);
-                        stream_pending_opens.remove(&stream_id);
-                        opened_streams.remove(&stream_id);
-                        stream_send_credit.remove(&stream_id);
-                        stream_pending_data.remove(&stream_id);
-                        stream_next_send_offset.remove(&stream_id);
-                        stream_sent_data.remove(&stream_id);
-                        stream_next_recv_offset.remove(&stream_id);
-                        stream_recv_pending.remove(&stream_id);
+                        cleanup_stream_state(
+                            stream_id,
+                            &mut opened_streams,
+                            &mut stream_send_credit,
+                            &mut stream_pending_data,
+                            &mut stream_pending_opens,
+                            &mut stream_next_send_offset,
+                            &mut stream_sent_data,
+                            &mut stream_next_recv_offset,
+                            &mut stream_recv_pending,
+                        );
                         stream_writers.remove(&stream_id);
                         #[cfg(unix)]
                         agent_writers.remove(&stream_id);
@@ -6140,6 +6150,28 @@ fn accept_stream_data(
     (writes, consumed, *expected)
 }
 
+#[allow(clippy::too_many_arguments)]
+fn cleanup_stream_state(
+    stream_id: u64,
+    opened_streams: &mut HashSet<u64>,
+    stream_send_credit: &mut HashMap<u64, usize>,
+    stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
+    stream_pending_opens: &mut HashMap<u64, PendingStreamOpen>,
+    stream_next_send_offset: &mut HashMap<u64, u64>,
+    stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
+    stream_next_recv_offset: &mut HashMap<u64, u64>,
+    stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
+) {
+    opened_streams.remove(&stream_id);
+    stream_send_credit.remove(&stream_id);
+    stream_pending_data.remove(&stream_id);
+    stream_pending_opens.remove(&stream_id);
+    stream_next_send_offset.remove(&stream_id);
+    stream_sent_data.remove(&stream_id);
+    stream_next_recv_offset.remove(&stream_id);
+    stream_recv_pending.remove(&stream_id);
+}
+
 fn ack_stream_data(
     stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
     stream_send_credit: &mut HashMap<u64, usize>,
@@ -7425,19 +7457,20 @@ mod tests {
         FRAME_GAP_RESYNC_AFTER_MS, FrameBuffer, LocalForward, MAX_PENDING_USER_INPUT_BYTES,
         POST_SUBMIT_ALL_INPUT_HOLD, PendingStreamOpen, PredictMode, Predictor,
         RESTART_STATUS_SCRIPT, RemoteForward, STARTUP_INPUT_HOLD, SshConfig, StartupGateMode,
-        StatusAction, auth_allows, cache_key, cache_server_prefix, clear_cached_credentials,
-        ensure_tui_safe_status_overlay, input_contains_focus_in, input_matches_escape,
-        is_local_status_target, is_resume_response_for_client, latest_release_download_url,
-        load_first_native_identity_with_prompt, parse_dynamic_forward, parse_escape_key,
-        parse_local_forward, parse_remote_forward, parse_ssh_config, post_submit_hold_duration,
-        queue_pending_user_input, raw_contains_host_table, recv_response_until, refresh_live_addr,
-        release_tag_download_url, release_tag_from_effective_url, release_version_from_tag,
-        render_status_clear, render_status_overlay, requested_env, resolved_startup_command,
-        retransmit_stream_opens, rewrite_forward_command, selected_predict_mode, selected_udp_host,
-        server_version_mismatch, should_flush_terminal_input_after_contact,
-        should_hold_during_startup_gate, should_hold_post_submit_input,
-        should_repaint_idle_alternate_screen, split_after_command_submit, ssh_destination_host,
-        ssh_username, ssh_with_user, startup_command, status_ssh_target, strip_stale_mouse_reports,
+        StatusAction, auth_allows, cache_key, cache_server_prefix, cleanup_stream_state,
+        clear_cached_credentials, ensure_tui_safe_status_overlay, input_contains_focus_in,
+        input_matches_escape, is_local_status_target, is_resume_response_for_client,
+        latest_release_download_url, load_first_native_identity_with_prompt, parse_dynamic_forward,
+        parse_escape_key, parse_local_forward, parse_remote_forward, parse_ssh_config,
+        post_submit_hold_duration, queue_pending_user_input, raw_contains_host_table,
+        recv_response_until, refresh_live_addr, release_tag_download_url,
+        release_tag_from_effective_url, release_version_from_tag, render_status_clear,
+        render_status_overlay, requested_env, resolved_startup_command, retransmit_stream_opens,
+        rewrite_forward_command, selected_predict_mode, selected_udp_host, server_version_mismatch,
+        should_flush_terminal_input_after_contact, should_hold_during_startup_gate,
+        should_hold_post_submit_input, should_repaint_idle_alternate_screen,
+        split_after_command_submit, ssh_destination_host, ssh_username, ssh_with_user,
+        startup_command, status_ssh_target, strip_stale_mouse_reports,
         strip_terminal_focus_reports, terminal_private_mode_transition, toml_bare_key_or_quoted,
         update_check_requested, update_version_status, upsert_managed_block, valid_forward_host,
         vscode_safe_alias,
@@ -7447,7 +7480,7 @@ mod tests {
     use dosh::protocol::{self, Frame, PacketKind};
     use dosh::udp::is_transient_udp_send_error;
     use ed25519_dalek::{SigningKey, VerifyingKey};
-    use std::collections::{HashMap, VecDeque};
+    use std::collections::{BTreeMap, HashMap, VecDeque};
     use std::fs;
     use std::time::{Duration, Instant};
 
@@ -7493,6 +7526,100 @@ mod tests {
         assert_eq!(vscode_safe_alias(":///"), "host");
     }
 
+    #[test]
+    fn cleanup_stream_state_removes_all_per_stream_state_only_for_target() {
+        let stream_id = 42;
+        let other_stream_id = 43;
+        let mut opened_streams = std::collections::HashSet::from([stream_id, other_stream_id]);
+        let mut stream_send_credit = HashMap::from([(stream_id, 12), (other_stream_id, 34)]);
+        let mut stream_pending_data = HashMap::from([
+            (stream_id, VecDeque::from([vec![1]])),
+            (other_stream_id, VecDeque::from([vec![2]])),
+        ]);
+        let mut stream_pending_opens = HashMap::from([
+            (
+                stream_id,
+                PendingStreamOpen {
+                    target_host: "127.0.0.1".to_string(),
+                    target_port: 1234,
+                    last_sent: Instant::now(),
+                    attempts: 1,
+                },
+            ),
+            (
+                other_stream_id,
+                PendingStreamOpen {
+                    target_host: "127.0.0.1".to_string(),
+                    target_port: 1235,
+                    last_sent: Instant::now(),
+                    attempts: 1,
+                },
+            ),
+        ]);
+        let mut stream_next_send_offset = HashMap::from([(stream_id, 55), (other_stream_id, 66)]);
+        let mut stream_sent_data = HashMap::from([
+            (
+                stream_id,
+                BTreeMap::from([(
+                    0,
+                    super::PendingStreamChunk {
+                        offset: 0,
+                        bytes: vec![3],
+                        last_sent: Instant::now(),
+                        attempts: 1,
+                    },
+                )]),
+            ),
+            (
+                other_stream_id,
+                BTreeMap::from([(
+                    0,
+                    super::PendingStreamChunk {
+                        offset: 0,
+                        bytes: vec![4],
+                        last_sent: Instant::now(),
+                        attempts: 1,
+                    },
+                )]),
+            ),
+        ]);
+        let mut stream_next_recv_offset = HashMap::from([(stream_id, 77), (other_stream_id, 88)]);
+        let mut stream_recv_pending = HashMap::from([
+            (stream_id, BTreeMap::from([(77, vec![5])])),
+            (other_stream_id, BTreeMap::from([(88, vec![6])])),
+        ]);
+
+        cleanup_stream_state(
+            stream_id,
+            &mut opened_streams,
+            &mut stream_send_credit,
+            &mut stream_pending_data,
+            &mut stream_pending_opens,
+            &mut stream_next_send_offset,
+            &mut stream_sent_data,
+            &mut stream_next_recv_offset,
+            &mut stream_recv_pending,
+        );
+
+        assert!(!opened_streams.contains(&stream_id));
+        assert!(!stream_send_credit.contains_key(&stream_id));
+        assert!(!stream_pending_data.contains_key(&stream_id));
+        assert!(!stream_pending_opens.contains_key(&stream_id));
+        assert!(!stream_next_send_offset.contains_key(&stream_id));
+        assert!(!stream_sent_data.contains_key(&stream_id));
+        assert!(!stream_next_recv_offset.contains_key(&stream_id));
+        assert!(!stream_recv_pending.contains_key(&stream_id));
+
+        assert!(opened_streams.contains(&other_stream_id));
+        assert!(stream_send_credit.contains_key(&other_stream_id));
+        assert!(stream_pending_data.contains_key(&other_stream_id));
+        assert!(stream_pending_opens.contains_key(&other_stream_id));
+        assert!(stream_next_send_offset.contains_key(&other_stream_id));
+        assert!(stream_sent_data.contains_key(&other_stream_id));
+        assert!(stream_next_recv_offset.contains_key(&other_stream_id));
+        assert!(stream_recv_pending.contains_key(&other_stream_id));
+    }
+
     #[test]
     fn managed_ssh_block_is_replaced_not_duplicated() {
         let dir = tempfile::tempdir().unwrap();

src/bin/dosh-server.rs

@@ -2133,11 +2133,7 @@ async fn handle_stream_open_reject(
     }
     client.endpoint = peer;
     client.last_seen = Instant::now();
-    client.stream_writers.remove(&reject.stream_id);
-    client.stream_pending_opens.remove(&reject.stream_id);
-    client.opened_streams.remove(&reject.stream_id);
-    client.stream_send_credit.remove(&reject.stream_id);
-    client.stream_pending_data.remove(&reject.stream_id);
+    cleanup_client_stream(client, reject.stream_id);
     Ok(())
 }
 
@@ -2192,15 +2188,7 @@ async fn handle_stream_close(
     // Connection migration on any authenticated, fresh packet (spec §11).
     client.endpoint = peer;
     client.last_seen = Instant::now();
-    client.stream_writers.remove(&close.stream_id);
-    client.stream_pending_opens.remove(&close.stream_id);
-    client.opened_streams.remove(&close.stream_id);
-    client.stream_send_credit.remove(&close.stream_id);
-    client.stream_pending_data.remove(&close.stream_id);
-    client.stream_next_send_offset.remove(&close.stream_id);
-    client.stream_sent_data.remove(&close.stream_id);
-    client.stream_next_recv_offset.remove(&close.stream_id);
-    client.stream_recv_pending.remove(&close.stream_id);
+    cleanup_client_stream(client, close.stream_id);
     Ok(())
 }
 
@@ -3280,6 +3268,18 @@ fn accept_stream_data(client: &mut ClientState, data: StreamData) -> (Vec<Vec<u8
     (writes, consumed, *expected)
 }
 
+fn cleanup_client_stream(client: &mut ClientState, stream_id: u64) {
+    client.stream_writers.remove(&stream_id);
+    client.stream_pending_opens.remove(&stream_id);
+    client.opened_streams.remove(&stream_id);
+    client.stream_send_credit.remove(&stream_id);
+    client.stream_pending_data.remove(&stream_id);
+    client.stream_next_send_offset.remove(&stream_id);
+    client.stream_sent_data.remove(&stream_id);
+    client.stream_next_recv_offset.remove(&stream_id);
+    client.stream_recv_pending.remove(&stream_id);
+}
+
 fn ack_stream_data(client: &mut ClientState, stream_id: u64, received_offset: u64) {
     let Some(sent) = client.stream_sent_data.get_mut(&stream_id) else {
         return;
@@ -3340,14 +3340,7 @@ async fn send_stream_close_to_client(
     {
         let mut locked = state.lock().expect("server state poisoned");
         if let Some(client) = locked.client_mut(&client_id) {
-            client.stream_writers.remove(&stream_id);
-            client.opened_streams.remove(&stream_id);
-            client.stream_send_credit.remove(&stream_id);
-            client.stream_pending_data.remove(&stream_id);
-            client.stream_next_send_offset.remove(&stream_id);
-            client.stream_sent_data.remove(&stream_id);
-            client.stream_next_recv_offset.remove(&stream_id);
-            client.stream_recv_pending.remove(&stream_id);
+            cleanup_client_stream(client, stream_id);
         }
     }
     let body = protocol::to_body(&StreamClose { stream_id })?;
@@ -4053,6 +4046,106 @@ mod tests {
         }
     }
 
+    #[test]
+    fn cleanup_client_stream_removes_all_per_stream_state_only_for_target() {
+        let stream_id = 42;
+        let other_stream_id = 43;
+        let mut client = test_client_state([7u8; 32]);
+        let (writer_tx, _writer_rx) = mpsc::channel(1);
+        client.stream_writers.insert(stream_id, writer_tx);
+        client.opened_streams.insert(stream_id);
+        client.opened_streams.insert(other_stream_id);
+        client.stream_send_credit.insert(stream_id, 12);
+        client.stream_send_credit.insert(other_stream_id, 34);
+        client
+            .stream_pending_data
+            .insert(stream_id, VecDeque::from([vec![1]]));
+        client
+            .stream_pending_data
+            .insert(other_stream_id, VecDeque::from([vec![2]]));
+        client.stream_pending_opens.insert(
+            stream_id,
+            PendingStreamOpen {
+                target_host: "127.0.0.1".to_string(),
+                target_port: 1234,
+                last_sent: Instant::now(),
+                attempts: 1,
+            },
+        );
+        client.stream_pending_opens.insert(
+            other_stream_id,
+            PendingStreamOpen {
+                target_host: "127.0.0.1".to_string(),
+                target_port: 1235,
+                last_sent: Instant::now(),
+                attempts: 1,
+            },
+        );
+        client.stream_next_send_offset.insert(stream_id, 55);
+        client.stream_next_send_offset.insert(other_stream_id, 66);
+        client.stream_sent_data.insert(
+            stream_id,
+            BTreeMap::from([(
+                0,
+                PendingStreamChunk {
+                    offset: 0,
+                    bytes: vec![3],
+                    last_sent: Instant::now(),
+                    attempts: 1,
+                },
+            )]),
+        );
+        client.stream_sent_data.insert(
+            other_stream_id,
+            BTreeMap::from([(
+                0,
+                PendingStreamChunk {
+                    offset: 0,
+                    bytes: vec![4],
+                    last_sent: Instant::now(),
+                    attempts: 1,
+                },
+            )]),
+        );
+        client.stream_next_recv_offset.insert(stream_id, 77);
+        client.stream_next_recv_offset.insert(other_stream_id, 88);
+        client
+            .stream_recv_pending
+            .insert(stream_id, BTreeMap::from([(77, vec![5])]));
+        client
+            .stream_recv_pending
+            .insert(other_stream_id, BTreeMap::from([(88, vec![6])]));
+
+        cleanup_client_stream(&mut client, stream_id);
+
+        assert!(!client.stream_writers.contains_key(&stream_id));
+        assert!(!client.opened_streams.contains(&stream_id));
+        assert!(!client.stream_send_credit.contains_key(&stream_id));
+        assert!(!client.stream_pending_data.contains_key(&stream_id));
+        assert!(!client.stream_pending_opens.contains_key(&stream_id));
+        assert!(!client.stream_next_send_offset.contains_key(&stream_id));
+        assert!(!client.stream_sent_data.contains_key(&stream_id));
+        assert!(!client.stream_next_recv_offset.contains_key(&stream_id));
+        assert!(!client.stream_recv_pending.contains_key(&stream_id));
+
+        assert!(client.opened_streams.contains(&other_stream_id));
+        assert!(client.stream_send_credit.contains_key(&other_stream_id));
+        assert!(client.stream_pending_data.contains_key(&other_stream_id));
+        assert!(client.stream_pending_opens.contains_key(&other_stream_id));
+        assert!(
+            client
+                .stream_next_send_offset
+                .contains_key(&other_stream_id)
+        );
+        assert!(client.stream_sent_data.contains_key(&other_stream_id));
+        assert!(
+            client
+                .stream_next_recv_offset
+                .contains_key(&other_stream_id)
+        );
+        assert!(client.stream_recv_pending.contains_key(&other_stream_id));
+    }
+
     #[test]
     fn client_index_stays_in_sync_with_session_clients() {
         let (pty_tx, _pty_rx) = mpsc::unbounded_channel();