Add Mosh parity hardening

src/bin/dosh-client.rs

@@ -110,6 +110,14 @@ struct Args {
     #[arg(long)]
     predict: bool,
     #[arg(long)]
+    predict_mode: Option<String>,
+    #[arg(short = 'a', long = "predict-always")]
+    predict_always: bool,
+    #[arg(short = 'n', long = "predict-never")]
+    predict_never: bool,
+    #[arg(long)]
+    escape_key: Option<String>,
+    #[arg(long)]
     ssh_port: Option<u16>,
     #[arg(long)]
     ssh_auth_command: Option<String>,
@@ -253,6 +261,8 @@ async fn main() -> Result<()> {
         || !dynamic_forwards.is_empty()
         || forward_agent;
     let predict = args.predict || host.predict.unwrap_or(config.predict);
+    let predict_mode = selected_predict_mode(&args, &config)?;
+    let escape_key = selected_escape_key(&args, &config)?;
     let session = select_session(args.session.as_deref(), args.new);
     let mode = if args.forward_only {
         "forward-only"
@@ -302,22 +312,17 @@ async fn main() -> Result<()> {
 
     let started = Instant::now();
     let mut resolved_ssh_config = None;
-    let target_udp_host = if let Some(host) = args
+    let requested_udp_host = args
         .dosh_host
         .clone()
         .or_else(|| host.dosh_host.clone())
-        .or_else(|| config.dosh_host.clone())
-    {
-        host
-    } else if args.local_auth {
-        "127.0.0.1".to_string()
+        .or_else(|| config.dosh_host.clone());
+    let target_udp_host = if args.local_auth {
+        requested_udp_host.unwrap_or_else(|| "127.0.0.1".to_string())
     } else {
         match ssh_config(&server, ssh_port) {
             Ok(parsed) => {
-                let hostname = parsed
-                    .hostname
-                    .clone()
-                    .unwrap_or_else(|| ssh_destination_host(&server));
+                let hostname = selected_udp_host(requested_udp_host.as_deref(), &server, &parsed)?;
                 resolved_ssh_config = Some(parsed);
                 hostname
             }
@@ -328,7 +333,11 @@ async fn main() -> Result<()> {
                     &format!("dosh could not resolve SSH config hostname, using {server}: {err:#}"),
                 );
                 resolved_ssh_config = Some(SshConfig::default());
-                ssh_destination_host(&server)
+                selected_udp_host(
+                    requested_udp_host.as_deref(),
+                    &server,
+                    &SshConfig::default(),
+                )?
             }
         }
     };
@@ -362,6 +371,8 @@ async fn main() -> Result<()> {
                         cred,
                         Some(frame),
                         predict,
+                        predict_mode,
+                        escape_key,
                         startup_command,
                         config.reconnect_timeout_secs,
                         Vec::new(),
@@ -398,6 +409,8 @@ async fn main() -> Result<()> {
                         cred,
                         Some(frame),
                         predict,
+                        predict_mode,
+                        escape_key,
                         startup_command,
                         config.reconnect_timeout_secs,
                         Vec::new(),
@@ -464,6 +477,8 @@ async fn main() -> Result<()> {
                     cred,
                     Some(frame),
                     predict,
+                    predict_mode,
+                    escape_key,
                     startup_command,
                     config.reconnect_timeout_secs,
                     local_forwards,
@@ -540,6 +555,8 @@ async fn main() -> Result<()> {
         cred,
         Some(first),
         predict,
+        predict_mode,
+        escape_key,
         startup_command,
         config.reconnect_timeout_secs,
         Vec::new(),
@@ -779,17 +796,19 @@ async fn run_doctor_for_host(
         Err(err) => println!("[warn] ssh fallback: {err:#}"),
     }
 
-    let udp_host = args
+    let requested_udp_host = args
         .dosh_host
         .clone()
         .or_else(|| host_config.dosh_host.clone())
-        .or_else(|| config.dosh_host.clone())
-        .unwrap_or_else(|| {
-            parsed_ssh_config
-                .hostname
-                .clone()
-                .unwrap_or_else(|| ssh_destination_host(&server))
-        });
+        .or_else(|| config.dosh_host.clone());
+    let udp_host =
+        match selected_udp_host(requested_udp_host.as_deref(), &server, &parsed_ssh_config) {
+            Ok(host) => host,
+            Err(err) => {
+                println!("[fail] native endpoint: {err:#}");
+                ssh_destination_host(&server)
+            }
+        };
     println!("[info] native endpoint: {udp_host}:{dosh_port}");
     match resolve_addr(&udp_host, dosh_port) {
         Ok(addr) => println!("[ok] udp resolve: {addr}"),
@@ -1717,6 +1736,30 @@ fn ssh_destination_host(server: &str) -> String {
         .to_string()
 }
 
+fn selected_udp_host(
+    explicit: Option<&str>,
+    server: &str,
+    ssh_config: &SshConfig,
+) -> Result<String> {
+    let ssh_host = || {
+        ssh_config
+            .hostname
+            .clone()
+            .unwrap_or_else(|| ssh_destination_host(server))
+    };
+    let Some(raw) = explicit.map(str::trim).filter(|value| !value.is_empty()) else {
+        return Ok(ssh_host());
+    };
+    match raw.to_ascii_lowercase().as_str() {
+        "ssh" | "auto" => Ok(ssh_host()),
+        "localhost" => Ok("127.0.0.1".to_string()),
+        "any" => Err(anyhow!(
+            "dosh_host={raw:?} is a server bind policy, not a client destination; use ssh/auto or a host/IP"
+        )),
+        _ => Ok(raw.to_string()),
+    }
+}
+
 fn ssh_username(server: &str) -> Option<String> {
     let (user, _) = server.rsplit_once('@')?;
     if user.is_empty() {
@@ -2356,7 +2399,13 @@ async fn bootstrap_attach(
         };
         Ok(protocol::from_body::<AttachOk>(&plain).ok())
     })
-    .await?;
+    .await
+    .with_context(|| {
+        format!(
+            "no Dosh bootstrap response from {}:{}; check UDP reachability and server firewall",
+            bootstrap.udp_host, bootstrap.udp_port
+        )
+    })?;
     let cred = CachedCredential {
         server: server_name.to_string(),
         session: ok.session.clone(),
@@ -2443,7 +2492,13 @@ async fn try_ticket_attach(
         };
         Ok(protocol::from_body::<AttachOk>(&plain).ok())
     })
-    .await?;
+    .await
+    .with_context(|| {
+        format!(
+            "cached ticket attach to {}:{} timed out; check UDP reachability or run `dosh recover {}`",
+            cached.udp_host, cached.udp_port, cached.server
+        )
+    })?;
     let frame = Frame {
         session: ok.session.clone(),
         output_seq: ok.initial_seq,
@@ -2500,7 +2555,13 @@ async fn try_resume_fresh_process(
         };
         Ok(protocol::from_body::<Frame>(&plain).ok())
     })
-    .await?;
+    .await
+    .with_context(|| {
+        format!(
+            "cached resume to {}:{} timed out; check UDP reachability or run `dosh recover {}`",
+            cached.udp_host, cached.udp_port, cached.server
+        )
+    })?;
     let mut next = cached.clone();
     next.last_rendered_seq = frame.output_seq;
     Ok((frame, next))
@@ -2549,7 +2610,13 @@ async fn try_live_resume(
         };
         Ok(protocol::from_body::<Frame>(&plain).ok())
     })
-    .await?;
+    .await
+    .with_context(|| {
+        format!(
+            "live resume to {}:{} timed out; still waiting for the server to become reachable",
+            cached.udp_host, cached.udp_port
+        )
+    })?;
     let mut next = cached.clone();
     next.last_rendered_seq = frame.output_seq;
     Ok((frame, next))
@@ -2584,6 +2651,8 @@ async fn run_terminal(
     mut cred: CachedCredential,
     first_frame: Option<Frame>,
     predict: bool,
+    predict_mode: PredictMode,
+    escape_key: Option<Vec<u8>>,
     startup_command: Option<Vec<u8>>,
     reconnect_timeout_secs: u64,
     local_forwards: Vec<LocalForward>,
@@ -2618,7 +2687,6 @@ async fn run_terminal(
     // env var wins for ad-hoc tuning; otherwise the client config's
     // `predict_mode` provides the persistent default. Predictions only run in a
     // read-write interactive session.
-    let predict_mode = resolve_predict_mode();
     let mut predictor = Predictor::with_mode(
         predict && cred.mode != "view-only" && !forward_only,
         predict_mode,
@@ -2703,7 +2771,7 @@ async fn run_terminal(
             stdin_msg = stdin_rx.recv() => {
                 match stdin_msg {
                     Some(bytes) => {
-                        if bytes == [0x1d] {
+                        if input_matches_escape(&bytes, escape_key.as_deref()) {
                             break;
                         }
                         if cred.mode != "view-only" && cred.mode != "forward-only" {
@@ -4444,17 +4512,81 @@ fn emit_overlay(_bytes: &[u8]) -> Result<()> {
     Ok(())
 }
 
-/// Resolve the prediction display policy. `DOSH_PREDICT_MODE` (off / experimental
-/// / always) overrides for quick experiments; otherwise the client config's
-/// `predict_mode` is used, defaulting to the adaptive (experimental) policy.
-fn resolve_predict_mode() -> PredictMode {
+/// Resolve the prediction display policy. `DOSH_PREDICT_MODE` (off /
+/// experimental / always) overrides for quick experiments; otherwise the
+/// already-selected CLI/config policy is used.
+fn selected_predict_mode(args: &Args, config: &dosh::config::ClientConfig) -> Result<PredictMode> {
+    if args.predict_always && args.predict_never {
+        return Err(anyhow!("--predict-always and --predict-never conflict"));
+    }
+    if args.predict_always {
+        return Ok(PredictMode::Always);
+    }
+    if args.predict_never {
+        return Ok(PredictMode::Off);
+    }
+    if let Some(mode) = args.predict_mode.as_deref() {
+        return Ok(PredictMode::parse(mode));
+    }
     if let Ok(env) = std::env::var("DOSH_PREDICT_MODE") {
-        return PredictMode::parse(&env);
+        return Ok(PredictMode::parse(&env));
     }
-    match load_client_config(None) {
-        Ok(cfg) => PredictMode::parse(&cfg.predict_mode),
-        Err(_) => PredictMode::Experimental,
+    Ok(PredictMode::parse(&config.predict_mode))
+}
+
+fn selected_escape_key(
+    args: &Args,
+    config: &dosh::config::ClientConfig,
+) -> Result<Option<Vec<u8>>> {
+    if let Some(raw) = args.escape_key.as_deref() {
+        return parse_escape_key(raw);
     }
+    if let Ok(raw) = std::env::var("DOSH_ESCAPE_KEY") {
+        return parse_escape_key(&raw);
+    }
+    parse_escape_key(&config.escape_key)
+}
+
+fn parse_escape_key(raw: &str) -> Result<Option<Vec<u8>>> {
+    let value = raw.trim();
+    if value.eq_ignore_ascii_case("none")
+        || value.eq_ignore_ascii_case("off")
+        || value.eq_ignore_ascii_case("disabled")
+    {
+        return Ok(None);
+    }
+    if let Some(rest) = value.strip_prefix("^") {
+        anyhow::ensure!(
+            rest.chars().count() == 1,
+            "invalid control escape key {raw:?}"
+        );
+        let ch = rest.as_bytes()[0];
+        let byte = match ch {
+            b'?' => 0x7f,
+            b'@'..=b'_' => ch - b'@',
+            b'a'..=b'z' => ch - b'a' + 1,
+            _ => return Err(anyhow!("invalid control escape key {raw:?}")),
+        };
+        return Ok(Some(vec![byte]));
+    }
+    if let Some(hex) = value.strip_prefix("0x") {
+        let byte = u8::from_str_radix(hex, 16)
+            .with_context(|| format!("invalid hex escape key {raw:?}"))?;
+        return Ok(Some(vec![byte]));
+    }
+    if value.len() == 1 {
+        return Ok(Some(value.as_bytes().to_vec()));
+    }
+    Err(anyhow!(
+        "escape key must be none, ^X, 0xNN, or one ASCII character"
+    ))
+}
+
+fn input_matches_escape(bytes: &[u8], escape_key: Option<&[u8]>) -> bool {
+    let Some(escape_key) = escape_key else {
+        return false;
+    };
+    !escape_key.is_empty() && bytes.windows(escape_key.len()).any(|w| w == escape_key)
 }
 
 /// Resolve whether the non-destructive disconnect status line is enabled.
@@ -4870,10 +5002,11 @@ mod tests {
         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,
+        input_matches_escape, 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, 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, ssh_destination_host,
         ssh_username, ssh_with_user, startup_command, toml_bare_key_or_quoted,
         update_check_requested, valid_forward_host,
     };
@@ -4905,6 +5038,10 @@ mod tests {
             forward_only: false,
             background: false,
             predict: false,
+            predict_mode: None,
+            predict_always: false,
+            predict_never: false,
+            escape_key: None,
             ssh_port: None,
             ssh_auth_command: None,
             ssh_key: None,
@@ -5046,6 +5183,31 @@ mod tests {
         assert_eq!(ssh_destination_host("alice@example.com"), "example.com");
     }
 
+    #[test]
+    fn selected_udp_host_supports_ssh_auto_and_localhost() {
+        let ssh = SshConfig {
+            hostname: Some("server.lan".to_string()),
+            ..SshConfig::default()
+        };
+        assert_eq!(
+            selected_udp_host(None, "alias", &ssh).unwrap(),
+            "server.lan"
+        );
+        assert_eq!(
+            selected_udp_host(Some("ssh"), "alias", &ssh).unwrap(),
+            "server.lan"
+        );
+        assert_eq!(
+            selected_udp_host(Some("auto"), "alias", &ssh).unwrap(),
+            "server.lan"
+        );
+        assert_eq!(
+            selected_udp_host(Some("localhost"), "alias", &ssh).unwrap(),
+            "127.0.0.1"
+        );
+        assert!(selected_udp_host(Some("any"), "alias", &ssh).is_err());
+    }
+
     #[test]
     fn parses_user_from_ssh_destination() {
         assert_eq!(ssh_username("alice@example.com"), Some("alice".to_string()));
@@ -5269,6 +5431,8 @@ mod tests {
         assert!(!Predictor::predictable(b"\n"));
         assert!(!Predictor::predictable(b"\x1b[A"));
         assert!(!Predictor::predictable("snowman: \u{2603}".as_bytes()));
+        assert!(!Predictor::predictable("wide: 你".as_bytes()));
+        assert!(!Predictor::predictable("combine: e\u{301}".as_bytes()));
     }
 
     #[test]
@@ -5446,6 +5610,42 @@ mod tests {
         assert_eq!(PredictMode::parse("anything"), PredictMode::Experimental);
     }
 
+    #[test]
+    fn cli_prediction_flags_override_config() {
+        let mut args = test_args("host", &[]);
+        let mut config = ClientConfig {
+            predict_mode: "experimental".to_string(),
+            ..ClientConfig::default()
+        };
+
+        args.predict_always = true;
+        assert_eq!(
+            selected_predict_mode(&args, &config).unwrap(),
+            PredictMode::Always
+        );
+
+        args.predict_always = false;
+        args.predict_never = true;
+        assert_eq!(
+            selected_predict_mode(&args, &config).unwrap(),
+            PredictMode::Off
+        );
+
+        args.predict_never = false;
+        args.predict_mode = Some("always".to_string());
+        assert_eq!(
+            selected_predict_mode(&args, &config).unwrap(),
+            PredictMode::Always
+        );
+
+        args.predict_mode = None;
+        config.predict_mode = "off".to_string();
+        assert_eq!(
+            selected_predict_mode(&args, &config).unwrap(),
+            PredictMode::Off
+        );
+    }
+
     #[test]
     fn startup_command_joins_args_for_remote_shell() {
         assert_eq!(startup_command(&[]), None);
@@ -5699,6 +5899,24 @@ mod tests {
         assert_eq!(pending_bytes, MAX_PENDING_USER_INPUT_BYTES);
     }
 
+    #[test]
+    fn escape_key_parser_accepts_mosh_style_controls() {
+        assert_eq!(parse_escape_key("^]").unwrap(), Some(vec![0x1d]));
+        assert_eq!(parse_escape_key("^C").unwrap(), Some(vec![0x03]));
+        assert_eq!(parse_escape_key("^?").unwrap(), Some(vec![0x7f]));
+        assert_eq!(parse_escape_key("0x1d").unwrap(), Some(vec![0x1d]));
+        assert_eq!(parse_escape_key("~").unwrap(), Some(vec![b'~']));
+        assert_eq!(parse_escape_key("none").unwrap(), None);
+        assert!(parse_escape_key("too-long").is_err());
+    }
+
+    #[test]
+    fn escape_key_matches_inside_stdin_chunk() {
+        assert!(input_matches_escape(b"abc\x1d", Some(&[0x1d])));
+        assert!(!input_matches_escape(b"abc", Some(&[0x1d])));
+        assert!(!input_matches_escape(b"abc\x1d", None));
+    }
+
     // --- Item 1: disconnect status line state machine ---
 
     /// The status line stays hidden while the link is fresh and only appears once