diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index d41ce3e..f3079bb 100644 --- a/src/bin/dosh-client.rs +++ b/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, + stream_send_credit: &mut HashMap, + stream_pending_data: &mut HashMap>>, + stream_pending_opens: &mut HashMap, + stream_next_send_offset: &mut HashMap, + stream_sent_data: &mut HashMap>, + stream_next_recv_offset: &mut HashMap, + stream_recv_pending: &mut HashMap>>, +) { + 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>, stream_send_credit: &mut HashMap, @@ -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(); diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 0bee9cb..9104348 100644 --- a/src/bin/dosh-server.rs +++ b/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