Add public comparison benchmarks and SSH import

src/bin/dosh-bench.rs

@@ -1,7 +1,11 @@
 use anyhow::{Context, Result, anyhow};
 use clap::Parser;
+use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem};
+use std::ffi::OsStr;
+use std::io::Read;
 use std::path::PathBuf;
 use std::process::Command;
+use std::sync::mpsc;
 use std::time::{Duration, Instant};
 
 #[derive(Debug, Parser)]
@@ -36,7 +40,17 @@ struct Args {
     #[arg(long)]
     no_cache: bool,
     #[arg(long)]
+    include_mosh: bool,
+    #[arg(long, default_value = "mosh")]
+    mosh: PathBuf,
+    #[arg(long, default_value = "mosh-server")]
+    mosh_server_command: String,
+    #[arg(long)]
+    mosh_port: Option<String>,
+    #[arg(long)]
     assert_ssh_plus_ms: Option<f64>,
+    #[arg(long)]
+    assert_mosh_minus_ms: Option<f64>,
 }
 
 fn main() -> Result<()> {
@@ -44,6 +58,7 @@ fn main() -> Result<()> {
     let client = args.client.clone().unwrap_or_else(default_client_path);
     let mut ssh_times = Vec::new();
     let mut dosh_times = Vec::new();
+    let mut mosh_times = Vec::new();
     let generated_control_path = if args.controlmaster {
         Some(std::env::temp_dir().join(format!("dosh-bench-control-{}", std::process::id())))
     } else {
@@ -97,6 +112,10 @@ fn main() -> Result<()> {
             cmd.arg(&args.server);
         }
         dosh_times.push(time_command(&mut cmd)?);
+
+        if args.include_mosh {
+            mosh_times.push(time_mosh_in_pty(&args)?);
+        }
     }
 
     if !ssh_times.is_empty() {
@@ -111,6 +130,13 @@ fn main() -> Result<()> {
         avg_ms(&dosh_times),
         dosh_times
     );
+    if !mosh_times.is_empty() {
+        println!(
+            "mosh_start_true_ms avg={:.2} samples={:?}",
+            avg_ms(&mosh_times),
+            mosh_times
+        );
+    }
     if let Some(margin) = args.assert_ssh_plus_ms {
         if ssh_times.is_empty() {
             return Err(anyhow!(
@@ -126,6 +152,21 @@ fn main() -> Result<()> {
         }
         println!("gate ok: dosh avg {dosh_avg:.2}ms <= ssh avg {ssh_avg:.2}ms + {margin:.2}ms");
     }
+    if let Some(margin) = args.assert_mosh_minus_ms {
+        if mosh_times.is_empty() {
+            return Err(anyhow!(
+                "--assert-mosh-minus-ms requires --include-mosh benchmark"
+            ));
+        }
+        let dosh_avg = avg_ms(&dosh_times);
+        let mosh_avg = avg_ms(&mosh_times);
+        if dosh_avg + margin > mosh_avg {
+            return Err(anyhow!(
+                "dosh attach avg {dosh_avg:.2}ms was not at least {margin:.2}ms faster than mosh avg {mosh_avg:.2}ms"
+            ));
+        }
+        println!("gate ok: dosh avg {dosh_avg:.2}ms + {margin:.2}ms <= mosh avg {mosh_avg:.2}ms");
+    }
     Ok(())
 }
 
@@ -143,6 +184,94 @@ fn add_ssh_options(cmd: &mut Command, args: &Args, control_path: Option<&PathBuf
     }
 }
 
+fn mosh_ssh_command(args: &Args) -> String {
+    let mut parts = vec![
+        "ssh".to_string(),
+        "-p".to_string(),
+        args.ssh_port.to_string(),
+    ];
+    if let Some(key) = &args.ssh_key {
+        parts.push("-i".to_string());
+        parts.push(shell_word(key.as_os_str()));
+    }
+    if let Some(known_hosts) = &args.ssh_known_hosts {
+        parts.push("-o".to_string());
+        parts.push(shell_word(OsStr::new(&format!(
+            "UserKnownHostsFile={}",
+            known_hosts.display()
+        ))));
+    }
+    parts.join(" ")
+}
+
+fn time_mosh_in_pty(args: &Args) -> Result<Duration> {
+    let pty_system = NativePtySystem::default();
+    let pair = pty_system
+        .openpty(PtySize {
+            rows: 24,
+            cols: 80,
+            pixel_width: 0,
+            pixel_height: 0,
+        })
+        .context("open pty for mosh benchmark")?;
+    let mut command_line = vec![
+        shell_word(args.mosh.as_os_str()),
+        shell_word(OsStr::new(&format!("--ssh={}", mosh_ssh_command(args)))),
+        shell_word(OsStr::new(&format!(
+            "--server={}",
+            args.mosh_server_command
+        ))),
+        "--no-init".to_string(),
+    ];
+    if let Some(port) = &args.mosh_port {
+        command_line.push("-p".to_string());
+        command_line.push(shell_word(OsStr::new(port)));
+    }
+    command_line.push(shell_word(OsStr::new(&args.server)));
+    command_line.push("--".to_string());
+    command_line.push("true".to_string());
+
+    let mut cmd = CommandBuilder::new("/bin/sh");
+    cmd.arg("-c");
+    cmd.arg(command_line.join(" "));
+
+    let start = Instant::now();
+    let mut child = pair
+        .slave
+        .spawn_command(cmd)
+        .context("spawn mosh benchmark in pty")?;
+    drop(pair.slave);
+
+    let mut reader = pair.master.try_clone_reader().context("clone pty reader")?;
+    let (tx, rx) = mpsc::channel();
+    std::thread::spawn(move || {
+        let mut output = Vec::new();
+        let _ = reader.read_to_end(&mut output);
+        let _ = tx.send(output);
+    });
+
+    {
+        let _writer = pair.master.take_writer().context("take pty writer")?;
+    }
+
+    let status = child.wait().context("wait for mosh benchmark")?;
+    let elapsed = start.elapsed();
+    drop(pair.master);
+    let output = rx.recv().unwrap_or_default();
+    if !status.success() {
+        return Err(anyhow!(
+            "mosh command failed with status {status:?}\noutput:\n{}",
+            String::from_utf8_lossy(&output)
+        ));
+    }
+    Ok(elapsed)
+}
+
+fn shell_word(value: &OsStr) -> String {
+    let value = value.to_string_lossy();
+    format!("'{}'", value.replace('\'', "'\\''"))
+}
+
 struct ControlMaster {
     server: String,
     ssh_port: u16,

src/bin/dosh-client.rs

@@ -18,6 +18,7 @@ use std::collections::BTreeMap;
 use std::fs;
 use std::io::{Read, Write};
 use std::net::{SocketAddr, ToSocketAddrs};
+use std::path::{Path, PathBuf};
 use std::process::{Command, Stdio};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 use tokio::net::UdpSocket;
@@ -49,11 +50,11 @@ struct Args {
     #[arg(long)]
     ssh_auth_command: Option<String>,
     #[arg(long)]
-    ssh_key: Option<std::path::PathBuf>,
+    ssh_key: Option<PathBuf>,
     #[arg(long)]
-    ssh_known_hosts: Option<std::path::PathBuf>,
+    ssh_known_hosts: Option<PathBuf>,
     #[arg(long)]
-    ssh_control_path: Option<std::path::PathBuf>,
+    ssh_control_path: Option<PathBuf>,
     #[arg(long)]
     dosh_port: Option<u16>,
     #[arg(long)]
@@ -84,6 +85,9 @@ async fn main() -> Result<()> {
     if args.server.as_deref() == Some("update") {
         return run_update(&config);
     }
+    if args.server.as_deref() == Some("import-ssh") {
+        return run_import_ssh(&config, &args.command);
+    }
     let requested_server = args.server.unwrap_or_else(|| config.server.clone());
     let hosts = load_hosts_config(None).unwrap_or_default();
     let host = hosts
@@ -105,7 +109,7 @@ async fn main() -> Result<()> {
         "read-write"
     }
     .to_string();
-    let ssh_port = args.ssh_port.or(config.ssh_port);
+    let ssh_port = args.ssh_port.or(host.ssh_port).or(config.ssh_port);
     let ssh_auth_command = args
         .ssh_auth_command
         .clone()
@@ -356,6 +360,77 @@ fn run_update(config: &dosh::config::ClientConfig) -> Result<()> {
     Ok(())
 }
 
+fn run_import_ssh(config: &dosh::config::ClientConfig, aliases: &[String]) -> Result<()> {
+    if aliases.is_empty() {
+        return Err(anyhow!("usage: dosh import-ssh <ssh-host> [ssh-host...]"));
+    }
+    let path = expand_tilde("~/.config/dosh/hosts.toml");
+    let mut raw = if path.exists() {
+        fs::read_to_string(&path).with_context(|| format!("read {}", path.display()))?
+    } else {
+        String::new()
+    };
+    if !raw.ends_with('\n') && !raw.is_empty() {
+        raw.push('\n');
+    }
+    for alias in aliases {
+        let ssh_config = ssh_config(alias, None)
+            .with_context(|| format!("read SSH config for {alias} with ssh -G"))?;
+        let udp_host = ssh_config
+            .hostname
+            .clone()
+            .unwrap_or_else(|| ssh_destination_host(alias));
+        if raw_contains_host_table(&raw, alias) {
+            eprintln!("dosh import-ssh: [{alias}] already exists, skipping");
+            continue;
+        }
+        raw.push('\n');
+        raw.push_str(&format!("[{}]\n", toml_bare_key_or_quoted(alias)));
+        raw.push_str(&format!("ssh = {}\n", toml_string(alias)));
+        raw.push_str(&format!("dosh_host = {}\n", toml_string(&udp_host)));
+        raw.push_str(&format!("port = {}\n", config.dosh_port));
+        if let Some(port) = ssh_config.port
+            && port != 22
+        {
+            raw.push_str(&format!("ssh_port = {port}\n"));
+        }
+        raw.push_str("predict = false\n");
+    }
+    if let Some(parent) = path.parent() {
+        fs::create_dir_all(parent)?;
+    }
+    fs::write(&path, raw).with_context(|| format!("write {}", path.display()))?;
+    eprintln!("dosh import-ssh: wrote {}", path.display());
+    Ok(())
+}
+
+fn raw_contains_host_table(raw: &str, alias: &str) -> bool {
+    let quoted = toml_string(alias);
+    raw.lines().any(|line| {
+        let line = line.trim();
+        line == format!("[{alias}]") || line == format!("[{quoted}]")
+    })
+}
+
+fn toml_string(value: &str) -> String {
+    let escaped = value
+        .replace('\\', "\\\\")
+        .replace('"', "\\\"")
+        .replace('\n', "\\n");
+    format!("\"{escaped}\"")
+}
+
+fn toml_bare_key_or_quoted(value: &str) -> String {
+    if value
+        .chars()
+        .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
+    {
+        value.to_string()
+    } else {
+        toml_string(value)
+    }
+}
+
 fn raw_base_from_repo(repo: &str) -> String {
     repo.trim_end_matches(".git").to_string()
 }
@@ -420,9 +495,9 @@ fn ssh_bootstrap(
     server: &str,
     ssh_port: Option<u16>,
     ssh_auth_command: &str,
-    ssh_key: Option<&std::path::Path>,
-    ssh_known_hosts: Option<&std::path::Path>,
-    ssh_control_path: Option<&std::path::Path>,
+    ssh_key: Option<&Path>,
+    ssh_known_hosts: Option<&Path>,
+    ssh_control_path: Option<&Path>,
     session: &str,
     mode: &str,
     cols: u16,
@@ -488,6 +563,18 @@ fn ssh_destination_host(server: &str) -> String {
 }
 
 fn ssh_config_hostname(server: &str, ssh_port: Option<u16>) -> Result<String> {
+    ssh_config(server, ssh_port)?
+        .hostname
+        .ok_or_else(|| anyhow!("ssh -G did not print hostname"))
+}
+
+#[derive(Debug, Clone, Default)]
+struct SshConfig {
+    hostname: Option<String>,
+    port: Option<u16>,
+}
+
+fn ssh_config(server: &str, ssh_port: Option<u16>) -> Result<SshConfig> {
     let mut command = Command::new("ssh");
     command.arg("-G");
     if let Some(ssh_port) = ssh_port {
@@ -504,21 +591,24 @@ fn ssh_config_hostname(server: &str, ssh_port: Option<u16>) -> Result<String> {
         ));
     }
     let raw = String::from_utf8(output.stdout)?;
-    parse_ssh_config_hostname(&raw).ok_or_else(|| anyhow!("ssh -G did not print hostname"))
+    Ok(parse_ssh_config(&raw))
 }
 
-fn parse_ssh_config_hostname(raw: &str) -> Option<String> {
-    raw.lines().find_map(|line| {
+fn parse_ssh_config(raw: &str) -> SshConfig {
+    let mut config = SshConfig::default();
+    for line in raw.lines() {
         let line = line.trim();
-        let (key, value) = line.split_once(char::is_whitespace)?;
-        if key.eq_ignore_ascii_case("hostname") {
-            let value = value.trim();
-            if !value.is_empty() {
-                return Some(value.to_string());
-            }
+        let Some((key, value)) = line.split_once(char::is_whitespace) else {
+            continue;
+        };
+        let value = value.trim();
+        if key.eq_ignore_ascii_case("hostname") && !value.is_empty() {
+            config.hostname = Some(value.to_string());
+        } else if key.eq_ignore_ascii_case("port") {
+            config.port = value.parse().ok();
         }
-        None
-    })
+    }
+    config
 }
 
 fn resolve_addr(host: &str, port: u16) -> Result<SocketAddr> {
@@ -1146,7 +1236,8 @@ const TERMINAL_CLEANUP: &[u8] = concat!(
 #[cfg(test)]
 mod tests {
     use super::{
-        FrameBuffer, Predictor, parse_ssh_config_hostname, ssh_destination_host, startup_command,
+        FrameBuffer, Predictor, parse_ssh_config, raw_contains_host_table, ssh_destination_host,
+        startup_command, toml_bare_key_or_quoted,
     };
     use dosh::protocol::Frame;
 
@@ -1154,7 +1245,7 @@ mod tests {
     fn parses_ssh_config_hostname() {
         let raw = "user palav\nhostname palav.dev\nport 22\n";
         assert_eq!(
-            parse_ssh_config_hostname(raw),
+            parse_ssh_config(raw).hostname,
             Some("palav.dev".to_string())
         );
     }
@@ -1163,11 +1254,29 @@ mod tests {
     fn parses_ssh_config_hostname_case_insensitively() {
         let raw = "HostName 192.0.2.10\n";
         assert_eq!(
-            parse_ssh_config_hostname(raw),
+            parse_ssh_config(raw).hostname,
             Some("192.0.2.10".to_string())
         );
     }
 
+    #[test]
+    fn parses_ssh_config_port() {
+        let parsed = parse_ssh_config("hostname example.com\nport 2222\n");
+        assert_eq!(parsed.hostname, Some("example.com".to_string()));
+        assert_eq!(parsed.port, Some(2222));
+    }
+
+    #[test]
+    fn detects_existing_imported_host_tables() {
+        assert!(raw_contains_host_table(
+            "[palav]\nssh = \"palav\"\n",
+            "palav"
+        ));
+        assert!(raw_contains_host_table("[\"home box\"]\n", "home box"));
+        assert_eq!(toml_bare_key_or_quoted("palav"), "palav");
+        assert_eq!(toml_bare_key_or_quoted("home box"), "\"home box\"");
+    }
+
     #[test]
     fn strips_user_from_fallback_udp_host() {
         assert_eq!(ssh_destination_host("palav@example.com"), "example.com");