diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index df7bfc2..777e9e1 100644 --- a/src/bin/dosh-client.rs +++ b/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> = HashMap::new(); let mut stream_next_recv_offset: HashMap = HashMap::new(); let mut stream_recv_pending: HashMap>> = HashMap::new(); + let mut pending_user_input: VecDeque> = VecDeque::new(); + let mut pending_user_input_bytes = 0usize; if let Some(frame) = first_frame { if !forward_only { render_frame(&frame)?; @@ -2704,8 +2707,42 @@ async fn run_terminal( break; } if cred.mode != "view-only" && cred.mode != "forward-only" { - predictor.observe_input(&bytes)?; - send_input(&socket, addr, &cred, &mut send_seq, bytes).await?; + 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>, + pending_bytes: &mut usize, + bytes: Vec, +) -> 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>, + 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, - RemoteForward, SshConfig, StatusAction, auth_allows, cache_key, cache_server_prefix, - clear_cached_credentials, ensure_tui_safe_status_overlay, latest_release_download_url, - load_first_native_identity_with_prompt, parse_dynamic_forward, parse_local_forward, - parse_remote_forward, parse_ssh_config, raw_contains_host_table, recv_response_until, - render_status_clear, render_status_overlay, requested_env, resolved_startup_command, - rewrite_forward_command, ssh_destination_host, ssh_username, ssh_with_user, - startup_command, toml_bare_key_or_quoted, update_check_requested, valid_forward_host, + DisconnectStatus, DynamicForward, FrameBuffer, LocalForward, MAX_PENDING_USER_INPUT_BYTES, + PredictMode, Predictor, RemoteForward, SshConfig, StatusAction, auth_allows, cache_key, + cache_server_prefix, clear_cached_credentials, ensure_tui_safe_status_overlay, + latest_release_download_url, load_first_native_identity_with_prompt, parse_dynamic_forward, + parse_local_forward, parse_remote_forward, parse_ssh_config, queue_pending_user_input, + raw_contains_host_table, recv_response_until, render_status_clear, render_status_overlay, + requested_env, resolved_startup_command, rewrite_forward_command, ssh_destination_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