Retransmit stream window adjustments

src/bin/dosh-client.rs

@@ -249,6 +249,14 @@ struct PendingStreamControl {
     attempts: u32,
 }
 
+#[derive(Clone)]
+struct PendingWindowAdjust {
+    received_offset: u64,
+    bytes: usize,
+    last_sent: Instant,
+    attempts: u32,
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum StartupGateMode {
     HoldAll,
@@ -5914,6 +5922,7 @@ async fn run_terminal(
     let mut stream_remote_eof_received: HashSet<u64> = HashSet::new();
     let mut stream_eof_retransmit: HashMap<u64, PendingStreamControl> = HashMap::new();
     let mut stream_close_retransmit: HashMap<u64, PendingStreamControl> = HashMap::new();
+    let mut stream_window_adjust_retransmit: HashMap<u64, PendingWindowAdjust> = HashMap::new();
     let mut stream_retired: HashSet<u64> = HashSet::new();
     let mut stream_retired_order: VecDeque<u64> = VecDeque::new();
     let mut pending_user_input: VecDeque<PendingUserInput> = VecDeque::new();
@@ -6853,6 +6862,7 @@ async fn run_terminal(
                             &mut stream_local_eof_sent,
                             &mut stream_remote_eof_received,
                             &mut stream_eof_retransmit,
+                            &mut stream_window_adjust_retransmit,
                             &mut stream_retired,
                             &mut stream_retired_order,
                         );
@@ -6895,6 +6905,17 @@ async fn run_terminal(
                             consumed,
                             received_offset,
                         ).await?;
+                        if consumed > 0 {
+                            stream_window_adjust_retransmit.insert(
+                                stream_id,
+                                PendingWindowAdjust {
+                                    received_offset,
+                                    bytes: consumed,
+                                    last_sent: Instant::now(),
+                                    attempts: 1,
+                                },
+                            );
+                        }
                     }
                     PacketKind::StreamWindowAdjust => {
                         let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
@@ -6964,6 +6985,7 @@ async fn run_terminal(
                                 &mut stream_local_eof_sent,
                                 &mut stream_remote_eof_received,
                                 &mut stream_eof_retransmit,
+                                &mut stream_window_adjust_retransmit,
                                 &mut stream_retired,
                                 &mut stream_retired_order,
                             );
@@ -6995,6 +7017,7 @@ async fn run_terminal(
                             &mut stream_local_eof_sent,
                             &mut stream_remote_eof_received,
                             &mut stream_eof_retransmit,
+                            &mut stream_window_adjust_retransmit,
                             &mut stream_retired,
                             &mut stream_retired_order,
                         );
@@ -7078,6 +7101,7 @@ async fn run_terminal(
                                 &mut stream_local_eof_sent,
                                 &mut stream_remote_eof_received,
                                 &mut stream_eof_retransmit,
+                                &mut stream_window_adjust_retransmit,
                                 &mut stream_retired,
                                 &mut stream_retired_order,
                             );
@@ -7101,6 +7125,7 @@ async fn run_terminal(
                             &mut stream_local_eof_sent,
                             &mut stream_remote_eof_received,
                             &mut stream_eof_retransmit,
+                            &mut stream_window_adjust_retransmit,
                             &mut stream_retired,
                             &mut stream_retired_order,
                         );
@@ -7292,6 +7317,14 @@ async fn run_terminal(
                     &mut stream_close_retransmit,
                     stream_retransmit_srtt,
                 ).await?;
+                retransmit_stream_window_adjusts(
+                    &socket,
+                    addr,
+                    &cred,
+                    &mut send_seq,
+                    &mut stream_window_adjust_retransmit,
+                    stream_retransmit_srtt,
+                ).await?;
             }
         }
     }
@@ -8376,6 +8409,47 @@ async fn retransmit_stream_closes(
     Ok(())
 }
 
