dosh/tests/hostile_network.rs

//! Hostile-network integration tests (Track B, spec milestone 5 / §16).
//!
//! These spin up a real `dosh-server` process bound to 127.0.0.1 on a free port
//! in a temp HOME, and drive the wire protocol directly from the test (mirroring
//! the `direct_attach` pattern in tests/integration_smoke.rs). Between the test
//! "client" and the server sits an in-process UDP relay/shim that can drop,
//! reorder, and duplicate datagrams, and can switch the client's source address
//! (by rebinding its upstream socket) mid-session.
//!
//! Assertions, mapped to §16 verification items:
//!   * "Stale encrypted packets after reconnect are ignored, not fatal"
//!     -> session survives loss/reorder; stale packets after resume don't kill it.
//!   * "Replayed transport packets are rejected" / "no double-apply"
//!     -> duplicated & replayed Input is applied at most once.
//!   * "Client IP/port change preserves the session"
//!     -> after the relay rebinds its upstream socket the session keeps working.
//!
//! Determinism: the relay's drop/reorder/dup behavior is driven by a fixed-seed
//! PRNG and by explicit one-shot toggles, never by wall-clock timing, so the
//! tests are reproducible and fast.
#![allow(
    clippy::collapsible_if,
    clippy::explicit_counter_loop,
    clippy::single_match
)]

use std::fs;
use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, UdpSocket};
use std::process::{Child, Command, Stdio};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::thread;
use std::time::{Duration, Instant};

use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use dosh::auth::{BootstrapResponse, build_bootstrap, load_or_create_server_secret};
use dosh::config::load_server_config;
use dosh::crypto;
use dosh::native::{self, ForwardingKind, ForwardingRequest, NativeAuthOk, NativeClientHello};
use dosh::protocol::{
    self, AttachOk, AttachReject, CLIENT_TO_SERVER, Frame, Header, Input, NativeAuthOkBody,
    NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind, ResumeRequest,
    SERVER_TO_CLIENT, StreamData, StreamOpen, StreamOpenOk, StreamWindowAdjust,
};
use ed25519_dalek::{SigningKey, VerifyingKey};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};

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}"
host_key = "{host_key}"
authorized_keys = ["{authorized_keys}"]
persist_sessions = false
"#,
            sessions = dir.path().join("sessions").display(),
            secret = dir.path().join("secret").display(),
            host_key = dir.path().join("host_key").display(),
            authorized_keys = dir.path().join("authorized_keys").display(),
        ),
    )
    .unwrap();
    config
}

fn authorize_ed25519_key(dir: &tempfile::TempDir, signing_key: &SigningKey) {
    let verifying = VerifyingKey::from(signing_key);
    let mut blob = Vec::new();
    write_ssh_string(&mut blob, b"ssh-ed25519");
    write_ssh_string(&mut blob, verifying.as_bytes());
    fs::write(
        dir.path().join("authorized_keys"),
        format!("ssh-ed25519 {} hostile-test\n", STANDARD.encode(blob)),
    )
    .unwrap();
}

fn write_ssh_string(out: &mut Vec<u8>, value: &[u8]) {
    out.extend_from_slice(&(value.len() as u32).to_be_bytes());
    out.extend_from_slice(value);
}

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
}

/// Knobs that control how the relay mishandles datagrams. All knobs default to
/// pass-through. Each is read on every forwarded packet.
struct RelayControls {
    /// Probability [0,100] that an upstream (client->server) packet is dropped.
    drop_c2s_percent: AtomicU64,
    /// Probability [0,100] that a downstream (server->client) packet is dropped.
    drop_s2c_percent: AtomicU64,
    /// Probability [0,100] that a packet (either direction) is duplicated.
    dup_percent: AtomicU64,
    /// When set, the relay holds one packet back and releases it after the next
    /// one passes, producing a 2-1 reorder. Used as a deterministic toggle.
    reorder_next: AtomicBool,
    /// Count of upstream packets the relay observed (for assertions).
    c2s_count: AtomicU64,
    /// Count of downstream packets the relay observed.
    s2c_count: AtomicU64,
    shutdown: AtomicBool,
}

impl RelayControls {
    fn new() -> Self {
        Self {
            drop_c2s_percent: AtomicU64::new(0),
            drop_s2c_percent: AtomicU64::new(0),
            dup_percent: AtomicU64::new(0),
            reorder_next: AtomicBool::new(false),
            c2s_count: AtomicU64::new(0),
            s2c_count: AtomicU64::new(0),
            shutdown: AtomicBool::new(false),
        }
    }
}

/// Command sent to the relay from the test thread.
enum RelayCmd {
    /// Rebind the upstream (toward-server) socket to a fresh local address,
    /// simulating a client NAT rebind / IP-port change. Replies the new addr.
    RebindUpstream(Sender<SocketAddr>),
}

/// A UDP relay sitting between the test client and the real server.
///
/// `front` is the address the test client sends to. The relay forwards each
/// client datagram to the server over `upstream`, remembering the client's
/// address so server replies can be returned. `upstream` can be replaced on
/// demand to simulate a client source-address change as the server observes it.
struct Relay {
    front_addr: SocketAddr,
    controls: Arc<RelayControls>,
    cmd_tx: Sender<RelayCmd>,
    handle: Option<thread::JoinHandle<()>>,
}

