diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index 5962639..9fa3677 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -67,6 +67,7 @@ use tokio::sync::mpsc; const STREAM_INITIAL_WINDOW: usize = 1024 * 1024; const STREAM_RETIRED_TOMBSTONES: usize = 16 * 1024; +const STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS: u32 = 8; const MAX_PENDING_USER_INPUT_BYTES: usize = 1024 * 1024; const STARTUP_INPUT_HOLD: Duration = Duration::from_millis(750); const POST_SUBMIT_ALL_INPUT_HOLD: Duration = Duration::from_millis(120); @@ -5912,6 +5913,7 @@ async fn run_terminal( let mut stream_local_eof_sent: HashSet = HashSet::new(); let mut stream_remote_eof_received: HashSet = HashSet::new(); let mut stream_eof_retransmit: HashMap = HashMap::new(); + let mut stream_close_retransmit: HashMap = HashMap::new(); let mut stream_retired: HashSet = HashSet::new(); let mut stream_retired_order: VecDeque = VecDeque::new(); let mut pending_user_input: VecDeque = VecDeque::new(); @@ -6979,6 +6981,7 @@ async fn run_terminal( }; last_packet_at = Instant::now(); pending_socks_replies.remove(&close.stream_id); + stream_close_retransmit.remove(&close.stream_id); retire_stream_state( close.stream_id, &mut opened_streams, @@ -7105,6 +7108,13 @@ async fn run_terminal( #[cfg(unix)] agent_writers.remove(&stream_id); send_stream_close(&socket, addr, &cred, &mut send_seq, stream_id).await?; + stream_close_retransmit.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now(), + attempts: 1, + }, + ); } None if forward_only => break, None => {} @@ -7274,6 +7284,14 @@ async fn run_terminal( &mut stream_eof_retransmit, stream_retransmit_srtt, ).await?; + retransmit_stream_closes( + &socket, + addr, + &cred, + &mut send_seq, + &mut stream_close_retransmit, + stream_retransmit_srtt, + ).await?; } } } @@ -8326,6 +8344,38 @@ async fn retransmit_stream_eofs( Ok(()) } +async fn retransmit_stream_closes( + socket: &UdpSocket, + addr: SocketAddr, + cred: &CachedCredential, + send_seq: &mut u64, + stream_close_retransmit: &mut HashMap, + stream_retransmit_srtt: Option, +) -> Result<()> { + let now = Instant::now(); + let retransmit_after = dosh::transport::adaptive_retransmit_after( + dosh::transport::DEFAULT_RETRANSMIT_AFTER, + stream_retransmit_srtt, + ); + let stream_ids: Vec = stream_close_retransmit.keys().copied().collect(); + for stream_id in stream_ids { + let Some(pending) = stream_close_retransmit.get_mut(&stream_id) else { + continue; + }; + if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { + stream_close_retransmit.remove(&stream_id); + continue; + } + if now.duration_since(pending.last_sent) < retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + send_stream_close(socket, addr, cred, send_seq, stream_id).await?; + } + Ok(()) +} + fn accept_stream_data( stream_next_recv_offset: &mut HashMap, stream_recv_pending: &mut HashMap>>, @@ -9862,32 +9912,34 @@ mod tests { LOCAL_SLEEP_REPAINT_RETRY_WINDOW, LocalForward, MAX_PENDING_USER_INPUT_BYTES, NativeIdentityContext, POST_RECONNECT_STALE_INPUT_GRACE, POST_SUBMIT_ALL_INPUT_HOLD, PendingStreamControl, PendingStreamOpen, PredictMode, Predictor, RESTART_STATUS_SCRIPT, - RemoteForward, STALE_TERMINAL_INPUT_AFTER, STARTUP_INPUT_HOLD, STREAM_INITIAL_WINDOW, - SshConfig, SshPathTokenContext, StartupGateMode, StatusAction, TERMINAL_SNAPSHOT_RESET, - UpdateOptions, UpdateRole, auth_allows, cache_key, cache_server_prefix, - cleanup_stream_state, clear_cached_credentials, effective_update_artifact_tag, - ensure_tui_safe_status_overlay, expand_ssh_path_tokens, 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, - native_proxy_udp_warning, newest_client_trace_path_from, parse_dynamic_forward, - parse_escape_key, parse_local_forward, parse_remote_forward, parse_single_remote_path, - parse_ssh_config, parse_trace_line, parse_trace_options, parse_trace_report_options, - parse_trace_summary, parse_update_options, post_submit_hold_duration, - queue_or_send_stream_data, queue_pending_user_input, queue_stale_pending_user_input, - raw_contains_host_table, recv_response_until, refresh_live_addr, release_tag_download_url, + RemoteForward, STALE_TERMINAL_INPUT_AFTER, STARTUP_INPUT_HOLD, + STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS, STREAM_INITIAL_WINDOW, SshConfig, + SshPathTokenContext, StartupGateMode, StatusAction, TERMINAL_SNAPSHOT_RESET, UpdateOptions, + UpdateRole, auth_allows, cache_key, cache_server_prefix, cleanup_stream_state, + clear_cached_credentials, effective_update_artifact_tag, ensure_tui_safe_status_overlay, + expand_ssh_path_tokens, 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, native_proxy_udp_warning, + newest_client_trace_path_from, parse_dynamic_forward, parse_escape_key, + parse_local_forward, parse_remote_forward, parse_single_remote_path, parse_ssh_config, + parse_trace_line, parse_trace_options, parse_trace_report_options, parse_trace_summary, + parse_update_options, post_submit_hold_duration, queue_or_send_stream_data, + queue_pending_user_input, queue_stale_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_frame_bytes, render_status_clear, render_status_overlay, requested_env, resolved_startup_command, - retire_stream_state, retransmit_stream_eofs, retransmit_stream_opens, - rewrite_forward_command, sanitize_trace_name, selected_predict_mode, selected_udp_host, - send_stream_eof, server_version_mismatch, should_flush_terminal_input_after_contact, - should_health_log_client_start, should_hold_during_startup_gate, - should_hold_post_submit_input, should_reconnect_before_input_for_local_sleep, - should_repaint_idle_terminal, should_strip_unowned_terminal_reports, - split_after_command_submit, split_trace_tokens, ssh_command_target, ssh_config_uses_proxy, - ssh_destination_host, ssh_username, ssh_with_user, startup_command, status_ssh_target, - strip_stale_mouse_reports, strip_terminal_focus_reports, strip_unowned_terminal_reports, - summarize_trace_file, summarize_trace_file_with_mode, terminal_private_mode_transition, - toml_bare_key_or_quoted, top_trace_events, trace_report_warnings, unix_update_script, + retire_stream_state, retransmit_stream_closes, retransmit_stream_eofs, + retransmit_stream_opens, rewrite_forward_command, sanitize_trace_name, + selected_predict_mode, selected_udp_host, send_stream_eof, server_version_mismatch, + should_flush_terminal_input_after_contact, should_health_log_client_start, + should_hold_during_startup_gate, should_hold_post_submit_input, + should_reconnect_before_input_for_local_sleep, should_repaint_idle_terminal, + should_strip_unowned_terminal_reports, split_after_command_submit, split_trace_tokens, + ssh_command_target, ssh_config_uses_proxy, ssh_destination_host, ssh_username, + ssh_with_user, startup_command, status_ssh_target, strip_stale_mouse_reports, + strip_terminal_focus_reports, strip_unowned_terminal_reports, summarize_trace_file, + summarize_trace_file_with_mode, terminal_private_mode_transition, toml_bare_key_or_quoted, + top_trace_events, trace_report_warnings, unix_update_script, update_binary_version_for_installer, update_installer_url, update_version_status, upsert_managed_block, valid_forward_host, vscode_safe_alias, wake_repaint_retry_deadline, windows_update_script, @@ -11346,6 +11398,102 @@ mod tests { assert!(!eof_retransmit.contains_key(&99)); } + #[tokio::test] + async fn client_retransmits_pending_stream_close() { + let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = receiver.local_addr().unwrap(); + let cred = CachedCredential { + server: "host".to_string(), + session: "term".to_string(), + mode: "forward-only".to_string(), + udp_host: "127.0.0.1".to_string(), + udp_port: addr.port(), + client_id: [1; 16], + session_key: [2; 32], + session_key_id: protocol::session_key_id(&[2; 32]), + attach_ticket: Vec::new(), + attach_ticket_psk: [3; 32], + last_rendered_seq: 0, + }; + let mut send_seq = 10; + let mut close_retransmit = HashMap::from([( + 44, + PendingStreamControl { + last_sent: Instant::now() - Duration::from_millis(250), + attempts: 1, + }, + )]); + + retransmit_stream_closes( + &sender, + addr, + &cred, + &mut send_seq, + &mut close_retransmit, + None, + ) + .await + .unwrap(); + + let mut buf = [0u8; 2048]; + let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf)) + .await + .unwrap() + .unwrap(); + let packet = protocol::decode(&buf[..n]).unwrap(); + assert_eq!(packet.header.kind, PacketKind::StreamClose); + assert_eq!(packet.header.seq, 10); + let plain = + protocol::decrypt_body(&packet, &cred.session_key, protocol::CLIENT_TO_SERVER).unwrap(); + let close: protocol::StreamClose = protocol::from_body(&plain).unwrap(); + assert_eq!(close.stream_id, 44); + assert_eq!(send_seq, 11); + assert_eq!(close_retransmit[&44].attempts, 2); + } + + #[tokio::test] + async fn client_expires_stream_close_retransmit_after_attempt_cap() { + let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = receiver.local_addr().unwrap(); + let cred = CachedCredential { + server: "host".to_string(), + session: "term".to_string(), + mode: "forward-only".to_string(), + udp_host: "127.0.0.1".to_string(), + udp_port: addr.port(), + client_id: [1; 16], + session_key: [2; 32], + session_key_id: protocol::session_key_id(&[2; 32]), + attach_ticket: Vec::new(), + attach_ticket_psk: [3; 32], + last_rendered_seq: 0, + }; + let mut send_seq = 10; + let mut close_retransmit = HashMap::from([( + 44, + PendingStreamControl { + last_sent: Instant::now() - Duration::from_millis(250), + attempts: STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS, + }, + )]); + + retransmit_stream_closes( + &sender, + addr, + &cred, + &mut send_seq, + &mut close_retransmit, + None, + ) + .await + .unwrap(); + + assert_eq!(send_seq, 10); + assert!(close_retransmit.is_empty()); + } + #[tokio::test] async fn client_stream_retransmit_uses_observed_rtt() { let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index d46b906..ad1929c 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -48,6 +48,7 @@ use tokio::sync::mpsc; const STREAM_INITIAL_WINDOW: usize = 1024 * 1024; const STREAM_RETIRED_TOMBSTONES: usize = 16 * 1024; +const STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS: u32 = 8; /// Persist a persistent session's screen to disk after this many bytes of output /// have accumulated since the last mirror. Keeps the atomic write off the @@ -436,6 +437,7 @@ struct ClientState { stream_local_eof_sent: HashSet, stream_remote_eof_received: HashSet, stream_eof_retransmit: HashMap, + stream_close_retransmit: HashMap, stream_retired: HashSet, stream_retired_order: VecDeque, stream_retransmit_srtt: Option, @@ -1281,6 +1283,7 @@ async fn handle_native_user_auth( stream_local_eof_sent: HashSet::new(), stream_remote_eof_received: HashSet::new(), stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), stream_retransmit_srtt: None, @@ -1434,6 +1437,7 @@ async fn handle_bootstrap_attach( stream_local_eof_sent: HashSet::new(), stream_remote_eof_received: HashSet::new(), stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), stream_retransmit_srtt: None, @@ -1587,6 +1591,7 @@ async fn handle_ticket_attach( stream_local_eof_sent: HashSet::new(), stream_remote_eof_received: HashSet::new(), stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), stream_retransmit_srtt: None, @@ -2479,6 +2484,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_close_retransmit.remove(&close.stream_id); retire_client_stream(client, close.stream_id); Ok(()) } @@ -3763,6 +3769,13 @@ 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) { retire_client_stream(client, stream_id); + client.stream_close_retransmit.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now(), + attempts: 1, + }, + ); } } let body = protocol::to_body(&StreamClose { stream_id })?; @@ -4092,6 +4105,34 @@ async fn retransmit_pending( )?; sends.push((client.endpoint, packet)); } + let stream_close_ids: Vec = + client.stream_close_retransmit.keys().copied().collect(); + for stream_id in stream_close_ids { + let Some(pending) = client.stream_close_retransmit.get_mut(&stream_id) else { + continue; + }; + if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { + client.stream_close_retransmit.remove(&stream_id); + continue; + } + if now.duration_since(pending.last_sent) < stream_retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + client.send_seq += 1; + let body = protocol::to_body(&StreamClose { stream_id })?; + let packet = protocol::encode_encrypted( + PacketKind::StreamClose, + *client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + sends.push((client.endpoint, packet)); + } let stream_ids: Vec = client.stream_sent_data.keys().copied().collect(); let mut retransmit_chunks = Vec::new(); for stream_id in stream_ids { @@ -4689,6 +4730,7 @@ mod tests { stream_local_eof_sent: HashSet::new(), stream_remote_eof_received: HashSet::new(), stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), stream_retransmit_srtt: None, @@ -5356,6 +5398,7 @@ mod tests { stream_local_eof_sent: HashSet::new(), stream_remote_eof_received: HashSet::new(), stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), stream_retransmit_srtt: None, @@ -5458,6 +5501,7 @@ mod tests { stream_local_eof_sent: HashSet::new(), stream_remote_eof_received: HashSet::new(), stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), stream_retransmit_srtt: None, @@ -5564,6 +5608,107 @@ mod tests { assert!(!client.stream_eof_retransmit.contains_key(&999)); } + #[tokio::test] + async fn pending_server_stream_close_is_retransmitted() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [19u8; 16]; + let session_key = [20u8; 32]; + let stream_id = 126; + let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); + let mut client = test_client_state(session_key); + client.endpoint = receiver.local_addr().unwrap(); + client.stream_close_retransmit.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now() - Duration::from_millis(250), + attempts: 1, + }, + ); + state.sessions.insert( + "test".to_string(), + Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, + clients: HashMap::from([(client_id, client)]), + output_seq: 0, + recent: VecDeque::new(), + empty_since: None, + holder_control: None, + persistent: false, + bytes_since_persist: 0, + last_persisted_seq: 0, + last_screen_persist_at: Instant::now(), + }, + ); + state.client_index.insert(client_id, "test".to_string()); + let state = Arc::new(Mutex::new(state)); + + retransmit_pending(&state, &sender).await.unwrap(); + + let mut buf = [0u8; 2048]; + let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf)) + .await + .unwrap() + .unwrap(); + let packet = protocol::decode(&buf[..n]).unwrap(); + assert_eq!(packet.header.kind, PacketKind::StreamClose); + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); + let close: StreamClose = protocol::from_body(&plain).unwrap(); + assert_eq!(close.stream_id, stream_id); + + let locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert_eq!(client.stream_close_retransmit[&stream_id].attempts, 2); + } + + #[tokio::test] + async fn pending_server_stream_close_expires_after_attempt_cap() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [21u8; 16]; + let session_key = [22u8; 32]; + let stream_id = 127; + let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); + let mut client = test_client_state(session_key); + client.endpoint = receiver.local_addr().unwrap(); + client.stream_close_retransmit.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now() - Duration::from_millis(250), + attempts: STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS, + }, + ); + state.sessions.insert( + "test".to_string(), + Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, + clients: HashMap::from([(client_id, client)]), + output_seq: 0, + recent: VecDeque::new(), + empty_since: None, + holder_control: None, + persistent: false, + bytes_since_persist: 0, + last_persisted_seq: 0, + last_screen_persist_at: Instant::now(), + }, + ); + state.client_index.insert(client_id, "test".to_string()); + let state = Arc::new(Mutex::new(state)); + + retransmit_pending(&state, &sender).await.unwrap(); + + let locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert!(!client.stream_close_retransmit.contains_key(&stream_id)); + } + #[tokio::test] async fn server_stream_send_splits_large_writes() { let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); @@ -5605,6 +5750,7 @@ mod tests { stream_local_eof_sent: HashSet::new(), stream_remote_eof_received: HashSet::new(), stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), stream_retransmit_srtt: None, @@ -5710,6 +5856,7 @@ mod tests { stream_local_eof_sent: HashSet::new(), stream_remote_eof_received: HashSet::new(), stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), stream_retired: HashSet::new(), stream_retired_order: VecDeque::new(), stream_retransmit_srtt: None,