Buffer input across reconnect

src/bin/dosh-client.rs

@@ -46,6 +46,7 @@ use tokio::net::{TcpListener, TcpStream, UdpSocket};
 use tokio::sync::mpsc;
 
 const STREAM_INITIAL_WINDOW: usize = 1024 * 1024;
+const MAX_PENDING_USER_INPUT_BYTES: usize = 1024 * 1024;
 
 /// Sentinel `target_host` the server uses on a server-initiated `StreamOpen` that
 /// represents an SSH-agent connection (rather than a TCP target). The client
@@ -2650,6 +2651,8 @@ async fn run_terminal(
     let mut stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>> = HashMap::new();
     let mut stream_next_recv_offset: HashMap<u64, u64> = HashMap::new();
     let mut stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>> = HashMap::new();
+    let mut pending_user_input: VecDeque<Vec<u8>> = VecDeque::new();
+    let mut pending_user_input_bytes = 0usize;
     if let Some(frame) = first_frame {
         if !forward_only {
             render_frame(&frame)?;
@@ -2704,10 +2707,44 @@ async fn run_terminal(
                             break;
                         }
                         if cred.mode != "view-only" && cred.mode != "forward-only" {
+                            if last_packet_at.elapsed() >= Duration::from_secs(2) {
+                                queue_pending_user_input(
+                                    &mut pending_user_input,
+                                    &mut pending_user_input_bytes,
+                                    bytes,
+                                )?;
+                                if let Some(frame) = reconnect(
+                                    &socket,
+                                    &mut cred,
+                                    &mut send_seq,
+                                    last_size,
+                                    &mut frame_buffer,
+                                    &mut predictor,
+                                )
+                                .await?
+                                {
+                                    if !forward_only {
+                                        render_frame(&frame)?;
+                                        predictor.observe_output(&frame.bytes);
+                                    }
+                                    last_packet_at = Instant::now();
+                                    flush_pending_user_input(
+                                        &socket,
+                                        addr,
+                                        &cred,
+                                        &mut send_seq,
+                                        &mut predictor,
+                                        &mut pending_user_input,
+                                        &mut pending_user_input_bytes,
+                                    )
+                                    .await?;
+                                }
+                            } else {
                                 predictor.observe_input(&bytes)?;
                                 send_input(&socket, addr, &cred, &mut send_seq, bytes).await?;
                             }
                         }
+                    }
                     None => break,
                 }
             }
@@ -2762,6 +2799,16 @@ async fn run_terminal(
                                     predictor.observe_output(&frame.bytes);
                                 }
                                 last_packet_at = Instant::now();
+                                flush_pending_user_input(
+                                    &socket,
+                                    addr,
+                                    &cred,
+                                    &mut send_seq,
+                                    &mut predictor,
+                                    &mut pending_user_input,
+                                    &mut pending_user_input_bytes,
+                                )
+                                .await?;
                             }
                             continue;
                         };
@@ -2846,6 +2893,16 @@ async fn run_terminal(
                                     predictor.observe_output(&frame.bytes);
                                 }
                                 last_packet_at = Instant::now();