impl Relay {
    fn spawn(server_port: u16, seed: u64) -> Self {
        let front = UdpSocket::bind("127.0.0.1:0").unwrap();
        front
            .set_read_timeout(Some(Duration::from_millis(20)))
            .unwrap();
        let front_addr = front.local_addr().unwrap();
        let server_addr: SocketAddr = format!("127.0.0.1:{server_port}").parse().unwrap();
        let controls = Arc::new(RelayControls::new());
        let (cmd_tx, cmd_rx) = channel::<RelayCmd>();

        let thread_controls = Arc::clone(&controls);
        let handle = thread::spawn(move || {
            relay_loop(front, server_addr, thread_controls, cmd_rx, seed);
        });

        Self {
            front_addr,
            controls,
            cmd_tx,
            handle: Some(handle),
        }
    }

    fn front_addr(&self) -> SocketAddr {
        self.front_addr
    }

    fn set_drop_c2s(&self, percent: u64) {
        self.controls
            .drop_c2s_percent
            .store(percent, Ordering::SeqCst);
    }

    fn set_drop_s2c(&self, percent: u64) {
        self.controls
            .drop_s2c_percent
            .store(percent, Ordering::SeqCst);
    }

    fn set_dup(&self, percent: u64) {
        self.controls.dup_percent.store(percent, Ordering::SeqCst);
    }

    fn arm_reorder(&self) {
        self.controls.reorder_next.store(true, Ordering::SeqCst);
    }

    fn clear_impairments(&self) {
        self.set_drop_c2s(0);
        self.set_drop_s2c(0);
        self.set_dup(0);
    }

    /// Rebind the relay's upstream socket; the server will see traffic from a
    /// new source address afterward. Returns the new upstream local address.
    fn rebind_upstream(&self) -> SocketAddr {
        let (tx, rx) = channel();
        self.cmd_tx.send(RelayCmd::RebindUpstream(tx)).unwrap();
        rx.recv_timeout(Duration::from_secs(2))
            .expect("relay rebind ack")
    }
}

impl Drop for Relay {
    fn drop(&mut self) {
        self.controls.shutdown.store(true, Ordering::SeqCst);
        if let Some(handle) = self.handle.take() {
            let _ = handle.join();
        }
    }
}

fn relay_loop(
    front: UdpSocket,
    server_addr: SocketAddr,
    controls: Arc<RelayControls>,
    cmd_rx: Receiver<RelayCmd>,
    seed: u64,
) {
    let mut upstream = new_upstream();
    let mut client_addr: Option<SocketAddr> = None;
    let mut rng = StdRng::seed_from_u64(seed);
    let mut held_c2s: Option<Vec<u8>> = None;
    let mut held_s2c: Option<Vec<u8>> = None;
    let mut buf = [0u8; 65535];

    loop {
        if controls.shutdown.load(Ordering::SeqCst) {
            return;
        }

        // Process any pending control commands.
        while let Ok(cmd) = cmd_rx.try_recv() {
            match cmd {
                RelayCmd::RebindUpstream(reply) => {
                    upstream = new_upstream();
                    let _ = reply.send(upstream.local_addr().unwrap());
                }
            }
        }

        // Client -> server.
        match front.recv_from(&mut buf) {
            Ok((n, src)) => {
                client_addr = Some(src);
                controls.c2s_count.fetch_add(1, Ordering::SeqCst);
                let packet = buf[..n].to_vec();
                let drop_pct = controls.drop_c2s_percent.load(Ordering::SeqCst);
                if drop_pct == 0 || rng.gen_range(0..100) >= drop_pct {
                    forward_with_effects(
                        &upstream,
                        server_addr,
                        packet,
                        &controls,
                        &mut rng,
                        &mut held_c2s,
                    );
                }
            }
            Err(ref e)
                if e.kind() == std::io::ErrorKind::WouldBlock
                    || e.kind() == std::io::ErrorKind::TimedOut => {}
            Err(_) => return,
        }

        // Server -> client.
        match upstream.recv_from(&mut buf) {
            Ok((n, _src)) => {
                controls.s2c_count.fetch_add(1, Ordering::SeqCst);
                if let Some(dst) = client_addr {
                    let packet = buf[..n].to_vec();
                    let drop_pct = controls.drop_s2c_percent.load(Ordering::SeqCst);
                    if drop_pct == 0 || rng.gen_range(0..100) >= drop_pct {
                        forward_with_effects(
                            &front,
                            dst,
                            packet,
                            &controls,
                            &mut rng,
                            &mut held_s2c,
                        );
                    }
                }
            }
            Err(ref e)
                if e.kind() == std::io::ErrorKind::WouldBlock
                    || e.kind() == std::io::ErrorKind::TimedOut => {}
            Err(_) => return,
        }
    }
}

fn new_upstream() -> UdpSocket {
    let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
    socket
        .set_read_timeout(Some(Duration::from_millis(20)))
        .unwrap();
    socket
}

