use anyhow::{Context, Result, anyhow}; use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use clap::{ArgAction, Parser}; use crossterm::terminal::{disable_raw_mode, enable_raw_mode, size}; use dosh::auth::{ BootstrapResponse, build_bootstrap, decode_bootstrap, load_or_create_server_secret, }; use dosh::config::{expand_tilde, load_client_config, load_hosts_config, load_server_config}; use dosh::crypto; use dosh::native::{ EnvVar, ForwardingKind, ForwardingRequest, KnownHostStatus, TrustResult, derive_native_session_key, generate_native_ephemeral, host_fingerprint, load_ed25519_identity, load_ed25519_identity_with_passphrase, parse_host_public_key_line, remove_trusted_host, sign_user_auth, trust_host, verify_known_host, verify_server_hello, }; use dosh::protocol::{ self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; use dosh::ssh_agent; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; use std::fs; use std::io::{Read, Write}; use std::net::{ Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener as StdTcpListener, TcpStream as StdTcpStream, ToSocketAddrs, }; use std::path::{Path, PathBuf}; use std::process::{Command, Stdio}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream, UdpSocket}; use tokio::sync::mpsc; const STREAM_INITIAL_WINDOW: usize = 1024 * 1024; /// Current terminal size, with a sane fallback. /// /// `crossterm::size()` returns `Err` when stdout is not a TTY (piped), but it /// can also return `Ok((0, 0))` for a real PTY that has not been sized yet /// (e.g. launched under `script` with no controlling window). Sending `0` to /// the server is meaningless and used to crash older servers, so treat any /// zero dimension the same as a missing size and fall back to 80x24. fn terminal_size() -> (u16, u16) { match size() { Ok((cols, rows)) if cols > 0 && rows > 0 => (cols, rows), _ => (80, 24), } } #[derive(Debug, Parser)] #[command(name = "dosh-client")] struct Args { #[arg()] server: Option, #[arg()] command: Vec, #[arg(long)] session: Option, #[arg(long)] new: bool, #[arg(long)] view_only: bool, #[arg(long)] local_auth: bool, #[arg(long)] auth: Option, #[arg(long)] no_cache: bool, #[arg(long)] remove: bool, #[arg(long)] replace: bool, #[arg(long)] attach_only: bool, #[arg(short = 'L', long = "local-forward")] local_forward: Vec, #[arg(short = 'R', long = "remote-forward")] remote_forward: Vec, #[arg(short = 'D', long = "dynamic-forward")] dynamic_forward: Vec, #[arg(short = 'N', long)] forward_only: bool, #[arg(short = 'f', long)] background: bool, #[arg(long)] predict: bool, #[arg(long)] ssh_port: Option, #[arg(long)] ssh_auth_command: Option, #[arg(long)] ssh_key: Option, #[arg(long)] ssh_known_hosts: Option, #[arg(long)] ssh_control_path: Option, #[arg(long)] dosh_port: Option, #[arg(long)] dosh_host: Option, #[arg(short, long, action = ArgAction::Count)] verbose: u8, } #[derive(Debug, Clone, Serialize, Deserialize)] struct CachedCredential { server: String, session: String, mode: String, udp_host: String, udp_port: u16, client_id: [u8; 16], session_key: [u8; 32], session_key_id: [u8; 16], attach_ticket: Vec, attach_ticket_psk: [u8; 32], last_rendered_seq: u64, } #[derive(Debug, Clone, PartialEq, Eq)] struct LocalForward { bind_host: String, listen_port: u16, target_host: String, target_port: u16, } #[derive(Debug, Clone, PartialEq, Eq)] struct RemoteForward { bind_host: String, listen_port: u16, target_host: String, target_port: u16, } #[derive(Debug, Clone, PartialEq, Eq)] struct DynamicForward { bind_host: String, listen_port: u16, } enum ForwardEvent { Open { stream_id: u64, target_host: String, target_port: u16, writer: tokio::net::tcp::OwnedWriteHalf, socks5_reply: bool, }, Data { stream_id: u64, bytes: Vec, }, Close { stream_id: u64, }, } #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { let args = Args::parse(); let config = load_client_config(None).unwrap_or_default(); 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); } if args.server.as_deref() == Some("trust") { return run_trust_command(&config, &args); } if args.server.as_deref() == Some("doctor") { return run_doctor_command(&config, &args).await; } let requested_server = args.server.clone().unwrap_or_else(|| config.server.clone()); let hosts = load_hosts_config(None).unwrap_or_default(); let host = hosts .hosts .get(&requested_server) .cloned() .unwrap_or_default(); let raw_server = host.ssh.clone().unwrap_or_else(|| requested_server.clone()); let server = ssh_with_user(&raw_server, host.user.as_deref()); let startup_command = startup_command(&args.command).or_else(|| { host.default_command .as_deref() .map(startup_command_from_string) }); let local_forwards = parse_local_forwards(&args.local_forward)?; let remote_forwards = parse_remote_forwards(&args.remote_forward)?; let dynamic_forwards = parse_dynamic_forwards(&args.dynamic_forward)?; let forwarding_requested = !local_forwards.is_empty() || !remote_forwards.is_empty() || !dynamic_forwards.is_empty(); let predict = args.predict || host.predict.unwrap_or(config.predict); let session = select_session(args.session.as_deref(), args.new); let mode = if args.forward_only { "forward-only" } else if args.view_only || config.view_only { "view-only" } else { "read-write" } .to_string(); let ssh_port = args.ssh_port.or(host.ssh_port).or(config.ssh_port); let ssh_auth_command = args .ssh_auth_command .clone() .or_else(|| config.ssh_auth_command.clone()) .unwrap_or_else(|| "~/.local/bin/dosh-auth".to_string()); if args .command .first() .is_some_and(|command| command == "sessions") { return run_sessions_command( &server, ssh_port, args.ssh_key.as_deref(), args.ssh_known_hosts.as_deref(), ); } let dosh_port = args.dosh_port.or(host.port).unwrap_or(config.dosh_port); let cache_path = cache_path(&config.credential_cache, &requested_server, &session, &mode); let (cols, rows) = terminal_size(); let auth_preference = args .auth .clone() .unwrap_or_else(|| config.auth_preference.clone()); if forwarding_requested && args.local_auth { return Err(anyhow!( "port forwarding requires native auth; remove --local-auth" )); } if forwarding_requested && !auth_allows(&auth_preference, "native") { return Err(anyhow!("port forwarding requires --auth native or auto")); } if args.background && std::env::var_os("DOSH_BACKGROUND_CHILD").is_none() { return spawn_background_forwarder(&args, forwarding_requested); } let allow_cache = !args.no_cache && !forwarding_requested; let started = Instant::now(); let mut resolved_ssh_config = None; let target_udp_host = if let Some(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() } else { match ssh_config(&server, ssh_port) { Ok(parsed) => { let hostname = parsed .hostname .clone() .unwrap_or_else(|| ssh_destination_host(&server)); resolved_ssh_config = Some(parsed); hostname } Err(err) => { log_debug( args.verbose, 2, &format!("dosh could not resolve SSH config hostname, using {server}: {err:#}"), ); resolved_ssh_config = Some(SshConfig::default()); ssh_destination_host(&server) } } }; let credential = if allow_cache { load_cache(&cache_path).ok() } else { None }; let cache_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref()); let socket = UdpSocket::bind("0.0.0.0:0").await?; if let Some(mut cached) = credential.clone() { cached.udp_host = target_udp_host.clone(); cached.udp_port = dosh_port; log_timing(args.verbose, "credential_lookup_end", started.elapsed()); if config.cache_attach_tickets && allow_cache { match try_ticket_attach(&socket, &cached, cols, rows, cache_requested_env.clone()).await { Ok((frame, cred)) => { log_timing(args.verbose, "udp_ticket_attach_ready", started.elapsed()); save_cache(&cache_path, &cred)?; if args.attach_only { render_frame(&frame)?; detach_once(&socket, &cred, 2).await?; return Ok(()); } return run_terminal( socket, cred, Some(frame), predict, startup_command, config.reconnect_timeout_secs, Vec::new(), Vec::new(), false, ) .await; } Err(err) => { log_debug( args.verbose, 2, &format!( "dosh ticket attach failed, falling back to SSH bootstrap: {err:#}" ), ); } } } else { match try_resume_fresh_process(&socket, &cached, cols, rows).await { Ok((frame, cred)) => { log_timing(args.verbose, "udp_resume_ready", started.elapsed()); if allow_cache { save_cache(&cache_path, &cred)?; } if args.attach_only { render_frame(&frame)?; detach_once(&socket, &cred, 2).await?; return Ok(()); } return run_terminal( socket, cred, Some(frame), predict, startup_command, config.reconnect_timeout_secs, Vec::new(), Vec::new(), false, ) .await; } Err(err) => { log_debug( args.verbose, 2, &format!("dosh resume failed, falling back to SSH bootstrap: {err:#}"), ); } } } } if !args.local_auth && auth_allows(&auth_preference, "native") { if resolved_ssh_config.is_none() { resolved_ssh_config = Some(ssh_config(&server, ssh_port).unwrap_or_default()); } let cold_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref()); let native_start = Instant::now(); match try_native_auth( &socket, &config, &requested_server, &server, ssh_port, &target_udp_host, dosh_port, args.ssh_key.as_deref(), &session, &mode, cols, rows, forwarding_requests(&local_forwards, &remote_forwards, &dynamic_forwards), resolved_ssh_config.as_ref(), cold_requested_env, ) .await { Ok((frame, cred)) => { log_timing(args.verbose, "native_auth", native_start.elapsed()); if allow_cache { save_cache(&cache_path, &cred)?; } log_timing(args.verbose, "terminal_ready", started.elapsed()); if args.attach_only { render_frame(&frame)?; detach_once(&socket, &cred, 2).await?; return Ok(()); } return run_terminal( socket, cred, Some(frame), predict, startup_command, config.reconnect_timeout_secs, local_forwards, dynamic_forwards, args.forward_only, ) .await; } Err(err) if !forwarding_requested && auth_allows(&auth_preference, "ssh") => { log_debug( args.verbose, 2, &format!("dosh native auth failed, falling back to SSH bootstrap: {err:#}"), ); } Err(err) => return Err(err.context("native auth failed")), } } if !args.local_auth && !auth_allows(&auth_preference, "ssh") { return Err(anyhow!( "auth preference {auth_preference:?} does not allow SSH fallback" )); } let bootstrap_start = Instant::now(); if !args.local_auth && resolved_ssh_config.is_none() { resolved_ssh_config = Some(ssh_config(&server, ssh_port).unwrap_or_default()); } let cold_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref()); let bootstrap = if args.local_auth { local_bootstrap(&session, &mode, cols, rows, dosh_port, target_udp_host)? } else { let mut bootstrap = ssh_bootstrap( &server, ssh_port, &ssh_auth_command, args.ssh_key.as_deref(), args.ssh_known_hosts.as_deref(), args.ssh_control_path.as_deref(), &session, &mode, cols, rows, )?; bootstrap.udp_host = target_udp_host; bootstrap.udp_port = dosh_port; bootstrap }; log_timing(args.verbose, "ssh_bootstrap", bootstrap_start.elapsed()); let (ok, mut cred) = bootstrap_attach(&socket, &server, &bootstrap, cols, rows, cold_requested_env).await?; cred.last_rendered_seq = ok.initial_seq; if allow_cache { save_cache(&cache_path, &cred)?; } log_timing(args.verbose, "terminal_ready", started.elapsed()); let first = Frame { session: ok.session, output_seq: ok.initial_seq, bytes: ok.snapshot, snapshot: true, closed: false, }; if args.attach_only { render_frame(&first)?; detach_once(&socket, &cred, 2).await?; return Ok(()); } run_terminal( socket, cred, Some(first), predict, startup_command, config.reconnect_timeout_secs, Vec::new(), Vec::new(), false, ) .await } fn run_trust_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> { let host = args .command .first() .ok_or_else(|| anyhow!("usage: dosh trust [--remove] [--replace] "))?; let path = expand_tilde(&config.known_hosts); if args.remove { if remove_trusted_host(&path, host)? { eprintln!("dosh trust: removed {host} from {}", path.display()); } else { eprintln!("dosh trust: {host} was not trusted in {}", path.display()); } return Ok(()); } let hosts = load_hosts_config(None).unwrap_or_default(); let host_config = hosts.hosts.get(host).cloned().unwrap_or_default(); let raw_server = host_config.ssh.clone().unwrap_or_else(|| host.to_string()); let server = ssh_with_user(&raw_server, host_config.user.as_deref()); let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port); let ssh_auth_command = args .ssh_auth_command .clone() .or_else(|| config.ssh_auth_command.clone()) .unwrap_or_else(|| "~/.local/bin/dosh-auth".to_string()); let public_key = fetch_host_key_over_ssh( &server, ssh_port, &ssh_auth_command, args.ssh_key.as_deref(), args.ssh_known_hosts.as_deref(), args.ssh_control_path.as_deref(), )?; match trust_host(&path, host, &public_key, "ssh", args.replace)? { TrustResult::AlreadyTrusted => { eprintln!( "dosh trust: {host} already trusted ({})", host_fingerprint(&public_key) ); } TrustResult::Trusted => { eprintln!( "dosh trust: trusted {host} {} in {}", host_fingerprint(&public_key), path.display() ); } } Ok(()) } async fn run_doctor_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> { let requested = args .command .first() .cloned() .ok_or_else(|| anyhow!("usage: dosh doctor "))?; let hosts = load_hosts_config(None).unwrap_or_default(); let host_config = hosts.hosts.get(&requested).cloned().unwrap_or_default(); let raw_server = host_config.ssh.clone().unwrap_or_else(|| requested.clone()); let server = ssh_with_user(&raw_server, host_config.user.as_deref()); let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port); let dosh_port = args .dosh_port .or(host_config.port) .unwrap_or(config.dosh_port); println!("Dosh doctor for {requested}"); let parsed_ssh_config = match ssh_config(&server, ssh_port) { Ok(parsed) => { let hostname = parsed .hostname .clone() .unwrap_or_else(|| ssh_destination_host(&server)); println!( "[ok] ssh config: host={} user={} port={} identities_only={} proxy={}", hostname, parsed.user.as_deref().unwrap_or(""), parsed.port.unwrap_or(ssh_port.unwrap_or(22)), parsed.identities_only, parsed .proxy_jump .as_deref() .or(parsed.proxy_command.as_deref()) .unwrap_or("") ); parsed } Err(err) => { println!("[fail] ssh config: {err:#}"); SshConfig::default() } }; match ssh_fallback_probe( &server, ssh_port, args.ssh_key.as_deref(), args.ssh_known_hosts.as_deref(), ) { Ok(()) => println!("[ok] ssh fallback: reachable"), Err(err) => println!("[warn] ssh fallback: {err:#}"), } let 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)) }); println!("[info] native endpoint: {udp_host}:{dosh_port}"); let socket = UdpSocket::bind("0.0.0.0:0").await?; match try_native_auth_check( &socket, config, &requested, &server, ssh_port, &udp_host, dosh_port, args.ssh_key.as_deref(), Some(&parsed_ssh_config), ) .await { Ok(check) => { println!("[ok] native udp: reachable"); println!( "[ok] server version: {}", if check.server_version.is_empty() { "" } else { &check.server_version } ); println!("[ok] known host: trusted"); println!("[ok] native auth: authorized as {}", check.requested_user); println!( "[ok] forwarding policy: tcp={} remote={} agent={}", check.allow_tcp_forwarding, check.allow_remote_forwarding, check.allow_agent_forwarding ); Ok(()) } Err(err) => { println!("[fail] native/auth check: {err:#}"); Err(err.context("dosh doctor failed")) } } } fn ssh_fallback_probe( server: &str, ssh_port: Option, ssh_key: Option<&Path>, ssh_known_hosts: Option<&Path>, ) -> Result<()> { let mut command = Command::new("ssh"); command.arg("-o").arg("BatchMode=yes"); command.arg("-o").arg("ConnectTimeout=5"); command.arg("-T"); if let Some(ssh_port) = ssh_port { command.arg("-p").arg(ssh_port.to_string()); } if let Some(key) = ssh_key { command.arg("-i").arg(key); } if let Some(known_hosts) = ssh_known_hosts { command .arg("-o") .arg(format!("UserKnownHostsFile={}", known_hosts.display())); } let output = command.arg(server).arg("true").output()?; anyhow::ensure!( output.status.success(), "{}", String::from_utf8_lossy(&output.stderr) ); Ok(()) } fn fetch_host_key_over_ssh( server: &str, ssh_port: Option, ssh_auth_command: &str, ssh_key: Option<&Path>, ssh_known_hosts: Option<&Path>, ssh_control_path: Option<&Path>, ) -> Result { let mut command = Command::new("ssh"); if let Some(ssh_port) = ssh_port { command.arg("-p").arg(ssh_port.to_string()); } command.arg("-T"); if let Some(key) = ssh_key { command.arg("-i").arg(key); } if let Some(known_hosts) = ssh_known_hosts { command .arg("-o") .arg(format!("UserKnownHostsFile={}", known_hosts.display())); } if let Some(control_path) = ssh_control_path { command.arg("-S").arg(control_path); } let output = command .arg(server) .arg(ssh_auth_command) .arg("--host-key") .output() .context("fetch Dosh host key over SSH")?; if !output.status.success() { return Err(anyhow!( "Dosh host-key fetch failed: {}", String::from_utf8_lossy(&output.stderr) )); } let raw = String::from_utf8(output.stdout)?; parse_host_public_key_line(&raw) } fn startup_command(args: &[String]) -> Option> { if args.is_empty() { return None; } let mut command = args.join(" "); command.push('\n'); Some(command.into_bytes()) } fn startup_command_from_string(command: &str) -> Vec { let mut command = command.to_string(); command.push('\n'); command.into_bytes() } fn spawn_background_forwarder(args: &Args, forwarding_requested: bool) -> Result<()> { if !forwarding_requested { return Err(anyhow!( "-f requires at least one -L, -R, or -D forwarding request" )); } if !args.forward_only { return Err(anyhow!("-f currently requires -N for forward-only mode")); } let listener = StdTcpListener::bind("127.0.0.1:0").context("bind background readiness socket")?; listener .set_nonblocking(true) .context("configure background readiness socket")?; let readiness_addr = listener.local_addr()?; let token = URL_SAFE_NO_PAD.encode(crypto::random_16()); let exe = std::env::current_exe().context("locate current dosh-client executable")?; let mut child = Command::new(exe) .args(std::env::args_os().skip(1)) .env("DOSH_BACKGROUND_CHILD", "1") .env("DOSH_BACKGROUND_READY_ADDR", readiness_addr.to_string()) .env("DOSH_BACKGROUND_READY_TOKEN", &token) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::null()) .spawn() .context("spawn background dosh forwarder")?; if let Some(path) = std::env::var_os("DOSH_BACKGROUND_PID_FILE") { fs::write(path, child.id().to_string()).context("write background pid file")?; } let deadline = Instant::now() + Duration::from_secs(30); loop { match listener.accept() { Ok((mut stream, _)) => { let mut raw = String::new(); stream.read_to_string(&mut raw).ok(); if raw.trim() == token { eprintln!("dosh background forwarding ready"); return Ok(()); } return Err(anyhow!("background forwarder sent invalid readiness token")); } Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {} Err(err) => return Err(err).context("wait for background readiness"), } if let Some(status) = child.try_wait()? { return Err(anyhow!( "background dosh forwarder exited before readiness: {status}" )); } if Instant::now() >= deadline { let _ = child.kill(); let _ = child.wait(); return Err(anyhow!("background dosh forwarder did not become ready")); } std::thread::sleep(Duration::from_millis(20)); } } fn signal_background_ready() { let Some(addr) = std::env::var_os("DOSH_BACKGROUND_READY_ADDR") else { return; }; let Some(token) = std::env::var_os("DOSH_BACKGROUND_READY_TOKEN") else { return; }; if let Ok(mut stream) = StdTcpStream::connect(addr.to_string_lossy().as_ref()) { let _ = stream.write_all(token.to_string_lossy().as_bytes()); let _ = stream.shutdown(std::net::Shutdown::Write); } } fn parse_local_forwards(raw: &[String]) -> Result> { raw.iter().map(|value| parse_local_forward(value)).collect() } fn parse_remote_forwards(raw: &[String]) -> Result> { raw.iter() .map(|value| parse_remote_forward(value)) .collect() } fn parse_dynamic_forwards(raw: &[String]) -> Result> { raw.iter() .map(|value| parse_dynamic_forward(value)) .collect() } fn parse_local_forward(raw: &str) -> Result { let parts = raw.split(':').collect::>(); let (bind_host, listen, target_host, target_port) = match parts.as_slice() { [listen, target_host, target_port] => { ("127.0.0.1".to_string(), *listen, *target_host, *target_port) } [bind_host, listen, target_host, target_port] => ( (*bind_host).to_string(), *listen, *target_host, *target_port, ), _ => { return Err(anyhow!( "invalid -L {raw:?}; expected listen_port:target_host:target_port or bind_host:listen_port:target_host:target_port" )); } }; let listen_port = listen .parse::() .with_context(|| format!("invalid local forward listen port in {raw:?}"))?; let target_port = target_port .parse::() .with_context(|| format!("invalid local forward target port in {raw:?}"))?; if target_host.is_empty() { return Err(anyhow!("invalid -L {raw:?}; target host cannot be empty")); } if bind_host.is_empty() { return Err(anyhow!("invalid -L {raw:?}; bind host cannot be empty")); } Ok(LocalForward { bind_host, listen_port, target_host: target_host.to_string(), target_port, }) } fn parse_remote_forward(raw: &str) -> Result { let local = parse_local_forward(raw)?; Ok(RemoteForward { bind_host: local.bind_host, listen_port: local.listen_port, target_host: local.target_host, target_port: local.target_port, }) } fn parse_dynamic_forward(raw: &str) -> Result { let parts = raw.split(':').collect::>(); let (bind_host, listen) = match parts.as_slice() { [listen] => ("127.0.0.1".to_string(), *listen), [bind_host, listen] => ((*bind_host).to_string(), *listen), _ => { return Err(anyhow!( "invalid -D {raw:?}; expected listen_port or bind_host:listen_port" )); } }; let listen_port = listen .parse::() .with_context(|| format!("invalid dynamic forward listen port in {raw:?}"))?; if bind_host.is_empty() { return Err(anyhow!("invalid -D {raw:?}; bind host cannot be empty")); } Ok(DynamicForward { bind_host, listen_port, }) } fn forwarding_requests( local_forwards: &[LocalForward], remote_forwards: &[RemoteForward], dynamic_forwards: &[DynamicForward], ) -> Vec { let mut requests = local_forwards .iter() .map(|forward| ForwardingRequest { kind: ForwardingKind::Local, bind_host: Some(forward.bind_host.clone()), listen_port: forward.listen_port, target_host: Some(forward.target_host.clone()), target_port: Some(forward.target_port), }) .collect::>(); requests.extend(remote_forwards.iter().map(|forward| ForwardingRequest { kind: ForwardingKind::Remote, bind_host: Some(forward.bind_host.clone()), listen_port: forward.listen_port, target_host: Some(forward.target_host.clone()), target_port: Some(forward.target_port), })); requests.extend(dynamic_forwards.iter().map(|forward| ForwardingRequest { kind: ForwardingKind::Dynamic, bind_host: Some(forward.bind_host.clone()), listen_port: forward.listen_port, target_host: None, target_port: None, })); requests } fn run_sessions_command( server: &str, ssh_port: Option, ssh_key: Option<&std::path::Path>, ssh_known_hosts: Option<&std::path::Path>, ) -> Result<()> { let mut command = Command::new("ssh"); if let Some(ssh_port) = ssh_port { command.arg("-p").arg(ssh_port.to_string()); } command.arg("-T"); if let Some(key) = ssh_key { command.arg("-i").arg(key); } if let Some(known_hosts) = ssh_known_hosts { command .arg("-o") .arg(format!("UserKnownHostsFile={}", known_hosts.display())); } let script = r#"printf 'tmux sessions:\n'; tmux ls 2>/dev/null || printf ' none\n'; printf '\nDosh service:\n'; systemctl --user --no-pager --plain status dosh-server.service 2>/dev/null | sed -n '1,8p' || pgrep -af dosh-server || printf ' not running\n'"#; let remote_command = format!("sh -c {}", shell_word(script)); let status = command .arg(server) .arg(remote_command) .status() .context("run remote sessions command")?; if !status.success() { return Err(anyhow!("sessions command failed with status {status}")); } Ok(()) } fn run_update(config: &dosh::config::ClientConfig) -> Result<()> { let repo = config .update_repo .clone() .unwrap_or_else(|| "https://git.palav.dev/Palav/dosh.git".to_string()); let raw_base = raw_base_from_repo(&repo); let installer = format!("{raw_base}/raw/branch/main/install.sh"); let update_cache = dirs::cache_dir() .unwrap_or_else(|| expand_tilde("~/.cache")) .join("dosh") .join("source"); let mut script = format!( "curl -fsSL {} | DOSH_REPO={} DOSH_PORT={} DOSH_UPDATE_CACHE={} DOSH_UPDATE_QUIET=1 sh -s -- client", shell_word(&installer), shell_word(&repo), config.update_port.unwrap_or(config.dosh_port), shell_word(&update_cache.display().to_string()) ); if config.server != "user@example.com" { script.push_str(&format!(" --server {}", shell_word(&config.server))); } if let Some(host) = &config.dosh_host { script.push_str(&format!(" --dosh-host {}", shell_word(host))); } let status = Command::new("sh") .arg("-c") .arg(script) .stdin(Stdio::null()) .status() .context("run dosh update installer")?; if !status.success() { return Err(anyhow!("dosh update failed with status {status}")); } 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...]")); } 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(user) = &ssh_config.user { raw.push_str(&format!("user = {}\n", toml_string(user))); } 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() } fn shell_word(value: &str) -> String { format!("'{}'", value.replace('\'', "'\\''")) } fn local_bootstrap( session: &str, mode: &str, cols: u16, rows: u16, port: u16, host: String, ) -> Result { let mut server_config = load_server_config(None)?; server_config.port = port; let secret = load_or_create_server_secret(&server_config)?; let user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string()); let nonce = crypto::random_12(); build_bootstrap( &server_config, &secret, user, session.to_string(), mode.to_string(), (cols, rows), nonce, host, ) } fn auth_allows(preference: &str, method: &str) -> bool { preference .split(',') .map(str::trim) .any(|entry| entry == "auto" || entry == method) } fn log_timing(verbose: u8, name: &str, elapsed: Duration) { if verbose >= 1 { eprintln!("dosh timing {name}={}ms", elapsed.as_millis()); } } fn log_debug(verbose: u8, level: u8, message: &str) { if verbose >= level { eprintln!("{message}"); } } fn select_session(requested: Option<&str>, _force_new: bool) -> String { if let Some(session) = requested { return session.to_string(); } unique_session_name() } fn unique_session_name() -> String { let millis = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_millis(); format!("term-{millis}-{}", std::process::id()) } fn ssh_bootstrap( server: &str, ssh_port: Option, ssh_auth_command: &str, ssh_key: Option<&Path>, ssh_known_hosts: Option<&Path>, ssh_control_path: Option<&Path>, session: &str, mode: &str, cols: u16, rows: u16, ) -> Result { let nonce = crypto::random_12(); let nonce_b64 = URL_SAFE_NO_PAD.encode(nonce); let size = format!("{cols}x{rows}"); let mut command = Command::new("ssh"); if let Some(ssh_port) = ssh_port { command.arg("-p").arg(ssh_port.to_string()); } command.arg("-T"); if let Some(key) = ssh_key { command.arg("-i").arg(key); } if let Some(known_hosts) = ssh_known_hosts { command .arg("-o") .arg(format!("UserKnownHostsFile={}", known_hosts.display())); } if let Some(control_path) = ssh_control_path { command.arg("-S").arg(control_path); } let output = command .arg(server) .arg(ssh_auth_command) .arg("--protocol") .arg("1") .arg(format!("--nonce={nonce_b64}")) .arg(format!("--session={session}")) .arg(format!("--mode={mode}")) .arg(format!("--size={size}")) .output() .context("run ssh dosh-auth")?; if !output.status.success() { return Err(anyhow!( "ssh bootstrap failed: {}", String::from_utf8_lossy(&output.stderr) )); } let raw = String::from_utf8(output.stdout)?; decode_bootstrap(&raw) } fn ssh_destination_host(server: &str) -> String { let without_user = server.rsplit_once('@').map_or(server, |(_, host)| host); let without_path = without_user .strip_prefix("ssh://") .unwrap_or(without_user) .split('/') .next() .unwrap_or(without_user); if let Some(stripped) = without_path.strip_prefix('[') { if let Some((host, _)) = stripped.split_once(']') { return host.to_string(); } } without_path .split_once(':') .map_or(without_path, |(host, _)| host) .to_string() } fn ssh_username(server: &str) -> Option { let (user, _) = server.rsplit_once('@')?; if user.is_empty() { None } else { Some(user.to_string()) } } fn ssh_with_user(server: &str, user: Option<&str>) -> String { let Some(user) = user else { return server.to_string(); }; if user.is_empty() || ssh_username(server).is_some() || server.starts_with("ssh://") { server.to_string() } else { format!("{user}@{server}") } } #[derive(Debug, Clone, Default)] struct SshConfig { hostname: Option, port: Option, user: Option, identity_files: Vec, user_known_hosts_file: Option, identities_only: bool, proxy_jump: Option, proxy_command: Option, send_env: Vec, set_env: Vec, } fn ssh_config(server: &str, ssh_port: Option) -> Result { let mut command = Command::new("ssh"); command.arg("-G"); if let Some(ssh_port) = ssh_port { command.arg("-p").arg(ssh_port.to_string()); } let output = command .arg(server) .output() .with_context(|| format!("run ssh -G {server}"))?; if !output.status.success() { return Err(anyhow!( "ssh -G failed: {}", String::from_utf8_lossy(&output.stderr) )); } let raw = String::from_utf8(output.stdout)?; Ok(parse_ssh_config(&raw)) } fn parse_ssh_config(raw: &str) -> SshConfig { let mut config = SshConfig::default(); for line in raw.lines() { let line = line.trim(); let Some((key, value)) = line.split_once(char::is_whitespace) else { continue; }; let value = value.trim(); if key.eq_ignore_ascii_case("hostname") && !empty_or_none(value) { config.hostname = Some(value.to_string()); } else if key.eq_ignore_ascii_case("port") { config.port = value.parse().ok(); } else if key.eq_ignore_ascii_case("user") && !empty_or_none(value) { config.user = Some(value.to_string()); } else if key.eq_ignore_ascii_case("identityfile") && !empty_or_none(value) { config.identity_files.push(value.to_string()); } else if key.eq_ignore_ascii_case("userknownhostsfile") && !empty_or_none(value) { config.user_known_hosts_file = Some(value.to_string()); } else if key.eq_ignore_ascii_case("identitiesonly") { config.identities_only = value.eq_ignore_ascii_case("yes"); } else if key.eq_ignore_ascii_case("proxyjump") && !empty_or_none(value) { config.proxy_jump = Some(value.to_string()); } else if key.eq_ignore_ascii_case("proxycommand") && !empty_or_none(value) { config.proxy_command = Some(value.to_string()); } else if key.eq_ignore_ascii_case("sendenv") && !empty_or_none(value) { config .send_env .extend(value.split_whitespace().map(ToString::to_string)); } else if key.eq_ignore_ascii_case("setenv") && !empty_or_none(value) { config.set_env.extend(parse_set_env_values(value)); } } config } fn parse_set_env_values(value: &str) -> Vec { value .split_whitespace() .filter_map(|entry| { let (name, value) = entry.split_once('=')?; if valid_env_name(name) && !value.as_bytes().contains(&0) { Some(EnvVar { name: name.to_string(), value: value.to_string(), }) } else { None } }) .collect() } fn empty_or_none(value: &str) -> bool { value.is_empty() || value.eq_ignore_ascii_case("none") } fn requested_env( config: &dosh::config::ClientConfig, host: &dosh::config::HostConfig, ssh_config: Option<&SshConfig>, ) -> Vec { let mut values = BTreeMap::new(); let mut patterns = host .send_env .clone() .unwrap_or_else(|| config.send_env.clone()); if let Some(ssh_config) = ssh_config { patterns.extend(ssh_config.send_env.clone()); } for (name, value) in std::env::vars() { if valid_env_name(&name) && patterns.iter().any(|pattern| glob_matches(pattern, &name)) { values.insert(name, value); } } for (name, value) in &config.set_env { if valid_env_name(name) && !value.as_bytes().contains(&0) { values.insert(name.clone(), value.clone()); } } for (name, value) in &host.set_env { if valid_env_name(name) && !value.as_bytes().contains(&0) { values.insert(name.clone(), value.clone()); } } if let Some(ssh_config) = ssh_config { for env in &ssh_config.set_env { values.insert(env.name.clone(), env.value.clone()); } } values .into_iter() .map(|(name, value)| EnvVar { name, value }) .collect() } fn valid_env_name(name: &str) -> bool { let mut chars = name.chars(); let Some(first) = chars.next() else { return false; }; if !(first == '_' || first.is_ascii_alphabetic()) { return false; } chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric()) } fn glob_matches(pattern: &str, value: &str) -> bool { let pattern = pattern.as_bytes(); let value = value.as_bytes(); let (mut p, mut v) = (0, 0); let mut star = None; let mut star_value = 0; while v < value.len() { if p < pattern.len() && (pattern[p] == b'?' || pattern[p] == value[v]) { p += 1; v += 1; } else if p < pattern.len() && pattern[p] == b'*' { star = Some(p); p += 1; star_value = v; } else if let Some(star_pos) = star { p = star_pos + 1; star_value += 1; v = star_value; } else { return false; } } while p < pattern.len() && pattern[p] == b'*' { p += 1; } p == pattern.len() } fn resolve_addr(host: &str, port: u16) -> Result { (host, port) .to_socket_addrs() .with_context(|| format!("resolve UDP target {host}:{port}"))? .next() .ok_or_else(|| anyhow!("no UDP address resolved for {host}:{port}")) } async fn try_native_auth( socket: &UdpSocket, config: &dosh::config::ClientConfig, requested_server: &str, server: &str, ssh_port: Option, udp_host: &str, udp_port: u16, cli_identity: Option<&Path>, session: &str, mode: &str, cols: u16, rows: u16, requested_forwardings: Vec, ssh_config_hint: Option<&SshConfig>, requested_env: Vec, ) -> Result<(Frame, CachedCredential)> { let addr = resolve_addr(udp_host, udp_port)?; let owned_ssh_config; let ssh_config = if let Some(ssh_config) = ssh_config_hint { ssh_config } else { owned_ssh_config = ssh_config(server, ssh_port).unwrap_or_default(); &owned_ssh_config }; let requested_user = ssh_username(server) .or(ssh_config.user.clone()) .unwrap_or_else(|| std::env::var("USER").unwrap_or_else(|_| "unknown".to_string())); let (client_secret, client_public) = generate_native_ephemeral(); let hello = dosh::native::NativeClientHello { protocol_version: dosh::native::NATIVE_PROTOCOL_VERSION, client_random: crypto::random_32(), client_ephemeral_public: client_public, requested_host: requested_server.to_string(), requested_user, requested_session: session.to_string(), requested_mode: mode.to_string(), terminal_size: (cols, rows), supported_aead: vec!["chacha20poly1305".to_string()], supported_user_key_algorithms: vec!["ssh-ed25519".to_string()], cached_host_key_fingerprint: None, attach_ticket_envelope: None, requested_env, }; let packet = protocol::encode_plain( PacketKind::NativeClientHello, [0u8; 16], 1, 0, &protocol::to_body(&NativeClientHelloBody { hello: hello.clone(), })?, )?; socket.send_to(&packet, addr).await?; let mut buf = vec![0u8; 65535]; let (n, _) = tokio::time::timeout( Duration::from_millis(config.native_auth_timeout_ms.max(1)), socket.recv_from(&mut buf), ) .await??; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeServerHello { if packet.header.kind == PacketKind::AttachReject { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("native auth rejected: {}", reject.reason)); } return Err(anyhow!("native auth received unexpected server response")); } let server_hello: NativeServerHelloBody = protocol::from_body(&packet.body)?; verify_server_hello(&hello, &server_hello.hello)?; let known_hosts = expand_tilde(&config.known_hosts); match verify_known_host(&known_hosts, requested_server, &server_hello.hello.host_key)? { KnownHostStatus::Trusted => {} KnownHostStatus::Unknown if config.trust_on_first_use => { trust_host( &known_hosts, requested_server, &server_hello.hello.host_key, "tofu", false, )?; } KnownHostStatus::Unknown => { return Err(anyhow!( "Dosh host key for {requested_server} is not trusted; run `dosh trust {requested_server}` first" )); } KnownHostStatus::Mismatch { expected, actual } => { return Err(anyhow!( "Dosh host key mismatch for {requested_server}: expected {expected}, got {actual}" )); } } let session_key = derive_native_session_key( &client_secret, server_hello.hello.server_ephemeral_public, &hello, &server_hello.hello, )?; let auth = sign_native_user_auth( config, cli_identity, &ssh_config, &hello, &server_hello.hello, requested_forwardings, )?; let mut pending_id = [0u8; 16]; pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]); let auth_body = protocol::to_body(&NativeUserAuthBody { auth })?; let packet = protocol::encode_encrypted( PacketKind::NativeUserAuth, pending_id, 2, 1, &session_key, CLIENT_TO_SERVER, &auth_body, )?; socket.send_to(&packet, addr).await?; let (n, _) = tokio::time::timeout( Duration::from_millis(config.native_auth_timeout_ms.max(1)), socket.recv_from(&mut buf), ) .await??; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeAuthOk { if packet.header.kind == PacketKind::AttachReject { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("native auth rejected: {}", reject.reason)); } return Err(anyhow!("native auth received unexpected auth response")); } let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT)?; let ok: NativeAuthOkBody = protocol::from_body(&plain)?; let frame = Frame { session: ok.ok.session.clone(), output_seq: ok.ok.initial_seq, bytes: ok.ok.snapshot.clone(), snapshot: true, closed: false, }; let cred = CachedCredential { server: requested_server.to_string(), session: ok.ok.session, mode: ok.ok.mode, udp_host: udp_host.to_string(), udp_port, client_id: ok.ok.client_id, session_key: ok.ok.session_key, session_key_id: ok.ok.session_key_id, attach_ticket: ok.ok.attach_ticket, attach_ticket_psk: ok.ok.attach_ticket_psk, last_rendered_seq: ok.ok.initial_seq, }; Ok((frame, cred)) } async fn try_native_auth_check( socket: &UdpSocket, config: &dosh::config::ClientConfig, requested_server: &str, server: &str, ssh_port: Option, udp_host: &str, udp_port: u16, cli_identity: Option<&Path>, ssh_config_hint: Option<&SshConfig>, ) -> Result { let addr = resolve_addr(udp_host, udp_port)?; let owned_ssh_config; let ssh_config = if let Some(ssh_config) = ssh_config_hint { ssh_config } else { owned_ssh_config = ssh_config(server, ssh_port).unwrap_or_default(); &owned_ssh_config }; let requested_user = ssh_username(server) .or(ssh_config.user.clone()) .unwrap_or_else(|| std::env::var("USER").unwrap_or_else(|_| "unknown".to_string())); let (client_secret, client_public) = generate_native_ephemeral(); let hello = dosh::native::NativeClientHello { protocol_version: dosh::native::NATIVE_PROTOCOL_VERSION, client_random: crypto::random_32(), client_ephemeral_public: client_public, requested_host: requested_server.to_string(), requested_user, requested_session: "doctor".to_string(), requested_mode: "doctor".to_string(), terminal_size: (80, 24), supported_aead: vec!["chacha20poly1305".to_string()], supported_user_key_algorithms: vec!["ssh-ed25519".to_string()], cached_host_key_fingerprint: None, attach_ticket_envelope: None, requested_env: Vec::new(), }; let packet = protocol::encode_plain( PacketKind::NativeClientHello, [0u8; 16], 1, 0, &protocol::to_body(&NativeClientHelloBody { hello: hello.clone(), })?, )?; socket.send_to(&packet, addr).await?; let mut buf = vec![0u8; 65535]; let (n, _) = tokio::time::timeout( Duration::from_millis(config.native_auth_timeout_ms.max(1)), socket.recv_from(&mut buf), ) .await??; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeServerHello { if packet.header.kind == PacketKind::AttachReject { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("native doctor rejected: {}", reject.reason)); } return Err(anyhow!("native doctor received unexpected server response")); } let server_hello: NativeServerHelloBody = protocol::from_body(&packet.body)?; verify_server_hello(&hello, &server_hello.hello)?; let known_hosts = expand_tilde(&config.known_hosts); match verify_known_host(&known_hosts, requested_server, &server_hello.hello.host_key)? { KnownHostStatus::Trusted => {} KnownHostStatus::Unknown if config.trust_on_first_use => { trust_host( &known_hosts, requested_server, &server_hello.hello.host_key, "tofu", false, )?; } KnownHostStatus::Unknown => { return Err(anyhow!( "Dosh host key for {requested_server} is not trusted; run `dosh trust {requested_server}` first" )); } KnownHostStatus::Mismatch { expected, actual } => { return Err(anyhow!( "Dosh host key mismatch for {requested_server}: expected {expected}, got {actual}" )); } } let session_key = derive_native_session_key( &client_secret, server_hello.hello.server_ephemeral_public, &hello, &server_hello.hello, )?; let auth = sign_native_user_auth( config, cli_identity, &ssh_config, &hello, &server_hello.hello, Vec::new(), )?; let mut pending_id = [0u8; 16]; pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]); let packet = protocol::encode_encrypted( PacketKind::NativeUserAuth, pending_id, 2, 1, &session_key, CLIENT_TO_SERVER, &protocol::to_body(&NativeUserAuthBody { auth })?, )?; socket.send_to(&packet, addr).await?; let (n, _) = tokio::time::timeout( Duration::from_millis(config.native_auth_timeout_ms.max(1)), socket.recv_from(&mut buf), ) .await??; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::NativeAuthCheckOk { if packet.header.kind == PacketKind::AttachReject { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("native doctor rejected: {}", reject.reason)); } return Err(anyhow!("native doctor received unexpected auth response")); } let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT)?; protocol::from_body(&plain) } fn sign_native_user_auth( config: &dosh::config::ClientConfig, cli_identity: Option<&Path>, ssh_config: &SshConfig, hello: &dosh::native::NativeClientHello, server_hello: &dosh::native::NativeServerHello, requested_forwardings: Vec, ) -> Result { let mut errors = Vec::new(); if config.use_ssh_agent && !ssh_config.identities_only { match ssh_agent::sign_user_auth_with_agent( hello, server_hello, requested_forwardings.clone(), ) { Ok(auth) => return Ok(auth), Err(err) => errors.push(format!("ssh-agent: {err:#}")), } } else if config.use_ssh_agent && ssh_config.identities_only { errors.push("ssh-agent: skipped because SSH config sets IdentitiesOnly=yes".to_string()); } match load_first_native_identity(config, cli_identity, ssh_config) { Ok(identity) => sign_user_auth(&identity, hello, server_hello, requested_forwardings), Err(err) => { errors.push(format!("identity files: {err:#}")); Err(anyhow!( "native auth found no usable key: {}", errors.join("; ") )) } } } fn load_first_native_identity( config: &dosh::config::ClientConfig, cli_identity: Option<&Path>, ssh_config: &SshConfig, ) -> Result { load_first_native_identity_with_prompt( config, cli_identity, ssh_config, prompt_identity_passphrase, ) } fn load_first_native_identity_with_prompt( config: &dosh::config::ClientConfig, cli_identity: Option<&Path>, ssh_config: &SshConfig, mut prompt: F, ) -> Result where F: FnMut(&Path) -> Result, { let mut paths = Vec::new(); if let Some(path) = cli_identity { push_identity_path(&mut paths, path.to_path_buf()); } for path in &ssh_config.identity_files { push_identity_path(&mut paths, expand_tilde(path)); } if !ssh_config.identities_only { for path in &config.identity_files { push_identity_path(&mut paths, expand_tilde(path)); } } let mut errors = Vec::new(); for path in paths { if !path.exists() { continue; } match load_ed25519_identity(&path) { Ok(identity) => return Ok(identity), Err(err) if err.to_string().contains(" is encrypted") => { match prompt(&path).and_then(|passphrase| { load_ed25519_identity_with_passphrase(&path, Some(&passphrase)) }) { Ok(identity) => return Ok(identity), Err(err) => errors.push(format!("{}: {err:#}", path.display())), } } Err(err) => errors.push(format!("{}: {err:#}", path.display())), } } if errors.is_empty() { Err(anyhow!("no usable ssh-ed25519 identity found")) } else { Err(anyhow!( "no usable ssh-ed25519 identity found: {}", errors.join("; ") )) } } fn push_identity_path(paths: &mut Vec, path: PathBuf) { if !paths.iter().any(|existing| existing == &path) { paths.push(path); } } fn prompt_identity_passphrase(path: &Path) -> Result { rpassword::prompt_password(format!("Enter passphrase for {}: ", path.display())) .with_context(|| format!("read passphrase for {}", path.display())) } async fn bootstrap_attach( socket: &UdpSocket, server_name: &str, bootstrap: &BootstrapResponse, cols: u16, rows: u16, requested_env: Vec, ) -> Result<(AttachOk, CachedCredential)> { let addr = resolve_addr(&bootstrap.udp_host, bootstrap.udp_port)?; let req = BootstrapAttachRequest { bootstrap: bootstrap.clone(), cols, rows, requested_env, }; let body = protocol::to_body(&req)?; let packet = protocol::encode_plain(PacketKind::BootstrapAttachRequest, [0u8; 16], 1, 0, &body)?; socket.send_to(&packet, addr).await?; let expected_key_id = protocol::session_key_id(&bootstrap.session_key); let ok = recv_response_until(socket, Duration::from_secs(5), |packet| { if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("attach rejected: {}", reject.reason)); } if packet.header.kind != PacketKind::AttachOk || packet.header.session_key_id != expected_key_id { return Ok(None); } let Ok(plain) = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT) else { return Ok(None); }; Ok(protocol::from_body::(&plain).ok()) }) .await?; let cred = CachedCredential { server: server_name.to_string(), session: ok.session.clone(), mode: ok.mode.clone(), udp_host: bootstrap.udp_host.clone(), udp_port: bootstrap.udp_port, client_id: ok.client_id, session_key: ok.session_key, session_key_id: ok.session_key_id, attach_ticket: bootstrap.attach_ticket.clone(), attach_ticket_psk: bootstrap.attach_ticket_psk, last_rendered_seq: ok.initial_seq, }; Ok((ok, cred)) } async fn try_ticket_attach( socket: &UdpSocket, cached: &CachedCredential, cols: u16, rows: u16, requested_env: Vec, ) -> Result<(Frame, CachedCredential)> { let addr = resolve_addr(&cached.udp_host, cached.udp_port)?; let client_nonce = crypto::random_12(); let request_key = crypto::hkdf32( &cached.attach_ticket_psk, &client_nonce, b"dosh/ticket-attach-request/v1", )?; let body = protocol::to_body(&TicketAttachBody { session: cached.session.clone(), mode: cached.mode.clone(), cols, rows, requested_env, })?; let ciphertext = crypto::seal( &request_key, &client_nonce, b"dosh-ticket-attach-request-v1", &body, )?; let envelope = TicketAttachEnvelope { ticket: cached.attach_ticket.clone(), client_nonce, ciphertext, }; let packet = protocol::encode_plain( PacketKind::TicketAttachRequest, [0u8; 16], 1, 0, &protocol::to_body(&envelope)?, )?; socket.send_to(&packet, addr).await?; let ok = recv_response_until(socket, Duration::from_millis(700), |packet| { if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("ticket attach rejected: {}", reject.reason)); } if packet.header.kind != PacketKind::AttachOk { return Ok(None); } let Ok(envelope) = protocol::from_body::(&packet.body) else { return Ok(None); }; let mut salt = Vec::with_capacity(24); salt.extend_from_slice(&client_nonce); salt.extend_from_slice(&envelope.server_nonce); let response_key = crypto::hkdf32( &cached.attach_ticket_psk, &salt, b"dosh/ticket-attach-ok/v1", )?; let Ok(plain) = crypto::open( &response_key, &envelope.server_nonce, b"dosh-ticket-attach-ok-v1", &envelope.ciphertext, ) else { return Ok(None); }; Ok(protocol::from_body::(&plain).ok()) }) .await?; let frame = Frame { session: ok.session.clone(), output_seq: ok.initial_seq, bytes: ok.snapshot.clone(), snapshot: true, closed: false, }; let mut next = cached.clone(); next.client_id = ok.client_id; next.session_key = ok.session_key; next.session_key_id = ok.session_key_id; next.last_rendered_seq = ok.initial_seq; Ok((frame, next)) } async fn try_resume_fresh_process( socket: &UdpSocket, cached: &CachedCredential, cols: u16, rows: u16, ) -> Result<(Frame, CachedCredential)> { let addr = resolve_addr(&cached.udp_host, cached.udp_port)?; let req = ResumeRequest { session: cached.session.clone(), last_rendered_seq: cached.last_rendered_seq, cols, rows, }; let body = protocol::to_body(&req)?; let packet = protocol::encode_encrypted( PacketKind::ResumeRequest, cached.client_id, 1, 0, &cached.session_key, CLIENT_TO_SERVER, &body, )?; socket.send_to(&packet, addr).await?; let frame = recv_response_until(socket, Duration::from_millis(700), |packet| { if packet.header.conn_id != cached.client_id { return Ok(None); } if packet.header.kind == PacketKind::AttachReject { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("resume rejected: {}", reject.reason)); } if packet.header.kind != PacketKind::ResumeOk { return Ok(None); } let Ok(plain) = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT) else { return Ok(None); }; Ok(protocol::from_body::(&plain).ok()) }) .await?; let mut next = cached.clone(); next.last_rendered_seq = frame.output_seq; Ok((frame, next)) } async fn try_live_resume( socket: &UdpSocket, cached: &CachedCredential, send_seq: &mut u64, cols: u16, rows: u16, ) -> Result<(Frame, CachedCredential)> { let addr = resolve_addr(&cached.udp_host, cached.udp_port)?; let req = ResumeRequest { session: cached.session.clone(), last_rendered_seq: cached.last_rendered_seq, cols, rows, }; let body = protocol::to_body(&req)?; let packet = protocol::encode_encrypted( PacketKind::ResumeRequest, cached.client_id, *send_seq, cached.last_rendered_seq, &cached.session_key, CLIENT_TO_SERVER, &body, )?; *send_seq += 1; socket.send_to(&packet, addr).await?; let frame = recv_response_until(socket, Duration::from_millis(700), |packet| { if packet.header.conn_id != cached.client_id { return Ok(None); } if packet.header.kind == PacketKind::AttachReject { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("resume rejected: {}", reject.reason)); } if packet.header.kind != PacketKind::ResumeOk { return Ok(None); } let Ok(plain) = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT) else { return Ok(None); }; Ok(protocol::from_body::(&plain).ok()) }) .await?; let mut next = cached.clone(); next.last_rendered_seq = frame.output_seq; Ok((frame, next)) } async fn recv_response_until(socket: &UdpSocket, wait: Duration, mut accept: F) -> Result where F: FnMut(protocol::Packet) -> Result>, { let started = Instant::now(); let mut buf = vec![0u8; 65535]; loop { let elapsed = started.elapsed(); if elapsed >= wait { return Err(anyhow!("timed out waiting for server response")); } let remaining = wait - elapsed; let (n, _) = tokio::time::timeout(remaining, socket.recv_from(&mut buf)) .await .context("timed out waiting for server response")??; let Ok(packet) = protocol::decode(&buf[..n]) else { continue; }; if let Some(response) = accept(packet)? { return Ok(response); } } } async fn run_terminal( socket: UdpSocket, mut cred: CachedCredential, first_frame: Option, predict: bool, startup_command: Option>, reconnect_timeout_secs: u64, local_forwards: Vec, dynamic_forwards: Vec, forward_only: bool, ) -> Result<()> { let _raw = if forward_only { None } else { Some(RawMode::enter()?) }; let addr = resolve_addr(&cred.udp_host, cred.udp_port)?; let mut send_seq = 2u64; let mut last_packet_at = Instant::now(); let mut status_tick = tokio::time::interval(Duration::from_secs(1)); let mut resize_tick = tokio::time::interval(Duration::from_millis(250)); let mut last_size = terminal_size(); let mut frame_buffer = FrameBuffer::default(); let mut predictor = Predictor::new(predict && cred.mode != "view-only" && !forward_only); let (forward_tx, mut forward_rx) = mpsc::channel::(1024); let _forward_keepalive = if local_forwards.is_empty() { Some(forward_tx.clone()) } else { None }; let remote_forward_tx = forward_tx.clone(); let stream_ids = Arc::new(AtomicU64::new(1)); start_local_forwards(&local_forwards, forward_tx.clone(), Arc::clone(&stream_ids)).await?; start_dynamic_forwards(&dynamic_forwards, forward_tx.clone(), stream_ids).await?; signal_background_ready(); let mut stream_writers: HashMap = HashMap::new(); let mut pending_socks_replies: HashSet = HashSet::new(); let mut opened_streams: HashSet = HashSet::new(); let mut stream_send_credit: HashMap = HashMap::new(); let mut stream_pending_data: HashMap>> = HashMap::new(); if let Some(frame) = first_frame { if !forward_only { render_frame(&frame)?; predictor.observe_output(&frame.bytes); } if frame.closed { return Ok(()); } cred.last_rendered_seq = frame.output_seq; send_ack(&socket, addr, &cred, &mut send_seq).await?; } if let Some(bytes) = startup_command && cred.mode != "view-only" && !forward_only { send_input(&socket, addr, &cred, &mut send_seq, bytes).await?; } let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::>(); let _stdin_keepalive = if forward_only { Some(stdin_tx) } else { std::thread::Builder::new() .name("dosh-stdin".to_string()) .spawn(move || { let mut stdin = std::io::stdin(); let mut buf = [0u8; 4096]; loop { match stdin.read(&mut buf) { Ok(0) => break, Ok(n) => { let _ = stdin_tx.send(buf[..n].to_vec()); } Err(_) => break, } } })?; None }; let mut recv_buf = vec![0u8; 65535]; loop { tokio::select! { biased; stdin_msg = stdin_rx.recv() => { match stdin_msg { Some(bytes) => { if bytes == [0x1d] { break; } if cred.mode != "view-only" && cred.mode != "forward-only" { predictor.observe_input(&bytes)?; send_input(&socket, addr, &cred, &mut send_seq, bytes).await?; } } None => break, } } _ = resize_tick.tick() => { if let Ok((cols, rows)) = size() && cols > 0 && rows > 0 && (cols, rows) != last_size { last_size = (cols, rows); if cred.mode != "view-only" && cred.mode != "forward-only" { send_resize(&socket, addr, &cred, &mut send_seq, cols, rows).await?; } } } recv = socket.recv_from(&mut recv_buf) => { let (n, _) = recv?; let Ok(packet) = protocol::decode(&recv_buf[..n]) else { continue; }; if packet.header.conn_id != [0u8; 16] && packet.header.conn_id != cred.client_id { continue; } match packet.header.kind { PacketKind::Frame | PacketKind::ResumeOk => { let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { 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(); } continue; }; let Ok(frame) = protocol::from_body::(&plain) else { continue; }; last_packet_at = Instant::now(); let frames = frame_buffer.accept(frame, &mut cred.last_rendered_seq); for frame in frames { predictor.clear_pending()?; if !forward_only { render_frame(&frame)?; predictor.observe_output(&frame.bytes); } if frame.closed { send_ack(&socket, addr, &cred, &mut send_seq).await?; return Ok(()); } } send_ack(&socket, addr, &cred, &mut send_seq).await?; } PacketKind::Pong => { last_packet_at = Instant::now(); } PacketKind::AttachReject => { let reject: AttachReject = protocol::from_body(&packet.body)?; if reject.reason == "unknown client" { 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(); } } else { return Err(anyhow!("server rejected session: {}", reject.reason)); } } PacketKind::StreamOpen => { let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { continue; }; let Ok(open) = protocol::from_body::(&plain) else { continue; }; last_packet_at = Instant::now(); match TcpStream::connect((open.target_host.as_str(), open.target_port)).await { Ok(stream) => { let (mut reader, writer) = stream.into_split(); stream_writers.insert(open.stream_id, writer); opened_streams.insert(open.stream_id); stream_send_credit.insert(open.stream_id, STREAM_INITIAL_WINDOW); send_stream_open_ok(&socket, addr, &cred, &mut send_seq, open.stream_id).await?; let forward_tx = remote_forward_tx.clone(); tokio::spawn(async move { let mut buf = [0u8; 16 * 1024]; loop { match reader.read(&mut buf).await { Ok(0) => break, Ok(n) => { if forward_tx .send(ForwardEvent::Data { stream_id: open.stream_id, bytes: buf[..n].to_vec(), }) .await .is_err() { return; } } Err(_) => break, } } let _ = forward_tx.send(ForwardEvent::Close { stream_id: open.stream_id, }).await; }); } Err(err) => { send_stream_open_reject( &socket, addr, &cred, &mut send_seq, open.stream_id, format!("connect {}:{}: {err}", open.target_host, open.target_port), ) .await?; } } } PacketKind::StreamOpenOk => { let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { continue; }; let Ok(ok) = protocol::from_body::(&plain) else { continue; }; last_packet_at = Instant::now(); opened_streams.insert(ok.stream_id); stream_send_credit.entry(ok.stream_id).or_insert(STREAM_INITIAL_WINDOW); if pending_socks_replies.remove(&ok.stream_id) && let Some(writer) = stream_writers.get_mut(&ok.stream_id) { let _ = writer.write_all(&[5, 0, 0, 1, 0, 0, 0, 0, 0, 0]).await; } flush_stream_pending_data( &socket, addr, &cred, &mut send_seq, ok.stream_id, &mut stream_send_credit, &mut stream_pending_data, ).await?; } PacketKind::StreamOpenReject => { let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { continue; }; let Ok(reject) = protocol::from_body::(&plain) else { continue; }; last_packet_at = Instant::now(); if pending_socks_replies.remove(&reject.stream_id) && let Some(writer) = stream_writers.get_mut(&reject.stream_id) { let _ = writer.write_all(&[5, 5, 0, 1, 0, 0, 0, 0, 0, 0]).await; } opened_streams.remove(&reject.stream_id); stream_send_credit.remove(&reject.stream_id); stream_pending_data.remove(&reject.stream_id); stream_writers.remove(&reject.stream_id); } PacketKind::StreamData => { let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { continue; }; let Ok(data) = protocol::from_body::(&plain) else { continue; }; last_packet_at = Instant::now(); let len = data.bytes.len(); if let Some(writer) = stream_writers.get_mut(&data.stream_id) { let _ = writer.write_all(&data.bytes).await; } // Always return flow-control credit so a stream whose local // writer is already gone cannot wedge the server's send window. send_stream_window_adjust( &socket, addr, &cred, &mut send_seq, data.stream_id, len, ).await?; } PacketKind::StreamWindowAdjust => { let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { continue; }; let Ok(adjust) = protocol::from_body::(&plain) else { continue; }; last_packet_at = Instant::now(); add_stream_credit(&mut stream_send_credit, adjust.stream_id, adjust.bytes); flush_stream_pending_data( &socket, addr, &cred, &mut send_seq, adjust.stream_id, &mut stream_send_credit, &mut stream_pending_data, ).await?; } PacketKind::StreamClose => { let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { continue; }; let Ok(close) = protocol::from_body::(&plain) else { continue; }; last_packet_at = Instant::now(); pending_socks_replies.remove(&close.stream_id); opened_streams.remove(&close.stream_id); stream_send_credit.remove(&close.stream_id); stream_pending_data.remove(&close.stream_id); stream_writers.remove(&close.stream_id); } _ => {} } } forward_event = forward_rx.recv() => { match forward_event { Some(ForwardEvent::Open { stream_id, target_host, target_port, writer, socks5_reply }) => { stream_writers.insert(stream_id, writer); stream_send_credit.insert(stream_id, STREAM_INITIAL_WINDOW); if socks5_reply { pending_socks_replies.insert(stream_id); } send_stream_open( &socket, addr, &cred, &mut send_seq, stream_id, target_host, target_port, ) .await?; } Some(ForwardEvent::Data { stream_id, bytes }) => { queue_or_send_stream_data( &socket, addr, &cred, &mut send_seq, stream_id, bytes, &opened_streams, &mut stream_send_credit, &mut stream_pending_data, ).await?; } Some(ForwardEvent::Close { stream_id }) => { pending_socks_replies.remove(&stream_id); opened_streams.remove(&stream_id); stream_send_credit.remove(&stream_id); stream_pending_data.remove(&stream_id); stream_writers.remove(&stream_id); send_stream_close(&socket, addr, &cred, &mut send_seq, stream_id).await?; } None if forward_only => break, None => {} } } _ = status_tick.tick() => { let stale = last_packet_at.elapsed(); if stale >= Duration::from_secs(reconnect_timeout_secs.max(1)) { 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(); } } else if stale >= Duration::from_secs(2) { let packet = protocol::encode_encrypted( PacketKind::Ping, cred.client_id, send_seq, cred.last_rendered_seq, &cred.session_key, CLIENT_TO_SERVER, b"", )?; send_seq += 1; socket.send_to(&packet, addr).await?; } } } } let _ = detach_once(&socket, &cred, send_seq).await; Ok(()) } async fn start_local_forwards( local_forwards: &[LocalForward], forward_tx: mpsc::Sender, stream_ids: Arc, ) -> Result<()> { for forward in local_forwards { let bind = format!("{}:{}", forward.bind_host, forward.listen_port); let listener = TcpListener::bind(&bind) .await .with_context(|| format!("bind local forward {bind}"))?; let forward = forward.clone(); let forward_tx = forward_tx.clone(); let stream_ids = Arc::clone(&stream_ids); tokio::spawn(async move { loop { let Ok((stream, _)) = listener.accept().await else { break; }; let stream_id = stream_ids.fetch_add(1, Ordering::Relaxed); let (mut reader, writer) = stream.into_split(); let _ = forward_tx .send(ForwardEvent::Open { stream_id, target_host: forward.target_host.clone(), target_port: forward.target_port, writer, socks5_reply: false, }) .await; let forward_tx = forward_tx.clone(); tokio::spawn(async move { let mut buf = [0u8; 16 * 1024]; loop { match reader.read(&mut buf).await { Ok(0) => break, Ok(n) => { if forward_tx .send(ForwardEvent::Data { stream_id, bytes: buf[..n].to_vec(), }) .await .is_err() { return; } } Err(_) => break, } } let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await; }); } }); } Ok(()) } async fn start_dynamic_forwards( dynamic_forwards: &[DynamicForward], forward_tx: mpsc::Sender, stream_ids: Arc, ) -> Result<()> { for forward in dynamic_forwards { let bind = format!("{}:{}", forward.bind_host, forward.listen_port); let listener = TcpListener::bind(&bind) .await .with_context(|| format!("bind dynamic forward {bind}"))?; let forward_tx = forward_tx.clone(); let stream_ids = Arc::clone(&stream_ids); tokio::spawn(async move { loop { let Ok((stream, _)) = listener.accept().await else { break; }; let forward_tx = forward_tx.clone(); let stream_ids = Arc::clone(&stream_ids); tokio::spawn(async move { let _ = handle_socks5_connect(stream, forward_tx, stream_ids).await; }); } }); } Ok(()) } async fn handle_socks5_connect( mut stream: TcpStream, forward_tx: mpsc::Sender, stream_ids: Arc, ) -> Result<()> { let mut hello = [0u8; 2]; stream.read_exact(&mut hello).await?; anyhow::ensure!(hello[0] == 5, "unsupported SOCKS version {}", hello[0]); let mut methods = vec![0u8; hello[1] as usize]; stream.read_exact(&mut methods).await?; if !methods.contains(&0) { stream.write_all(&[5, 0xff]).await?; return Ok(()); } stream.write_all(&[5, 0]).await?; let mut header = [0u8; 4]; stream.read_exact(&mut header).await?; anyhow::ensure!( header[0] == 5, "unsupported SOCKS request version {}", header[0] ); if header[1] != 1 { stream.write_all(&[5, 7, 0, 1, 0, 0, 0, 0, 0, 0]).await?; return Ok(()); } let target_host = match header[3] { 1 => { let mut raw = [0u8; 4]; stream.read_exact(&mut raw).await?; Ipv4Addr::from(raw).to_string() } 3 => { let mut len = [0u8; 1]; stream.read_exact(&mut len).await?; let mut raw = vec![0u8; len[0] as usize]; stream.read_exact(&mut raw).await?; String::from_utf8(raw).context("SOCKS domain is not valid UTF-8")? } 4 => { let mut raw = [0u8; 16]; stream.read_exact(&mut raw).await?; Ipv6Addr::from(raw).to_string() } atyp => { stream.write_all(&[5, 8, 0, 1, 0, 0, 0, 0, 0, 0]).await?; return Err(anyhow!("unsupported SOCKS address type {atyp}")); } }; let mut port = [0u8; 2]; stream.read_exact(&mut port).await?; let target_port = u16::from_be_bytes(port); let stream_id = stream_ids.fetch_add(1, Ordering::Relaxed); let (mut reader, writer) = stream.into_split(); forward_tx .send(ForwardEvent::Open { stream_id, target_host, target_port, writer, socks5_reply: true, }) .await?; tokio::spawn(async move { let mut buf = [0u8; 16 * 1024]; loop { match reader.read(&mut buf).await { Ok(0) => break, Ok(n) => { if forward_tx .send(ForwardEvent::Data { stream_id, bytes: buf[..n].to_vec(), }) .await .is_err() { return; } } Err(_) => break, } } let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await; }); Ok(()) } async fn reconnect( socket: &UdpSocket, cred: &mut CachedCredential, send_seq: &mut u64, size: (u16, u16), frame_buffer: &mut FrameBuffer, predictor: &mut Predictor, ) -> Result> { if let Ok((frame, next)) = try_live_resume(socket, cred, send_seq, size.0, size.1).await { *cred = next; frame_buffer.clear(); predictor.reset(); cred.last_rendered_seq = frame.output_seq; send_ack( socket, resolve_addr(&cred.udp_host, cred.udp_port)?, cred, send_seq, ) .await?; return Ok(Some(frame)); } let Ok((frame, next)) = try_ticket_attach(socket, cred, size.0, size.1, Vec::new()).await else { return Ok(None); }; *cred = next; *send_seq = 2; frame_buffer.clear(); predictor.reset(); cred.last_rendered_seq = frame.output_seq; send_ack( socket, resolve_addr(&cred.udp_host, cred.udp_port)?, cred, send_seq, ) .await?; Ok(Some(frame)) } async fn send_input( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, bytes: Vec, ) -> Result<()> { let body = protocol::to_body(&Input { bytes })?; let packet = protocol::encode_encrypted( PacketKind::Input, cred.client_id, *send_seq, cred.last_rendered_seq, &cred.session_key, CLIENT_TO_SERVER, &body, )?; *send_seq += 1; socket.send_to(&packet, addr).await?; Ok(()) } async fn send_stream_open( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, stream_id: u64, target_host: String, target_port: u16, ) -> Result<()> { let body = protocol::to_body(&StreamOpen { stream_id, target_host, target_port, })?; send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamOpen, &body).await } async fn send_stream_open_ok( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, stream_id: u64, ) -> Result<()> { let body = protocol::to_body(&StreamOpenOk { stream_id })?; send_stream_packet( socket, addr, cred, send_seq, PacketKind::StreamOpenOk, &body, ) .await } async fn send_stream_open_reject( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, stream_id: u64, reason: String, ) -> Result<()> { let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?; send_stream_packet( socket, addr, cred, send_seq, PacketKind::StreamOpenReject, &body, ) .await } async fn send_stream_data( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, stream_id: u64, bytes: Vec, ) -> Result<()> { let body = protocol::to_body(&StreamData { stream_id, bytes })?; send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamData, &body).await } async fn queue_or_send_stream_data( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, stream_id: u64, bytes: Vec, opened_streams: &HashSet, stream_send_credit: &mut HashMap, stream_pending_data: &mut HashMap>>, ) -> Result<()> { if !opened_streams.contains(&stream_id) || stream_send_credit.get(&stream_id).copied().unwrap_or(0) < bytes.len() || stream_pending_data .get(&stream_id) .is_some_and(|pending| !pending.is_empty()) { stream_pending_data .entry(stream_id) .or_default() .push_back(bytes); return Ok(()); } *stream_send_credit.entry(stream_id).or_default() -= bytes.len(); send_stream_data(socket, addr, cred, send_seq, stream_id, bytes).await } async fn flush_stream_pending_data( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, stream_id: u64, stream_send_credit: &mut HashMap, stream_pending_data: &mut HashMap>>, ) -> Result<()> { let Some(pending) = stream_pending_data.get_mut(&stream_id) else { return Ok(()); }; loop { let Some(bytes) = pending.front() else { break; }; let credit = stream_send_credit.get(&stream_id).copied().unwrap_or(0); if credit < bytes.len() { break; } let bytes = pending.pop_front().expect("pending front exists"); *stream_send_credit.entry(stream_id).or_default() -= bytes.len(); send_stream_data(socket, addr, cred, send_seq, stream_id, bytes).await?; } if pending.is_empty() { stream_pending_data.remove(&stream_id); } Ok(()) } fn add_stream_credit(stream_send_credit: &mut HashMap, stream_id: u64, bytes: u32) { let credit = stream_send_credit.entry(stream_id).or_default(); *credit = credit.saturating_add(bytes as usize); } async fn send_stream_window_adjust( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, stream_id: u64, bytes: usize, ) -> Result<()> { let body = protocol::to_body(&StreamWindowAdjust { stream_id, bytes: bytes.min(u32::MAX as usize) as u32, })?; send_stream_packet( socket, addr, cred, send_seq, PacketKind::StreamWindowAdjust, &body, ) .await } async fn send_stream_close( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, stream_id: u64, ) -> Result<()> { let body = protocol::to_body(&StreamClose { stream_id })?; send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamClose, &body).await } async fn send_stream_packet( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, kind: PacketKind, body: &[u8], ) -> Result<()> { let packet = protocol::encode_encrypted( kind, cred.client_id, *send_seq, cred.last_rendered_seq, &cred.session_key, CLIENT_TO_SERVER, body, )?; *send_seq += 1; socket.send_to(&packet, addr).await?; Ok(()) } struct Predictor { enabled: bool, alternate_screen: bool, pending: Vec, } impl Predictor { fn new(enabled: bool) -> Self { Self { enabled, alternate_screen: false, pending: Vec::new(), } } fn reset(&mut self) { self.pending.clear(); self.alternate_screen = false; } fn observe_input(&mut self, bytes: &[u8]) -> Result<()> { if !self.enabled || self.alternate_screen { return Ok(()); } if !Self::predictable(bytes) { self.clear_pending()?; return Ok(()); } self.pending.extend_from_slice(bytes); self.draw_pending() } fn observe_output(&mut self, bytes: &[u8]) { if contains_bytes(bytes, b"\x1b[?1049h") || contains_bytes(bytes, b"\x1b[?1047h") || contains_bytes(bytes, b"\x1b[?47h") { self.alternate_screen = true; self.pending.clear(); } if contains_bytes(bytes, b"\x1b[?1049l") || contains_bytes(bytes, b"\x1b[?1047l") || contains_bytes(bytes, b"\x1b[?47l") { self.alternate_screen = false; self.pending.clear(); } } fn predictable(bytes: &[u8]) -> bool { !bytes.is_empty() && bytes.iter().all(|byte| matches!(byte, 0x20..=0x7e)) } fn draw_pending(&self) -> Result<()> { if self.pending.is_empty() { return Ok(()); } let mut stdout = std::io::stdout(); stdout.write_all(b"\x1b7\x1b[4m")?; stdout.write_all(&self.pending)?; stdout.write_all(b"\x1b[24m\x1b8")?; stdout.flush()?; Ok(()) } fn clear_pending(&mut self) -> Result<()> { if self.pending.is_empty() { return Ok(()); } let mut stdout = std::io::stdout(); stdout.write_all(b"\x1b7")?; for _ in 0..self.pending.len() { stdout.write_all(b" ")?; } stdout.write_all(b"\x1b8")?; stdout.flush()?; self.pending.clear(); Ok(()) } } fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool { haystack .windows(needle.len()) .any(|window| window == needle) } #[derive(Default)] struct FrameBuffer { pending: BTreeMap, } impl FrameBuffer { fn clear(&mut self) { self.pending.clear(); } fn accept(&mut self, frame: Frame, last_rendered_seq: &mut u64) -> Vec { if frame.snapshot { if frame.output_seq < *last_rendered_seq { return Vec::new(); } self.pending.clear(); *last_rendered_seq = frame.output_seq; return vec![frame]; } if frame.output_seq <= *last_rendered_seq { return Vec::new(); } self.pending.entry(frame.output_seq).or_insert(frame); let mut ready = Vec::new(); while let Some(frame) = self.pending.remove(&(*last_rendered_seq + 1)) { *last_rendered_seq = frame.output_seq; ready.push(frame); } ready } } async fn send_resize( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, cols: u16, rows: u16, ) -> Result<()> { let body = protocol::to_body(&Resize { cols, rows })?; let packet = protocol::encode_encrypted( PacketKind::Resize, cred.client_id, *send_seq, cred.last_rendered_seq, &cred.session_key, CLIENT_TO_SERVER, &body, )?; *send_seq += 1; socket.send_to(&packet, addr).await?; Ok(()) } async fn send_ack( socket: &UdpSocket, addr: SocketAddr, cred: &CachedCredential, send_seq: &mut u64, ) -> Result<()> { let packet = protocol::encode_encrypted( PacketKind::Ack, cred.client_id, *send_seq, cred.last_rendered_seq, &cred.session_key, CLIENT_TO_SERVER, b"", )?; *send_seq += 1; socket.send_to(&packet, addr).await?; Ok(()) } async fn detach_once(socket: &UdpSocket, cred: &CachedCredential, seq: u64) -> Result<()> { let addr = resolve_addr(&cred.udp_host, cred.udp_port)?; let packet = protocol::encode_encrypted( PacketKind::Detach, cred.client_id, seq, cred.last_rendered_seq, &cred.session_key, CLIENT_TO_SERVER, b"", )?; socket.send_to(&packet, addr).await?; Ok(()) } fn render_frame(frame: &Frame) -> Result<()> { let mut stdout = std::io::stdout(); stdout.write_all(&frame.bytes)?; stdout.flush()?; Ok(()) } fn cache_path(root: &str, server: &str, session: &str, mode: &str) -> std::path::PathBuf { let safe = format!("{server}_{session}_{mode}") .chars() .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' }) .collect::(); expand_tilde(root).join(format!("{safe}.bin")) } fn load_cache(path: &std::path::Path) -> Result { let raw = fs::read(path)?; Ok(bincode::deserialize(&raw)?) } fn save_cache(path: &std::path::Path, cred: &CachedCredential) -> Result<()> { if let Some(parent) = path.parent() { fs::create_dir_all(parent)?; } fs::write(path, bincode::serialize(cred)?)?; Ok(()) } struct RawMode; impl RawMode { fn enter() -> Result { enable_raw_mode()?; Ok(Self) } } impl Drop for RawMode { fn drop(&mut self) { let mut stdout = std::io::stdout(); let _ = stdout.write_all(TERMINAL_CLEANUP); let _ = stdout.flush(); let _ = disable_raw_mode(); } } const TERMINAL_CLEANUP: &[u8] = concat!( "\x1b[?1l", "\x1b[0m", "\x1b[?25h", "\x1b[?1003l", "\x1b[?1002l", "\x1b[?1001l", "\x1b[?1000l", "\x1b[?1015l", "\x1b[?1006l", "\x1b[?1005l", "\x1b[?2004l", "\x1b[?1049l" ) .as_bytes(); #[cfg(test)] mod tests { use super::{ DynamicForward, FrameBuffer, LocalForward, Predictor, RemoteForward, SshConfig, auth_allows, 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, requested_env, ssh_destination_host, ssh_username, ssh_with_user, startup_command, toml_bare_key_or_quoted, }; use dosh::config::{ClientConfig, HostConfig}; use dosh::native::EnvVar; use dosh::protocol::{self, Frame, PacketKind}; use ed25519_dalek::{SigningKey, VerifyingKey}; #[test] fn parses_ssh_config_hostname() { let raw = "user palav\nhostname palav.dev\nport 22\n"; assert_eq!( parse_ssh_config(raw).hostname, Some("palav.dev".to_string()) ); } #[test] fn parses_ssh_config_hostname_case_insensitively() { let raw = "HostName 192.0.2.10\n"; assert_eq!( 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 parses_ssh_config_user_and_identity_files() { let parsed = parse_ssh_config( "user palav\nidentityfile ~/.ssh/id_ed25519\nidentityfile ~/.ssh/work\n", ); assert_eq!(parsed.user, Some("palav".to_string())); assert_eq!( parsed.identity_files, vec!["~/.ssh/id_ed25519".to_string(), "~/.ssh/work".to_string()] ); } #[test] fn parses_ssh_config_native_parity_options() { let parsed = parse_ssh_config( "userknownhostsfile ~/.ssh/known_hosts ~/.ssh/work_hosts\n\ identitiesonly yes\n\ proxyjump bastion\n\ proxycommand ssh bastion -W %h:%p\n\ identityfile none\n\ sendenv LANG LC_*\n\ setenv DOSH_MODE=native DOSH_COLOR=true\n", ); assert_eq!( parsed.user_known_hosts_file, Some("~/.ssh/known_hosts ~/.ssh/work_hosts".to_string()) ); assert!(parsed.identities_only); assert_eq!(parsed.proxy_jump, Some("bastion".to_string())); assert_eq!( parsed.proxy_command, Some("ssh bastion -W %h:%p".to_string()) ); assert_eq!( parsed.send_env, vec!["LANG".to_string(), "LC_*".to_string()] ); assert_eq!( parsed.set_env, vec![ EnvVar { name: "DOSH_MODE".to_string(), value: "native".to_string() }, EnvVar { name: "DOSH_COLOR".to_string(), value: "true".to_string() } ] ); assert!(parsed.identity_files.is_empty()); } #[test] fn requested_env_merges_config_host_and_ssh_setenv() { let mut config = ClientConfig::default(); config.send_env.clear(); config.set_env.insert("DOSH_MODE".into(), "client".into()); config.set_env.insert("DOSH_KEEP".into(), "yes".into()); let mut host = HostConfig::default(); host.set_env.insert("DOSH_MODE".into(), "host".into()); let ssh_config = SshConfig { set_env: vec![EnvVar { name: "DOSH_SSH".to_string(), value: "true".to_string(), }], ..SshConfig::default() }; assert_eq!( requested_env(&config, &host, Some(&ssh_config)), vec![ EnvVar { name: "DOSH_KEEP".to_string(), value: "yes".to_string() }, EnvVar { name: "DOSH_MODE".to_string(), value: "host".to_string() }, EnvVar { name: "DOSH_SSH".to_string(), value: "true".to_string() } ] ); } #[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"); } #[test] fn parses_user_from_ssh_destination() { assert_eq!(ssh_username("palav@example.com"), Some("palav".to_string())); assert_eq!(ssh_username("example.com"), None); } #[test] fn applies_host_config_user_to_plain_ssh_destination() { assert_eq!( ssh_with_user("example.com", Some("palav")), "palav@example.com" ); assert_eq!( ssh_with_user("root@example.com", Some("palav")), "root@example.com" ); assert_eq!( ssh_with_user("ssh://example.com", Some("palav")), "ssh://example.com" ); } #[test] fn auth_preference_accepts_auto_or_named_method() { assert!(auth_allows("native,ssh", "native")); assert!(auth_allows("native,ssh", "ssh")); assert!(auth_allows("auto", "native")); assert!(!auth_allows("native", "ssh")); } #[test] fn encrypted_identity_file_prompts_and_loads() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("id_ed25519"); let signing = SigningKey::from_bytes(&[61u8; 32]); write_encrypted_identity(&path, &signing, "let-me-in"); let mut config = ClientConfig::default(); config.identity_files.clear(); let loaded = load_first_native_identity_with_prompt( &config, Some(&path), &SshConfig::default(), |_| Ok("let-me-in".to_string()), ) .unwrap(); assert_eq!( VerifyingKey::from(&loaded).to_bytes(), VerifyingKey::from(&signing).to_bytes() ); } #[test] fn encrypted_identity_file_reports_wrong_passphrase() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("id_ed25519"); let signing = SigningKey::from_bytes(&[62u8; 32]); write_encrypted_identity(&path, &signing, "correct"); let mut config = ClientConfig::default(); config.identity_files.clear(); let err = load_first_native_identity_with_prompt( &config, Some(&path), &SshConfig::default(), |_| Ok("wrong".to_string()), ) .unwrap_err(); assert!(err.to_string().contains("no usable ssh-ed25519 identity")); } #[test] fn identities_only_skips_dosh_config_identity_fallback() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("id_ed25519"); let signing = SigningKey::from_bytes(&[63u8; 32]); write_identity(&path, &signing); let mut config = ClientConfig::default(); config.identity_files = vec![path.display().to_string()]; let ssh_config = SshConfig { identities_only: true, ..SshConfig::default() }; let err = load_first_native_identity_with_prompt(&config, None, &ssh_config, |_| { unreachable!("unencrypted key should not prompt") }) .unwrap_err(); assert!(err.to_string().contains("no usable ssh-ed25519 identity")); } #[test] fn identities_only_still_allows_ssh_config_identity_files() { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("id_ed25519"); let signing = SigningKey::from_bytes(&[64u8; 32]); write_identity(&path, &signing); let mut config = ClientConfig::default(); config.identity_files.clear(); let ssh_config = SshConfig { identity_files: vec![path.display().to_string()], identities_only: true, ..SshConfig::default() }; let loaded = load_first_native_identity_with_prompt(&config, None, &ssh_config, |_| { unreachable!("unencrypted key should not prompt") }) .unwrap(); assert_eq!( VerifyingKey::from(&loaded).to_bytes(), VerifyingKey::from(&signing).to_bytes() ); } #[tokio::test] async fn recv_response_until_ignores_unmatched_packets() { let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); let addr = receiver.local_addr().unwrap(); let stale = protocol::encode_plain(PacketKind::Ping, [1u8; 16], 1, 0, b"stale").unwrap(); let valid = protocol::encode_plain(PacketKind::Pong, [2u8; 16], 2, 0, b"valid").unwrap(); sender.send_to(&stale, addr).await.unwrap(); sender.send_to(&valid, addr).await.unwrap(); let packet = recv_response_until(&receiver, std::time::Duration::from_secs(1), |packet| { if packet.header.kind == PacketKind::Pong { Ok(Some(packet)) } else { Ok(None) } }) .await .unwrap(); assert_eq!(packet.header.kind, PacketKind::Pong); assert_eq!(packet.body, b"valid"); } fn write_identity(path: &std::path::Path, signing: &SigningKey) { let keypair = ssh_key::private::Ed25519Keypair::from(signing); let private = ssh_key::PrivateKey::new(ssh_key::private::KeypairData::from(keypair), "test").unwrap(); private .write_openssh_file(path, ssh_key::LineEnding::LF) .unwrap(); } fn write_encrypted_identity(path: &std::path::Path, signing: &SigningKey, passphrase: &str) { let keypair = ssh_key::private::Ed25519Keypair::from(signing); let private = ssh_key::PrivateKey::new(ssh_key::private::KeypairData::from(keypair), "test").unwrap(); let encrypted = private.encrypt(&mut rand::rngs::OsRng, passphrase).unwrap(); encrypted .write_openssh_file(path, ssh_key::LineEnding::LF) .unwrap(); } fn test_frame(seq: u64, snapshot: bool) -> Frame { Frame { session: "test".to_string(), output_seq: seq, bytes: format!("frame-{seq}").into_bytes(), snapshot, closed: false, } } #[test] fn frame_buffer_renders_only_contiguous_frames() { let mut buffer = FrameBuffer::default(); let mut last = 10; assert!(buffer.accept(test_frame(12, false), &mut last).is_empty()); assert_eq!(last, 10); let ready = buffer.accept(test_frame(11, false), &mut last); assert_eq!(ready.len(), 2); assert_eq!(ready[0].output_seq, 11); assert_eq!(ready[1].output_seq, 12); assert_eq!(last, 12); } #[test] fn frame_buffer_ignores_duplicate_frames() { let mut buffer = FrameBuffer::default(); let mut last = 4; assert_eq!(buffer.accept(test_frame(5, false), &mut last).len(), 1); assert!(buffer.accept(test_frame(5, false), &mut last).is_empty()); assert_eq!(last, 5); } #[test] fn frame_buffer_accepts_snapshot_as_complete_state() { let mut buffer = FrameBuffer::default(); let mut last = 10; assert!(buffer.accept(test_frame(12, false), &mut last).is_empty()); let ready = buffer.accept(test_frame(20, true), &mut last); assert_eq!(ready.len(), 1); assert_eq!(ready[0].output_seq, 20); assert_eq!(last, 20); assert!(buffer.accept(test_frame(12, false), &mut last).is_empty()); } #[test] fn frame_buffer_ignores_stale_snapshot() { let mut buffer = FrameBuffer::default(); let mut last = 20; assert!(buffer.accept(test_frame(19, true), &mut last).is_empty()); assert_eq!(last, 20); } #[test] fn predictor_only_accepts_plain_printable_input() { assert!(Predictor::predictable(b"hello world")); assert!(!Predictor::predictable(b"")); assert!(!Predictor::predictable(b"\n")); assert!(!Predictor::predictable(b"\x1b[A")); assert!(!Predictor::predictable("snowman: \u{2603}".as_bytes())); } #[test] fn predictor_tracks_alternate_screen_state() { let mut predictor = Predictor::new(true); predictor.observe_output(b"\x1b[?1049h"); assert!(predictor.alternate_screen); predictor.observe_output(b"\x1b[?1049l"); assert!(!predictor.alternate_screen); } #[test] fn startup_command_joins_args_for_remote_shell() { assert_eq!(startup_command(&[]), None); assert_eq!(startup_command(&["tm".to_string()]), Some(b"tm\n".to_vec())); assert_eq!( startup_command(&["tm".to_string(), "work".to_string()]), Some(b"tm work\n".to_vec()) ); } #[test] fn parses_local_forward_with_default_bind() { assert_eq!( parse_local_forward("8080:127.0.0.1:80").unwrap(), LocalForward { bind_host: "127.0.0.1".to_string(), listen_port: 8080, target_host: "127.0.0.1".to_string(), target_port: 80, } ); } #[test] fn parses_local_forward_with_explicit_bind() { assert_eq!( parse_local_forward("0.0.0.0:8080:localhost:8000").unwrap(), LocalForward { bind_host: "0.0.0.0".to_string(), listen_port: 8080, target_host: "localhost".to_string(), target_port: 8000, } ); } #[test] fn parses_remote_forward_with_default_bind() { assert_eq!( parse_remote_forward("9090:127.0.0.1:5432").unwrap(), RemoteForward { bind_host: "127.0.0.1".to_string(), listen_port: 9090, target_host: "127.0.0.1".to_string(), target_port: 5432, } ); } #[test] fn parses_dynamic_forward_with_default_bind() { assert_eq!( parse_dynamic_forward("1080").unwrap(), DynamicForward { bind_host: "127.0.0.1".to_string(), listen_port: 1080, } ); } }