Harden reconnect and stream open reliability

src/bin/dosh-client.rs

@@ -203,6 +203,14 @@ struct PendingStreamChunk {
     attempts: u32,
 }
 
+#[derive(Clone)]
+struct PendingStreamOpen {
+    target_host: String,
+    target_port: u16,
+    last_sent: Instant,
+    attempts: u32,
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum StartupGateMode {
     HoldAll,
@@ -2667,7 +2675,7 @@ async fn try_resume_fresh_process(
     )?;
     socket.send_to(&packet, addr).await?;
     let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
-        if packet.header.conn_id != cached.client_id {
+        if !is_resume_response_for_client(&packet, cached.client_id) {
             return Ok(None);
         }
         if packet.header.kind == PacketKind::AttachReject {
@@ -2722,7 +2730,7 @@ async fn try_live_resume(
     *send_seq += 1;
     socket.send_to(&packet, addr).await?;
     let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
-        if packet.header.conn_id != cached.client_id {
+        if !is_resume_response_for_client(&packet, cached.client_id) {
             return Ok(None);
         }
         if packet.header.kind == PacketKind::AttachReject {
@@ -2750,6 +2758,11 @@ async fn try_live_resume(
     Ok((frame, next))
 }
 
+fn is_resume_response_for_client(packet: &protocol::Packet, client_id: [u8; 16]) -> bool {
+    packet.header.conn_id == client_id
+        || (packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16])
+}
+
 async fn recv_response_until<T, F>(socket: &UdpSocket, wait: Duration, mut accept: F) -> Result<T>
 where
     F: FnMut(protocol::Packet) -> Result<Option<T>>,
@@ -2844,6 +2857,7 @@ async fn run_terminal(
     let mut opened_streams: HashSet<u64> = HashSet::new();
     let mut stream_send_credit: HashMap<u64, usize> = HashMap::new();
     let mut stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>> = HashMap::new();
+    let mut stream_pending_opens: HashMap<u64, PendingStreamOpen> = HashMap::new();
     let mut stream_next_send_offset: HashMap<u64, u64> = HashMap::new();
     let mut stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>> = HashMap::new();
     let mut stream_next_recv_offset: HashMap<u64, u64> = HashMap::new();
@@ -3195,6 +3209,17 @@ async fn run_terminal(
                             continue;
                         };
                         last_packet_at = Instant::now();
+                        if opened_streams.contains(&open.stream_id) {
+                            send_stream_open_ok(
+                                &socket,
+                                addr,
+                                &cred,
+                                &mut send_seq,
+                                open.stream_id,
+                            )
+                            .await?;
+                            continue;
+                        }
                         // SSH-agent stream: the server is asking us to splice this
                         // stream into our LOCAL ssh-agent. Only honor it when we
                         // actually opted in (agent_sock is Some); otherwise reject,
@@ -3343,6 +3368,7 @@ async fn run_terminal(
                             continue;
                         };
                         last_packet_at = Instant::now();
+                        stream_pending_opens.remove(&ok.stream_id);
                         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);
@@ -3372,6 +3398,7 @@ async fn run_terminal(
                             continue;
                         };
                         last_packet_at = Instant::now();
+                        stream_pending_opens.remove(&reject.stream_id);
                         if pending_socks_replies.remove(&reject.stream_id)
                             && let Some(writer) = stream_writers.get_mut(&reject.stream_id)
                         {
@@ -3456,6 +3483,7 @@ 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);
@@ -3480,6 +3508,12 @@ async fn run_terminal(
                         if socks5_reply {
                             pending_socks_replies.insert(stream_id);
                         }
+                        stream_pending_opens.insert(stream_id, PendingStreamOpen {
+                            target_host: target_host.clone(),
+                            target_port,
+                            last_sent: Instant::now(),
+                            attempts: 1,
+                        });
                         send_stream_open(
                             &socket,
                             addr,
@@ -3508,6 +3542,7 @@ 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);
@@ -3603,6 +3638,13 @@ async fn run_terminal(
                     &mut send_seq,
                     &mut stream_sent_data,
                 ).await?;
+                retransmit_stream_opens(
+                    &socket,
+                    addr,
+                    &cred,
+                    &mut send_seq,
+                    &mut stream_pending_opens,
+                ).await?;
             }
         }
     }
@@ -4135,6 +4177,38 @@ async fn retransmit_stream_data(
     Ok(())
 }
 
+async fn retransmit_stream_opens(
+    socket: &UdpSocket,
+    addr: SocketAddr,
+    cred: &CachedCredential,
+    send_seq: &mut u64,
+    stream_pending_opens: &mut HashMap<u64, PendingStreamOpen>,
+) -> Result<()> {
+    let now = Instant::now();
+    let mut retransmit = Vec::new();
+    for (stream_id, pending) in stream_pending_opens.iter_mut() {
+        if now.duration_since(pending.last_sent) < Duration::from_millis(200) {
+            continue;
+        }
+        pending.last_sent = now;
+        pending.attempts = pending.attempts.saturating_add(1);
+        retransmit.push((*stream_id, pending.target_host.clone(), pending.target_port));
+    }
+    for (stream_id, target_host, target_port) in retransmit {
+        send_stream_open(
+            socket,
+            addr,
+            cred,
+            send_seq,
+            stream_id,
+            target_host,
+            target_port,
+        )
+        .await?;
+    }
+    Ok(())
+}
+
 fn accept_stream_data(
     stream_next_recv_offset: &mut HashMap<u64, u64>,
     stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
@@ -5347,25 +5421,26 @@ const TERMINAL_CLEANUP: &[u8] = concat!(
 mod tests {
     use super::{
         CachedCredential, DisconnectStatus, DynamicForward, FrameBuffer, LocalForward,
-        MAX_PENDING_USER_INPUT_BYTES, POST_SUBMIT_ALL_INPUT_HOLD, 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_matches_escape, is_local_status_target,
-        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, render_status_clear, render_status_overlay,
-        requested_env, resolved_startup_command, rewrite_forward_command, selected_predict_mode,
-        selected_udp_host, server_version_mismatch, should_hold_during_startup_gate,
-        should_hold_post_submit_input, split_after_command_submit, ssh_destination_host,
-        ssh_username, ssh_with_user, startup_command, status_ssh_target, toml_bare_key_or_quoted,
-        update_check_requested, valid_forward_host,
+        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_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,
+        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_hold_during_startup_gate, should_hold_post_submit_input,
+        split_after_command_submit, ssh_destination_host, ssh_username, ssh_with_user,
+        startup_command, status_ssh_target, toml_bare_key_or_quoted, update_check_requested,
+        valid_forward_host,
     };
     use dosh::config::{ClientConfig, CommandExtension, HostConfig};
     use dosh::native::EnvVar;
     use dosh::protocol::{self, Frame, PacketKind};
     use ed25519_dalek::{SigningKey, VerifyingKey};
-    use std::collections::VecDeque;
+    use std::collections::{HashMap, VecDeque};
     use std::fs;
     use std::time::Duration;
 
@@ -6083,6 +6158,92 @@ mod tests {
         assert_eq!(addr.port(), 50001);
     }
 
+    #[test]
+    fn resume_accepts_unknown_client_reject_from_legacy_zero_id_server() {
+        let client_id = [7u8; 16];
+        let matching_reject = protocol::Packet {
+            header: protocol::Header {
+                kind: PacketKind::AttachReject,
+                flags: 0,
+                conn_id: client_id,
+                seq: 0,
+                ack: 0,
+                session_key_id: [0u8; 16],
+                body_len: 0,
+            },
+            body: Vec::new(),
+        };
+        let zero_id_reject = protocol::Packet {
+            header: protocol::Header {
+                conn_id: [0u8; 16],
+                ..matching_reject.header.clone()
+            },
+            body: Vec::new(),
+        };
+        let unrelated_frame = protocol::Packet {
+            header: protocol::Header {
+                kind: PacketKind::Frame,
+                conn_id: [9u8; 16],
+                ..matching_reject.header.clone()
+            },
+            body: Vec::new(),
+        };
+
+        assert!(is_resume_response_for_client(&matching_reject, client_id));
+        assert!(is_resume_response_for_client(&zero_id_reject, client_id));
+        assert!(!is_resume_response_for_client(&unrelated_frame, client_id));
+    }
+
+    #[tokio::test]
+    async fn pending_stream_open_is_retransmitted_until_answered() {
+        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 pending = HashMap::from([(
+            44,
+            PendingStreamOpen {
+                target_host: "127.0.0.1".to_string(),
+                target_port: 8080,
+                last_sent: std::time::Instant::now() - Duration::from_millis(250),
+                attempts: 1,
+            },
+        )]);
+
+        retransmit_stream_opens(&sender, addr, &cred, &mut send_seq, &mut pending)
+            .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::StreamOpen);
+        assert_eq!(packet.header.seq, 10);
+        let plain =
+            protocol::decrypt_body(&packet, &cred.session_key, protocol::CLIENT_TO_SERVER).unwrap();
+        let open: protocol::StreamOpen = protocol::from_body(&plain).unwrap();
+        assert_eq!(open.stream_id, 44);
+        assert_eq!(open.target_port, 8080);
+        assert_eq!(send_seq, 11);
+        assert_eq!(pending[&44].attempts, 2);
+    }
+
     #[test]
     fn update_check_accepts_only_check_flag() {
         assert!(!update_check_requested(&[]).unwrap());