/// Forward a packet to `dst` over `out`, applying duplication and reorder.
fn forward_with_effects(
    out: &UdpSocket,
    dst: SocketAddr,
    packet: Vec<u8>,
    controls: &RelayControls,
    rng: &mut StdRng,
    held: &mut Option<Vec<u8>>,
) {
    // Reorder: if armed, hold this packet and release the previously held one
    // afterward (so two consecutive packets swap order).
    if controls.reorder_next.swap(false, Ordering::SeqCst) {
        if let Some(prev) = held.take() {
            let _ = out.send_to(&packet, dst);
            let _ = out.send_to(&prev, dst);
            return;
        }
        *held = Some(packet);
        return;
    }
    if let Some(prev) = held.take() {
        let _ = out.send_to(&prev, dst);
    }

    let _ = out.send_to(&packet, dst);

    let dup_pct = controls.dup_percent.load(Ordering::SeqCst);
    if dup_pct > 0 && rng.gen_range(0..100) < dup_pct {
        let _ = out.send_to(&packet, dst);
    }
}

/// Build a bootstrap and attach through the relay, returning the client socket,
/// the bootstrap (for the session key), and the AttachOk.
fn attach_through_relay(
    config: &std::path::Path,
    relay: &Relay,
) -> (UdpSocket, 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(),
        "read-write".to_string(),
        (80, 24),
        crypto::random_12(),
        "127.0.0.1".to_string(),
    )
    .unwrap();
    let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
    socket
        .set_read_timeout(Some(Duration::from_millis(200)))
        .unwrap();
    let req = protocol::BootstrapAttachRequest {
        bootstrap: bootstrap.clone(),
        cols: 80,
        rows: 24,
        requested_env: Vec::new(),
    };
    let packet = protocol::encode_plain(
        PacketKind::BootstrapAttachRequest,
        [0u8; 16],
        1,
        0,
        &protocol::to_body(&req).unwrap(),
    )
    .unwrap();

    // Retry the attach request a few times in case the relay drops it.
    let deadline = Instant::now() + Duration::from_secs(5);
    loop {
        socket.send_to(&packet, relay.front_addr()).unwrap();
        let mut buf = [0u8; 65535];
        match socket.recv_from(&mut buf) {
            Ok((n, _)) => {
                if let Ok(decoded) = protocol::decode(&buf[..n]) {
                    if decoded.header.kind == PacketKind::AttachOk {
                        let plain = protocol::decrypt_body(
                            &decoded,
                            &bootstrap.session_key,
                            SERVER_TO_CLIENT,
                        )
                        .unwrap();
                        let ok: AttachOk = protocol::from_body(&plain).unwrap();
                        return (socket, bootstrap, ok);
                    }
                }
            }
            Err(_) => {}
        }
        if Instant::now() > deadline {
            panic!("attach through relay timed out");
        }
    }
}

struct NativeAttached {
    socket: UdpSocket,
    ok: NativeAuthOk,
}

fn native_attach_through_relay(
    relay: &Relay,
    signing_key: &SigningKey,
    requested_forwardings: Vec<ForwardingRequest>,
) -> NativeAttached {
    let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
    socket
        .set_read_timeout(Some(Duration::from_millis(500)))
        .unwrap();
    let (client_secret, client_public) = native::generate_native_ephemeral();
    let hello = NativeClientHello {
        protocol_version: native::NATIVE_PROTOCOL_VERSION,
        client_random: crypto::random_32(),
        client_ephemeral_public: client_public,
        requested_host: "127.0.0.1".to_string(),
        requested_user: std::env::var("USER").unwrap_or_else(|_| "unknown".to_string()),
        requested_session: "default".to_string(),
        requested_mode: "read-write".to_string(),
        terminal_size: (80, 24),
        supported_aead: vec!["chacha20poly1305".to_string()],
        supported_user_key_algorithms: native::supported_user_key_algorithms(),
        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(),
        })
        .unwrap(),
    )
    .unwrap();
    let hello_deadline = Instant::now() + Duration::from_secs(5);
    let server_hello = loop {
        socket.send_to(&packet, relay.front_addr()).unwrap();
        let mut buf = [0u8; 65535];
        match socket.recv_from(&mut buf) {
            Ok((n, _)) => {
                let packet = protocol::decode(&buf[..n]).unwrap();
                match packet.header.kind {
                    PacketKind::NativeServerHello => {
                        let body: NativeServerHelloBody =
                            protocol::from_body(&packet.body).unwrap();
                        native::verify_server_hello(&hello, &body.hello).unwrap();
                        break body.hello;
                    }
                    PacketKind::AttachReject => {
                        let reject: AttachReject = protocol::from_body(&packet.body).unwrap();
                        panic!("native hello rejected: {}", reject.reason);
                    }
                    kind => panic!("unexpected native hello response: {kind:?}"),
                }
            }
            Err(_) if Instant::now() < hello_deadline => {}
            Err(err) => panic!("native server hello timed out: {err}"),
        }
    };

    let session_key = native::derive_native_session_key(
        &client_secret,
        server_hello.server_ephemeral_public,
        &hello,
        &server_hello,
    )
    .unwrap();
    let auth =
        native::sign_user_auth(signing_key, &hello, &server_hello, requested_forwardings).unwrap();
    let mut pending_id = [0u8; 16];
    pending_id.copy_from_slice(&server_hello.auth_challenge[..16]);
    let auth_packet = protocol::encode_encrypted(
        PacketKind::NativeUserAuth,
        pending_id,
        2,
        1,
        &session_key,
        CLIENT_TO_SERVER,
        &protocol::to_body(&NativeUserAuthBody { auth }).unwrap(),
    )
    .unwrap();
    let auth_deadline = Instant::now() + Duration::from_secs(5);
    let ok = loop {
        socket.send_to(&auth_packet, relay.front_addr()).unwrap();
        let mut buf = [0u8; 65535];
        match socket.recv_from(&mut buf) {
            Ok((n, _)) => {
                let packet = protocol::decode(&buf[..n]).unwrap();
                match packet.header.kind {
                    PacketKind::NativeAuthOk => {
                        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT)
                            .unwrap();
                        let body: NativeAuthOkBody = protocol::from_body(&plain).unwrap();
                        break body.ok;
                    }
                    PacketKind::AttachReject => {
                        let reject: AttachReject = protocol::from_body(&packet.body).unwrap();
                        panic!("native auth rejected: {}", reject.reason);
                    }
                    kind => panic!("unexpected native auth response: {kind:?}"),
                }
            }
            Err(_) if Instant::now() < auth_deadline => {}
            Err(err) => panic!("native auth ok timed out: {err}"),
        }
    };
    NativeAttached { socket, ok }
}

