dosh/src/bin/dosh-client.rs

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::collections::BTreeMap;
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<String>,
    #[arg(long)]
    session: Option<String>,
    #[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)]
    predict: bool,
    #[arg(long)]
    ssh_port: Option<u16>,
    #[arg(long)]
    ssh_auth_command: Option<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>,
    #[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<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();
    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 =
        if let Some(host) = args.dosh_host.clone().or_else(|| config.dosh_host.clone()) {
            host
        } else if args.local_auth {
            "127.0.0.1".to_string()
        } else {
            ssh_config_hostname(&server, ssh_port).unwrap_or_else(|err| {
                log_debug(
                    args.verbose,
                    2,
                    &format!("dosh could not resolve SSH config hostname, using {server}: {err:#}"),
                );
                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), args.predict).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), args.predict).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), args.predict).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<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 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<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");
    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_config_hostname(server: &str, ssh_port: Option<u16>) -> Result<String> {
    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)?;
    parse_ssh_config_hostname(&raw).ok_or_else(|| anyhow!("ssh -G did not print hostname"))
}

fn parse_ssh_config_hostname(raw: &str) -> Option<String> {
    raw.lines().find_map(|line| {
        let line = line.trim();
        let (key, value) = line.split_once(char::is_whitespace)?;
        if key.eq_ignore_ascii_case("hostname") {
            let value = value.trim();
            if !value.is_empty() {
                return Some(value.to_string());
            }
        }
        None
    })
}

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 {
        if packet.header.kind == PacketKind::AttachReject {
            let reject: AttachReject = protocol::from_body(&packet.body)?;
            return Err(anyhow!("attach rejected: {}", reject.reason));
        }
        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 {
        if packet.header.kind == PacketKind::AttachReject {
            let reject: AttachReject = protocol::from_body(&packet.body)?;
            return Err(anyhow!("ticket attach rejected: {}", reject.reason));
        }
        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 {
        if packet.header.kind == PacketKind::AttachReject {
            let reject: AttachReject = protocol::from_body(&packet.body)?;
            return Err(anyhow!("resume rejected: {}", reject.reason));
        }
        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>,
    predict: bool,
) -> 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_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 frame_buffer = FrameBuffer::default();
    let mut predictor = Predictor::new(predict && cred.mode != "view-only");
    if let Some(frame) = first_frame {
        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?;
    }

    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! {
            stdin_msg = stdin_rx.recv() => {
                match stdin_msg {
                    Some(bytes) => {
                        if bytes == [0x1d] {
                            break;
                        }
                        if cred.mode != "view-only" {
                            predictor.observe_input(&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();
                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)?;
                        let frames = frame_buffer.accept(frame, &mut cred.last_rendered_seq);
                        for frame in frames {
                            predictor.clear_pending()?;
                            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 => {}
                    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) {
                    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(())
}

struct Predictor {
    enabled: bool,
    alternate_screen: bool,
    pending: Vec<u8>,
}

impl Predictor {
    fn new(enabled: bool) -> Self {
        Self {
            enabled,
            alternate_screen: false,
            pending: Vec::new(),
        }
    }

    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<u64, Frame>,
}

impl FrameBuffer {
    fn accept(&mut self, frame: Frame, last_rendered_seq: &mut u64) -> Vec<Frame> {
        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::<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 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::{FrameBuffer, Predictor, parse_ssh_config_hostname, ssh_destination_host};
    use dosh::protocol::Frame;

    #[test]
    fn parses_ssh_config_hostname() {
        let raw = "user palav\nhostname palav.dev\nport 22\n";
        assert_eq!(
            parse_ssh_config_hostname(raw),
            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_hostname(raw),
            Some("192.0.2.10".to_string())
        );
    }

    #[test]
    fn strips_user_from_fallback_udp_host() {
        assert_eq!(ssh_destination_host("palav@example.com"), "example.com");
    }

    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);
    }
}