dosh/tests/integration_smoke.rs

use std::fs;
use std::net::UdpSocket;
use std::process::{Child, Command, Stdio};
use std::thread;
use std::time::Duration;

use dosh::auth::{build_bootstrap, load_or_create_server_secret};
use dosh::config::load_server_config;
use dosh::crypto;
use dosh::protocol::{
    self, AttachOk, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, PacketKind, Resize,
    ResumeRequest, SERVER_TO_CLIENT,
};

fn free_udp_port() -> u16 {
    let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
    socket.local_addr().unwrap().port()
}

fn write_server_config(dir: &tempfile::TempDir, port: u16) -> std::path::PathBuf {
    let config_dir = dir.path().join(".config/dosh");
    fs::create_dir_all(&config_dir).unwrap();
    let config = config_dir.join("server.toml");
    fs::write(
        &config,
        format!(
            r#"
port = {port}
bind = "127.0.0.1"
scrollback = 5000
auth_ttl_secs = 30
attach_ticket_ttl_secs = 3600
allow_attach_tickets = true
client_timeout_secs = 30
retransmit_window = 256
default_input_mode = "read-write"
prewarm_sessions = ["default"]
create_on_attach = true
shell = "/bin/sh"
sessions_dir = "{sessions}"
secret_path = "{secret}"
"#,
            sessions = dir.path().join("sessions").display(),
            secret = dir.path().join("secret").display(),
        ),
    )
    .unwrap();
    config
}

fn start_server(dir: &tempfile::TempDir, config: &std::path::Path) -> Child {
    let server = env!("CARGO_BIN_EXE_dosh-server");
    let child = Command::new(server)
        .arg("serve")
        .arg("--config")
        .arg(config)
        .env("HOME", dir.path())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .unwrap();
    thread::sleep(Duration::from_millis(500));
    child
}

fn attach_once(dir: &tempfile::TempDir, port: u16, cache: bool) -> std::process::Output {
    let client = env!("CARGO_BIN_EXE_dosh-client");
    let mut cmd = Command::new(client);
    cmd.arg("--local-auth")
        .arg("--attach-only")
        .arg("--dosh-port")
        .arg(port.to_string())
        .arg("local")
        .env("HOME", dir.path());
    if !cache {
        cmd.arg("--no-cache");
    }
    cmd.output().unwrap()
}

fn direct_attach(
    config: &std::path::Path,
    port: u16,
    mode: &str,
) -> (std::net::UdpSocket, dosh::auth::BootstrapResponse, AttachOk) {
    let config = load_server_config(Some(config.to_path_buf())).unwrap();
    let secret = load_or_create_server_secret(&config).unwrap();
    let bootstrap = build_bootstrap(
        &config,
        &secret,
        "tester".to_string(),
        "default".to_string(),
        mode.to_string(),
        (80, 24),
        crypto::random_12(),
        "127.0.0.1".to_string(),
    )
    .unwrap();
    let socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
    socket
        .set_read_timeout(Some(Duration::from_secs(2)))
        .unwrap();
    let req = BootstrapAttachRequest {
        bootstrap: bootstrap.clone(),
        cols: 80,
        rows: 24,
    };
    let packet = protocol::encode_plain(
        PacketKind::BootstrapAttachRequest,
        [0u8; 16],
        1,
        0,
        &protocol::to_body(&req).unwrap(),
    )
    .unwrap();
    socket
        .send_to(&packet, format!("127.0.0.1:{port}"))
        .unwrap();
    let mut buf = [0u8; 65535];
    let (n, _) = socket.recv_from(&mut buf).unwrap();
    let packet = protocol::decode(&buf[..n]).unwrap();
    assert_eq!(packet.header.kind, PacketKind::AttachOk);
    let plain = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT).unwrap();
    let ok: AttachOk = protocol::from_body(&plain).unwrap();
    (socket, bootstrap, ok)
}

fn send_encrypted(
    socket: &UdpSocket,
    port: u16,
    kind: PacketKind,
    client_id: [u8; 16],
    seq: u64,
    ack: u64,
    key: &[u8; 32],
    body: &[u8],
) {
    let packet =
        protocol::encode_encrypted(kind, client_id, seq, ack, key, CLIENT_TO_SERVER, body).unwrap();
    socket
        .send_to(&packet, format!("127.0.0.1:{port}"))
        .unwrap();
}

