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_server_config}; use dosh::crypto; use dosh::protocol::{ self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, PacketKind, Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; use serde::{Deserialize, Serialize}; use std::fs; use std::io::{Read, Write}; use std::net::{SocketAddr, ToSocketAddrs}; use std::process::{Command, Stdio}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::net::UdpSocket; use tokio::sync::mpsc; #[derive(Debug, Parser)] #[command(name = "dosh-client")] struct Args { #[arg()] server: Option, #[arg(long)] session: Option, #[arg(long)] new: bool, #[arg(long)] view_only: bool, #[arg(long)] local_auth: bool, #[arg(long)] no_cache: bool, #[arg(long)] attach_only: 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, } #[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); } let server = args.server.unwrap_or(config.server); let session = select_session(args.session.as_deref(), args.new); let mode = if args.view_only || config.view_only { "view-only" } else { "read-write" } .to_string(); let ssh_port = args.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 dosh_port = args.dosh_port.unwrap_or(config.dosh_port); let cache_path = cache_path(&config.credential_cache, &server, &session, &mode); let (cols, rows) = size().unwrap_or((80, 24)); let started = Instant::now(); let target_udp_host = args .dosh_host .clone() .or_else(|| config.dosh_host.clone()) .unwrap_or_else(|| { if args.local_auth { "127.0.0.1".to_string() } else { ssh_destination_host(&server) } }); let credential = if !args.no_cache { load_cache(&cache_path).ok() } else { None }; 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()); match try_resume(&socket, &cached, cols, rows).await { Ok((frame, cred)) => { log_timing(args.verbose, "udp_resume_ready", started.elapsed()); if !args.no_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)).await; } Err(err) => { log_debug( args.verbose, 2, &format!("dosh resume failed, trying ticket attach before SSH: {err:#}"), ); if config.cache_attach_tickets && !args.no_cache { match try_ticket_attach(&socket, &cached, cols, rows).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)).await; } Err(err) => { log_debug( args.verbose, 2, &format!( "dosh ticket attach failed, falling back to SSH bootstrap: {err:#}" ), ); } } } } } } let bootstrap_start = Instant::now(); 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).await?; cred.last_rendered_seq = ok.initial_seq; if !args.no_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)).await } 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={} 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 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 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<&std::path::Path>, ssh_known_hosts: Option<&std::path::Path>, ssh_control_path: Option<&std::path::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 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 bootstrap_attach( socket: &UdpSocket, server_name: &str, bootstrap: &BootstrapResponse, cols: u16, rows: u16, ) -> Result<(AttachOk, CachedCredential)> { let addr = resolve_addr(&bootstrap.udp_host, bootstrap.udp_port)?; let req = BootstrapAttachRequest { bootstrap: bootstrap.clone(), cols, rows, }; 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 mut buf = vec![0u8; 65535]; let (n, _) = tokio::time::timeout(Duration::from_secs(5), socket.recv_from(&mut buf)).await??; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::AttachOk { return Err(anyhow!("attach rejected or unexpected response")); } let plain = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT)?; let ok: AttachOk = protocol::from_body(&plain)?; 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, ) -> 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, })?; 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 mut buf = vec![0u8; 65535]; let (n, _) = tokio::time::timeout(Duration::from_millis(700), socket.recv_from(&mut buf)).await??; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::AttachOk { return Err(anyhow!("ticket attach rejected")); } let envelope: TicketAttachOkEnvelope = protocol::from_body(&packet.body)?; 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 plain = crypto::open( &response_key, &envelope.server_nonce, b"dosh-ticket-attach-ok-v1", &envelope.ciphertext, )?; let ok: AttachOk = protocol::from_body(&plain)?; 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( 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 mut buf = vec![0u8; 65535]; let (n, _) = tokio::time::timeout(Duration::from_millis(700), socket.recv_from(&mut buf)).await??; let packet = protocol::decode(&buf[..n])?; if packet.header.kind != PacketKind::ResumeOk { return Err(anyhow!("resume rejected")); } let plain = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)?; let frame: Frame = protocol::from_body(&plain)?; let mut next = cached.clone(); next.last_rendered_seq = frame.output_seq; Ok((frame, next)) } async fn run_terminal( socket: UdpSocket, mut cred: CachedCredential, first_frame: Option, ) -> Result<()> { let _raw = 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_visible = false; 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 = size().unwrap_or((80, 24)); let mut input_normalizer = InputNormalizer::default(); if let Some(frame) = first_frame { render_frame(&frame)?; if frame.closed { return Ok(()); } cred.last_rendered_seq = frame.output_seq; send_ack(&socket, addr, &cred, &mut send_seq).await?; } let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::>(); 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, } } })?; let mut recv_buf = vec![0u8; 65535]; loop { tokio::select! { stdin_msg = stdin_rx.recv() => { match stdin_msg { Some(bytes) => { if bytes == [0x1d] { break; } if cred.mode != "view-only" { let bytes = input_normalizer.normalize(&bytes); 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?; } } None => break, } } _ = resize_tick.tick() => { if let Ok((cols, rows)) = size() && (cols, rows) != last_size { last_size = (cols, rows); if cred.mode != "view-only" { send_resize(&socket, addr, &cred, &mut send_seq, cols, rows).await?; } } } recv = socket.recv_from(&mut recv_buf) => { let (n, _) = recv?; last_packet_at = Instant::now(); if status_visible { clear_status_bar()?; status_visible = false; } let packet = protocol::decode(&recv_buf[..n])?; match packet.header.kind { PacketKind::Frame | PacketKind::ResumeOk => { let plain = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT)?; let frame: Frame = protocol::from_body(&plain)?; render_frame(&frame)?; cred.last_rendered_seq = frame.output_seq; send_ack(&socket, addr, &cred, &mut send_seq).await?; if frame.closed { break; } } PacketKind::Pong => {} PacketKind::AttachReject => { let reject: AttachReject = protocol::from_body(&packet.body)?; return Err(anyhow!("server rejected session: {}", reject.reason)); } _ => {} } } _ = status_tick.tick() => { let stale = last_packet_at.elapsed(); if stale >= Duration::from_secs(2) { draw_status_bar(&format!( "DOSH disconnected for {}s | session {} | trying UDP", stale.as_secs(), cred.session ))?; status_visible = true; 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(()) } #[derive(Default)] struct InputNormalizer { state: InputNormalizerState, } #[derive(Default)] enum InputNormalizerState { #[default] Ground, Esc, Ss3, } impl InputNormalizer { fn normalize(&mut self, bytes: &[u8]) -> Vec { let mut out = Vec::with_capacity(bytes.len()); for &byte in bytes { match self.state { InputNormalizerState::Ground => { out.push(byte); if byte == 0x1b { self.state = InputNormalizerState::Esc; } } InputNormalizerState::Esc => { if byte == b'O' { self.state = InputNormalizerState::Ss3; } else { out.push(byte); self.state = InputNormalizerState::Ground; } } InputNormalizerState::Ss3 => { if matches!(byte, b'A'..=b'D') { out.push(b'['); out.push(byte); } else { out.push(b'O'); out.push(byte); } self.state = InputNormalizerState::Ground; } } } out } } 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 draw_status_bar(message: &str) -> Result<()> { let mut stdout = std::io::stdout(); write!( stdout, "\x1b7\x1b[1;1H\x1b[44;97m\x1b[K {} \x1b[0m\x1b8", message )?; stdout.flush()?; Ok(()) } fn clear_status_bar() -> Result<()> { let mut stdout = std::io::stdout(); stdout.write_all(b"\x1b7\x1b[1;1H\x1b[K\x1b8")?; 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()?; let mut stdout = std::io::stdout(); stdout.write_all(TERMINAL_CLEANUP)?; stdout.flush()?; 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::InputNormalizer; #[test] fn normalizes_ss3_cursor_keys_to_csi() { let mut normalizer = InputNormalizer::default(); assert_eq!(normalizer.normalize(b"\x1bOA"), b"\x1b[A"); assert_eq!(normalizer.normalize(b"\x1bOB"), b"\x1b[B"); assert_eq!(normalizer.normalize(b"\x1bOC"), b"\x1b[C"); assert_eq!(normalizer.normalize(b"\x1bOD"), b"\x1b[D"); } #[test] fn preserves_split_ss3_cursor_state_between_reads() { let mut normalizer = InputNormalizer::default(); assert_eq!(normalizer.normalize(b"\x1b"), b"\x1b"); assert_eq!(normalizer.normalize(b"O"), b""); assert_eq!(normalizer.normalize(b"B"), b"[B"); } #[test] fn leaves_non_cursor_ss3_sequences_raw() { let mut normalizer = InputNormalizer::default(); assert_eq!(normalizer.normalize(b"\x1bOP"), b"\x1bOP"); } }