fn send_input(
    socket: &UdpSocket,
    relay: &Relay,
    client_id: [u8; 16],
    seq: u64,
    ack: u64,
    key: &[u8; 32],
    text: &[u8],
) {
    let input = Input {
        bytes: text.to_vec(),
    };
    let packet = protocol::encode_encrypted(
        PacketKind::Input,
        client_id,
        seq,
        ack,
        key,
        CLIENT_TO_SERVER,
        &protocol::to_body(&input).unwrap(),
    )
    .unwrap();
    socket.send_to(&packet, relay.front_addr()).unwrap();
}

fn send_raw(socket: &UdpSocket, relay: &Relay, packet: &[u8]) {
    socket.send_to(packet, relay.front_addr()).unwrap();
}

fn send_stream_packet(
    socket: &UdpSocket,
    relay: &Relay,
    ok: &NativeAuthOk,
    kind: PacketKind,
    seq: u64,
    ack: u64,
    body: Vec<u8>,
) {
    let packet = protocol::encode_encrypted(
        kind,
        ok.client_id,
        seq,
        ack,
        &ok.session_key,
        CLIENT_TO_SERVER,
        &body,
    )
    .unwrap();
    socket.send_to(&packet, relay.front_addr()).unwrap();
}

fn recv_stream_packet<T: serde::de::DeserializeOwned>(
    socket: &UdpSocket,
    key: &[u8; 32],
    kind: PacketKind,
) -> Option<(Header, T)> {
    let mut buf = [0u8; 65535];
    let (n, _) = socket.recv_from(&mut buf).ok()?;
    let packet = protocol::decode(&buf[..n]).ok()?;
    if packet.header.kind != kind {
        return None;
    }
    let plain = protocol::decrypt_body(&packet, key, SERVER_TO_CLIENT).ok()?;
    let body = protocol::from_body(&plain).ok()?;
    Some((packet.header, body))
}

fn wait_for_stream_open_ok(
    socket: &UdpSocket,
    key: &[u8; 32],
    stream_id: u64,
    millis: u64,
) -> Option<Header> {
    let deadline = Instant::now() + Duration::from_millis(millis);
    while Instant::now() < deadline {
        if let Some((header, ok)) =
            recv_stream_packet::<StreamOpenOk>(socket, key, PacketKind::StreamOpenOk)
        {
            if ok.stream_id == stream_id {
                return Some(header);
            }
        }
    }
    None
}

fn wait_for_stream_data(
    socket: &UdpSocket,
    key: &[u8; 32],
    stream_id: u64,
    millis: u64,
) -> Option<(Header, StreamData)> {
    let deadline = Instant::now() + Duration::from_millis(millis);
    while Instant::now() < deadline {
        if let Some((header, data)) =
            recv_stream_packet::<StreamData>(socket, key, PacketKind::StreamData)
        {
            if data.stream_id == stream_id {
                return Some((header, data));
            }
        }
    }
    None
}

fn start_tcp_collector() -> (u16, Receiver<Vec<u8>>) {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let port = listener.local_addr().unwrap().port();
    let (tx, rx) = channel();
    thread::spawn(move || {
        let Ok((mut stream, _)) = listener.accept() else {
            return;
        };
        stream
            .set_read_timeout(Some(Duration::from_millis(100)))
            .unwrap();
        let deadline = Instant::now() + Duration::from_secs(5);
        let mut buf = [0u8; 1024];
        while Instant::now() < deadline {
            match stream.read(&mut buf) {
                Ok(0) => return,
                Ok(n) => {
                    let _ = tx.send(buf[..n].to_vec());
                }
                Err(ref err)
                    if err.kind() == std::io::ErrorKind::WouldBlock
                        || err.kind() == std::io::ErrorKind::TimedOut => {}
                Err(_) => return,
            }
        }
    });
    (port, rx)
}

