diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index ac2a60f..66f4e34 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -48,6 +48,7 @@ use tokio::sync::mpsc; const STREAM_INITIAL_WINDOW: usize = 1024 * 1024; 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); /// Sentinel `target_host` the server uses on a server-initiated `StreamOpen` that /// represents an SSH-agent connection (rather than a TCP target). The client @@ -202,6 +203,12 @@ struct PendingStreamChunk { attempts: u32, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StartupGateMode { + HoldAll, + HoldControl, +} + #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { let mut args = Args::parse(); @@ -2844,7 +2851,7 @@ async fn run_terminal( let mut pending_user_input: VecDeque> = VecDeque::new(); let mut pending_user_input_bytes = 0usize; let mut startup_input_hold_until: Option = None; - let mut hold_all_startup_input = false; + let mut startup_gate_mode = StartupGateMode::HoldControl; if let Some(frame) = first_frame { if !forward_only { render_frame(&frame)?; @@ -2862,7 +2869,7 @@ async fn run_terminal( { send_input(&socket, addr, &cred, &mut send_seq, bytes).await?; startup_input_hold_until = Some(Instant::now() + STARTUP_INPUT_HOLD); - hold_all_startup_input = true; + startup_gate_mode = StartupGateMode::HoldAll; } let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::>(); @@ -2905,7 +2912,7 @@ async fn run_terminal( if !predictor.alternate_screen && Instant::now() < deadline && should_hold_during_startup_gate( - hold_all_startup_input, + startup_gate_mode, !pending_user_input.is_empty(), &bytes, ) @@ -2918,7 +2925,7 @@ async fn run_terminal( continue; } startup_input_hold_until = None; - hold_all_startup_input = false; + startup_gate_mode = StartupGateMode::HoldControl; flush_pending_user_input( &socket, addr, @@ -2936,9 +2943,16 @@ async fn run_terminal( { predictor.observe_input(&send_now)?; send_input(&socket, addr, &cred, &mut send_seq, send_now).await?; - startup_input_hold_until = - Some(Instant::now() + STARTUP_INPUT_HOLD); - hold_all_startup_input = false; + startup_input_hold_until = Some( + Instant::now() + + post_submit_hold_duration(&hold_for_later), + ); + startup_gate_mode = if should_hold_post_submit_input(&hold_for_later) + { + StartupGateMode::HoldControl + } else { + StartupGateMode::HoldAll + }; if should_hold_post_submit_input(&hold_for_later) { queue_pending_user_input( &mut pending_user_input, @@ -3081,7 +3095,7 @@ async fn run_terminal( &mut send_seq, &mut predictor, (&mut pending_user_input, &mut pending_user_input_bytes), - (&mut startup_input_hold_until, &mut hold_all_startup_input), + (&mut startup_input_hold_until, &mut startup_gate_mode), ) .await?; } @@ -3561,7 +3575,7 @@ async fn run_terminal( &mut send_seq, &mut predictor, (&mut pending_user_input, &mut pending_user_input_bytes), - (&mut startup_input_hold_until, &mut hold_all_startup_input), + (&mut startup_input_hold_until, &mut startup_gate_mode), ) .await?; } @@ -3846,8 +3860,20 @@ fn should_hold_post_submit_input(bytes: &[u8]) -> bool { matches!(bytes.first(), Some(0x1b | 0x9b)) } -fn should_hold_during_startup_gate(hold_all: bool, already_queued: bool, bytes: &[u8]) -> bool { - hold_all || already_queued || should_hold_post_submit_input(bytes) +fn post_submit_hold_duration(followup: &[u8]) -> Duration { + if should_hold_post_submit_input(followup) { + STARTUP_INPUT_HOLD + } else { + POST_SUBMIT_ALL_INPUT_HOLD + } +} + +fn should_hold_during_startup_gate( + mode: StartupGateMode, + already_queued: bool, + bytes: &[u8], +) -> bool { + mode == StartupGateMode::HoldAll || already_queued || should_hold_post_submit_input(bytes) } async fn flush_pending_user_input( @@ -3874,9 +3900,9 @@ async fn flush_startup_input_if_ready( send_seq: &mut u64, predictor: &mut Predictor, pending: (&mut VecDeque>, &mut usize), - hold: (&mut Option, &mut bool), + hold: (&mut Option, &mut StartupGateMode), ) -> Result<()> { - let (startup_hold_until, hold_all_startup_input) = hold; + let (startup_hold_until, startup_gate_mode) = hold; let Some(deadline) = *startup_hold_until else { return Ok(()); }; @@ -3884,7 +3910,7 @@ async fn flush_startup_input_if_ready( return Ok(()); } *startup_hold_until = None; - *hold_all_startup_input = false; + *startup_gate_mode = StartupGateMode::HoldControl; let (pending, pending_bytes) = pending; flush_pending_user_input( socket, @@ -5312,11 +5338,12 @@ const TERMINAL_CLEANUP: &[u8] = concat!( mod tests { use super::{ DisconnectStatus, DynamicForward, FrameBuffer, LocalForward, MAX_PENDING_USER_INPUT_BYTES, - PredictMode, Predictor, RESTART_STATUS_SCRIPT, RemoteForward, SshConfig, 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_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, render_status_clear, render_status_overlay, requested_env, resolved_startup_command, rewrite_forward_command, selected_predict_mode, selected_udp_host, server_version_mismatch, @@ -5997,10 +6024,32 @@ mod tests { #[test] fn startup_gate_keeps_later_input_behind_queued_control_input() { - assert!(should_hold_during_startup_gate(false, false, b"\x1b[B")); - assert!(should_hold_during_startup_gate(false, true, b"x")); - assert!(should_hold_during_startup_gate(true, false, b"x")); - assert!(!should_hold_during_startup_gate(false, false, b"x")); + assert!(should_hold_during_startup_gate( + StartupGateMode::HoldControl, + false, + b"\x1b[B" + )); + assert!(should_hold_during_startup_gate( + StartupGateMode::HoldControl, + true, + b"x" + )); + assert!(should_hold_during_startup_gate( + StartupGateMode::HoldAll, + false, + b"x" + )); + assert!(!should_hold_during_startup_gate( + StartupGateMode::HoldControl, + false, + b"x" + )); + } + + #[test] + fn post_submit_hold_duration_has_short_printable_phase() { + assert_eq!(post_submit_hold_duration(b"\x1b[B"), STARTUP_INPUT_HOLD); + assert_eq!(post_submit_hold_duration(b"x"), POST_SUBMIT_ALL_INPUT_HOLD); } #[test]