fn recv_frame(socket: &UdpSocket, key: &[u8; 32]) -> Option<(protocol::Header, Frame)> {
    let mut buf = [0u8; 65535];
    let (n, _) = socket.recv_from(&mut buf).ok()?;
    let packet = protocol::decode(&buf[..n]).ok()?;
    match packet.header.kind {
        PacketKind::Frame | PacketKind::ResumeOk => {
            let plain = protocol::decrypt_body(&packet, key, SERVER_TO_CLIENT).ok()?;
            let frame: Frame = protocol::from_body(&plain).ok()?;
            Some((packet.header, frame))
        }
        _ => None,
    }
}

fn collect_frame_text(socket: &UdpSocket, key: &[u8; 32], millis: u64) -> String {
    socket
        .set_read_timeout(Some(Duration::from_millis(100)))
        .unwrap();
    let deadline = std::time::Instant::now() + Duration::from_millis(millis);
    let mut text = String::new();
    while std::time::Instant::now() < deadline {
        if let Some((_header, frame)) = recv_frame(socket, key) {
            text.push_str(&String::from_utf8_lossy(&frame.bytes));
        }
    }
    socket
        .set_read_timeout(Some(Duration::from_secs(2)))
        .unwrap();
    text
}

#[test]
fn local_attach_only_smoke() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let mut server = start_server(&dir, &config);
    let output = attach_once(&dir, port, false);
    let _ = server.kill();
    let _ = server.wait();
    assert!(
        output.status.success(),
        "stderr={}",
        String::from_utf8_lossy(&output.stderr)
    );
    assert!(
        String::from_utf8_lossy(&output.stderr).contains("terminal_ready"),
        "stderr={}",
        String::from_utf8_lossy(&output.stderr)
    );
}

#[test]
fn ticket_attach_after_server_restart_smoke() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let mut server = start_server(&dir, &config);
    let first = attach_once(&dir, port, true);
    let _ = server.kill();
    let _ = server.wait();
    assert!(
        first.status.success(),
        "stderr={}",
        String::from_utf8_lossy(&first.stderr)
    );

    let mut server = start_server(&dir, &config);
    let second = attach_once(&dir, port, true);
    let _ = server.kill();
    let _ = server.wait();
    let stderr = String::from_utf8_lossy(&second.stderr);
    assert!(second.status.success(), "stderr={stderr}");
    assert!(
        stderr.contains("udp_ticket_attach_ready"),
        "stderr={stderr}"
    );
}

#[test]
fn view_only_input_is_rejected_by_server() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let mut server = start_server(&dir, &config);

    let (view_socket, view_bootstrap, view_ok) = direct_attach(&config, port, "view-only");
    let input = Input {
        bytes: b"echo DOSH_VIEW_ONLY_BUG\n".to_vec(),
    };
    let packet = protocol::encode_encrypted(
        PacketKind::Input,
        view_ok.client_id,
        2,
        0,
        &view_bootstrap.session_key,
        CLIENT_TO_SERVER,
        &protocol::to_body(&input).unwrap(),
    )
    .unwrap();
    view_socket
        .send_to(&packet, format!("127.0.0.1:{port}"))
        .unwrap();
    thread::sleep(Duration::from_millis(250));

    let (_rw_socket, _rw_bootstrap, rw_ok) = direct_attach(&config, port, "read-write");
    let snapshot = String::from_utf8_lossy(&rw_ok.snapshot);

    let _ = server.kill();
    let _ = server.wait();

    assert!(
        !snapshot.contains("DOSH_VIEW_ONLY_BUG"),
        "view-only input reached PTY; snapshot={snapshot:?}"
    );
}

#[test]
fn multiple_clients_share_one_session_screen() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let mut server = start_server(&dir, &config);
    let (writer_socket, writer_bootstrap, writer_ok) = direct_attach(&config, port, "read-write");
    let (reader_socket, reader_bootstrap, _reader_ok) = direct_attach(&config, port, "read-write");

    let input = Input {
        bytes: b"printf DOSH_MULTI_CLIENT\\n\n".to_vec(),
    };
    send_encrypted(
        &writer_socket,
        port,
        PacketKind::Input,
        writer_ok.client_id,
        2,
        0,
        &writer_bootstrap.session_key,
        &protocol::to_body(&input).unwrap(),
    );
    let text = collect_frame_text(&reader_socket, &reader_bootstrap.session_key, 2000);

    let _ = server.kill();
    let _ = server.wait();

    assert!(
        text.contains("DOSH_MULTI_CLIENT"),
        "expected second client to see first client's output, got {text:?}"
    );
}