fn start_tcp_echo() -> u16 {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let port = listener.local_addr().unwrap().port();
    thread::spawn(move || {
        let Ok((mut stream, _)) = listener.accept() else {
            return;
        };
        stream
            .set_read_timeout(Some(Duration::from_millis(100)))
            .unwrap();
        let deadline = Instant::now() + Duration::from_secs(5);
        let mut buf = [0u8; 1024];
        while Instant::now() < deadline {
            match stream.read(&mut buf) {
                Ok(0) => return,
                Ok(n) => {
                    let _ = stream.write_all(&buf[..n]);
                }
                Err(ref err)
                    if err.kind() == std::io::ErrorKind::WouldBlock
                        || err.kind() == std::io::ErrorKind::TimedOut => {}
                Err(_) => return,
            }
        }
    });
    port
}

fn collect_tcp(rx: &Receiver<Vec<u8>>, millis: u64) -> Vec<u8> {
    let deadline = Instant::now() + Duration::from_millis(millis);
    let mut out = Vec::new();
    while Instant::now() < deadline {
        match rx.recv_timeout(Duration::from_millis(50)) {
            Ok(chunk) => out.extend_from_slice(&chunk),
            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
            Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break,
        }
    }
    out
}

fn recv_frame(socket: &UdpSocket, key: &[u8; 32]) -> Option<(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,
    }
}

/// Collect terminal output text for up to `millis`, returning all decoded frame
/// bytes concatenated.
fn collect_text(socket: &UdpSocket, key: &[u8; 32], millis: u64) -> String {
    let prev = socket.read_timeout().unwrap();
    socket
        .set_read_timeout(Some(Duration::from_millis(100)))
        .unwrap();
    let deadline = Instant::now() + Duration::from_millis(millis);
    let mut text = String::new();
    while Instant::now() < deadline {
        if let Some((_h, frame)) = recv_frame(socket, key) {
            text.push_str(&String::from_utf8_lossy(&frame.bytes));
        }
    }
    socket.set_read_timeout(prev).unwrap();
    text
}

/// Wait until terminal output containing `needle` is observed, retrying for up
/// to `millis`. Returns true if seen.
fn wait_for_text(socket: &UdpSocket, key: &[u8; 32], needle: &str, millis: u64) -> bool {
    let deadline = Instant::now() + Duration::from_millis(millis);
    let mut acc = String::new();
    while Instant::now() < deadline {
        acc.push_str(&collect_text(socket, key, 200));
        if acc.contains(needle) {
            return true;
        }
    }
    acc.contains(needle)
}

fn wait_for_text_with_ack(
    socket: &UdpSocket,
    relay: &Relay,
    client_id: [u8; 16],
    seq: &mut u64,
    key: &[u8; 32],
    needle: &str,
    millis: u64,
) -> bool {
    let prev = socket.read_timeout().unwrap();
    socket
        .set_read_timeout(Some(Duration::from_millis(100)))
        .unwrap();
    let deadline = Instant::now() + Duration::from_millis(millis);
    let mut acc = String::new();
    while Instant::now() < deadline {
        if let Some((_header, frame)) = recv_frame(socket, key) {
            acc.push_str(&String::from_utf8_lossy(&frame.bytes));
            let ack = protocol::encode_encrypted(
                PacketKind::Ack,
                client_id,
                *seq,
                frame.output_seq,
                key,
                CLIENT_TO_SERVER,
                b"",
            )
            .unwrap();
            *seq += 1;
            socket.send_to(&ack, relay.front_addr()).unwrap();
            if acc.contains(needle) {
                socket.set_read_timeout(prev).unwrap();
                return true;
            }
        }
    }
    socket.set_read_timeout(prev).unwrap();
    acc.contains(needle)
}

#[test]
fn session_survives_packet_loss_and_reorder() {
    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 relay = Relay::spawn(port, 0x1055_u64 ^ 0x1111);

    let (socket, bootstrap, ok) = attach_through_relay(&config, &relay);

    // Introduce 40% loss in both directions and frequent duplication, plus
    // reorder on the input flight. The server retransmits unacked frames and
    // the client retransmits input, so the command must still land.
    relay.set_drop_c2s(40);
    relay.set_drop_s2c(40);
    relay.set_dup(30);

    let mut seq = 2u64;
    let mut seen = false;
    // Send the same logical command several times with monotonically rising
    // sequence numbers (as a real client retransmitting would), interleaving a
    // reorder toggle, until the output is observed despite the lossy link.
    for attempt in 0..20 {
        if attempt % 3 == 0 {
            relay.arm_reorder();
        }
        send_input(
            &socket,
            &relay,
            ok.client_id,
            seq,
            0,
            &bootstrap.session_key,
            b"printf DOSH_LOSSY_OK\\n\n",
        );
        seq += 1;
        if wait_for_text(&socket, &bootstrap.session_key, "DOSH_LOSSY_OK", 400) {
            seen = true;
            break;
        }
    }

    relay.clear_impairments();
    drop(relay);
    let _ = server.kill();
    let _ = server.wait();

    assert!(
        seen,
        "terminal output never arrived across a lossy/reordering link"
    );
}