+async fn retransmit_stream_window_adjusts(
+    socket: &UdpSocket,
+    addr: SocketAddr,
+    cred: &CachedCredential,
+    send_seq: &mut u64,
+    stream_window_adjust_retransmit: &mut HashMap<u64, PendingWindowAdjust>,
+    stream_retransmit_srtt: Option<Duration>,
+) -> 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<u64> = stream_window_adjust_retransmit.keys().copied().collect();
+    for stream_id in stream_ids {
+        let Some(pending) = stream_window_adjust_retransmit.get_mut(&stream_id) else {
+            continue;
+        };
+        if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS {
+            stream_window_adjust_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_window_adjust(
+            socket,
+            addr,
+            cred,
+            send_seq,
+            stream_id,
+            pending.bytes,
+            pending.received_offset,
+        )
+        .await?;
+    }
+    Ok(())
+}
+
 fn accept_stream_data(
     stream_next_recv_offset: &mut HashMap<u64, u64>,
     stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
@@ -8405,6 +8479,7 @@ fn cleanup_stream_state(
     stream_local_eof_sent: &mut HashSet<u64>,
     stream_remote_eof_received: &mut HashSet<u64>,
     stream_eof_retransmit: &mut HashMap<u64, PendingStreamControl>,
+    stream_window_adjust_retransmit: &mut HashMap<u64, PendingWindowAdjust>,
 ) {
     opened_streams.remove(&stream_id);
     stream_send_credit.remove(&stream_id);
@@ -8417,6 +8492,7 @@ fn cleanup_stream_state(
     stream_local_eof_sent.remove(&stream_id);
     stream_remote_eof_received.remove(&stream_id);
     stream_eof_retransmit.remove(&stream_id);
+    stream_window_adjust_retransmit.remove(&stream_id);
 }
 
 fn stream_eof_complete(
@@ -8441,6 +8517,7 @@ fn stream_state_exists(
     stream_local_eof_sent: &HashSet<u64>,
     stream_remote_eof_received: &HashSet<u64>,
     stream_eof_retransmit: &HashMap<u64, PendingStreamControl>,
+    stream_window_adjust_retransmit: &HashMap<u64, PendingWindowAdjust>,
 ) -> bool {
     opened_streams.contains(&stream_id)
         || stream_send_credit.contains_key(&stream_id)
@@ -8453,6 +8530,7 @@ fn stream_state_exists(
         || stream_local_eof_sent.contains(&stream_id)
         || stream_remote_eof_received.contains(&stream_id)
         || stream_eof_retransmit.contains_key(&stream_id)
+        || stream_window_adjust_retransmit.contains_key(&stream_id)
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -8469,6 +8547,7 @@ fn retire_stream_state(
     stream_local_eof_sent: &mut HashSet<u64>,
     stream_remote_eof_received: &mut HashSet<u64>,
     stream_eof_retransmit: &mut HashMap<u64, PendingStreamControl>,
+    stream_window_adjust_retransmit: &mut HashMap<u64, PendingWindowAdjust>,
     stream_retired: &mut HashSet<u64>,
     stream_retired_order: &mut VecDeque<u64>,
 ) {
@@ -8485,6 +8564,7 @@ fn retire_stream_state(
         stream_local_eof_sent,
         stream_remote_eof_received,
         stream_eof_retransmit,
+        stream_window_adjust_retransmit,
     );
     cleanup_stream_state(
         stream_id,
@@ -8499,6 +8579,7 @@ fn retire_stream_state(
         stream_local_eof_sent,
         stream_remote_eof_received,
         stream_eof_retransmit,
+        stream_window_adjust_retransmit,
     );
     if !had_state || STREAM_RETIRED_TOMBSTONES == 0 {
         return;
@@ -9911,8 +9992,8 @@ mod tests {
         FRAME_GAP_RESYNC_AFTER_MS, FrameBuffer, LOCAL_SLEEP_REPAINT_AFTER,
         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,
+        PendingStreamControl, PendingStreamOpen, PendingWindowAdjust, PredictMode, Predictor,
+        RESTART_STATUS_SCRIPT, 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,
@@ -9929,17 +10010,17 @@ mod tests {
         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_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,
+        retransmit_stream_opens, retransmit_stream_window_adjusts, 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,
@@ -10075,6 +10156,26 @@ mod tests {
                 },
             ),
         ]);
+        let mut stream_window_adjust_retransmit = HashMap::from([
+            (
+                stream_id,
+                PendingWindowAdjust {
+                    received_offset: 77,
+                    bytes: 5,
+                    last_sent: Instant::now(),
+                    attempts: 1,
+                },
+            ),
+            (
+                other_stream_id,
+                PendingWindowAdjust {
+                    received_offset: 88,
+                    bytes: 6,
+                    last_sent: Instant::now(),
+                    attempts: 1,
+                },
+            ),
+        ]);
 
         cleanup_stream_state(
             stream_id,
@@ -10089,6 +10190,7 @@ mod tests {
             &mut stream_local_eof_sent,
             &mut stream_remote_eof_received,
             &mut stream_eof_retransmit,
+            &mut stream_window_adjust_retransmit,
         );
 
         assert!(!opened_streams.contains(&stream_id));
@@ -10102,6 +10204,7 @@ mod tests {
         assert!(!stream_local_eof_sent.contains(&stream_id));
         assert!(!stream_remote_eof_received.contains(&stream_id));
         assert!(!stream_eof_retransmit.contains_key(&stream_id));
+        assert!(!stream_window_adjust_retransmit.contains_key(&stream_id));
 
         assert!(opened_streams.contains(&other_stream_id));
         assert!(stream_send_credit.contains_key(&other_stream_id));
@@ -10114,6 +10217,7 @@ mod tests {
         assert!(stream_local_eof_sent.contains(&other_stream_id));
         assert!(stream_remote_eof_received.contains(&other_stream_id));
         assert!(stream_eof_retransmit.contains_key(&other_stream_id));
+        assert!(stream_window_adjust_retransmit.contains_key(&other_stream_id));
     }
 
     #[test]
@@ -10130,6 +10234,7 @@ mod tests {
         let mut stream_local_eof_sent = HashSet::new();
         let mut stream_remote_eof_received = HashSet::new();
         let mut stream_eof_retransmit = HashMap::new();
+        let mut stream_window_adjust_retransmit = HashMap::new();
         let mut stream_retired = HashSet::new();
         let mut stream_retired_order = VecDeque::new();
 
@@ -10146,6 +10251,7 @@ mod tests {
             &mut stream_local_eof_sent,
             &mut stream_remote_eof_received,
             &mut stream_eof_retransmit,
+            &mut stream_window_adjust_retransmit,
             &mut stream_retired,
             &mut stream_retired_order,
         );
@@ -10162,6 +10268,7 @@ mod tests {
             &mut stream_local_eof_sent,
             &mut stream_remote_eof_received,
             &mut stream_eof_retransmit,
+            &mut stream_window_adjust_retransmit,
             &mut stream_retired,
             &mut stream_retired_order,
         );
@@ -11494,6 +11601,94 @@ mod tests {
         assert!(close_retransmit.is_empty());
     }
 
+    #[tokio::test]
+    async fn client_retransmits_pending_stream_window_adjust() {
+        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,
+            PendingWindowAdjust {
+                received_offset: 1234,
+                bytes: 99,
+                last_sent: Instant::now() - Duration::from_millis(250),
+                attempts: 1,
+            },
+        )]);
+
+        retransmit_stream_window_adjusts(&sender, addr, &cred, &mut send_seq, &mut pending, 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::StreamWindowAdjust);
+        assert_eq!(packet.header.seq, 10);
+        let plain =
+            protocol::decrypt_body(&packet, &cred.session_key, protocol::CLIENT_TO_SERVER).unwrap();
+        let adjust: protocol::StreamWindowAdjust = protocol::from_body(&plain).unwrap();
+        assert_eq!(adjust.stream_id, 44);
+        assert_eq!(adjust.received_offset, 1234);
+        assert_eq!(adjust.bytes, 99);
+        assert_eq!(send_seq, 11);
+        assert_eq!(pending[&44].attempts, 2);
+    }
+
+    #[tokio::test]
+    async fn client_expires_stream_window_adjust_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 pending = HashMap::from([(
+            44,
+            PendingWindowAdjust {
+                received_offset: 1234,
+                bytes: 99,
+                last_sent: Instant::now() - Duration::from_millis(250),
+                attempts: STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS,
+            },
+        )]);
+
+        retransmit_stream_window_adjusts(&sender, addr, &cred, &mut send_seq, &mut pending, None)
+            .await
+            .unwrap();
+
+        assert_eq!(send_seq, 10);
+        assert!(pending.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();