#[test]
fn server_retransmits_unacked_output_frames() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let mut server = start_server(&dir, &config);
    let (socket, bootstrap, ok) = direct_attach(&config, port, "read-write");

    let input = Input {
        bytes: b"printf DOSH_RETRANSMIT\\n\n".to_vec(),
    };
    send_encrypted(
        &socket,
        port,
        PacketKind::Input,
        ok.client_id,
        2,
        0,
        &bootstrap.session_key,
        &protocol::to_body(&input).unwrap(),
    );

    let mut seen = Vec::new();
    let mut duplicate = None;
    let deadline = std::time::Instant::now() + Duration::from_secs(3);
    while std::time::Instant::now() < deadline {
        if let Some((header, frame)) = recv_frame(&socket, &bootstrap.session_key) {
            if String::from_utf8_lossy(&frame.bytes).contains("DOSH_RETRANSMIT") {
                if seen
                    .iter()
                    .any(|(_, output_seq)| *output_seq == frame.output_seq)
                {
                    duplicate = Some((header.seq, frame.output_seq));
                    break;
                }
                seen.push((header.seq, frame.output_seq));
            }
        }
    }

    let _ = server.kill();
    let _ = server.wait();

    assert!(
        duplicate.is_some(),
        "expected retransmitted same output seq, seen={seen:?}"
    );
}

#[test]
fn resize_updates_pty_size() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let mut server = start_server(&dir, &config);
    let (socket, bootstrap, ok) = direct_attach(&config, port, "read-write");

    let resize = Resize {
        cols: 100,
        rows: 10,
    };
    send_encrypted(
        &socket,
        port,
        PacketKind::Resize,
        ok.client_id,
        2,
        0,
        &bootstrap.session_key,
        &protocol::to_body(&resize).unwrap(),
    );
    thread::sleep(Duration::from_millis(100));
    let input = Input {
        bytes: b"stty size\n".to_vec(),
    };
    send_encrypted(
        &socket,
        port,
        PacketKind::Input,
        ok.client_id,
        3,
        0,
        &bootstrap.session_key,
        &protocol::to_body(&input).unwrap(),
    );
    let text = collect_frame_text(&socket, &bootstrap.session_key, 2000);

    let _ = server.kill();
    let _ = server.wait();

    assert!(
        text.contains("10 100"),
        "expected resized pty to report 10 100, got {text:?}"
    );
}

#[test]
fn resume_updates_udp_endpoint_for_roaming() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let mut server = start_server(&dir, &config);
    let (_old_socket, bootstrap, ok) = direct_attach(&config, port, "read-write");
    let new_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
    new_socket
        .set_read_timeout(Some(Duration::from_secs(2)))
        .unwrap();

    let resume = ResumeRequest {
        session: "default".to_string(),
        last_rendered_seq: ok.initial_seq,
        cols: 80,
        rows: 24,
    };
    send_encrypted(
        &new_socket,
        port,
        PacketKind::ResumeRequest,
        ok.client_id,
        1,
        0,
        &bootstrap.session_key,
        &protocol::to_body(&resume).unwrap(),
    );
    let (_header, resume_frame) =
        recv_frame(&new_socket, &bootstrap.session_key).expect("resume response on new socket");
    assert!(resume_frame.snapshot);

    let input = Input {
        bytes: b"printf DOSH_ROAM\\n\n".to_vec(),
    };
    send_encrypted(
        &new_socket,
        port,
        PacketKind::Input,
        ok.client_id,
        2,
        resume_frame.output_seq,
        &bootstrap.session_key,
        &protocol::to_body(&input).unwrap(),
    );
    let text = collect_frame_text(&new_socket, &bootstrap.session_key, 2000);

    let _ = server.kill();
    let _ = server.wait();

    assert!(
        text.contains("DOSH_ROAM"),
        "expected output on resumed socket, got {text:?}"
    );
}