+                                flush_pending_user_input(
+                                    &socket,
+                                    addr,
+                                    &cred,
+                                    &mut send_seq,
+                                    &mut predictor,
+                                    &mut pending_user_input,
+                                    &mut pending_user_input_bytes,
+                                )
+                                .await?;
                             }
                         } else {
                             return Err(anyhow!("server rejected session: {}", reject.reason));
@@ -3206,6 +3263,16 @@ async fn run_terminal(
                             predictor.observe_output(&frame.bytes);
                         }
                         last_packet_at = Instant::now();
+                        flush_pending_user_input(
+                            &socket,
+                            addr,
+                            &cred,
+                            &mut send_seq,
+                            &mut predictor,
+                            &mut pending_user_input,
+                            &mut pending_user_input_bytes,
+                        )
+                        .await?;
                     }
                 } else if stale >= Duration::from_secs(2) {
                     let packet = protocol::encode_encrypted(
@@ -3469,6 +3536,40 @@ async fn reconnect(
     Ok(Some(frame))
 }
 
+fn queue_pending_user_input(
+    pending: &mut VecDeque<Vec<u8>>,
+    pending_bytes: &mut usize,
+    bytes: Vec<u8>,
+) -> Result<()> {
+    let next = pending_bytes
+        .checked_add(bytes.len())
+        .ok_or_else(|| anyhow!("pending input buffer overflow"))?;
+    anyhow::ensure!(
+        next <= MAX_PENDING_USER_INPUT_BYTES,
+        "pending input buffer full while disconnected"
+    );
+    *pending_bytes = next;
+    pending.push_back(bytes);
+    Ok(())
+}
+
+async fn flush_pending_user_input(
+    socket: &UdpSocket,
+    addr: SocketAddr,
+    cred: &CachedCredential,
+    send_seq: &mut u64,
+    predictor: &mut Predictor,
+    pending: &mut VecDeque<Vec<u8>>,
+    pending_bytes: &mut usize,
+) -> Result<()> {
+    while let Some(bytes) = pending.pop_front() {
+        *pending_bytes = pending_bytes.saturating_sub(bytes.len());
+        predictor.observe_input(&bytes)?;
+        send_input(socket, addr, cred, send_seq, bytes).await?;
+    }
+    Ok(())
+}
+
 async fn send_input(
     socket: &UdpSocket,
     addr: SocketAddr,
@@ -4766,19 +4867,21 @@ const TERMINAL_CLEANUP: &[u8] = concat!(
 #[cfg(test)]
 mod tests {
     use super::{
-        DisconnectStatus, DynamicForward, FrameBuffer, LocalForward, PredictMode, Predictor,
+        DisconnectStatus, DynamicForward, FrameBuffer, LocalForward, MAX_PENDING_USER_INPUT_BYTES,
-        RemoteForward, SshConfig, StatusAction, auth_allows, cache_key, cache_server_prefix,
+        PredictMode, Predictor, RemoteForward, SshConfig, StatusAction, auth_allows, cache_key,
-        clear_cached_credentials, ensure_tui_safe_status_overlay, latest_release_download_url,
+        cache_server_prefix, clear_cached_credentials, ensure_tui_safe_status_overlay,
-        load_first_native_identity_with_prompt, parse_dynamic_forward, parse_local_forward,
+        latest_release_download_url, load_first_native_identity_with_prompt, parse_dynamic_forward,
-        parse_remote_forward, parse_ssh_config, raw_contains_host_table, recv_response_until,
+        parse_local_forward, parse_remote_forward, parse_ssh_config, queue_pending_user_input,
-        render_status_clear, render_status_overlay, requested_env, resolved_startup_command,
+        raw_contains_host_table, recv_response_until, render_status_clear, render_status_overlay,
-        rewrite_forward_command, ssh_destination_host, ssh_username, ssh_with_user,
+        requested_env, resolved_startup_command, rewrite_forward_command, ssh_destination_host,
-        startup_command, toml_bare_key_or_quoted, update_check_requested, valid_forward_host,
+        ssh_username, ssh_with_user, startup_command, 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::fs;
     use std::time::Duration;
 
@@ -5579,6 +5682,23 @@ mod tests {
         assert!(similar.exists());
     }
 
+    #[test]
+    fn pending_user_input_is_bounded_and_ordered() {
+        let mut pending = VecDeque::new();
+        let mut pending_bytes = 0usize;
+        queue_pending_user_input(&mut pending, &mut pending_bytes, b"abc".to_vec()).unwrap();
+        queue_pending_user_input(&mut pending, &mut pending_bytes, b"def".to_vec()).unwrap();
+        assert_eq!(pending_bytes, 6);
+        assert_eq!(pending.pop_front().unwrap(), b"abc");
+        assert_eq!(pending.pop_front().unwrap(), b"def");
+
+        let mut pending = VecDeque::new();
+        let mut pending_bytes = MAX_PENDING_USER_INPUT_BYTES;
+        assert!(queue_pending_user_input(&mut pending, &mut pending_bytes, b"x".to_vec()).is_err());
+        assert!(pending.is_empty());
+        assert_eq!(pending_bytes, MAX_PENDING_USER_INPUT_BYTES);
+    }
+
     // --- Item 1: disconnect status line state machine ---
 
     /// The status line stays hidden while the link is fresh and only appears once