#[test]
#[ignore = "30-minute hostile-network TUI soak; run with `DOSH_BADNET_SOAK_SECONDS=1800 cargo test --test hostile_network bad_network_tui_work_soak_30m -- --ignored --nocapture`"]
fn bad_network_tui_work_soak_30m() {
    let soak_secs = std::env::var("DOSH_BADNET_SOAK_SECONDS")
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .unwrap_or(30 * 60);
    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 relay = Relay::spawn(port, 0xBADC0DEu64);
    let (socket, bootstrap, ok) = attach_through_relay(&config, &relay);

    relay.set_drop_c2s(15);
    relay.set_drop_s2c(20);
    relay.set_dup(10);

    let deadline = Instant::now() + Duration::from_secs(soak_secs);
    let mut seq = 2u64;
    let mut iteration = 0u64;
    while Instant::now() < deadline {
        if iteration % 2 == 0 {
            relay.arm_reorder();
        }
        if iteration > 0 && iteration % 7 == 0 {
            let _ = relay.rebind_upstream();
        }
        if iteration > 0 && iteration % 11 == 0 {
            relay.set_drop_s2c(100);
            thread::sleep(Duration::from_millis(750));
            relay.set_drop_s2c(20);
        }

        let marker = format!("DOSH_BADNET_TUI_{iteration}");
        let command = format!(
            "stty -echo; \
             printf '\\033[?1049h\\033[?2026h\\033[?25l'; \
             for i in 1 2 3 4 5 6; do \
               printf '\\033[%s;4H{} ⠀⠁⠃⠇⡇⣇⣧⣷⣿ █▇▆▅▄▃▂▁ %s\\033[0m' \"$i\" \"$i\"; \
             done; \
             printf '\\033[?25h\\033[?2026l\\033[?1049l'\n",
            marker
        );

        let step_deadline = Instant::now() + Duration::from_secs(8);
        let mut seen = false;
        let mut attempts = 0;
        while Instant::now() < step_deadline && attempts < 12 {
            send_input(
                &socket,
                &relay,
                ok.client_id,
                seq,
                0,
                &bootstrap.session_key,
                command.as_bytes(),
            );
            seq += 1;
            attempts += 1;
            if wait_for_text_with_ack(
                &socket,
                &relay,
                ok.client_id,
                &mut seq,
                &bootstrap.session_key,
                &marker,
                500,
            ) {
                seen = true;
                break;
            }
        }
        assert!(
            seen,
            "TUI marker {marker} did not arrive during hostile-network soak"
        );

        iteration += 1;
        thread::sleep(Duration::from_millis(500));
    }

    relay.clear_impairments();
    drop(relay);
    let _ = server.kill();
    let _ = server.wait();
}

#[test]
fn duplicated_and_replayed_input_is_applied_at_most_once() {
    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 relay = Relay::spawn(port, 0xD0D0u64);

    let (socket, bootstrap, ok) = attach_through_relay(&config, &relay);

    // Append a fixed token to a file once per *distinct* delivered input. We use
    // `>>` so every time the server's PTY actually executes the command, a new
    // line is appended. Replay protection must ensure the duplicate/replayed
    // packet at the same sequence number is NOT re-applied.
    let marker = dir.path().join("dup_marker");
    let cmd = format!("printf x >> {}\n", marker.display());

    // Build one encrypted Input packet at a fixed sequence and send it many
    // times verbatim (a true replay: identical bytes, identical seq/nonce).
    let input = Input {
        bytes: cmd.into_bytes(),
    };
    let replayed = protocol::encode_encrypted(
        PacketKind::Input,
        ok.client_id,
        2,
        0,
        &bootstrap.session_key,
        CLIENT_TO_SERVER,
        &protocol::to_body(&input).unwrap(),
    )
    .unwrap();

    // Also have the relay duplicate everything, to stack duplication on top of
    // our explicit replays.
    relay.set_dup(100);
    for _ in 0..12 {
        send_raw(&socket, &relay, &replayed);
        thread::sleep(Duration::from_millis(40));
    }
    relay.set_dup(0);

    // Give the PTY time to run and flush.
    thread::sleep(Duration::from_millis(800));

    // Drive a fence command so we know the PTY has processed at least up to here
    // before we read the marker file.
    send_input(
        &socket,
        &relay,
        ok.client_id,
        3,
        0,
        &bootstrap.session_key,
        b"printf DUP_FENCE\\n\n",
    );
    let _ = wait_for_text(&socket, &bootstrap.session_key, "DUP_FENCE", 2000);
    thread::sleep(Duration::from_millis(300));

    let count = fs::read(&marker).map(|b| b.len()).unwrap_or(0);

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

    // The replayed/duplicated identical packet must apply at most once. If
    // replay protection were broken we would see many 'x' bytes.
    assert!(
        count <= 1,
        "replayed/duplicated input was applied {count} times (expected at most 1)"
    );
}

