Initial Dosh implementation

src/bin/dosh-client.rs

@@ -0,0 +1,617 @@
+use anyhow::{Context, Result, anyhow};
+use base64::Engine;
+use base64::engine::general_purpose::URL_SAFE_NO_PAD;
+use clap::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, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, PacketKind,
+    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;
+use std::time::{Duration, Instant};
+use tokio::net::UdpSocket;
+use tokio::sync::mpsc;
+
+#[derive(Debug, Parser)]
+#[command(name = "dosh-client")]
+struct Args {
+    #[arg()]
+    server: Option<String>,
+    #[arg(long, default_value = "default")]
+    session: String,
+    #[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<u16>,
+    #[arg(long, default_value = "dosh-auth")]
+    ssh_auth_command: String,
+    #[arg(long)]
+    ssh_key: Option<std::path::PathBuf>,
+    #[arg(long)]
+    ssh_known_hosts: Option<std::path::PathBuf>,
+    #[arg(long)]
+    ssh_control_path: Option<std::path::PathBuf>,
+    #[arg(long)]
+    dosh_port: Option<u16>,
+    #[arg(long)]
+    dosh_host: Option<String>,
+}
+
+#[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<u8>,
+    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();
+    let server = args.server.unwrap_or(config.server);
+    let session = args.session;
+    let mode = if args.view_only || config.view_only {
+        "view-only"
+    } else {
+        "read-write"
+    }
+    .to_string();
+    let ssh_port = args.ssh_port.unwrap_or(config.ssh_port);
+    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;
+        eprintln!(
+            "dosh timing credential_lookup_end={}ms",
+            started.elapsed().as_millis()
+        );
+        match try_resume(&socket, &cached, cols, rows).await {
+            Ok((frame, cred)) => {
+                eprintln!(
+                    "dosh timing udp_resume_ready={}ms",
+                    started.elapsed().as_millis()
+                );
+                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) => {
+                eprintln!("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)) => {
+                            eprintln!(
+                                "dosh timing udp_ticket_attach_ready={}ms",
+                                started.elapsed().as_millis()
+                            );
+                            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) => {
+                            eprintln!(
+                                "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,
+            &args.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
+    };
+    eprintln!(
+        "dosh timing ssh_bootstrap={}ms",
+        bootstrap_start.elapsed().as_millis()
+    );
+
+    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)?;
+    }
+    eprintln!(
+        "dosh timing terminal_ready={}ms",
+        started.elapsed().as_millis()
+    );
+    let first = Frame {
+        session: ok.session,
+        output_seq: ok.initial_seq,
+        bytes: ok.snapshot,
+        snapshot: true,
+    };
+    if args.attach_only {
+        render_frame(&first)?;
+        detach_once(&socket, &cred, 2).await?;
+        return Ok(());
+    }
+    run_terminal(socket, cred, Some(first)).await
+}
+
+fn local_bootstrap(
+    session: &str,
+    mode: &str,
+    cols: u16,
+    rows: u16,
+    port: u16,
+    host: String,
+) -> Result<BootstrapResponse> {
+    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 ssh_bootstrap(
+    server: &str,
+    ssh_port: u16,
+    ssh_auth_command: &str,
+    ssh_key: Option<&std::path::Path>,
+    ssh_known_hosts: Option<&std::path::Path>,
+    ssh_control_path: Option<&std::path::Path>,
+    session: &str,
+    mode: &str,
+    cols: u16,
+    rows: u16,
+) -> Result<BootstrapResponse> {
+    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");
+    command.arg("-p").arg(ssh_port.to_string()).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("--nonce")
+        .arg(nonce_b64)
+        .arg("--session")
+        .arg(session)
+        .arg("--mode")
+        .arg(mode)
+        .arg("--size")
+        .arg(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<SocketAddr> {
+    (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,
+    };
+    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<Frame>,
+) -> Result<()> {
+    let _raw = RawMode::enter()?;
+    let addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
+    let mut send_seq = 2u64;
+    if let Some(frame) = first_frame {
+        render_frame(&frame)?;
+        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::<Vec<u8>>();
+    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! {
+            Some(bytes) = stdin_rx.recv() => {
+                if bytes == [0x1d] {
+                    break;
+                }
+                if cred.mode != "view-only" {
+                    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?;
+                }
+            }
+            recv = socket.recv_from(&mut recv_buf) => {
+                let (n, _) = recv?;
+                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?;
+                    }
+                    PacketKind::Pong => {}
+                    _ => {}
+                }
+            }
+        }
+    }
+    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::<String>();
+    expand_tilde(root).join(format!("{safe}.bin"))
+}
+
+fn load_cache(path: &std::path::Path) -> Result<CachedCredential> {
+    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<Self> {
+        enable_raw_mode()?;
+        Ok(Self)
+    }
+}
+
+impl Drop for RawMode {
+    fn drop(&mut self) {
+        let _ = disable_raw_mode();
+    }
+}