#[test]
fn forwarded_stream_data_survives_reorder_and_rejects_replay() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let signing_key = SigningKey::from_bytes(&[42u8; 32]);
    authorize_ed25519_key(&dir, &signing_key);
    let (target_port, tcp_rx) = start_tcp_collector();
    let mut server = start_server(&dir, &config);
    let relay = Relay::spawn(port, 0xF0E0D0u64);

    let attached = native_attach_through_relay(
        &relay,
        &signing_key,
        vec![ForwardingRequest {
            kind: ForwardingKind::Local,
            bind_host: Some("127.0.0.1".to_string()),
            listen_port: 0,
            target_host: Some("127.0.0.1".to_string()),
            target_port: Some(target_port),
        }],
    );
    let stream_id = 44;
    let open = StreamOpen {
        stream_id,
        target_host: "127.0.0.1".to_string(),
        target_port,
    };
    send_stream_packet(
        &attached.socket,
        &relay,
        &attached.ok,
        PacketKind::StreamOpen,
        3,
        attached.ok.initial_seq,
        protocol::to_body(&open).unwrap(),
    );
    let open_ok =
        wait_for_stream_open_ok(&attached.socket, &attached.ok.session_key, stream_id, 3000)
            .expect("stream did not open");

    // Swap the first two StreamData packets in flight. The server's replay
    // window must accept both out-of-order packets and write both to the TCP
    // target exactly once.
    relay.arm_reorder();
    send_stream_packet(
        &attached.socket,
        &relay,
        &attached.ok,
        PacketKind::StreamData,
        4,
        open_ok.seq,
        protocol::to_body(&StreamData {
            stream_id,
            offset: 0,
            bytes: b"A".to_vec(),
        })
        .unwrap(),
    );
    relay.arm_reorder();
    send_stream_packet(
        &attached.socket,
        &relay,
        &attached.ok,
        PacketKind::StreamData,
        5,
        open_ok.seq,
        protocol::to_body(&StreamData {
            stream_id,
            offset: 1,
            bytes: b"B".to_vec(),
        })
        .unwrap(),
    );

    let mut seen = collect_tcp(&tcp_rx, 1500);
    assert!(
        seen.starts_with(b"AB"),
        "reordered stream bytes were not delivered once in-order to TCP target: {:?}",
        String::from_utf8_lossy(&seen)
    );

    // Now send one encrypted StreamData packet repeatedly, with relay-level
    // duplication enabled too. This is a true replay: same sequence, same nonce,
    // same ciphertext. It must be applied to the TCP target at most once.
    let replayed = protocol::encode_encrypted(
        PacketKind::StreamData,
        attached.ok.client_id,
        6,
        open_ok.seq,
        &attached.ok.session_key,
        CLIENT_TO_SERVER,
        &protocol::to_body(&StreamData {
            stream_id,
            offset: 2,
            bytes: b"X".to_vec(),
        })
        .unwrap(),
    )
    .unwrap();
    relay.set_dup(100);
    for _ in 0..12 {
        send_raw(&attached.socket, &relay, &replayed);
        thread::sleep(Duration::from_millis(25));
    }
    relay.set_dup(0);

    seen.extend(collect_tcp(&tcp_rx, 1500));
    let replay_count = seen.iter().filter(|byte| **byte == b'X').count();

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

    assert!(
        replay_count <= 1,
        "replayed/duplicated StreamData was applied {replay_count} times: {:?}",
        String::from_utf8_lossy(&seen)
    );
}

#[test]
fn forwarded_stream_data_recovers_after_server_to_client_loss() {
    let dir = tempfile::tempdir().unwrap();
    let port = free_udp_port();
    let config = write_server_config(&dir, port);
    let signing_key = SigningKey::from_bytes(&[43u8; 32]);
    authorize_ed25519_key(&dir, &signing_key);
    let target_port = start_tcp_echo();
    let mut server = start_server(&dir, &config);
    let relay = Relay::spawn(port, 0xA11CEu64);

    let attached = native_attach_through_relay(
        &relay,
        &signing_key,
        vec![ForwardingRequest {
            kind: ForwardingKind::Local,
            bind_host: Some("127.0.0.1".to_string()),
            listen_port: 0,
            target_host: Some("127.0.0.1".to_string()),
            target_port: Some(target_port),
        }],
    );
    let stream_id = 55;
    let open = StreamOpen {
        stream_id,
        target_host: "127.0.0.1".to_string(),
        target_port,
    };
    send_stream_packet(
        &attached.socket,
        &relay,
        &attached.ok,
        PacketKind::StreamOpen,
        3,
        attached.ok.initial_seq,
        protocol::to_body(&open).unwrap(),
    );
    let open_ok =
        wait_for_stream_open_ok(&attached.socket, &attached.ok.session_key, stream_id, 3000)
            .expect("stream did not open");

    // Drop the first server->client echo and at least one retransmit tick. Once
    // the link clears, the unacked stream chunk must be re-encrypted with a new
    // packet sequence and delivered at the same stream offset.
    relay.set_drop_s2c(100);
    send_stream_packet(
        &attached.socket,
        &relay,
        &attached.ok,
        PacketKind::StreamData,
        4,
        open_ok.seq,
        protocol::to_body(&StreamData {
            stream_id,
            offset: 0,
            bytes: b"PING".to_vec(),
        })
        .unwrap(),
    );
    thread::sleep(Duration::from_millis(300));
    relay.set_drop_s2c(0);

    let (_header, data) =
        wait_for_stream_data(&attached.socket, &attached.ok.session_key, stream_id, 3000)
            .expect("lost server->client stream data was not retransmitted");
    assert_eq!(data.offset, 0);
    assert_eq!(data.bytes, b"PING");

    send_stream_packet(
        &attached.socket,
        &relay,
        &attached.ok,
        PacketKind::StreamWindowAdjust,
        5,
        open_ok.seq,
        protocol::to_body(&StreamWindowAdjust {
            stream_id,
            received_offset: 4,
            bytes: 4,
        })
        .unwrap(),
    );

    drop(relay);
    let _ = server.kill();
    let _ = server.wait();
}

#[test]
fn stale_packets_after_resume_are_ignored_not_fatal() {
    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 relay = Relay::spawn(port, 0x57A1u64);

    let (old_socket, bootstrap, ok) = attach_through_relay(&config, &relay);

    // Send an initial command on the original socket and confirm it lands.
    send_input(
        &old_socket,
        &relay,
        ok.client_id,
        2,
        0,
        &bootstrap.session_key,
        b"printf DOSH_STALE_BEFORE\\n\n",
    );
    assert!(
        wait_for_text(
            &old_socket,
            &bootstrap.session_key,
            "DOSH_STALE_BEFORE",
            3000
        ),
        "initial command did not land before resume"
    );

    // Simulate a reconnect from a new socket via a ResumeRequest (roaming),
    // mirroring tests/integration_smoke.rs::resume_updates_udp_endpoint.
    let new_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
    new_socket
        .set_read_timeout(Some(Duration::from_millis(200)))
        .unwrap();
    let resume = ResumeRequest {
        session: "default".to_string(),
        last_rendered_seq: ok.initial_seq,
        cols: 80,
        rows: 24,
    };
    let resume_packet = protocol::encode_encrypted(
        PacketKind::ResumeRequest,
        ok.client_id,
        100,
        0,
        &bootstrap.session_key,
        CLIENT_TO_SERVER,
        &protocol::to_body(&resume).unwrap(),
    )
    .unwrap();
    let mut resumed_seq = None;
    let deadline = Instant::now() + Duration::from_secs(5);
    while Instant::now() < deadline {
        new_socket
            .send_to(&resume_packet, relay.front_addr())
            .unwrap();
        if let Some((_h, frame)) = recv_frame(&new_socket, &bootstrap.session_key) {
            if frame.snapshot {
                resumed_seq = Some(frame.output_seq);
                break;
            }
        }
    }
    let resumed_seq = resumed_seq.expect("resume snapshot never arrived");

    // Now replay a STALE packet from the OLD socket with a low sequence number
    // (already-seen / out of the replay window). This must NOT terminate the
    // session.
    let stale = protocol::encode_encrypted(
        PacketKind::Input,
        ok.client_id,
        2, // old, already-consumed sequence
        0,
        &bootstrap.session_key,
        CLIENT_TO_SERVER,
        &protocol::to_body(&Input {
            bytes: b"printf DOSH_STALE_REPLAY\\n\n".to_vec(),
        })
        .unwrap(),
    )
    .unwrap();
    for _ in 0..5 {
        old_socket.send_to(&stale, relay.front_addr()).unwrap();
        thread::sleep(Duration::from_millis(30));
    }

    // The session must remain alive on the resumed socket: a fresh command with
    // a higher sequence still produces output.
    send_input(
        &new_socket,
        &relay,
        ok.client_id,
        101,
        resumed_seq,
        &bootstrap.session_key,
        b"printf DOSH_STALE_AFTER\\n\n",
    );
    let alive = wait_for_text(
        &new_socket,
        &bootstrap.session_key,
        "DOSH_STALE_AFTER",
        3000,
    );

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

    assert!(
        alive,
        "session was killed by stale packets after reconnect (should be ignored, not fatal)"
    );
}

#[test]
fn client_source_address_change_preserves_session() {
    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 relay = Relay::spawn(port, 0xADD12u64);

    let (socket, bootstrap, ok) = attach_through_relay(&config, &relay);

    // Confirm the session works before the address change.
    send_input(
        &socket,
        &relay,
        ok.client_id,
        2,
        0,
        &bootstrap.session_key,
        b"printf DOSH_ADDR_BEFORE\\n\n",
    );
    assert!(
        wait_for_text(&socket, &bootstrap.session_key, "DOSH_ADDR_BEFORE", 3000),
        "command before source-address change did not land"
    );

    // Switch the client's source address as the server sees it by rebinding the
    // relay's upstream socket (spec §11: "Connection migration must be accepted
    // after any valid encrypted packet from a new source address").
    let new_addr = relay.rebind_upstream();
    assert_ne!(new_addr.port(), 0);

    // The next valid encrypted packet now arrives from a new source address.
    // The session must keep working without a re-handshake.
    let mut migrated = false;
    let mut seq = 3u64;
    for _ in 0..12 {
        send_input(
            &socket,
            &relay,
            ok.client_id,
            seq,
            0,
            &bootstrap.session_key,
            b"printf DOSH_ADDR_AFTER\\n\n",
        );
        seq += 1;
        if wait_for_text(&socket, &bootstrap.session_key, "DOSH_ADDR_AFTER", 500) {
            migrated = true;
            break;
        }
    }

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

    assert!(
        migrated,
        "session did not survive a client source-address change (connection migration)"
    );
}