dosh/src/bin/dosh-server.rs

use anyhow::{Context, Result, anyhow, bail};
use clap::{Parser, Subcommand};
use dosh::auth::{
    build_attach_ticket, build_bootstrap, encode_bootstrap, load_or_create_server_secret, now_secs,
    open_attach_ticket, verify_bootstrap,
};
use dosh::config::{ServerConfig, expand_tilde, load_server_config};
use dosh::crypto;
use dosh::exec_service::{
    EXEC_STREAM_SENTINEL, ExecResponse, FrameDecoder as ExecFrameDecoder,
    encode_response as encode_exec_response,
};
use dosh::file_transfer::{
    CHUNK_SIZE, FILE_STREAM_SENTINEL, FileEntry, FileKind, FileMeta, FileRequest, FileResponse,
    FrameDecoder, clean_remote_path, encode_response,
};
use dosh::native::{
    EnvVar, ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello,
    derive_native_session_key, generate_native_ephemeral, host_public_key, host_public_key_line,
    load_or_create_host_key, sign_server_hello, verify_native_user_auth_from_config,
};
use dosh::persist;
use dosh::protocol::{
    self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
    NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody,
    NativeUserAuthBody, PacketKind, ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT,
    StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
    TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
};
use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session};
use dosh::udp::is_transient_udp_send_error;
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::fs;
use std::io::{Read, Seek, SeekFrom, Write};
use std::net::{IpAddr, SocketAddr};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::UnixStream as StdUnixStream;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket, UnixListener};
use tokio::process::Command as TokioCommand;
use tokio::sync::mpsc;

const STREAM_INITIAL_WINDOW: usize = 1024 * 1024;
const STREAM_RETIRED_TOMBSTONES: usize = 16 * 1024;

/// Persist a persistent session's screen to disk after this many bytes of output
/// have accumulated since the last mirror. Keeps the atomic write off the
/// per-packet hot path while bounding how much fresh output a crash can lose.
const SCREEN_PERSIST_BYTES: usize = 4096;
const SCREEN_PERSIST_MAX_AGE: Duration = Duration::from_millis(250);
const IMPLICIT_EMPTY_SESSION_GRACE_SECS: u64 = 300;

/// Sentinel `target_host` sent to the client on a server-initiated `StreamOpen`
/// that carries an SSH-agent connection (the client splices it into its local
/// `SSH_AUTH_SOCK` rather than dialing TCP). Must match the client's constant.
const AGENT_STREAM_SENTINEL: &str = "@dosh-agent";

/// Monotonic counter for unique agent proxy socket filenames within this process.
static AGENT_SOCK_COUNTER: AtomicU64 = AtomicU64::new(0);
/// How long a pending native handshake (ClientHello awaiting UserAuth) is kept
/// before it is swept by the cleanup task.
const NATIVE_HANDSHAKE_TTL_SECS: u64 = 30;

#[derive(Debug, Parser)]
#[command(
    name = "dosh-server",
    version = dosh::build_info::VERSION,
    long_version = dosh::build_info::LONG_VERSION
)]
struct Args {
    #[command(subcommand)]
    command: Command,
}

#[derive(Debug, Subcommand)]
enum Command {
    Serve {
        #[arg(long)]
        config: Option<std::path::PathBuf>,
    },
    Auth {
        #[arg(long, default_value_t = 1)]
        protocol: u8,
        #[arg(long)]
        nonce: Option<String>,
        #[arg(long)]
        host_key: bool,
        #[arg(long, default_value = "default")]
        session: String,
        #[arg(long, default_value = "read-write")]
        mode: String,
        #[arg(long, default_value = "80x24")]
        size: String,
        #[arg(long, default_value = "dev")]
        client_version: String,
        #[arg(long)]
        udp_host: Option<String>,
    },
    /// Internal: run as a per-session shell holder so the shell survives a
    /// server restart. Spawned by the server itself; not meant to be run by
    /// hand. Daemonizes, owns the PTY, and serves the master fd over a socket.
    #[command(hide = true)]
    Hold {
        #[arg(long)]
        runtime_dir: PathBuf,
        #[arg(long)]
        session: String,
        #[arg(long)]
        shell: String,
        #[arg(long)]
        cols: u16,
        #[arg(long)]
        rows: u16,
        /// `NAME=VALUE` pairs for the shell environment. Repeatable.
        #[arg(long = "env")]
        env: Vec<String>,
    },
}

fn main() -> Result<()> {
    let args = Args::parse();
    // The holder is handled BEFORE any tokio runtime exists: it forks/setsids to
    // daemonize, and forking out of a live async runtime is best avoided. It also
    // never uses tokio, so it has no reason to pay for one.
    if let Command::Hold {
        runtime_dir,
        session,
        shell,
        cols,
        rows,
        env,
    } = &args.command
    {
        let env = parse_env_pairs(env)?;
        return persist::run_holder(runtime_dir, session, shell, *cols, *rows, &env);
    }
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("build tokio runtime")?;
    runtime.block_on(async move { run_command(args.command).await })
}

async fn run_command(command: Command) -> Result<()> {
    match command {
        Command::Serve { config } => serve(config).await,
        Command::Hold { .. } => unreachable!("Hold is handled before the runtime starts"),
        Command::Auth {
            protocol,
            nonce,
            host_key,
            session,
            mode,
            size,
            client_version: _,
            udp_host,
        } => {
            anyhow::ensure!(protocol == 1, "unsupported protocol {protocol}");
            let config = load_server_config(None)?;
            if host_key {
                let signing_key = load_or_create_host_key(&config)?;
                println!("{}", host_public_key_line(&host_public_key(&signing_key)));
                return Ok(());
            }
            let secret = load_or_create_server_secret(&config)?;
            let nonce = parse_nonce(nonce.as_deref().context("--nonce is required")?)?;
            let size = parse_size(&size)?;
            let user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
            let udp_host = udp_host.unwrap_or_else(|| "127.0.0.1".to_string());
            let resp =
                build_bootstrap(&config, &secret, user, session, mode, size, nonce, udp_host)?;
            println!("{}", encode_bootstrap(&resp)?);
            Ok(())
        }
    }
}

async fn serve(config_path: Option<std::path::PathBuf>) -> Result<()> {
    let config = load_server_config(config_path)?;
    let secret = load_or_create_server_secret(&config)?;
    let bind = format!("{}:{}", config.bind, config.port);
    let socket = Arc::new(
        UdpSocket::bind(&bind)
            .await
            .with_context(|| format!("bind {bind}"))?,
    );
    eprintln!("dosh-server listening on {bind}");

    let (pty_tx, mut pty_rx) = mpsc::unbounded_channel();
    let state = Arc::new(Mutex::new(ServerState::new(
        config.clone(),
        secret,
        pty_tx.clone(),
    )));
    {
        let mut locked = state.lock().expect("server state poisoned");
        // Re-adopt any persistent sessions whose holders survived a previous
        // server's exit, BEFORE prewarming, so a prewarmed name reattaches to
        // its existing shell instead of spawning a duplicate.
        if config.persist_sessions {
            locked.readopt_persistent_sessions();
        }
        for session in config.prewarm_sessions.clone() {
            locked.ensure_session(&session, 80, 24, "read-write", &[])?;
        }
    }

    let output_state = Arc::clone(&state);
    let output_socket = Arc::clone(&socket);
    tokio::spawn(async move {
        while let Some(output) = pty_rx.recv().await {
            if let Err(err) = broadcast_output(&output_state, &output_socket, output).await {
                eprintln!("broadcast error: {err:#}");
            }
        }
    });

    let retransmit_state = Arc::clone(&state);
    let retransmit_socket = Arc::clone(&socket);
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(Duration::from_millis(100));
        loop {
            interval.tick().await;
            if let Err(err) = retransmit_pending(&retransmit_state, &retransmit_socket).await {
                eprintln!("retransmit error: {err:#}");
            }
            if let Err(err) = maybe_rekey_clients(&retransmit_state, &retransmit_socket).await {
                eprintln!("rekey error: {err:#}");
            }
        }
    });

    let cleanup_state = Arc::clone(&state);
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(Duration::from_secs(5));
        loop {
            interval.tick().await;
            cleanup_disconnected_clients(&cleanup_state);
        }
    });

    // Periodically mirror persistent sessions' screens to disk, bounding how much
    // recent output a sudden crash can lose regardless of throughput (the
    // per-packet path only persists every few KB). Cheap: one atomic write per
    // persistent session per tick, and only when its screen changed.
    if config.persist_sessions {
        let flush_state = Arc::clone(&state);
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(SCREEN_PERSIST_MAX_AGE);
            loop {
                interval.tick().await;
                flush_persistent_screens(&flush_state);
            }
        });
    }

    let mut buf = vec![0u8; 65535];
    loop {
        let (n, peer) = socket.recv_from(&mut buf).await?;
        if let Err(err) = handle_packet(&state, &socket, peer, &buf[..n]).await {
            eprintln!("packet from {peer}: {err:#}");
        }
    }
}

struct ServerState {
    config: ServerConfig,
    secret: [u8; 32],
    pty_tx: mpsc::UnboundedSender<PtyOutput>,
    sessions: HashMap<String, Session>,
    pending_native: HashMap<[u8; 16], PendingNativeAuth>,
    next_server_stream_id: u64,
    /// O(1) reverse index from a live client's `conn_id` to its session name, so
    /// per-packet handlers don't linear-scan every session's client map. Kept in
    /// lockstep with `Session::clients` on every insert and remove.
    client_index: HashMap<[u8; 16], String>,
    /// Per-source-IP token bucket for native ClientHello flood protection,
    /// checked before any X25519/Ed25519 work.
    native_auth_limiter: NativeAuthRateLimiter,
}

/// Per-source-IP token bucket throttling native auth ClientHellos.
///
/// Threat model item: a malicious unauthenticated client flooding auth attempts.
/// We charge one token per ClientHello *before* doing any X25519/Ed25519 crypto,
/// so a flood cannot burn server CPU. The bucket refills continuously at
/// `per_minute / 60` tokens per second up to a burst of `per_minute`, matching
/// the `native_auth_rate_limit_per_minute` value advertised in `ServerHello`.
struct NativeAuthRateLimiter {
    per_minute: u32,
    buckets: HashMap<IpAddr, TokenBucket>,
}

#[derive(Clone, Copy)]
struct TokenBucket {
    tokens: f64,
    last_refill: Instant,
}

impl NativeAuthRateLimiter {
    fn new(per_minute: u32) -> Self {
        Self {
            per_minute,
            buckets: HashMap::new(),
        }
    }

    /// Try to spend one token for `ip` at `now`. Returns `Ok(remaining)` whole
    /// tokens left when allowed, or `Err(())` when the bucket is empty. A
    /// `per_minute` of 0 disables native auth entirely (no tokens ever).
    fn check(&mut self, ip: IpAddr, now: Instant) -> std::result::Result<u32, ()> {
        if self.per_minute == 0 {
            return Err(());
        }
        let capacity = self.per_minute as f64;
        let refill_per_sec = capacity / 60.0;
        let bucket = self.buckets.entry(ip).or_insert(TokenBucket {
            tokens: capacity,
            last_refill: now,
        });
        let elapsed = now
            .saturating_duration_since(bucket.last_refill)
            .as_secs_f64();
        bucket.tokens = (bucket.tokens + elapsed * refill_per_sec).min(capacity);
        bucket.last_refill = now;
        if bucket.tokens < 1.0 {
            return Err(());
        }
        bucket.tokens -= 1.0;
        Ok(bucket.tokens as u32)
    }

    /// Drop buckets that have fully refilled, so the map can't grow unbounded
    /// under a spoofed-source-IP flood. Called from the periodic cleanup task.
    fn evict_full(&mut self, now: Instant) {
        if self.per_minute == 0 {
            self.buckets.clear();
            return;
        }
        let capacity = self.per_minute as f64;
        let refill_per_sec = capacity / 60.0;
        self.buckets.retain(|_, bucket| {
            let elapsed = now
                .saturating_duration_since(bucket.last_refill)
                .as_secs_f64();
            (bucket.tokens + elapsed * refill_per_sec) < capacity
        });
    }
}

struct Session {
    pty: Option<PtyHandle>,
    parser: vt100::Parser,
    clients: HashMap<[u8; 16], ClientState>,
    output_seq: u64,
    recent: VecDeque<Vec<u8>>,
    /// When the session last dropped to zero clients, used to reap abandoned
    /// sessions (and their shells) after a grace period. `None` while at least
    /// one client is attached.
    empty_since: Option<Instant>,
    /// Open control connection to this session's holder process, when the shell
    /// runs persistently out-of-process. Dropping it (server exit) leaves the
    /// holder running; sending a shutdown byte over it reaps the holder. `None`
    /// for non-persistent (in-process) sessions.
    holder_control: Option<StdUnixStream>,
    /// Whether this session's shell lives in a holder process (persistent).
    persistent: bool,
    /// Bytes of session output since the screen was last mirrored to disk, used
    /// to throttle the (atomic) screen-persistence writes.
    bytes_since_persist: usize,
    /// `output_seq` at the last successful screen mirror, so the periodic flush
    /// skips sessions whose screen has not changed.
    last_persisted_seq: u64,
    /// Last time this session's screen was considered for a mirror. Combined
    /// with the byte threshold so low-throughput prompt/TUI updates persist
    /// quickly without writing on every packet in high-throughput bursts.
    last_screen_persist_at: Instant,
}

#[derive(Clone)]
struct ClientState {
    endpoint: SocketAddr,
    mode: String,
    session_key: [u8; 32],
    last_acked: u64,
    replay: ReplayWindow,
    send_seq: u64,
    cols: u16,
    rows: u16,
    last_seen: Instant,
    pending: VecDeque<PendingFrame>,
    allowed_forwardings: Vec<ForwardingRequest>,
    stream_writers: HashMap<u64, mpsc::Sender<Vec<u8>>>,
    opened_streams: HashSet<u64>,
    stream_send_credit: HashMap<u64, usize>,
    stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
    stream_pending_opens: HashMap<u64, PendingStreamOpen>,
    stream_next_send_offset: HashMap<u64, u64>,
    stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
    stream_next_recv_offset: HashMap<u64, u64>,
    stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
    stream_retired: HashSet<u64>,
    stream_retired_order: VecDeque<u64>,
    /// Current transport key epoch (0 = original handshake key). Bumped on rekey.
    epoch: u64,
    /// When the current epoch began, for the wall-clock rekey trigger.
    epoch_started: Instant,
    /// Packets seen/sent in the current epoch, for the packet-count trigger.
    epoch_packets: u64,
    /// The previous epoch's key, retained briefly so in-flight pre-rekey packets
    /// still decrypt (matched via `session_key_id`) instead of dropping fatally.
    previous_session_key: Option<[u8; 32]>,
}

#[derive(Clone)]
struct PendingFrame {
    output_seq: u64,
    packet: Vec<u8>,
    last_sent: Instant,
    attempts: u8,
}

#[derive(Clone)]
struct PendingStreamChunk {
    offset: u64,
    bytes: Vec<u8>,
    last_sent: Instant,
    attempts: u32,
}

#[derive(Clone)]
struct PendingStreamOpen {
    target_host: String,
    target_port: u16,
    last_sent: Instant,
    attempts: u32,
}

#[derive(Clone)]
struct PendingNativeAuth {
    client: dosh::native::NativeClientHello,
    server: NativeServerHello,
    session_key: [u8; 32],
    peer: SocketAddr,
    created_at: Instant,
}

impl ServerState {
    fn new(
        config: ServerConfig,
        secret: [u8; 32],
        pty_tx: mpsc::UnboundedSender<PtyOutput>,
    ) -> Self {
        let per_minute = config.native_auth_rate_limit_per_minute;
        Self {
            config,
            secret,
            pty_tx,
            sessions: HashMap::new(),
            pending_native: HashMap::new(),
            next_server_stream_id: 1u64 << 63,
            client_index: HashMap::new(),
            native_auth_limiter: NativeAuthRateLimiter::new(per_minute),
        }
    }

    fn ensure_session(
        &mut self,
        name: &str,
        cols: u16,
        rows: u16,
        mode: &str,
        requested_env: &[EnvVar],
    ) -> Result<()> {
        let env = accepted_env(&self.config, requested_env)?;
        // Whether this existing session needs a (re)spawned PTY. Computed without
        // holding a mutable borrow of `self.sessions` so we can call the
        // `&self`-borrowing spawn helper.
        let needs_pty = self
            .sessions
            .get(name)
            .map(|session| session.pty.is_none() && mode_uses_pty(mode));
        if let Some(needs_pty) = needs_pty {
            if needs_pty {
                let (pty, control, persistent) = self.spawn_session_pty(name, cols, rows, &env)?;
                if let Some(session) = self.sessions.get_mut(name) {
                    session.pty = Some(pty);
                    session.holder_control = control;
                    session.persistent = persistent;
                }
            }
            return Ok(());
        }
        let (pty, control, persistent) = if mode_uses_pty(mode) {
            let (pty, control, persistent) = self.spawn_session_pty(name, cols, rows, &env)?;
            (Some(pty), control, persistent)
        } else {
            (None, None, false)
        };
        self.sessions.insert(
            name.to_string(),
            Session {
                pty,
                parser: vt100::Parser::new(rows.max(1), cols.max(1), self.config.scrollback),
                clients: HashMap::new(),
                output_seq: 0,
                recent: VecDeque::with_capacity(self.config.scrollback),
                empty_since: None,
                holder_control: control,
                persistent,
                bytes_since_persist: 0,
                last_persisted_seq: 0,
                last_screen_persist_at: Instant::now() - SCREEN_PERSIST_MAX_AGE,
            },
        );
        Ok(())
    }

    /// Whether a session should be backed by a persistent holder.
    ///
    /// With persistence enabled, every interactive session gets a detached holder
    /// so an update/restart of `dosh-server` behaves like a transient network
    /// break. Named sessions can sit empty for the normal client timeout.
    /// Generated one-off sessions are also persisted, but cleanup gives them a
    /// shorter empty grace period so invisible abandoned holders do not linger
    /// indefinitely.
    fn should_persist_session(&self, _name: &str) -> bool {
        self.config.persist_sessions
    }

    /// Spawn (or, in the persistent case, spawn-and-adopt) a PTY for a session.
    ///
    /// With `persist_sessions` on, the shell is launched in a detached holder
    /// process whose PTY master fd is passed back via SCM_RIGHTS, so the shell
    /// outlives this server. If anything in that path fails we degrade
    /// gracefully to the original in-process spawn (a shell that dies with the
    /// server), so persistence never makes a session un-attachable.
    fn spawn_session_pty(
        &self,
        name: &str,
        cols: u16,
        rows: u16,
        env: &[(String, String)],
    ) -> Result<(PtyHandle, Option<StdUnixStream>, bool)> {
        let cols = cols.max(1);
        let rows = rows.max(1);
        if self.should_persist_session(name) {
            match self.spawn_persistent_pty(name, cols, rows, env) {
                Ok(pair) => return Ok((pair.0, Some(pair.1), true)),
                Err(err) => {
                    eprintln!(
                        "persistence unavailable for session {name}, using in-process shell: {err:#}"
                    );
                }
            }
        }
        let pty = spawn_pty_session(
            name.to_string(),
            &self.config.shell,
            cols,
            rows,
            env,
            self.pty_tx.clone(),
        )?;
        Ok((pty, None, false))
    }

    /// Spawn a fresh holder for `name` and adopt its PTY master fd.
    fn spawn_persistent_pty(
        &self,
        name: &str,
        cols: u16,
        rows: u16,
        env: &[(String, String)],
    ) -> Result<(PtyHandle, StdUnixStream)> {
        let sessions_dir = self.sessions_dir();
        persist::spawn_holder(&sessions_dir, name, &self.config.shell, cols, rows, env)?;
        self.adopt_persistent_pty(name)
    }

    /// Connect to an existing holder for `name` and build an adopted PtyHandle
    /// from the master fd it hands back. Used both right after spawning a holder
    /// and when re-adopting a holder left behind by a previous server.
    fn adopt_persistent_pty(&self, name: &str) -> Result<(PtyHandle, StdUnixStream)> {
        let sessions_dir = self.sessions_dir();
        let (fd, control) = persist::adopt_holder(&sessions_dir, name)?;
        let pty = match adopt_pty_from_fd(name.to_string(), fd, self.pty_tx.clone()) {
            Ok(pty) => pty,
            Err(err) => {
                // We own `fd`; close it so the failed adopt doesn't leak it.
                drop(persist::fd_into_file(fd));
                return Err(err);
            }
        };
        Ok((pty, control))
    }

    fn sessions_dir(&self) -> PathBuf {
        expand_tilde(&self.config.sessions_dir)
    }

    /// On startup, re-adopt persistent sessions whose holder processes outlived a
    /// previous server. For each live holder we reconnect, receive the PTY master
    /// fd, restore the persisted screen/scrollback, and rebuild the in-memory
    /// `Session` so a reattaching client lands on the same shell with its screen
    /// intact. A holder we cannot adopt is left for the cleanup pass / fresh
    /// re-create, never crashing startup.
    fn readopt_persistent_sessions(&mut self) {
        let sessions_dir = self.sessions_dir();
        let holders = persist::scan_existing_holders(&sessions_dir);
        for meta in holders {
            let name = meta.session.clone();
            if self.sessions.contains_key(&name) {
                continue;
            }
            let (pty, control) = match self.adopt_persistent_pty(&name) {
                Ok(pair) => pair,
                Err(err) => {
                    eprintln!("could not re-adopt session {name}: {err:#}");
                    persist::remove_runtime_dir(&sessions_dir, &name);
                    continue;
                }
            };
            // Restore screen + scrollback geometry/seq if we saved one. The
            // restored snapshot is replayed into a fresh parser so a reattaching
            // client (and `screen_snapshot`) sees the pre-restart screen.
            let saved = persist::load_screen(&sessions_dir, &name);
            let (cols, rows) = saved
                .as_ref()
                .map(|s| (s.cols.max(1), s.rows.max(1)))
                .unwrap_or((80, 24));
            let mut parser = vt100::Parser::new(rows, cols, self.config.scrollback);
            let mut output_seq = 0;
            if let Some(saved) = saved {
                parser.process(&saved.snapshot);
                output_seq = saved.output_seq;
            }
            self.sessions.insert(
                name.clone(),
                Session {
                    pty: Some(pty),
                    parser,
                    clients: HashMap::new(),
                    output_seq,
                    recent: VecDeque::with_capacity(self.config.scrollback),
                    empty_since: Some(Instant::now()),
                    holder_control: Some(control),
                    persistent: true,
                    bytes_since_persist: 0,
                    last_persisted_seq: output_seq,
                    last_screen_persist_at: Instant::now(),
                },
            );
            eprintln!(
                "re-adopted persistent session {name} (shell pid {})",
                meta.shell_pid
            );
        }
    }

    /// Insert a client into `session_name` and record it in the O(1) index. The
    /// session must already exist (callers `ensure_session` first).
    fn insert_client(&mut self, session_name: &str, client_id: [u8; 16], client: ClientState) {
        if let Some(session) = self.sessions.get_mut(session_name) {
            session.clients.insert(client_id, client);
            session.empty_since = None;
            self.client_index
                .insert(client_id, session_name.to_string());
        }
    }

    /// Remove a client from whichever session holds it, keeping the index in
    /// sync. Returns true when a client was actually removed.
    fn remove_client_everywhere(&mut self, client_id: &[u8; 16]) -> bool {
        if let Some(session_name) = self.client_index.remove(client_id)
            && let Some(session) = self.sessions.get_mut(&session_name)
        {
            return session.clients.remove(client_id).is_some();
        }
        false
    }

    /// O(1) resolve of a live client's session key and session name via the
    /// index, replacing the former O(sessions) linear scan.
    fn lookup_client(&self, client_id: &[u8; 16]) -> Option<([u8; 32], String)> {
        let session_name = self.client_index.get(client_id)?;
        let session = self.sessions.get(session_name)?;
        let client = session.clients.get(client_id)?;
        Some((client.session_key, session_name.clone()))
    }

    /// Resolve the decryption key for an inbound packet, honoring rekey grace.
    ///
    /// During the brief window after a rekey, in-flight packets sealed under the
    /// previous epoch's key must still decrypt instead of being dropped as fatal
    /// (spec: "stale encrypted packets ... ignored, not fatal"). We pick whichever
    /// of the current or retained-previous key matches the packet's
    /// `session_key_id`. If neither key id matches, the client is on a stale
    /// epoch and should be forced through the normal reconnect path instead of
    /// waiting for a silent timeout.
    fn lookup_client_key_for(
        &self,
        client_id: &[u8; 16],
        session_key_id: &[u8; 16],
    ) -> Option<([u8; 32], String)> {
        let session_name = self.client_index.get(client_id)?;
        let session = self.sessions.get(session_name)?;
        let client = session.clients.get(client_id)?;
        if protocol::session_key_id(&client.session_key) == *session_key_id {
            return Some((client.session_key, session_name.clone()));
        }
        if let Some(previous) = client.previous_session_key
            && protocol::session_key_id(&previous) == *session_key_id
        {
            return Some((previous, session_name.clone()));
        }
        None
    }

    /// O(1) mutable access to a live client via the conn_id index.
    fn client_mut(&mut self, client_id: &[u8; 16]) -> Option<&mut ClientState> {
        let session_name = self.client_index.get(client_id)?;
        self.sessions
            .get_mut(session_name)?
            .clients
            .get_mut(client_id)
    }
}

fn mode_uses_pty(mode: &str) -> bool {
    mode != "forward-only"
}

/// Resize the PTY and terminal parser for a session.
///
/// `cols`/`rows` come straight off the wire from the client, where a terminal
/// that has not been sized yet (or a buggy/hostile client) can report `0`.
/// `vt100::Parser::set_size(0, _)` panics with a subtract overflow and would
/// take down the whole single-threaded daemon, so clamp to at least `1` here,
/// the same way `Parser::new` is constructed in `ensure_session`.
fn apply_terminal_size(
    pty: Option<&PtyHandle>,
    parser: &mut vt100::Parser,
    cols: u16,
    rows: u16,
) -> Result<()> {
    let cols = cols.max(1);
    let rows = rows.max(1);
    pty.context("terminal session has no pty")?
        .resize(cols, rows)?;
    parser.set_size(rows, cols);
    Ok(())
}

fn mode_allows_terminal_updates(mode: &str) -> bool {
    mode != "view-only" && mode != "forward-only"
}

fn accepted_env(config: &ServerConfig, requested: &[EnvVar]) -> Result<Vec<(String, String)>> {
    let mut env = Vec::new();
    for entry in requested {
        anyhow::ensure!(
            valid_env_name(&entry.name),
            "invalid environment variable name {:?}",
            entry.name
        );
        anyhow::ensure!(
            !entry.value.as_bytes().contains(&0),
            "environment variable {:?} contains NUL",
            entry.name
        );
        if config
            .accept_env
            .iter()
            .any(|pattern| glob_matches(pattern, &entry.name))
        {
            env.push((entry.name.clone(), entry.value.clone()));
        }
    }
    Ok(env)
}

fn valid_env_name(name: &str) -> bool {
    let mut chars = name.chars();
    let Some(first) = chars.next() else {
        return false;
    };
    if !(first == '_' || first.is_ascii_alphabetic()) {
        return false;
    }
    chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
}

fn glob_matches(pattern: &str, value: &str) -> bool {
    let pattern = pattern.as_bytes();
    let value = value.as_bytes();
    let (mut p, mut v) = (0, 0);
    let mut star = None;
    let mut star_value = 0;
    while v < value.len() {
        if p < pattern.len() && (pattern[p] == b'?' || pattern[p] == value[v]) {
            p += 1;
            v += 1;
        } else if p < pattern.len() && pattern[p] == b'*' {
            star = Some(p);
            p += 1;
            star_value = v;
        } else if let Some(star_pos) = star {
            p = star_pos + 1;
            star_value += 1;
            v = star_value;
        } else {
            return false;
        }
    }
    while p < pattern.len() && pattern[p] == b'*' {
        p += 1;
    }
    p == pattern.len()
}

async fn handle_packet(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    raw: &[u8],
) -> Result<()> {
    // A peer running a different wire protocol version cannot be decoded at all,
    // because the header VERSION byte is part of the AEAD AAD and the framing.
    // Rather than drop it (which the peer sees as a silent timeout), answer with
    // a clear, named version-mismatch reject so the user is told to upgrade.
    if let Some(foreign) = protocol::peek_foreign_wire_version(raw) {
        let _ = foreign;
        return send_reject(socket, peer, protocol::VERSION_MISMATCH_REASON).await;
    }
    let packet = protocol::decode(raw)?;
    match packet.header.kind {
        PacketKind::NativeClientHello => {
            handle_native_client_hello(state, socket, peer, packet.body).await
        }
        PacketKind::NativeUserAuth => handle_native_user_auth(state, socket, peer, &packet).await,
        PacketKind::BootstrapAttachRequest => {
            handle_bootstrap_attach(state, socket, peer, packet.body).await
        }
        PacketKind::TicketAttachRequest => {
            handle_ticket_attach(state, socket, peer, packet.body).await
        }
        PacketKind::ResumeRequest => handle_resume(state, socket, peer, &packet).await,
        PacketKind::Input
        | PacketKind::Resize
        | PacketKind::Ping
        | PacketKind::Ack
        | PacketKind::StreamOpen
        | PacketKind::StreamOpenOk
        | PacketKind::StreamOpenReject
        | PacketKind::StreamData
        | PacketKind::StreamClose
        | PacketKind::StreamEof
        | PacketKind::StreamWindowAdjust
        | PacketKind::RekeyAck
        | PacketKind::Detach
            if find_client_decrypt_key(state, &packet.header).is_err() =>
        {
            send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client").await
        }
        PacketKind::Input => handle_input(state, peer, &packet).await,
        PacketKind::Resize => handle_resize(state, peer, &packet).await,
        PacketKind::Ping => handle_ping(state, socket, peer, &packet).await,
        PacketKind::Ack => handle_ack(state, peer, &packet).await,
        PacketKind::RekeyAck => handle_rekey_ack(state, peer, &packet).await,
        PacketKind::Detach => handle_detach(state, &packet).await,
        PacketKind::StreamOpen => handle_stream_open(state, socket, peer, &packet).await,
        PacketKind::StreamOpenOk => handle_stream_open_ok(state, socket, peer, &packet).await,
        PacketKind::StreamOpenReject => handle_stream_open_reject(state, peer, &packet).await,
        PacketKind::StreamData => handle_stream_data(state, socket, peer, &packet).await,
        PacketKind::StreamWindowAdjust => {
            handle_stream_window_adjust(state, socket, peer, &packet).await
        }
        PacketKind::StreamClose => handle_stream_close(state, peer, &packet).await,
        _ => Ok(()),
    }
}

async fn handle_native_client_hello(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    body: Vec<u8>,
) -> Result<()> {
    let req: NativeClientHelloBody = protocol::from_body(&body)?;
    // Defense in depth against a peer that shares our wire VERSION but advertises
    // a different native handshake protocol_version (the spec's negotiation
    // point). Reject with a named, actionable reason before any crypto.
    if let Err(err) =
        dosh::native::check_native_protocol_version(req.hello.protocol_version, "client")
    {
        return send_reject(socket, peer, &err.to_string()).await;
    }
    // Rate-limit native auth *before* touching the host key or any X25519/Ed25519
    // work, so a flood of ClientHellos from one source cannot burn server CPU.
    // A separate, non-crypto reject lets a legitimate client back off cleanly.
    let rate_limit_remaining = {
        let mut locked = state.lock().expect("server state poisoned");
        locked.native_auth_limiter.check(peer.ip(), Instant::now())
    };
    let rate_limit_remaining = match rate_limit_remaining {
        Ok(remaining) => remaining,
        Err(()) => {
            return send_reject(socket, peer, "native auth rate limit exceeded").await;
        }
    };
    let built: Result<([u8; 16], NativeServerHello)> = {
        let mut locked = state.lock().expect("server state poisoned");
        if !locked.config.native_auth {
            Err(anyhow!("native auth disabled"))
        } else if !req
            .hello
            .supported_aead
            .iter()
            .any(|algorithm| algorithm == "chacha20poly1305")
        {
            Err(anyhow!("native auth requires chacha20poly1305"))
        } else if !req
            .hello
            .supported_user_key_algorithms
            .iter()
            .any(|algorithm| dosh::native::is_supported_user_signature_algorithm(algorithm))
        {
            Err(anyhow!(
                "native auth requires a supported user key algorithm"
            ))
        } else {
            let current_user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
            if req.hello.requested_user != current_user {
                Err(anyhow!("native auth user mismatch"))
            } else {
                let host_signing = load_or_create_host_key(&locked.config)?;
                let (server_secret, server_public) = generate_native_ephemeral();
                let mut hello = NativeServerHello {
                    protocol_version: dosh::native::NATIVE_PROTOCOL_VERSION,
                    server_random: crypto::random_32(),
                    server_ephemeral_public: server_public,
                    host_key: host_public_key(&host_signing),
                    chosen_aead: "chacha20poly1305".to_string(),
                    server_key_epoch: 1,
                    auth_challenge: crypto::random_32(),
                    rate_limit_remaining: Some(rate_limit_remaining),
                    host_signature: Vec::new(),
                };
                sign_server_hello(&host_signing, &req.hello, &mut hello)?;
                let session_key = derive_native_session_key(
                    &server_secret,
                    req.hello.client_ephemeral_public,
                    &req.hello,
                    &hello,
                )?;
                let mut pending_id = [0u8; 16];
                pending_id.copy_from_slice(&hello.auth_challenge[..16]);
                locked.pending_native.insert(
                    pending_id,
                    PendingNativeAuth {
                        client: req.hello,
                        server: hello.clone(),
                        session_key,
                        peer,
                        created_at: Instant::now(),
                    },
                );
                Ok((pending_id, hello))
            }
        }
    };
    let (pending_id, hello) = match built {
        Ok(built) => built,
        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
    };
    let body = protocol::to_body(&NativeServerHelloBody { hello })?;
    let out = protocol::encode_plain(PacketKind::NativeServerHello, pending_id, 1, 0, &body)?;
    let _ = send_udp(socket, &out, peer).await?;
    Ok(())
}

async fn handle_native_user_auth(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let pending = {
        let locked = state.lock().expect("server state poisoned");
        locked.pending_native.get(&packet.header.conn_id).cloned()
    };
    let Some(pending) = pending else {
        return send_reject_to_client(socket, peer, packet.header.conn_id, "unknown native auth")
            .await;
    };
    if pending.peer != peer {
        return send_reject_to_client(
            socket,
            peer,
            packet.header.conn_id,
            "native auth peer changed",
        )
        .await;
    }
    if pending.created_at.elapsed() > Duration::from_secs(30) {
        state
            .lock()
            .expect("server state poisoned")
            .pending_native
            .remove(&packet.header.conn_id);
        return send_reject_to_client(socket, peer, packet.header.conn_id, "native auth expired")
            .await;
    }
    let body = match protocol::decrypt_body(packet, &pending.session_key, CLIENT_TO_SERVER) {
        Ok(body) => body,
        Err(_) => {
            return send_reject_to_client(
                socket,
                peer,
                packet.header.conn_id,
                "native auth decrypt failed",
            )
            .await;
        }
    };
    let req: NativeUserAuthBody = match protocol::from_body(&body) {
        Ok(req) => req,
        Err(_) => {
            return send_reject_to_client(
                socket,
                peer,
                packet.header.conn_id,
                "invalid native auth body",
            )
            .await;
        }
    };
    if pending.client.requested_mode == "doctor" {
        let check_result = {
            let locked = state.lock().expect("server state poisoned");
            verify_native_user_auth_from_config(
                &locked.config,
                &pending.client,
                &pending.server,
                &req.auth,
                Some(peer.ip()),
            )
            .context("verify native user auth")
            .map(|_| NativeAuthCheckOkBody {
                requested_user: pending.client.requested_user.clone(),
                native_auth_enabled: locked.config.native_auth,
                allow_tcp_forwarding: locked.config.allow_tcp_forwarding,
                allow_remote_forwarding: locked.config.allow_remote_forwarding,
                allow_agent_forwarding: locked.config.allow_agent_forwarding,
                allow_file_transfer: locked.config.allow_file_transfer,
                allow_exec_command: locked.config.allow_exec_command,
                policy_flags: Vec::new(),
                server_version: env!("CARGO_PKG_VERSION").to_string(),
            })
        };
        let check = match check_result {
            Ok(check) => check,
            Err(err) => {
                return send_reject_to_client(
                    socket,
                    peer,
                    packet.header.conn_id,
                    &format!("{err:#}"),
                )
                .await;
            }
        };
        state
            .lock()
            .expect("server state poisoned")
            .pending_native
            .remove(&packet.header.conn_id);
        let body = protocol::to_body(&check)?;
        let out = protocol::encode_encrypted(
            PacketKind::NativeAuthCheckOk,
            packet.header.conn_id,
            1,
            packet.header.seq,
            &pending.session_key,
            SERVER_TO_CLIENT,
            &body,
        )?;
        let _ = send_udp(socket, &out, peer).await?;
        return Ok(());
    }

    // SSH-agent forwarding (opt-in + policy-gated). When the client requested it
    // AND the server allows it, bind a per-session proxy unix socket now so its
    // path can be exported as the remote shell's SSH_AUTH_SOCK before the PTY is
    // spawned. The accept loop is started after the client_id exists below.
    // SECURITY: only ever reached on explicit client opt-in and with
    // `allow_agent_forwarding` enabled; the socket is private to this user.
    let agent_allowed = {
        let locked = state.lock().expect("server state poisoned");
        locked.config.allow_agent_forwarding
    };
    let agent_listener =
        if agent_allowed && agent_forwarding_requested(&req.auth.requested_forwardings) {
            match bind_agent_proxy_socket() {
                Ok(pair) => Some(pair),
                Err(err) => {
                    eprintln!("agent forwarding setup failed, continuing without it: {err:#}");
                    None
                }
            }
        } else {
            None
        };
    // Env handed to ensure_session: base requested env plus SSH_AUTH_SOCK when
    // agent forwarding is active. Note: only applies to a freshly spawned shell;
    // an already-running (prewarmed/attached) session keeps its existing env, the
    // same constraint ssh/mosh have.
    let mut session_env = pending.client.requested_env.clone();
    if let Some((_, ref path)) = agent_listener {
        session_env.retain(|e| e.name != "SSH_AUTH_SOCK");
        session_env.push(EnvVar {
            name: "SSH_AUTH_SOCK".to_string(),
            value: path.to_string_lossy().to_string(),
        });
    }

    struct NativeAttachResult {
        client_id: [u8; 16],
        key_id: [u8; 16],
        attach_ticket: Vec<u8>,
        attach_ticket_psk: [u8; 32],
        session_name: String,
        mode: String,
        output_seq: u64,
        snapshot: Vec<u8>,
        requested_forwardings: Vec<ForwardingRequest>,
    }

    let attached: Result<NativeAttachResult> = {
        let mut locked = state.lock().expect("server state poisoned");
        verify_native_user_auth_from_config(
            &locked.config,
            &pending.client,
            &pending.server,
            &req.auth,
            Some(peer.ip()),
        )
        .context("verify native user auth")
        .and_then(|_| {
            let session_name = pending.client.requested_session.clone();
            let mode = pending.client.requested_mode.clone();
            let cols = pending.client.terminal_size.0;
            let rows = pending.client.terminal_size.1;
            if !locked.sessions.contains_key(&session_name) && !locked.config.create_on_attach {
                return Err(anyhow!("session does not exist"));
            }
            locked.ensure_session(&session_name, cols, rows, &mode, &session_env)?;
            let session_key = pending.session_key;
            let key_id = protocol::session_key_id(&session_key);
            let attach_ticket_psk = crypto::random_32();
            let issued_at = now_secs()?;
            let attach_ticket = build_attach_ticket(
                &locked.secret,
                crypto::sha256(&locked.secret),
                1,
                pending.client.requested_user.clone(),
                session_name.clone(),
                mode.clone(),
                issued_at,
                issued_at + locked.config.attach_ticket_ttl_secs,
                &attach_ticket_psk,
            )?;
            let session = locked
                .sessions
                .get_mut(&session_name)
                .expect("session exists");
            if mode_allows_terminal_updates(&mode) {
                apply_terminal_size(session.pty.as_ref(), &mut session.parser, cols, rows)?;
            }
            let client_id = crypto::random_16();
            let snapshot = screen_snapshot(session.parser.screen());
            let output_seq = session.output_seq;
            locked.insert_client(
                &session_name,
                client_id,
                ClientState {
                    endpoint: peer,
                    mode: mode.clone(),
                    session_key,
                    last_acked: output_seq,
                    replay: ReplayWindow::default(),
                    send_seq: 1,
                    cols,
                    rows,
                    last_seen: Instant::now(),
                    pending: VecDeque::new(),
                    allowed_forwardings: req.auth.requested_forwardings.clone(),
                    stream_writers: HashMap::new(),
                    opened_streams: HashSet::new(),
                    stream_send_credit: HashMap::new(),
                    stream_pending_data: HashMap::new(),
                    stream_pending_opens: HashMap::new(),
                    stream_next_send_offset: HashMap::new(),
                    stream_sent_data: HashMap::new(),
                    stream_next_recv_offset: HashMap::new(),
                    stream_recv_pending: HashMap::new(),
                    stream_retired: HashSet::new(),
                    stream_retired_order: VecDeque::new(),
                    epoch: 0,
                    epoch_started: Instant::now(),
                    epoch_packets: 0,
                    previous_session_key: None,
                },
            );
            Ok(NativeAttachResult {
                client_id,
                key_id,
                attach_ticket,
                attach_ticket_psk,
                session_name,
                mode,
                output_seq,
                snapshot,
                requested_forwardings: req.auth.requested_forwardings.clone(),
            })
        })
    };

    let attached = match attached {
        Ok(attached) => attached,
        Err(err) => {
            // Auth/session setup failed: clean up the agent socket we bound early.
            if let Some((_, path)) = agent_listener {
                let _ = std::fs::remove_file(&path);
            }
            return send_reject_to_client(socket, peer, packet.header.conn_id, &format!("{err:#}"))
                .await;
        }
    };
    let client_id = attached.client_id;
    if let Err(err) = start_remote_forwards(
        state,
        socket,
        client_id,
        attached.requested_forwardings.clone(),
        pending.client.requested_session.clone(),
    )
    .await
    {
        if let Some((_, path)) = agent_listener {
            let _ = std::fs::remove_file(&path);
        }
        remove_client(state, client_id);
        return send_reject_to_client(socket, peer, packet.header.conn_id, &err.to_string()).await;
    }
    // Now that the client exists, start accepting forwarded agent connections.
    if let Some((listener, path)) = agent_listener {
        spawn_agent_forward(state, socket, client_id, listener, path);
    }
    state
        .lock()
        .expect("server state poisoned")
        .pending_native
        .remove(&packet.header.conn_id);

    let ok = NativeAuthOk {
        client_id,
        session: attached.session_name,
        mode: attached.mode,
        session_key: pending.session_key,
        session_key_id: attached.key_id,
        attach_ticket: attached.attach_ticket,
        attach_ticket_psk: attached.attach_ticket_psk,
        initial_seq: attached.output_seq,
        snapshot: attached.snapshot,
        policy_flags: Vec::new(),
    };
    let body = protocol::to_body(&NativeAuthOkBody { ok })?;
    let out = protocol::encode_encrypted(
        PacketKind::NativeAuthOk,
        client_id,
        1,
        packet.header.seq,
        &pending.session_key,
        SERVER_TO_CLIENT,
        &body,
    )?;
    let _ = send_udp(socket, &out, peer).await?;
    Ok(())
}

async fn handle_bootstrap_attach(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    body: Vec<u8>,
) -> Result<()> {
    let req: BootstrapAttachRequest = protocol::from_body(&body)?;
    type BootstrapAttachResult = ([u8; 16], [u8; 32], [u8; 16], String, String, u64, Vec<u8>);
    let attached: Result<BootstrapAttachResult> = {
        let mut locked = state.lock().expect("server state poisoned");
        if !verify_bootstrap(&req.bootstrap, &locked.secret)? {
            Err(anyhow!("invalid or expired bootstrap"))
        } else if !locked.sessions.contains_key(&req.bootstrap.session)
            && !locked.config.create_on_attach
        {
            Err(anyhow!("session does not exist"))
        } else {
            locked.ensure_session(
                &req.bootstrap.session,
                req.cols,
                req.rows,
                &req.bootstrap.mode,
                &req.requested_env,
            )?;
            let session = locked
                .sessions
                .get_mut(&req.bootstrap.session)
                .expect("session exists");
            if mode_allows_terminal_updates(&req.bootstrap.mode) {
                apply_terminal_size(
                    session.pty.as_ref(),
                    &mut session.parser,
                    req.cols,
                    req.rows,
                )?;
            }
            let client_id = crypto::random_16();
            let snapshot = screen_snapshot(session.parser.screen());
            let output_seq = session.output_seq;
            let bootstrap_session = req.bootstrap.session.clone();
            locked.insert_client(
                &bootstrap_session,
                client_id,
                ClientState {
                    endpoint: peer,
                    mode: req.bootstrap.mode.clone(),
                    session_key: req.bootstrap.session_key,
                    last_acked: output_seq,
                    replay: ReplayWindow::default(),
                    send_seq: 1,
                    cols: req.cols,
                    rows: req.rows,
                    last_seen: Instant::now(),
                    pending: VecDeque::new(),
                    allowed_forwardings: Vec::new(),
                    stream_writers: HashMap::new(),
                    opened_streams: HashSet::new(),
                    stream_send_credit: HashMap::new(),
                    stream_pending_data: HashMap::new(),
                    stream_pending_opens: HashMap::new(),
                    stream_next_send_offset: HashMap::new(),
                    stream_sent_data: HashMap::new(),
                    stream_next_recv_offset: HashMap::new(),
                    stream_recv_pending: HashMap::new(),
                    stream_retired: HashSet::new(),
                    stream_retired_order: VecDeque::new(),
                    epoch: 0,
                    epoch_started: Instant::now(),
                    epoch_packets: 0,
                    previous_session_key: None,
                },
            );
            Ok((
                client_id,
                req.bootstrap.session_key,
                req.bootstrap.session_key_id,
                req.bootstrap.session.clone(),
                req.bootstrap.mode.clone(),
                output_seq,
                snapshot,
            ))
        }
    };
    let (client_id, key, key_id, session_name, mode, output_seq, snapshot) = match attached {
        Ok(attached) => attached,
        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
    };
    let ok = AttachOk {
        client_id,
        session: session_name,
        mode,
        session_key: key,
        session_key_id: key_id,
        initial_seq: output_seq,
        snapshot,
    };
    let body = protocol::to_body(&ok)?;
    let out = protocol::encode_encrypted(
        PacketKind::AttachOk,
        client_id,
        1,
        0,
        &key,
        SERVER_TO_CLIENT,
        &body,
    )?;
    let _ = send_udp(socket, &out, peer).await?;
    Ok(())
}

async fn handle_ticket_attach(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    body: Vec<u8>,
) -> Result<()> {
    let env: TicketAttachEnvelope = protocol::from_body(&body)?;
    let opened = {
        let locked = state.lock().expect("server state poisoned");
        if !locked.config.allow_attach_tickets {
            Err(anyhow!("attach tickets disabled"))
        } else {
            let ticket = open_attach_ticket(&locked.secret, &env.ticket)?;
            let request_key = crypto::hkdf32(
                &ticket.psk,
                &env.client_nonce,
                b"dosh/ticket-attach-request/v1",
            )?;
            let request_plain = crypto::open(
                &request_key,
                &env.client_nonce,
                b"dosh-ticket-attach-request-v1",
                &env.ciphertext,
            )?;
            Ok((ticket, request_plain))
        }
    };
    let (ticket, request_plain) = match opened {
        Ok(opened) => opened,
        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
    };
    let req: TicketAttachBody = protocol::from_body(&request_plain)?;
    if req.session != ticket.session || req.mode != ticket.mode {
        return send_reject(socket, peer, "ticket scope mismatch").await;
    }

    let session_key = crypto::random_32();
    let session_key_id = protocol::session_key_id(&session_key);
    let attached = {
        let mut locked = state.lock().expect("server state poisoned");
        if !locked.sessions.contains_key(&req.session) && !locked.config.create_on_attach {
            Err(anyhow!("session does not exist"))
        } else {
            locked.ensure_session(
                &req.session,
                req.cols,
                req.rows,
                &req.mode,
                &req.requested_env,
            )?;
            let session = locked
                .sessions
                .get_mut(&req.session)
                .expect("session exists");
            if mode_allows_terminal_updates(&req.mode) {
                apply_terminal_size(
                    session.pty.as_ref(),
                    &mut session.parser,
                    req.cols,
                    req.rows,
                )?;
            }
            let client_id = crypto::random_16();
            let snapshot = screen_snapshot(session.parser.screen());
            let output_seq = session.output_seq;
            let ticket_session = req.session.clone();
            locked.insert_client(
                &ticket_session,
                client_id,
                ClientState {
                    endpoint: peer,
                    mode: req.mode.clone(),
                    session_key,
                    last_acked: output_seq,
                    replay: ReplayWindow::default(),
                    send_seq: 1,
                    cols: req.cols,
                    rows: req.rows,
                    last_seen: Instant::now(),
                    pending: VecDeque::new(),
                    allowed_forwardings: Vec::new(),
                    stream_writers: HashMap::new(),
                    opened_streams: HashSet::new(),
                    stream_send_credit: HashMap::new(),
                    stream_pending_data: HashMap::new(),
                    stream_pending_opens: HashMap::new(),
                    stream_next_send_offset: HashMap::new(),
                    stream_sent_data: HashMap::new(),
                    stream_next_recv_offset: HashMap::new(),
                    stream_recv_pending: HashMap::new(),
                    stream_retired: HashSet::new(),
                    stream_retired_order: VecDeque::new(),
                    epoch: 0,
                    epoch_started: Instant::now(),
                    epoch_packets: 0,
                    previous_session_key: None,
                },
            );
            Ok((client_id, output_seq, snapshot))
        }
    };
    let (client_id, output_seq, snapshot) = match attached {
        Ok(attached) => attached,
        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
    };

    let ok = AttachOk {
        client_id,
        session: req.session,
        mode: req.mode,
        session_key,
        session_key_id,
        initial_seq: output_seq,
        snapshot,
    };
    let ok_plain = protocol::to_body(&ok)?;
    let server_nonce = crypto::random_12();
    let mut salt = Vec::with_capacity(24);
    salt.extend_from_slice(&env.client_nonce);
    salt.extend_from_slice(&server_nonce);
    let response_key = crypto::hkdf32(&ticket.psk, &salt, b"dosh/ticket-attach-ok/v1")?;
    let ciphertext = crypto::seal(
        &response_key,
        &server_nonce,
        b"dosh-ticket-attach-ok-v1",
        &ok_plain,
    )?;
    let envelope = TicketAttachOkEnvelope {
        server_nonce,
        ciphertext,
    };
    let body = protocol::to_body(&envelope)?;
    let out = protocol::encode_plain(PacketKind::AttachOk, client_id, 1, 0, &body)?;
    let _ = send_udp(socket, &out, peer).await?;
    Ok(())
}

async fn send_reject(socket: &UdpSocket, peer: SocketAddr, reason: &str) -> Result<()> {
    send_reject_to_client(socket, peer, [0u8; 16], reason).await
}

async fn send_reject_to_client(
    socket: &UdpSocket,
    peer: SocketAddr,
    client_id: [u8; 16],
    reason: &str,
) -> Result<()> {
    let body = protocol::to_body(&AttachReject {
        reason: reason.to_string(),
    })?;
    let out = protocol::encode_plain(PacketKind::AttachReject, client_id, 0, 0, &body)?;
    let _ = send_udp(socket, &out, peer).await?;
    Ok(())
}

async fn handle_resume(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = match find_client_decrypt_key(state, &packet.header) {
        Ok(found) => found,
        Err(_) => {
            return send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client")
                .await;
        }
    };
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let req: ResumeRequest = protocol::from_body(&body)?;
    let (send_seq, output_seq, snapshot) = {
        let mut locked = state.lock().expect("server state poisoned");
        let session = locked
            .sessions
            .get_mut(&req.session)
            .ok_or_else(|| anyhow!("unknown session"))?;
        let client = session
            .clients
            .get_mut(&packet.header.conn_id)
            .ok_or_else(|| anyhow!("unknown client"))?;
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        client.endpoint = peer;
        client.last_acked = req.last_rendered_seq;
        client.cols = req.cols;
        client.rows = req.rows;
        client.last_seen = Instant::now();
        client.send_seq += 1;
        if mode_allows_terminal_updates(&client.mode) {
            apply_terminal_size(
                session.pty.as_ref(),
                &mut session.parser,
                req.cols,
                req.rows,
            )?;
        }
        let snapshot = screen_snapshot(session.parser.screen());
        (client.send_seq, session.output_seq, snapshot)
    };
    let frame = Frame {
        session: session_name,
        output_seq,
        bytes: snapshot,
        snapshot: true,
        closed: false,
    };
    let body = protocol::to_body(&frame)?;
    let out = protocol::encode_encrypted(
        PacketKind::ResumeOk,
        packet.header.conn_id,
        send_seq,
        packet.header.seq,
        &key,
        SERVER_TO_CLIENT,
        &body,
    )?;
    let _ = send_udp(socket, &out, peer).await?;
    Ok(())
}

async fn handle_input(
    state: &Arc<Mutex<ServerState>>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let input: Input = protocol::from_body(&body)?;
    let mut locked = state.lock().expect("server state poisoned");
    let session = locked
        .sessions
        .get_mut(&session_name)
        .ok_or_else(|| anyhow!("unknown session"))?;
    let client = session
        .clients
        .get_mut(&packet.header.conn_id)
        .ok_or_else(|| anyhow!("unknown client"))?;
    if !client.replay.accept(packet.header.seq) {
        return Ok(());
    }
    if client.endpoint != peer {
        client.endpoint = peer;
    }
    client.last_seen = Instant::now();
    if !mode_allows_terminal_updates(&client.mode) {
        return Ok(());
    }
    session
        .pty
        .as_ref()
        .context("terminal session has no pty")?
        .write_all(&input.bytes)?;
    Ok(())
}

async fn handle_resize(
    state: &Arc<Mutex<ServerState>>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let resize: Resize = protocol::from_body(&body)?;
    let mut locked = state.lock().expect("server state poisoned");
    let session = locked
        .sessions
        .get_mut(&session_name)
        .ok_or_else(|| anyhow!("unknown session"))?;
    let client = session
        .clients
        .get_mut(&packet.header.conn_id)
        .ok_or_else(|| anyhow!("unknown client"))?;
    if !client.replay.accept(packet.header.seq) {
        return Ok(());
    }
    // Connection migration happens for any authenticated, fresh packet (spec §11),
    // independent of the session mode; a view-only client roams too.
    client.endpoint = peer;
    client.last_seen = Instant::now();
    if mode_allows_terminal_updates(&client.mode) {
        client.cols = resize.cols;
        client.rows = resize.rows;
        apply_terminal_size(
            session.pty.as_ref(),
            &mut session.parser,
            resize.cols,
            resize.rows,
        )?;
    }
    Ok(())
}

async fn handle_ping(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, _) = find_client_decrypt_key(state, &packet.header)?;
    // Authenticate the ping (verify its AEAD tag) BEFORE acting on it. Pings carry
    // an empty plaintext but still a tag, so this rejects a spoofed ping that
    // could otherwise replay-poison or, worse, hijack `endpoint` migration.
    let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let seq = {
        let mut locked = state.lock().expect("server state poisoned");
        let client = locked
            .client_mut(&packet.header.conn_id)
            .ok_or_else(|| anyhow!("unknown client"))?;
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        // Connection migration: any authenticated, fresh packet from a new source
        // address roams the session there (spec §11). Gated on replay-accept so a
        // captured old packet can't redirect the stream.
        client.endpoint = peer;
        client.last_seen = Instant::now();
        client.send_seq += 1;
        client.send_seq
    };
    let out = protocol::encode_encrypted(
        PacketKind::Pong,
        packet.header.conn_id,
        seq,
        packet.header.seq,
        &key,
        SERVER_TO_CLIENT,
        b"",
    )?;
    let _ = send_udp(socket, &out, peer).await?;
    Ok(())
}

async fn handle_detach(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
    let (key, _) = find_client_decrypt_key(state, &packet.header)?;
    let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let mut locked = state.lock().expect("server state poisoned");
    if let Some(client) = locked.client_mut(&packet.header.conn_id)
        && !client.replay.accept(packet.header.seq)
    {
        return Ok(());
    }
    locked.remove_client_everywhere(&packet.header.conn_id);
    Ok(())
}

fn remove_client(state: &Arc<Mutex<ServerState>>, client_id: [u8; 16]) {
    let mut locked = state.lock().expect("server state poisoned");
    locked.remove_client_everywhere(&client_id);
}

async fn send_udp(socket: &UdpSocket, packet: &[u8], peer: SocketAddr) -> Result<bool> {
    match socket.send_to(packet, peer).await {
        Ok(_) => Ok(true),
        Err(err) if is_transient_udp_send_error(&err) => Ok(false),
        Err(err) => Err(err).with_context(|| format!("send UDP packet to {peer}")),
    }
}

async fn handle_ack(
    state: &Arc<Mutex<ServerState>>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, _) = find_client_decrypt_key(state, &packet.header)?;
    let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let mut locked = state.lock().expect("server state poisoned");
    if let Some(client) = locked.client_mut(&packet.header.conn_id) {
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        // Connection migration on any authenticated, fresh packet (spec §11).
        client.endpoint = peer;
        client.last_seen = Instant::now();
        client.last_acked = packet.header.ack;
        while client
            .pending
            .front()
            .is_some_and(|pending| pending.output_seq <= packet.header.ack)
        {
            client.pending.pop_front();
        }
    }
    Ok(())
}

/// The client confirms it has switched to the new epoch key. Sealed under the
/// new (now-current) key, so decrypting it authenticates the confirmation. We
/// retire the previous-epoch key immediately, ending the grace early.
async fn handle_rekey_ack(
    state: &Arc<Mutex<ServerState>>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, _) = find_client_decrypt_key(state, &packet.header)?;
    let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let mut locked = state.lock().expect("server state poisoned");
    if let Some(client) = locked.client_mut(&packet.header.conn_id) {
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        client.endpoint = peer;
        client.last_seen = Instant::now();
        // Only retire the old key if the ack was sealed under the *current* key,
        // confirming the client really is on the new epoch.
        if protocol::session_key_id(&client.session_key) == packet.header.session_key_id {
            client.previous_session_key = None;
        }
    }
    Ok(())
}

async fn handle_stream_open(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let open: StreamOpen = protocol::from_body(&body)?;
    enum StreamOpenCheck {
        Duplicate,
        New(Result<()>),
    }
    let check = {
        let mut locked = state.lock().expect("server state poisoned");
        let config = locked.config.clone();
        let session = locked
            .sessions
            .get_mut(&session_name)
            .ok_or_else(|| anyhow!("unknown session"))?;
        let client = session
            .clients
            .get_mut(&packet.header.conn_id)
            .ok_or_else(|| anyhow!("unknown client"))?;
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        client.endpoint = peer;
        client.last_seen = Instant::now();
        if client.stream_retired.contains(&open.stream_id) {
            StreamOpenCheck::New(Err(anyhow!("stream {} is closed", open.stream_id)))
        } else if client.opened_streams.contains(&open.stream_id) {
            StreamOpenCheck::Duplicate
        } else {
            StreamOpenCheck::New(stream_open_allowed(&config, client, &open))
        }
    };
    match check {
        StreamOpenCheck::Duplicate => {
            return send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await;
        }
        StreamOpenCheck::New(Ok(())) => {}
        StreamOpenCheck::New(Err(err)) => {
            return send_stream_open_reject(
                state,
                socket,
                packet.header.conn_id,
                open.stream_id,
                err.to_string(),
            )
            .await;
        }
    }

    if open.target_host == FILE_STREAM_SENTINEL {
        let (writer_tx, writer_rx) = mpsc::channel::<Vec<u8>>(1024);
        register_opened_stream(
            state,
            &session_name,
            packet.header.conn_id,
            open.stream_id,
            writer_tx,
        )?;
        send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;

        let state = Arc::clone(state);
        let socket = Arc::clone(socket);
        let client_id = packet.header.conn_id;
        let stream_id = open.stream_id;
        tokio::spawn(async move {
            if let Err(err) = run_file_stream_service(
                state.clone(),
                socket.clone(),
                client_id,
                stream_id,
                writer_rx,
            )
            .await
            {
                let _ = send_file_response_to_client(
                    &state,
                    &socket,
                    client_id,
                    stream_id,
                    FileResponse::Error {
                        message: err.to_string(),
                    },
                )
                .await;
            }
            let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
        });
        return Ok(());
    }

    if open.target_host == EXEC_STREAM_SENTINEL {
        let (writer_tx, writer_rx) = mpsc::channel::<Vec<u8>>(1024);
        register_opened_stream(
            state,
            &session_name,
            packet.header.conn_id,
            open.stream_id,
            writer_tx,
        )?;
        send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;

        let state = Arc::clone(state);
        let socket = Arc::clone(socket);
        let client_id = packet.header.conn_id;
        let stream_id = open.stream_id;
        tokio::spawn(async move {
            if let Err(err) = run_exec_stream_service(
                state.clone(),
                socket.clone(),
                client_id,
                stream_id,
                writer_rx,
            )
            .await
            {
                let _ = send_exec_response_to_client(
                    &state,
                    &socket,
                    client_id,
                    stream_id,
                    ExecResponse::Error {
                        message: err.to_string(),
                    },
                )
                .await;
            }
            let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
        });
        return Ok(());
    }

    let target = format!("{}:{}", open.target_host, open.target_port);
    let stream = match TcpStream::connect((open.target_host.as_str(), open.target_port)).await {
        Ok(stream) => stream,
        Err(err) => {
            return send_stream_open_reject(
                state,
                socket,
                packet.header.conn_id,
                open.stream_id,
                format!("connect {target}: {err}"),
            )
            .await;
        }
    };
    let (mut reader, mut writer) = stream.into_split();
    let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
    register_opened_stream(
        state,
        &session_name,
        packet.header.conn_id,
        open.stream_id,
        writer_tx,
    )?;

    tokio::spawn(async move {
        while let Some(bytes) = writer_rx.recv().await {
            if writer.write_all(&bytes).await.is_err() {
                break;
            }
        }
    });

    send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;

    let state = Arc::clone(state);
    let socket = Arc::clone(socket);
    let client_id = packet.header.conn_id;
    let stream_id = open.stream_id;
    tokio::spawn(async move {
        let mut buf = [0u8; 16 * 1024];
        loop {
            match reader.read(&mut buf).await {
                Ok(0) => break,
                Ok(n) => {
                    if let Err(err) = send_stream_data_to_client(
                        &state,
                        &socket,
                        client_id,
                        stream_id,
                        buf[..n].to_vec(),
                    )
                    .await
                    {
                        eprintln!("stream send error: {err:#}");
                        break;
                    }
                }
                Err(_) => break,
            }
        }
        let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
    });

    Ok(())
}

fn register_opened_stream(
    state: &Arc<Mutex<ServerState>>,
    session_name: &str,
    client_id: [u8; 16],
    stream_id: u64,
    writer_tx: mpsc::Sender<Vec<u8>>,
) -> Result<()> {
    let mut locked = state.lock().expect("server state poisoned");
    let session = locked
        .sessions
        .get_mut(session_name)
        .ok_or_else(|| anyhow!("unknown session"))?;
    let client = session
        .clients
        .get_mut(&client_id)
        .ok_or_else(|| anyhow!("unknown client"))?;
    client.stream_writers.insert(stream_id, writer_tx);
    client.opened_streams.insert(stream_id);
    client
        .stream_send_credit
        .insert(stream_id, STREAM_INITIAL_WINDOW);
    client.stream_next_send_offset.entry(stream_id).or_insert(0);
    client.stream_next_recv_offset.entry(stream_id).or_insert(0);
    Ok(())
}

async fn handle_stream_data(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let data: StreamData = protocol::from_body(&body)?;
    let stream_id = data.stream_id;
    let (writer, writes, consumed, received_offset) = {
        let mut locked = state.lock().expect("server state poisoned");
        let session = locked
            .sessions
            .get_mut(&session_name)
            .ok_or_else(|| anyhow!("unknown session"))?;
        let client = session
            .clients
            .get_mut(&packet.header.conn_id)
            .ok_or_else(|| anyhow!("unknown client"))?;
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        client.endpoint = peer;
        client.last_seen = Instant::now();
        if !client.opened_streams.contains(&data.stream_id) {
            return Ok(());
        }
        let writer = client.stream_writers.get(&data.stream_id).cloned();
        let (writes, consumed, received_offset) = accept_stream_data(client, data);
        (writer, writes, consumed, received_offset)
    };
    if let Some(writer) = writer {
        for bytes in writes {
            let _ = writer.send(bytes).await;
        }
    }
    // Always return flow-control credit for live streams, even when the stream's
    // writer is already gone. The peer debited its send window for these bytes,
    // so skipping the adjust would wedge the stream at a permanent credit deficit.
    send_stream_window_adjust_to_client(
        state,
        socket,
        packet.header.conn_id,
        stream_id,
        consumed,
        received_offset,
    )
    .await?;
    Ok(())
}

async fn handle_stream_open_ok(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let ok: StreamOpenOk = protocol::from_body(&body)?;
    {
        let mut locked = state.lock().expect("server state poisoned");
        let session = locked
            .sessions
            .get_mut(&session_name)
            .ok_or_else(|| anyhow!("unknown session"))?;
        let client = session
            .clients
            .get_mut(&packet.header.conn_id)
            .ok_or_else(|| anyhow!("unknown client"))?;
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        client.endpoint = peer;
        client.last_seen = Instant::now();
        if client.stream_pending_opens.remove(&ok.stream_id).is_none() {
            return Ok(());
        }
        client.opened_streams.insert(ok.stream_id);
        client
            .stream_send_credit
            .entry(ok.stream_id)
            .or_insert(STREAM_INITIAL_WINDOW);
    }
    flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, ok.stream_id).await
}

async fn handle_stream_open_reject(
    state: &Arc<Mutex<ServerState>>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let reject: StreamOpenReject = protocol::from_body(&body)?;
    let mut locked = state.lock().expect("server state poisoned");
    let session = locked
        .sessions
        .get_mut(&session_name)
        .ok_or_else(|| anyhow!("unknown session"))?;
    let client = session
        .clients
        .get_mut(&packet.header.conn_id)
        .ok_or_else(|| anyhow!("unknown client"))?;
    if !client.replay.accept(packet.header.seq) {
        return Ok(());
    }
    client.endpoint = peer;
    client.last_seen = Instant::now();
    retire_client_stream(client, reject.stream_id);
    Ok(())
}

async fn handle_stream_window_adjust(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let adjust: StreamWindowAdjust = protocol::from_body(&body)?;
    {
        let mut locked = state.lock().expect("server state poisoned");
        let session = locked
            .sessions
            .get_mut(&session_name)
            .ok_or_else(|| anyhow!("unknown session"))?;
        let client = session
            .clients
            .get_mut(&packet.header.conn_id)
            .ok_or_else(|| anyhow!("unknown client"))?;
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        client.endpoint = peer;
        client.last_seen = Instant::now();
        if !client.opened_streams.contains(&adjust.stream_id) {
            return Ok(());
        }
        ack_stream_data(client, adjust.stream_id, adjust.received_offset);
    }
    flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, adjust.stream_id)
        .await
}

async fn handle_stream_close(
    state: &Arc<Mutex<ServerState>>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let close: StreamClose = protocol::from_body(&body)?;
    let mut locked = state.lock().expect("server state poisoned");
    let Some(session) = locked.sessions.get_mut(&session_name) else {
        return Ok(());
    };
    let Some(client) = session.clients.get_mut(&packet.header.conn_id) else {
        return Ok(());
    };
    if !client.replay.accept(packet.header.seq) {
        return Ok(());
    }
    // Connection migration on any authenticated, fresh packet (spec §11).
    client.endpoint = peer;
    client.last_seen = Instant::now();
    retire_client_stream(client, close.stream_id);
    Ok(())
}

async fn start_remote_forwards(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    requested_forwardings: Vec<ForwardingRequest>,
    session_name: String,
) -> Result<()> {
    let remote_forwards = requested_forwardings
        .into_iter()
        .filter(|forward| forward.kind == ForwardingKind::Remote)
        .collect::<Vec<_>>();
    if remote_forwards.is_empty() {
        return Ok(());
    }
    let config = {
        let locked = state.lock().expect("server state poisoned");
        locked.config.clone()
    };
    anyhow::ensure!(
        config.allow_remote_forwarding,
        "remote forwarding disabled on server"
    );
    for forward in remote_forwards {
        let bind_host = forward
            .bind_host
            .clone()
            .unwrap_or_else(|| "127.0.0.1".to_string());
        remote_bind_allowed(&config, &bind_host)?;
        let listen_port = forward.listen_port;
        let target_host = forward
            .target_host
            .clone()
            .ok_or_else(|| anyhow!("remote forward target host is required"))?;
        let target_port = forward
            .target_port
            .ok_or_else(|| anyhow!("remote forward target port is required"))?;
        let bind = format!("{bind_host}:{listen_port}");
        let listener = TcpListener::bind(&bind)
            .await
            .with_context(|| format!("bind remote forward {bind}"))?;
        let state = Arc::clone(state);
        let socket = Arc::clone(socket);
        let session_name = session_name.clone();
        tokio::spawn(async move {
            loop {
                let Ok((stream, _)) = listener.accept().await else {
                    break;
                };
                let stream_id = allocate_server_stream_id(&state);
                let (mut reader, mut writer) = stream.into_split();
                let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
                {
                    let mut locked = state.lock().expect("server state poisoned");
                    if let Some(session) = locked.sessions.get_mut(&session_name)
                        && let Some(client) = session.clients.get_mut(&client_id)
                    {
                        client.stream_writers.insert(stream_id, writer_tx);
                        client
                            .stream_send_credit
                            .insert(stream_id, STREAM_INITIAL_WINDOW);
                    } else {
                        break;
                    }
                }
                tokio::spawn(async move {
                    while let Some(bytes) = writer_rx.recv().await {
                        if writer.write_all(&bytes).await.is_err() {
                            break;
                        }
                    }
                });
                if let Err(err) = send_stream_open_to_client(
                    &state,
                    &socket,
                    client_id,
                    stream_id,
                    target_host.clone(),
                    target_port,
                )
                .await
                {
                    eprintln!("remote forward open error: {err:#}");
                    break;
                }
                let state = Arc::clone(&state);
                let socket = Arc::clone(&socket);
                tokio::spawn(async move {
                    let mut buf = [0u8; 16 * 1024];
                    loop {
                        match reader.read(&mut buf).await {
                            Ok(0) => break,
                            Ok(n) => {
                                if let Err(err) = send_stream_data_to_client(
                                    &state,
                                    &socket,
                                    client_id,
                                    stream_id,
                                    buf[..n].to_vec(),
                                )
                                .await
                                {
                                    eprintln!("remote forward data error: {err:#}");
                                    break;
                                }
                            }
                            Err(_) => break,
                        }
                    }
                    let _ =
                        send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
                });
            }
        });
    }
    Ok(())
}

fn allocate_server_stream_id(state: &Arc<Mutex<ServerState>>) -> u64 {
    let mut locked = state.lock().expect("server state poisoned");
    let stream_id = locked.next_server_stream_id;
    locked.next_server_stream_id = locked.next_server_stream_id.wrapping_add(1).max(1u64 << 63);
    stream_id
}

/// Whether the client requested SSH-agent forwarding during auth.
fn agent_forwarding_requested(forwardings: &[ForwardingRequest]) -> bool {
    forwardings.iter().any(|f| f.kind == ForwardingKind::Agent)
}

/// Bind the per-session SSH-agent proxy unix socket, returning its `UnixListener`
/// and path. SECURITY: the socket lives in a process-private directory created
/// 0700 and the socket itself is chmod 0600, so only this user can connect — the
/// forwarded agent is never world-reachable. Called only when the client opted in
/// AND `allow_agent_forwarding` is set; the path is then exported as the remote
/// session's `SSH_AUTH_SOCK`.
fn bind_agent_proxy_socket() -> Result<(UnixListener, PathBuf)> {
    let dir = std::env::temp_dir().join(format!("dosh-agent-{}", std::process::id()));
    std::fs::create_dir_all(&dir).with_context(|| format!("create agent dir {}", dir.display()))?;
    std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700))
        .with_context(|| format!("chmod agent dir {}", dir.display()))?;
    let n = AGENT_SOCK_COUNTER.fetch_add(1, Ordering::Relaxed);
    let path = dir.join(format!("agent.{}.{n}.sock", std::process::id()));
    // Stale socket from a crashed prior run with the same name: remove first.
    let _ = std::fs::remove_file(&path);
    let listener = UnixListener::bind(&path)
        .with_context(|| format!("bind agent socket {}", path.display()))?;
    std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))
        .with_context(|| format!("chmod agent socket {}", path.display()))?;
    Ok((listener, path))
}

/// Run the accept loop for a forwarded SSH-agent: each connection from a remote
/// process to the proxy socket is tunneled to the client (which splices it into
/// the user's real ssh-agent) over a Dosh stream using the agent sentinel target.
/// The socket file is removed when the loop ends (client gone / session closed).
fn spawn_agent_forward(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    listener: UnixListener,
    socket_path: PathBuf,
) {
    let state = Arc::clone(state);
    let socket = Arc::clone(socket);
    tokio::spawn(async move {
        loop {
            let Ok((stream, _)) = listener.accept().await else {
                break;
            };
            // Stop accepting once the client has gone away.
            {
                let locked = state.lock().expect("server state poisoned");
                if locked.lookup_client(&client_id).is_none() {
                    break;
                }
            }
            let stream_id = allocate_server_stream_id(&state);
            let (mut reader, mut writer) = stream.into_split();
            let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
            {
                let mut locked = state.lock().expect("server state poisoned");
                if let Some(client) = locked.client_mut(&client_id) {
                    client.stream_writers.insert(stream_id, writer_tx);
                    client
                        .stream_send_credit
                        .insert(stream_id, STREAM_INITIAL_WINDOW);
                } else {
                    break;
                }
            }
            tokio::spawn(async move {
                while let Some(bytes) = writer_rx.recv().await {
                    if writer.write_all(&bytes).await.is_err() {
                        break;
                    }
                }
            });
            if let Err(err) = send_stream_open_to_client(
                &state,
                &socket,
                client_id,
                stream_id,
                AGENT_STREAM_SENTINEL.to_string(),
                0,
            )
            .await
            {
                eprintln!("agent forward open error: {err:#}");
                break;
            }
            let state = Arc::clone(&state);
            let socket = Arc::clone(&socket);
            tokio::spawn(async move {
                let mut buf = [0u8; 16 * 1024];
                loop {
                    match reader.read(&mut buf).await {
                        Ok(0) => break,
                        Ok(n) => {
                            if let Err(err) = send_stream_data_to_client(
                                &state,
                                &socket,
                                client_id,
                                stream_id,
                                buf[..n].to_vec(),
                            )
                            .await
                            {
                                eprintln!("agent forward data error: {err:#}");
                                break;
                            }
                        }
                        Err(_) => break,
                    }
                }
                let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
            });
        }
        let _ = std::fs::remove_file(&socket_path);
    });
}

fn is_loopback_bind(host: &str) -> bool {
    matches!(host, "127.0.0.1" | "localhost" | "::1")
}

fn remote_bind_allowed(config: &ServerConfig, host: &str) -> Result<()> {
    if !is_loopback_bind(host) {
        anyhow::ensure!(
            config.allow_remote_non_loopback_bind,
            "remote forwarding non-loopback bind disabled"
        );
    }
    Ok(())
}

fn stream_open_allowed(
    config: &ServerConfig,
    client: &ClientState,
    open: &StreamOpen,
) -> Result<()> {
    if open.target_host == FILE_STREAM_SENTINEL {
        anyhow::ensure!(config.allow_file_transfer, "file transfer disabled");
        let allowed = client.allowed_forwardings.iter().any(|forward| {
            forward.kind == ForwardingKind::Local
                && forward.target_host.as_deref() == Some(FILE_STREAM_SENTINEL)
                && forward.target_port == Some(0)
        });
        anyhow::ensure!(allowed, "file transfer was not requested during auth");
        return Ok(());
    }
    if open.target_host == EXEC_STREAM_SENTINEL {
        anyhow::ensure!(config.allow_exec_command, "exec command disabled");
        let allowed = client.allowed_forwardings.iter().any(|forward| {
            forward.kind == ForwardingKind::Local
                && forward.target_host.as_deref() == Some(EXEC_STREAM_SENTINEL)
                && forward.target_port == Some(0)
        });
        anyhow::ensure!(allowed, "exec command was not requested during auth");
        return Ok(());
    }
    anyhow::ensure!(config.allow_tcp_forwarding, "TCP forwarding disabled");
    let allowed = client
        .allowed_forwardings
        .iter()
        .any(|forward| match forward.kind {
            ForwardingKind::Local => {
                forward.target_host.as_deref() == Some(open.target_host.as_str())
                    && forward.target_port == Some(open.target_port)
            }
            ForwardingKind::Dynamic => true,
            ForwardingKind::Remote => false,
            // Agent forwarding never authorizes a client-initiated TCP open.
            ForwardingKind::Agent => false,
        });
    anyhow::ensure!(
        allowed,
        "stream target {}:{} was not requested during auth",
        open.target_host,
        open.target_port
    );
    Ok(())
}

struct UploadState {
    final_path: PathBuf,
    temp_path: Option<PathBuf>,
    file: fs::File,
    hasher: Sha256,
    written: u64,
    expected_size: u64,
    modified_secs: Option<u64>,
    atomic: bool,
}

async fn run_file_stream_service(
    state: Arc<Mutex<ServerState>>,
    socket: Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    mut rx: mpsc::Receiver<Vec<u8>>,
) -> Result<()> {
    let home = std::env::var_os("HOME")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("."));
    let mut decoder = FrameDecoder::default();
    let mut upload: Option<UploadState> = None;
    while let Some(bytes) = rx.recv().await {
        for frame in decoder.push(&bytes)? {
            let request = match dosh::file_transfer::decode_request(&frame) {
                Ok(request) => request,
                Err(err) => {
                    send_file_response_to_client(
                        &state,
                        &socket,
                        client_id,
                        stream_id,
                        FileResponse::Error {
                            message: err.to_string(),
                        },
                    )
                    .await?;
                    continue;
                }
            };
            if let Err(err) = handle_file_request(
                &state,
                &socket,
                client_id,
                stream_id,
                &home,
                &mut upload,
                request,
            )
            .await
            {
                send_file_response_to_client(
                    &state,
                    &socket,
                    client_id,
                    stream_id,
                    FileResponse::Error {
                        message: err.to_string(),
                    },
                )
                .await?;
            }
        }
    }
    Ok(())
}

async fn handle_file_request(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    home: &Path,
    upload: &mut Option<UploadState>,
    request: FileRequest,
) -> Result<()> {
    match request {
        FileRequest::Stat { path } => {
            let path = clean_remote_path(&path, home)?;
            let meta = file_meta(&path, home)?;
            send_file_response_to_client(
                state,
                socket,
                client_id,
                stream_id,
                FileResponse::Stat { meta },
            )
            .await
        }
        FileRequest::List { path } => {
            let path = clean_remote_path(&path, home)?;
            let mut entries = Vec::new();
            for entry in fs::read_dir(&path).with_context(|| format!("list {}", path.display()))? {
                let entry = entry?;
                let name = entry.file_name().to_string_lossy().to_string();
                entries.push(FileEntry {
                    name,
                    meta: file_meta(&entry.path(), home)?,
                });
            }
            entries.sort_by(|a, b| a.name.cmp(&b.name));
            send_file_response_to_client(
                state,
                socket,
                client_id,
                stream_id,
                FileResponse::List { entries },
            )
            .await
        }
        FileRequest::Mkdir { path, mode } => {
            let path = clean_remote_path(&path, home)?;
            fs::create_dir_all(&path).with_context(|| format!("create {}", path.display()))?;
            if let Some(mode) = mode {
                set_mode(&path, mode)?;
            }
            send_file_response_to_client(state, socket, client_id, stream_id, FileResponse::Ok)
                .await
        }
        FileRequest::Remove { path, recursive } => {
            let path = clean_remote_path(&path, home)?;
            let metadata =
                fs::symlink_metadata(&path).with_context(|| format!("stat {}", path.display()))?;
            if metadata.is_dir() {
                anyhow::ensure!(
                    recursive,
                    "{} is a directory; pass -r to remove directories",
                    path.display()
                );
                fs::remove_dir_all(&path).with_context(|| format!("remove {}", path.display()))?;
            } else {
                fs::remove_file(&path).with_context(|| format!("remove {}", path.display()))?;
            }
            send_file_response_to_client(state, socket, client_id, stream_id, FileResponse::Ok)
                .await
        }
        FileRequest::PutStart {
            path,
            size,
            mode,
            modified_secs,
            overwrite,
            resume,
        } => {
            if upload.is_some() {
                bail!("upload already in progress");
            }
            let final_path = clean_remote_path(&path, home)?;
            if let Some(parent) = final_path.parent() {
                fs::create_dir_all(parent)
                    .with_context(|| format!("create {}", parent.display()))?;
            }
            if final_path.exists() && !overwrite && !resume {
                bail!("destination exists: {}", final_path.display());
            }
            let mut hasher = Sha256::new();
            let (file, temp_path, written, atomic) = if resume && final_path.exists() {
                let mut file = fs::OpenOptions::new()
                    .read(true)
                    .append(true)
                    .open(&final_path)
                    .with_context(|| format!("open {}", final_path.display()))?;
                let written = file.metadata()?.len().min(size);
                file.set_len(written)
                    .with_context(|| format!("truncate {}", final_path.display()))?;
                if written > 0 {
                    let mut prefix = fs::File::open(&final_path)?;
                    hash_exact_prefix(&mut prefix, written, &mut hasher)?;
                }
                file.seek(SeekFrom::Start(written))?;
                (file, None, written, false)
            } else {
                let temp_path = upload_temp_path(&final_path, stream_id);
                let file = fs::OpenOptions::new()
                    .create(true)
                    .truncate(true)
                    .write(true)
                    .open(&temp_path)
                    .with_context(|| format!("create {}", temp_path.display()))?;
                (file, Some(temp_path), 0, true)
            };
            *upload = Some(UploadState {
                final_path,
                temp_path,
                file,
                hasher,
                written,
                expected_size: size,
                modified_secs,
                atomic,
            });
            if let Some(mode) = mode {
                let target = upload
                    .as_ref()
                    .and_then(|state| state.temp_path.as_ref())
                    .unwrap_or_else(|| &upload.as_ref().unwrap().final_path)
                    .clone();
                set_mode(&target, mode)?;
            }
            send_file_response_to_client(
                state,
                socket,
                client_id,
                stream_id,
                FileResponse::Resume { offset: written },
            )
            .await
        }
        FileRequest::PutChunk { offset, bytes } => {
            let Some(active) = upload.as_mut() else {
                bail!("no upload in progress");
            };
            anyhow::ensure!(
                offset == active.written,
                "upload offset mismatch: got {offset}, expected {}",
                active.written
            );
            active.file.write_all(&bytes)?;
            active.hasher.update(&bytes);
            active.written = active.written.saturating_add(bytes.len() as u64);
            Ok(())
        }
        FileRequest::PutFinish { sha256 } => {
            let Some(mut active) = upload.take() else {
                bail!("no upload in progress");
            };
            active.file.flush()?;
            active.file.sync_all()?;
            anyhow::ensure!(
                active.written == active.expected_size,
                "upload size mismatch: got {}, expected {}",
                active.written,
                active.expected_size
            );
            let digest: [u8; 32] = active.hasher.finalize().into();
            if digest != sha256 {
                if let Some(temp_path) = &active.temp_path {
                    let _ = fs::remove_file(temp_path);
                }
                bail!("upload checksum mismatch");
            }
            drop(active.file);
            if active.atomic
                && let Some(temp_path) = active.temp_path.take()
            {
                fs::rename(&temp_path, &active.final_path).with_context(|| {
                    format!(
                        "rename {} to {}",
                        temp_path.display(),
                        active.final_path.display()
                    )
                })?;
            }
            if let Some(modified_secs) = active.modified_secs {
                filetime::set_file_mtime(
                    &active.final_path,
                    filetime::FileTime::from_unix_time(modified_secs as i64, 0),
                )
                .with_context(|| format!("set mtime {}", active.final_path.display()))?;
            }
            send_file_response_to_client(
                state,
                socket,
                client_id,
                stream_id,
                FileResponse::Done {
                    sha256: digest,
                    bytes: active.written,
                },
            )
            .await
        }
        FileRequest::Get { path, offset } => {
            let path = clean_remote_path(&path, home)?;
            let meta = file_meta(&path, home)?;
            anyhow::ensure!(
                meta.kind == FileKind::File,
                "remote path is not a file: {}",
                meta.path
            );
            let mut file =
                fs::File::open(&path).with_context(|| format!("open {}", path.display()))?;
            file.seek(SeekFrom::Start(offset))?;
            send_file_response_to_client(
                state,
                socket,
                client_id,
                stream_id,
                FileResponse::Start {
                    meta: meta.clone(),
                    offset,
                },
            )
            .await?;
            let mut hasher = Sha256::new();
            if offset > 0 {
                let mut prefix = fs::File::open(&path)?;
                hash_exact_prefix(&mut prefix, offset, &mut hasher)?;
            }
            let mut sent = offset;
            let mut buf = vec![0u8; CHUNK_SIZE];
            loop {
                let n = file.read(&mut buf)?;
                if n == 0 {
                    break;
                }
                hasher.update(&buf[..n]);
                send_file_response_to_client(
                    state,
                    socket,
                    client_id,
                    stream_id,
                    FileResponse::Chunk {
                        offset: sent,
                        bytes: buf[..n].to_vec(),
                    },
                )
                .await?;
                sent = sent.saturating_add(n as u64);
            }
            let digest: [u8; 32] = hasher.finalize().into();
            send_file_response_to_client(
                state,
                socket,
                client_id,
                stream_id,
                FileResponse::Done {
                    sha256: digest,
                    bytes: sent,
                },
            )
            .await
        }
    }
}

async fn send_file_response_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    response: FileResponse,
) -> Result<()> {
    send_stream_data_to_client(
        state,
        socket,
        client_id,
        stream_id,
        encode_response(&response)?,
    )
    .await
}

async fn run_exec_stream_service(
    state: Arc<Mutex<ServerState>>,
    socket: Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    mut rx: mpsc::Receiver<Vec<u8>>,
) -> Result<()> {
    let mut decoder = ExecFrameDecoder::default();
    while let Some(bytes) = rx.recv().await {
        for frame in decoder.push(&bytes)? {
            let request = match dosh::exec_service::decode_request(&frame) {
                Ok(request) => request,
                Err(err) => {
                    send_exec_response_to_client(
                        &state,
                        &socket,
                        client_id,
                        stream_id,
                        ExecResponse::Error {
                            message: err.to_string(),
                        },
                    )
                    .await?;
                    continue;
                }
            };
            let output = TokioCommand::new("sh")
                .arg("-lc")
                .arg(&request.command)
                .stdin(Stdio::null())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped())
                .output()
                .await
                .with_context(|| format!("run command {:?}", request.command))?;
            for chunk in output.stdout.chunks(CHUNK_SIZE) {
                send_exec_response_to_client(
                    &state,
                    &socket,
                    client_id,
                    stream_id,
                    ExecResponse::Stdout(chunk.to_vec()),
                )
                .await?;
            }
            for chunk in output.stderr.chunks(CHUNK_SIZE) {
                send_exec_response_to_client(
                    &state,
                    &socket,
                    client_id,
                    stream_id,
                    ExecResponse::Stderr(chunk.to_vec()),
                )
                .await?;
            }
            send_exec_response_to_client(
                &state,
                &socket,
                client_id,
                stream_id,
                ExecResponse::Exit {
                    code: output.status.code(),
                },
            )
            .await?;
        }
    }
    Ok(())
}

async fn send_exec_response_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    response: ExecResponse,
) -> Result<()> {
    send_stream_data_to_client(
        state,
        socket,
        client_id,
        stream_id,
        encode_exec_response(&response)?,
    )
    .await
}

fn file_meta(path: &Path, home: &Path) -> Result<FileMeta> {
    let metadata =
        fs::symlink_metadata(path).with_context(|| format!("stat {}", path.display()))?;
    let file_type = metadata.file_type();
    let kind = if file_type.is_file() {
        FileKind::File
    } else if file_type.is_dir() {
        FileKind::Directory
    } else if file_type.is_symlink() {
        FileKind::Symlink
    } else {
        FileKind::Other
    };
    let modified_secs = metadata
        .modified()
        .ok()
        .and_then(|mtime| mtime.duration_since(std::time::UNIX_EPOCH).ok())
        .map(|duration| duration.as_secs());
    let display_path = path
        .strip_prefix(home)
        .ok()
        .and_then(|relative| relative.to_str())
        .filter(|relative| !relative.is_empty())
        .map(|relative| format!("~/{relative}"))
        .unwrap_or_else(|| path.display().to_string());
    Ok(FileMeta {
        path: display_path,
        kind,
        len: metadata.len(),
        mode: Some(metadata.permissions().mode() & 0o7777),
        modified_secs,
    })
}

fn set_mode(path: &Path, mode: u32) -> Result<()> {
    let mut permissions = fs::metadata(path)?.permissions();
    permissions.set_mode(mode & 0o7777);
    fs::set_permissions(path, permissions).with_context(|| format!("chmod {}", path.display()))
}

fn upload_temp_path(final_path: &Path, stream_id: u64) -> PathBuf {
    let name = final_path
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or("upload");
    final_path.with_file_name(format!(
        ".{name}.dosh-upload-{}-{stream_id}",
        std::process::id()
    ))
}

fn hash_exact_prefix(file: &mut fs::File, bytes: u64, hasher: &mut Sha256) -> Result<()> {
    file.seek(SeekFrom::Start(0))?;
    let mut remaining = bytes;
    let mut buf = vec![0u8; CHUNK_SIZE];
    while remaining > 0 {
        let want = remaining.min(buf.len() as u64) as usize;
        let n = file.read(&mut buf[..want])?;
        if n == 0 {
            bail!("file ended while hashing resume prefix");
        }
        hasher.update(&buf[..n]);
        remaining -= n as u64;
    }
    Ok(())
}

async fn send_stream_open_ok(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
) -> Result<()> {
    let body = protocol::to_body(&StreamOpenOk { stream_id })?;
    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenOk, body).await
}

async fn send_stream_open_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    target_host: String,
    target_port: u16,
) -> Result<()> {
    let body = protocol::to_body(&StreamOpen {
        stream_id,
        target_host: target_host.clone(),
        target_port,
    })?;
    let send = {
        let mut locked = state.lock().expect("server state poisoned");
        let mut found = None;
        if let Some(client) = locked.client_mut(&client_id) {
            client.send_seq += 1;
            let packet = protocol::encode_encrypted(
                PacketKind::StreamOpen,
                client_id,
                client.send_seq,
                client.last_acked,
                &client.session_key,
                SERVER_TO_CLIENT,
                &body,
            )?;
            client.stream_pending_opens.insert(
                stream_id,
                PendingStreamOpen {
                    target_host,
                    target_port,
                    last_sent: Instant::now(),
                    attempts: 1,
                },
            );
            found = Some((client.endpoint, packet));
        }
        found
    };
    if let Some((endpoint, packet)) = send {
        let _ = send_udp(socket, &packet, endpoint).await?;
    }
    Ok(())
}

async fn send_stream_open_reject(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    reason: String,
) -> Result<()> {
    let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?;
    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenReject, body).await
}

async fn send_stream_data_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    bytes: Vec<u8>,
) -> Result<()> {
    queue_or_send_stream_data_to_client(state, socket, client_id, stream_id, bytes).await
}

async fn queue_or_send_stream_data_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    bytes: Vec<u8>,
) -> Result<()> {
    let send = {
        let mut locked = state.lock().expect("server state poisoned");
        let mut send = None;
        if let Some(client) = locked.client_mut(&client_id) {
            if !client.opened_streams.contains(&stream_id)
                || client
                    .stream_send_credit
                    .get(&stream_id)
                    .copied()
                    .unwrap_or(0)
                    < bytes.len()
                || client
                    .stream_pending_data
                    .get(&stream_id)
                    .is_some_and(|pending| !pending.is_empty())
            {
                client
                    .stream_pending_data
                    .entry(stream_id)
                    .or_default()
                    .push_back(bytes);
                return Ok(());
            }
            *client.stream_send_credit.entry(stream_id).or_default() -= bytes.len();
            let offset = *client.stream_next_send_offset.entry(stream_id).or_default();
            client
                .stream_next_send_offset
                .insert(stream_id, offset.saturating_add(bytes.len() as u64));
            client.send_seq += 1;
            let body = protocol::to_body(&StreamData {
                stream_id,
                offset,
                bytes: bytes.clone(),
            })?;
            let packet = protocol::encode_encrypted(
                PacketKind::StreamData,
                client_id,
                client.send_seq,
                client.last_acked,
                &client.session_key,
                SERVER_TO_CLIENT,
                &body,
            )?;
            client
                .stream_sent_data
                .entry(stream_id)
                .or_default()
                .insert(
                    offset,
                    PendingStreamChunk {
                        offset,
                        bytes,
                        last_sent: Instant::now(),
                        attempts: 1,
                    },
                );
            send = Some((client.endpoint, packet));
        }
        send
    };
    if let Some((endpoint, packet)) = send {
        let _ = send_udp(socket, &packet, endpoint).await?;
    }
    Ok(())
}

async fn flush_stream_pending_data_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
) -> Result<()> {
    loop {
        let send = {
            let mut locked = state.lock().expect("server state poisoned");
            let mut send = None;
            if let Some(client) = locked.client_mut(&client_id) {
                let Some(pending) = client.stream_pending_data.get_mut(&stream_id) else {
                    return Ok(());
                };
                let Some(bytes) = pending.front() else {
                    client.stream_pending_data.remove(&stream_id);
                    return Ok(());
                };
                let credit = client
                    .stream_send_credit
                    .get(&stream_id)
                    .copied()
                    .unwrap_or(0);
                if credit < bytes.len() {
                    return Ok(());
                }
                let bytes = pending.pop_front().expect("pending front exists");
                if pending.is_empty() {
                    client.stream_pending_data.remove(&stream_id);
                }
                *client.stream_send_credit.entry(stream_id).or_default() -= bytes.len();
                let offset = *client.stream_next_send_offset.entry(stream_id).or_default();
                client
                    .stream_next_send_offset
                    .insert(stream_id, offset.saturating_add(bytes.len() as u64));
                client.send_seq += 1;
                let body = protocol::to_body(&StreamData {
                    stream_id,
                    offset,
                    bytes: bytes.clone(),
                })?;
                let packet = protocol::encode_encrypted(
                    PacketKind::StreamData,
                    client_id,
                    client.send_seq,
                    client.last_acked,
                    &client.session_key,
                    SERVER_TO_CLIENT,
                    &body,
                )?;
                client
                    .stream_sent_data
                    .entry(stream_id)
                    .or_default()
                    .insert(
                        offset,
                        PendingStreamChunk {
                            offset,
                            bytes,
                            last_sent: Instant::now(),
                            attempts: 1,
                        },
                    );
                send = Some((client.endpoint, packet));
            }
            send
        };
        let Some((endpoint, packet)) = send else {
            return Ok(());
        };
        let _ = send_udp(socket, &packet, endpoint).await?;
    }
}

fn accept_stream_data(client: &mut ClientState, data: StreamData) -> (Vec<Vec<u8>>, usize, u64) {
    let stream_id = data.stream_id;
    let expected = client.stream_next_recv_offset.entry(stream_id).or_insert(0);
    if data.offset < *expected {
        return (Vec::new(), 0, *expected);
    }
    if data.offset > *expected {
        client
            .stream_recv_pending
            .entry(stream_id)
            .or_default()
            .entry(data.offset)
            .or_insert(data.bytes);
        return (Vec::new(), 0, *expected);
    }

    let mut writes = vec![data.bytes];
    let mut consumed = writes[0].len();
    *expected = expected.saturating_add(consumed as u64);

    while let Some(bytes) = client
        .stream_recv_pending
        .entry(stream_id)
        .or_default()
        .remove(expected)
    {
        consumed += bytes.len();
        *expected = expected.saturating_add(bytes.len() as u64);
        writes.push(bytes);
    }
    if client
        .stream_recv_pending
        .get(&stream_id)
        .is_some_and(BTreeMap::is_empty)
    {
        client.stream_recv_pending.remove(&stream_id);
    }
    (writes, consumed, *expected)
}

fn cleanup_client_stream(client: &mut ClientState, stream_id: u64) {
    client.stream_writers.remove(&stream_id);
    client.stream_pending_opens.remove(&stream_id);
    client.opened_streams.remove(&stream_id);
    client.stream_send_credit.remove(&stream_id);
    client.stream_pending_data.remove(&stream_id);
    client.stream_next_send_offset.remove(&stream_id);
    client.stream_sent_data.remove(&stream_id);
    client.stream_next_recv_offset.remove(&stream_id);
    client.stream_recv_pending.remove(&stream_id);
}

fn client_has_stream_state(client: &ClientState, stream_id: u64) -> bool {
    client.stream_writers.contains_key(&stream_id)
        || client.stream_pending_opens.contains_key(&stream_id)
        || client.opened_streams.contains(&stream_id)
        || client.stream_send_credit.contains_key(&stream_id)
        || client.stream_pending_data.contains_key(&stream_id)
        || client.stream_next_send_offset.contains_key(&stream_id)
        || client.stream_sent_data.contains_key(&stream_id)
        || client.stream_next_recv_offset.contains_key(&stream_id)
        || client.stream_recv_pending.contains_key(&stream_id)
}

fn retire_client_stream(client: &mut ClientState, stream_id: u64) {
    let had_state = client_has_stream_state(client, stream_id);
    cleanup_client_stream(client, stream_id);
    if !had_state || STREAM_RETIRED_TOMBSTONES == 0 {
        return;
    }
    if client.stream_retired.insert(stream_id) {
        client.stream_retired_order.push_back(stream_id);
    }
    while client.stream_retired_order.len() > STREAM_RETIRED_TOMBSTONES {
        if let Some(expired) = client.stream_retired_order.pop_front() {
            client.stream_retired.remove(&expired);
        }
    }
}

fn ack_stream_data(client: &mut ClientState, stream_id: u64, received_offset: u64) {
    let Some(sent) = client.stream_sent_data.get_mut(&stream_id) else {
        return;
    };
    let acked_offsets: Vec<u64> = sent
        .iter()
        .filter_map(|(offset, chunk)| {
            let end = chunk.offset.saturating_add(chunk.bytes.len() as u64);
            (end <= received_offset).then_some(*offset)
        })
        .collect();
    let mut acked_bytes = 0usize;
    for offset in acked_offsets {
        if let Some(chunk) = sent.remove(&offset) {
            acked_bytes = acked_bytes.saturating_add(chunk.bytes.len());
        }
    }
    if sent.is_empty() {
        client.stream_sent_data.remove(&stream_id);
    }
    add_stream_credit(&mut client.stream_send_credit, stream_id, acked_bytes);
}

fn add_stream_credit(stream_send_credit: &mut HashMap<u64, usize>, stream_id: u64, bytes: usize) {
    let credit = stream_send_credit.entry(stream_id).or_default();
    *credit = credit.saturating_add(bytes).min(STREAM_INITIAL_WINDOW);
}

async fn send_stream_window_adjust_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    bytes: usize,
    received_offset: u64,
) -> Result<()> {
    let body = protocol::to_body(&StreamWindowAdjust {
        stream_id,
        received_offset,
        bytes: bytes.min(u32::MAX as usize) as u32,
    })?;
    send_stream_packet_to_client(
        state,
        socket,
        client_id,
        PacketKind::StreamWindowAdjust,
        body,
    )
    .await
}

async fn send_stream_close_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
) -> Result<()> {
    {
        let mut locked = state.lock().expect("server state poisoned");
        if let Some(client) = locked.client_mut(&client_id) {
            retire_client_stream(client, stream_id);
        }
    }
    let body = protocol::to_body(&StreamClose { stream_id })?;
    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamClose, body).await
}

async fn send_stream_packet_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    kind: PacketKind,
    body: Vec<u8>,
) -> Result<()> {
    let send = {
        let mut locked = state.lock().expect("server state poisoned");
        let mut found = None;
        if let Some(client) = locked.client_mut(&client_id) {
            client.send_seq += 1;
            let packet = protocol::encode_encrypted(
                kind,
                client_id,
                client.send_seq,
                client.last_acked,
                &client.session_key,
                SERVER_TO_CLIENT,
                &body,
            )?;
            found = Some((client.endpoint, packet));
        }
        found
    };
    if let Some((endpoint, packet)) = send {
        let _ = send_udp(socket, &packet, endpoint).await?;
    }
    Ok(())
}

async fn broadcast_output(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    output: PtyOutput,
) -> Result<()> {
    if output.exited {
        return broadcast_session_exit(state, socket, output.session).await;
    }

    let mut persist_screen: Option<(String, u16, u16, u64, Vec<u8>)> = None;
    let sends = {
        let mut locked = state.lock().expect("server state poisoned");
        let scrollback = locked.config.scrollback;
        let retransmit_window = locked.config.retransmit_window.max(1);
        let session = locked
            .sessions
            .get_mut(&output.session)
            .ok_or_else(|| anyhow!("unknown session"))?;
        session.parser.process(&output.bytes);
        session.output_seq += 1;
        let output_seq = session.output_seq;
        session.recent.push_back(output.bytes.clone());
        while session.recent.len() > scrollback {
            session.recent.pop_front();
        }
        // For a persistent session, mirror the screen to disk periodically so a
        // post-restart reattach repaints. Throttled by accumulated bytes or a
        // short wall-clock age so low-throughput prompt/TUI changes persist
        // quickly while large bursts still avoid per-packet disk writes. The
        // actual file write happens after the lock is dropped.
        if session.persistent {
            session.bytes_since_persist = session
                .bytes_since_persist
                .saturating_add(output.bytes.len());
            let now = Instant::now();
            if session.bytes_since_persist >= SCREEN_PERSIST_BYTES
                || now.duration_since(session.last_screen_persist_at) >= SCREEN_PERSIST_MAX_AGE
            {
                session.bytes_since_persist = 0;
                session.last_screen_persist_at = now;
                let screen = session.parser.screen();
                persist_screen = Some((
                    output.session.clone(),
                    screen.size().1,
                    screen.size().0,
                    output_seq,
                    screen_snapshot(screen),
                ));
            }
        }
        let mut sends = Vec::new();
        for (client_id, client) in session.clients.iter_mut() {
            client.send_seq += 1;
            // Count traffic toward the packet-count rekey trigger (spec §11).
            client.epoch_packets = client.epoch_packets.saturating_add(1);
            // Live terminal bytes are never rewritten into vt100 snapshots. A
            // snapshot is useful for attach/resume, but substituting it during a
            // running TUI can lose graph/background-cell details that apps like
            // btop rely on. If retransmit history is full, prune only history.
            while client.pending.len() >= retransmit_window {
                client.pending.pop_front();
            }
            let frame = Frame {
                session: output.session.clone(),
                output_seq,
                bytes: output.bytes.clone(),
                snapshot: false,
                closed: false,
            };
            let body = protocol::to_body(&frame)?;
            let packet = protocol::encode_encrypted(
                PacketKind::Frame,
                *client_id,
                client.send_seq,
                client.last_acked,
                &client.session_key,
                SERVER_TO_CLIENT,
                &body,
            )?;
            client.pending.push_back(PendingFrame {
                output_seq,
                packet: packet.clone(),
                last_sent: Instant::now(),
                attempts: 0,
            });
            sends.push((client.endpoint, packet));
        }
        sends
    };
    if let Some((name, cols, rows, output_seq, snapshot)) = persist_screen {
        let sessions_dir = {
            let locked = state.lock().expect("server state poisoned");
            locked.sessions_dir()
        };
        if let Err(err) =
            persist::save_screen(&sessions_dir, &name, cols, rows, output_seq, &snapshot)
        {
            eprintln!("screen persist for {name} failed: {err:#}");
        } else {
            mark_screen_persisted(state, &name, output_seq);
        }
    }
    for (endpoint, packet) in sends {
        let _ = send_udp(socket, &packet, endpoint).await?;
    }
    Ok(())
}

async fn broadcast_session_exit(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    session_name: String,
) -> Result<()> {
    let (sends, persistent) = {
        let mut locked = state.lock().expect("server state poisoned");
        let Some(mut session) = locked.sessions.remove(&session_name) else {
            return Ok(());
        };
        let persistent = session.persistent;
        session.output_seq += 1;
        let output_seq = session.output_seq;
        let mut sends = Vec::new();
        for (client_id, mut client) in session.clients.drain() {
            locked.client_index.remove(&client_id);
            client.send_seq += 1;
            let frame = Frame {
                session: session_name.clone(),
                output_seq,
                bytes: b"\r\n[dosh session exited]\r\n".to_vec(),
                snapshot: false,
                closed: true,
            };
            let body = protocol::to_body(&frame)?;
            let packet = protocol::encode_encrypted(
                PacketKind::Frame,
                client_id,
                client.send_seq,
                client.last_acked,
                &client.session_key,
                SERVER_TO_CLIENT,
                &body,
            )?;
            sends.push((client.endpoint, packet));
        }
        (sends, persistent)
    };
    if persistent {
        // The shell exited; its holder cleans up its own runtime dir on the same
        // event, but remove ours too so a re-adopt scan never sees a dead entry.
        let sessions_dir = {
            let locked = state.lock().expect("server state poisoned");
            locked.sessions_dir()
        };
        persist::remove_runtime_dir(&sessions_dir, &session_name);
    }
    for (endpoint, packet) in sends {
        let _ = send_udp(socket, &packet, endpoint).await?;
    }
    Ok(())
}

fn screen_snapshot(screen: &vt100::Screen) -> Vec<u8> {
    let mut bytes = Vec::new();
    if screen.alternate_screen() {
        bytes.extend_from_slice(b"\x1b[?1049h");
        bytes.extend_from_slice(b"\x1b[H\x1b[2J");
    } else {
        bytes.extend_from_slice(b"\x1b[?1049l");
    }
    bytes.extend(screen.state_formatted());
    bytes
}

async fn retransmit_pending(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
) -> Result<()> {
    let sends = {
        let mut locked = state.lock().expect("server state poisoned");
        let now = Instant::now();
        let mut sends = Vec::new();
        for session in locked.sessions.values_mut() {
            for (client_id, client) in session.clients.iter_mut() {
                for pending in client.pending.iter_mut() {
                    if pending.output_seq <= client.last_acked {
                        continue;
                    }
                    if now.duration_since(pending.last_sent) >= Duration::from_millis(200) {
                        pending.last_sent = now;
                        pending.attempts = pending.attempts.saturating_add(1);
                        sends.push((client.endpoint, pending.packet.clone()));
                    }
                }
                let pending_open_ids: Vec<u64> =
                    client.stream_pending_opens.keys().copied().collect();
                let mut retransmit_opens = Vec::new();
                for stream_id in pending_open_ids {
                    let Some(open) = client.stream_pending_opens.get_mut(&stream_id) else {
                        continue;
                    };
                    if now.duration_since(open.last_sent) < Duration::from_millis(200) {
                        continue;
                    }
                    open.last_sent = now;
                    open.attempts = open.attempts.saturating_add(1);
                    retransmit_opens.push((stream_id, open.target_host.clone(), open.target_port));
                }
                for (stream_id, target_host, target_port) in retransmit_opens {
                    client.send_seq += 1;
                    let body = protocol::to_body(&StreamOpen {
                        stream_id,
                        target_host,
                        target_port,
                    })?;
                    let packet = protocol::encode_encrypted(
                        PacketKind::StreamOpen,
                        *client_id,
                        client.send_seq,
                        client.last_acked,
                        &client.session_key,
                        SERVER_TO_CLIENT,
                        &body,
                    )?;
                    sends.push((client.endpoint, packet));
                }
                let stream_ids: Vec<u64> = client.stream_sent_data.keys().copied().collect();
                let mut retransmit_chunks = Vec::new();
                for stream_id in stream_ids {
                    let Some(chunks) = client.stream_sent_data.get_mut(&stream_id) else {
                        continue;
                    };
                    for chunk in chunks.values_mut() {
                        if now.duration_since(chunk.last_sent) < Duration::from_millis(200) {
                            continue;
                        }
                        chunk.last_sent = now;
                        chunk.attempts = chunk.attempts.saturating_add(1);
                        retransmit_chunks.push((stream_id, chunk.offset, chunk.bytes.clone()));
                    }
                }
                for (stream_id, offset, bytes) in retransmit_chunks {
                    client.send_seq += 1;
                    let body = protocol::to_body(&StreamData {
                        stream_id,
                        offset,
                        bytes,
                    })?;
                    let packet = protocol::encode_encrypted(
                        PacketKind::StreamData,
                        *client_id,
                        client.send_seq,
                        client.last_acked,
                        &client.session_key,
                        SERVER_TO_CLIENT,
                        &body,
                    )?;
                    sends.push((client.endpoint, packet));
                }
            }
        }
        sends
    };
    for (endpoint, packet) in sends {
        let _ = send_udp(socket, &packet, endpoint).await?;
    }
    Ok(())
}

/// How long the previous epoch's key is retained after a rekey so in-flight
/// pre-rekey packets still decrypt instead of being dropped as fatal.
const REKEY_PREVIOUS_KEY_GRACE_SECS: u64 = 5;

/// Rotate transport traffic keys for any client past its rekey threshold (spec
/// §11). Triggers on either a packet count or a wall-clock interval, configured
/// via `rekey_after_packets` / `rekey_after_secs`. The new key is derived from
/// fresh server material independent of the handshake keys
/// (see [`dosh::native::derive_rekey_session_key`]); the server installs it
/// immediately, retains the old key for a short grace, and ships the fresh
/// material to the client in a `Rekey` packet sealed under the *current* key.
async fn maybe_rekey_clients(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
) -> Result<()> {
    let sends = {
        let mut locked = state.lock().expect("server state poisoned");
        let after_packets = locked.config.rekey_after_packets;
        let after_secs = locked.config.rekey_after_secs;
        if after_packets == 0 && after_secs == 0 {
            return Ok(());
        }
        let now = Instant::now();
        let grace = Duration::from_secs(REKEY_PREVIOUS_KEY_GRACE_SECS);
        let mut sends = Vec::new();
        for session in locked.sessions.values_mut() {
            for (client_id, client) in session.clients.iter_mut() {
                // Expire the retained previous-epoch key once its grace is over.
                if client.previous_session_key.is_some()
                    && now.duration_since(client.epoch_started) >= grace
                {
                    client.previous_session_key = None;
                }
                let by_packets = after_packets != 0 && client.epoch_packets >= after_packets;
                let by_time = after_secs != 0
                    && now.duration_since(client.epoch_started).as_secs() >= after_secs;
                if !(by_packets || by_time) {
                    continue;
                }
                // Don't start a new rekey while the previous epoch's grace is
                // still active, so we never juggle three live keys at once.
                if client.previous_session_key.is_some() {
                    continue;
                }
                let previous_key = client.session_key;
                let previous_key_id = protocol::session_key_id(&previous_key);
                let new_epoch = client.epoch.saturating_add(1);
                let rekey_material = crypto::random_32();
                let new_key = match dosh::native::derive_rekey_session_key(
                    &previous_key,
                    &rekey_material,
                    &previous_key_id,
                    new_epoch,
                ) {
                    Ok(key) => key,
                    Err(err) => {
                        eprintln!("rekey derive error: {err:#}");
                        continue;
                    }
                };
                let new_key_id = protocol::session_key_id(&new_key);
                let rekey = protocol::Rekey {
                    epoch: new_epoch,
                    rekey_material,
                    new_session_key_id: new_key_id,
                };
                let body = protocol::to_body(&rekey)?;
                client.send_seq += 1;
                // The Rekey itself is sealed under the CURRENT key so the client
                // can decrypt it before switching.
                let packet = protocol::encode_encrypted(
                    PacketKind::Rekey,
                    *client_id,
                    client.send_seq,
                    client.last_acked,
                    &previous_key,
                    SERVER_TO_CLIENT,
                    &body,
                )?;
                // Install the new key, keep the old one for the grace window.
                client.previous_session_key = Some(previous_key);
                client.session_key = new_key;
                client.epoch = new_epoch;
                client.epoch_started = now;
                client.epoch_packets = 0;
                sends.push((client.endpoint, packet));
            }
        }
        sends
    };
    for (endpoint, packet) in sends {
        let _ = send_udp(socket, &packet, endpoint).await?;
    }
    Ok(())
}

/// Mirror every persistent session's current screen to disk if it changed since
/// the last mirror. Runs on a slow timer so a crash loses at most a couple of
/// seconds of screen state even for a low-throughput interactive session that
/// never crosses the per-packet byte threshold. Screen bytes are captured under
/// the lock; the (atomic) file writes happen after it is released.
fn flush_persistent_screens(state: &Arc<Mutex<ServerState>>) {
    let (sessions_dir, to_write) = {
        let mut locked = state.lock().expect("server state poisoned");
        if !locked.config.persist_sessions {
            return;
        }
        let sessions_dir = locked.sessions_dir();
        let mut to_write: Vec<(String, u16, u16, u64, Vec<u8>)> = Vec::new();
        for (name, session) in locked.sessions.iter_mut() {
            if !session.persistent || session.output_seq == session.last_persisted_seq {
                continue;
            }
            session.bytes_since_persist = 0;
            session.last_screen_persist_at = Instant::now();
            let screen = session.parser.screen();
            to_write.push((
                name.clone(),
                screen.size().1,
                screen.size().0,
                session.output_seq,
                screen_snapshot(screen),
            ));
        }
        (sessions_dir, to_write)
    };
    for (name, cols, rows, output_seq, snapshot) in to_write {
        if let Err(err) =
            persist::save_screen(&sessions_dir, &name, cols, rows, output_seq, &snapshot)
        {
            eprintln!("screen flush for {name} failed: {err:#}");
        } else {
            mark_screen_persisted(state, &name, output_seq);
        }
    }
}

fn mark_screen_persisted(state: &Arc<Mutex<ServerState>>, name: &str, output_seq: u64) {
    let mut locked = state.lock().expect("server state poisoned");
    if let Some(session) = locked.sessions.get_mut(name) {
        session.last_persisted_seq = session.last_persisted_seq.max(output_seq);
    }
}

fn cleanup_disconnected_clients(state: &Arc<Mutex<ServerState>>) {
    // Sessions removed here are dropped after the lock is released so their
    // shells are killed (PtyHandle::drop) without holding up the event loop.
    // For a persistent session, dropping the handle only DETACHES (the shell
    // lives in the holder), so we also tell the holder to shut down — otherwise
    // a truly-abandoned shell would linger forever.
    let sessions_dir = {
        let locked = state.lock().expect("server state poisoned");
        locked.sessions_dir()
    };
    let reaped: Vec<(String, Session)> = {
        let mut locked = state.lock().expect("server state poisoned");
        let timeout = Duration::from_secs(locked.config.client_timeout_secs.max(1));
        let now = Instant::now();
        let prewarm: HashSet<String> = locked.config.prewarm_sessions.iter().cloned().collect();
        // Collect timed-out clients first so we can purge them from the O(1)
        // conn_id index too, keeping it in lockstep with `Session::clients`.
        let mut timed_out: Vec<[u8; 16]> = Vec::new();
        for session in locked.sessions.values_mut() {
            session.clients.retain(|client_id, client| {
                let keep = now.duration_since(client.last_seen) <= timeout;
                if !keep {
                    timed_out.push(*client_id);
                }
                keep
            });
            if session.clients.is_empty() {
                session.empty_since.get_or_insert(now);
            } else {
                session.empty_since = None;
            }
        }
        for client_id in timed_out {
            locked.client_index.remove(&client_id);
        }
        // Evict half-finished native handshakes. Each unauthenticated ClientHello
        // inserts a PendingNativeAuth that is otherwise only removed by a matching
        // UserAuth, so without this a flood of hellos grows the map without bound.
        let handshake_ttl = Duration::from_secs(NATIVE_HANDSHAKE_TTL_SECS);
        locked
            .pending_native
            .retain(|_, pending| now.duration_since(pending.created_at) < handshake_ttl);
        // Drop fully-refilled rate-limit buckets so a spoofed-source flood cannot
        // grow the map without bound.
        locked.native_auth_limiter.evict_full(now);
        // Reap sessions that have been clientless past the grace period. Prewarmed
        // sessions stay hot on purpose, even with no clients attached.
        let to_reap: Vec<String> = locked
            .sessions
            .iter()
            .filter(|(name, session)| {
                !prewarm.contains(name.as_str())
                    && session.empty_since.is_some_and(|since| {
                        now.duration_since(since) >= empty_session_timeout(name, timeout)
                    })
            })
            .map(|(name, _)| name.clone())
            .collect();
        to_reap
            .into_iter()
            .filter_map(|name| locked.sessions.remove(&name).map(|session| (name, session)))
            .collect()
    };
    for (name, mut session) in reaped {
        if session.persistent {
            // Ask the holder to terminate the shell and exit, then clean its
            // runtime dir, so the abandoned shell does not survive forever.
            persist::request_shutdown(&sessions_dir, &name, session.holder_control.as_mut());
        }
        drop(session);
    }
}

fn empty_session_timeout(name: &str, configured_timeout: Duration) -> Duration {
    if protocol::is_implicit_session_name(name) {
        configured_timeout.min(Duration::from_secs(IMPLICIT_EMPTY_SESSION_GRACE_SECS))
    } else {
        configured_timeout
    }
}

/// Select the client key (current or retained previous epoch) matching the
/// packet header's `session_key_id`, so a packet sealed under the pre-rekey key
/// during the grace window still decrypts while an expired/stale epoch is
/// treated as an unknown client and triggers a client reconnect.
fn find_client_decrypt_key(
    state: &Arc<Mutex<ServerState>>,
    header: &protocol::Header,
) -> Result<([u8; 32], String)> {
    let locked = state.lock().expect("server state poisoned");
    locked
        .lookup_client_key_for(&header.conn_id, &header.session_key_id)
        .ok_or_else(|| anyhow!("unknown client"))
}

fn parse_nonce(raw: &str) -> Result<[u8; 12]> {
    let bytes = base64::Engine::decode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, raw)
        .context("decode nonce")?;
    anyhow::ensure!(bytes.len() == 12, "nonce must decode to 12 bytes");
    let mut out = [0u8; 12];
    out.copy_from_slice(&bytes);
    Ok(out)
}

fn parse_size(raw: &str) -> Result<(u16, u16)> {
    let (cols, rows) = raw.split_once('x').context("size must be COLSxROWS")?;
    Ok((cols.parse()?, rows.parse()?))
}

/// Parse `--env NAME=VALUE` arguments passed to the holder back into pairs. The
/// value may itself contain `=`, so split only on the first.
fn parse_env_pairs(raw: &[String]) -> Result<Vec<(String, String)>> {
    raw.iter()
        .map(|entry| {
            let (name, value) = entry
                .split_once('=')
                .with_context(|| format!("env entry {entry:?} must be NAME=VALUE"))?;
            Ok((name.to_string(), value.to_string()))
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn persists_terminal_sessions_when_enabled() {
        let (pty_tx, _rx) = mpsc::unbounded_channel();
        let config = ServerConfig {
            persist_sessions: true,
            prewarm_sessions: vec!["default".to_string()],
            ..ServerConfig::default()
        };
        let state = ServerState::new(config, [0u8; 32], pty_tx);
        assert!(state.should_persist_session("default")); // prewarmed
        assert!(state.should_persist_session("work")); // explicitly named
        assert!(state.should_persist_session("term-1781470634216-76685")); // one-off update-survivable terminal
    }

    #[test]
    fn persist_disabled_never_persists() {
        let (pty_tx, _rx) = mpsc::unbounded_channel();
        let config = ServerConfig {
            persist_sessions: false,
            prewarm_sessions: vec!["default".to_string()],
            ..ServerConfig::default()
        };
        let state = ServerState::new(config, [0u8; 32], pty_tx);
        assert!(!state.should_persist_session("default"));
        assert!(!state.should_persist_session("work"));
    }

    #[test]
    fn native_auth_rate_limiter_throttles_per_ip_and_refills() {
        let mut limiter = NativeAuthRateLimiter::new(3);
        let ip_a: IpAddr = "10.0.0.1".parse().unwrap();
        let ip_b: IpAddr = "10.0.0.2".parse().unwrap();
        let t0 = Instant::now();

        // Burst of 3 allowed, 4th rejected for ip_a.
        assert!(limiter.check(ip_a, t0).is_ok());
        assert!(limiter.check(ip_a, t0).is_ok());
        assert!(limiter.check(ip_a, t0).is_ok());
        assert!(limiter.check(ip_a, t0).is_err());

        // A different source IP has its own independent bucket.
        assert!(limiter.check(ip_b, t0).is_ok());

        // After 20s at 3/min (one token every 20s), ip_a gets exactly one back.
        let t1 = t0 + Duration::from_secs(20);
        assert!(limiter.check(ip_a, t1).is_ok());
        assert!(limiter.check(ip_a, t1).is_err());
    }

    #[test]
    fn native_auth_rate_limiter_zero_per_minute_blocks_everything() {
        let mut limiter = NativeAuthRateLimiter::new(0);
        assert!(
            limiter
                .check("127.0.0.1".parse().unwrap(), Instant::now())
                .is_err()
        );
    }

    #[test]
    fn native_auth_rate_limiter_evicts_refilled_buckets() {
        let mut limiter = NativeAuthRateLimiter::new(60);
        let ip: IpAddr = "10.0.0.9".parse().unwrap();
        let t0 = Instant::now();
        assert!(limiter.check(ip, t0).is_ok());
        assert_eq!(limiter.buckets.len(), 1);
        // 60/min = one token per second; after 60s the bucket is full again.
        limiter.evict_full(t0 + Duration::from_secs(60));
        assert!(limiter.buckets.is_empty());
    }

    #[test]
    fn default_client_timeout_supports_long_sleep() {
        assert_eq!(ServerConfig::default().client_timeout_secs, 2_592_000);
    }

    #[tokio::test]
    async fn unknown_resume_reject_keeps_client_id() {
        let (pty_tx, _rx) = mpsc::unbounded_channel();
        let state = Arc::new(Mutex::new(ServerState::new(
            ServerConfig::default(),
            [0u8; 32],
            pty_tx,
        )));
        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let client_id = [9u8; 16];
        let packet = protocol::Packet {
            header: protocol::Header {
                kind: PacketKind::ResumeRequest,
                flags: 1,
                conn_id: client_id,
                seq: 77,
                ack: 0,
                session_key_id: [8u8; 16],
                body_len: 0,
            },
            body: Vec::new(),
        };

        handle_resume(&state, &sender, receiver.local_addr().unwrap(), &packet)
            .await
            .unwrap();

        let mut buf = [0u8; 1024];
        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
            .await
            .unwrap()
            .unwrap();
        let reject = protocol::decode(&buf[..n]).unwrap();
        assert_eq!(reject.header.kind, PacketKind::AttachReject);
        assert_eq!(reject.header.conn_id, client_id);
        let body: AttachReject = protocol::from_body(&reject.body).unwrap();
        assert_eq!(body.reason, "unknown client");
    }

    fn test_client_state(session_key: [u8; 32]) -> ClientState {
        ClientState {
            endpoint: "127.0.0.1:9".parse().unwrap(),
            mode: "forward-only".to_string(),
            session_key,
            last_acked: 0,
            replay: ReplayWindow::default(),
            send_seq: 1,
            cols: 80,
            rows: 24,
            last_seen: Instant::now(),
            pending: VecDeque::new(),
            allowed_forwardings: Vec::new(),
            stream_writers: HashMap::new(),
            opened_streams: HashSet::new(),
            stream_send_credit: HashMap::new(),
            stream_pending_data: HashMap::new(),
            stream_pending_opens: HashMap::new(),
            stream_next_send_offset: HashMap::new(),
            stream_sent_data: HashMap::new(),
            stream_next_recv_offset: HashMap::new(),
            stream_recv_pending: HashMap::new(),
            stream_retired: HashSet::new(),
            stream_retired_order: VecDeque::new(),
            epoch: 0,
            epoch_started: Instant::now(),
            epoch_packets: 0,
            previous_session_key: None,
        }
    }

    #[test]
    fn cleanup_client_stream_removes_all_per_stream_state_only_for_target() {
        let stream_id = 42;
        let other_stream_id = 43;
        let mut client = test_client_state([7u8; 32]);
        let (writer_tx, _writer_rx) = mpsc::channel(1);
        client.stream_writers.insert(stream_id, writer_tx);
        client.opened_streams.insert(stream_id);
        client.opened_streams.insert(other_stream_id);
        client.stream_send_credit.insert(stream_id, 12);
        client.stream_send_credit.insert(other_stream_id, 34);
        client
            .stream_pending_data
            .insert(stream_id, VecDeque::from([vec![1]]));
        client
            .stream_pending_data
            .insert(other_stream_id, VecDeque::from([vec![2]]));
        client.stream_pending_opens.insert(
            stream_id,
            PendingStreamOpen {
                target_host: "127.0.0.1".to_string(),
                target_port: 1234,
                last_sent: Instant::now(),
                attempts: 1,
            },
        );
        client.stream_pending_opens.insert(
            other_stream_id,
            PendingStreamOpen {
                target_host: "127.0.0.1".to_string(),
                target_port: 1235,
                last_sent: Instant::now(),
                attempts: 1,
            },
        );
        client.stream_next_send_offset.insert(stream_id, 55);
        client.stream_next_send_offset.insert(other_stream_id, 66);
        client.stream_sent_data.insert(
            stream_id,
            BTreeMap::from([(
                0,
                PendingStreamChunk {
                    offset: 0,
                    bytes: vec![3],
                    last_sent: Instant::now(),
                    attempts: 1,
                },
            )]),
        );
        client.stream_sent_data.insert(
            other_stream_id,
            BTreeMap::from([(
                0,
                PendingStreamChunk {
                    offset: 0,
                    bytes: vec![4],
                    last_sent: Instant::now(),
                    attempts: 1,
                },
            )]),
        );
        client.stream_next_recv_offset.insert(stream_id, 77);
        client.stream_next_recv_offset.insert(other_stream_id, 88);
        client
            .stream_recv_pending
            .insert(stream_id, BTreeMap::from([(77, vec![5])]));
        client
            .stream_recv_pending
            .insert(other_stream_id, BTreeMap::from([(88, vec![6])]));

        cleanup_client_stream(&mut client, stream_id);

        assert!(!client.stream_writers.contains_key(&stream_id));
        assert!(!client.opened_streams.contains(&stream_id));
        assert!(!client.stream_send_credit.contains_key(&stream_id));
        assert!(!client.stream_pending_data.contains_key(&stream_id));
        assert!(!client.stream_pending_opens.contains_key(&stream_id));
        assert!(!client.stream_next_send_offset.contains_key(&stream_id));
        assert!(!client.stream_sent_data.contains_key(&stream_id));
        assert!(!client.stream_next_recv_offset.contains_key(&stream_id));
        assert!(!client.stream_recv_pending.contains_key(&stream_id));

        assert!(client.opened_streams.contains(&other_stream_id));
        assert!(client.stream_send_credit.contains_key(&other_stream_id));
        assert!(client.stream_pending_data.contains_key(&other_stream_id));
        assert!(client.stream_pending_opens.contains_key(&other_stream_id));
        assert!(
            client
                .stream_next_send_offset
                .contains_key(&other_stream_id)
        );
        assert!(client.stream_sent_data.contains_key(&other_stream_id));
        assert!(
            client
                .stream_next_recv_offset
                .contains_key(&other_stream_id)
        );
        assert!(client.stream_recv_pending.contains_key(&other_stream_id));
    }

    #[test]
    fn retire_client_stream_tombstones_only_real_stream_state() {
        let stream_id = 42;
        let mut client = test_client_state([8u8; 32]);
        client.opened_streams.insert(stream_id);
        client
            .stream_send_credit
            .insert(stream_id, STREAM_INITIAL_WINDOW);

        retire_client_stream(&mut client, stream_id);
        retire_client_stream(&mut client, 99);

        assert!(!client.opened_streams.contains(&stream_id));
        assert!(!client.stream_send_credit.contains_key(&stream_id));
        assert!(client.stream_retired.contains(&stream_id));
        assert!(!client.stream_retired.contains(&99));
        assert_eq!(client.stream_retired_order, VecDeque::from([stream_id]));
    }

    #[test]
    fn client_index_stays_in_sync_with_session_clients() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
        state
            .ensure_session("work", 80, 24, "forward-only", &[])
            .unwrap();
        let client_id = [5u8; 16];
        let key = [6u8; 32];

        // Inserting through the helper records the index for O(1) lookup.
        state.insert_client("work", client_id, test_client_state(key));
        assert_eq!(
            state.client_index.get(&client_id).map(String::as_str),
            Some("work")
        );
        assert_eq!(
            state.lookup_client(&client_id),
            Some((key, "work".to_string()))
        );
        assert!(state.client_mut(&client_id).is_some());
        state.sessions.get_mut("work").unwrap().empty_since = Some(Instant::now());
        state.insert_client("work", [7u8; 16], test_client_state([9u8; 32]));
        assert!(
            state.sessions["work"].empty_since.is_none(),
            "reattach must clear empty-since so cleanup cannot reap a live session"
        );

        // Removing purges both the session map and the index.
        assert!(state.remove_client_everywhere(&client_id));
        assert!(!state.client_index.contains_key(&client_id));
        assert!(state.lookup_client(&client_id).is_none());
        assert!(state.client_mut(&client_id).is_none());
        assert!(!state.sessions["work"].clients.contains_key(&client_id));
        // Removing an absent client is a no-op.
        assert!(!state.remove_client_everywhere(&client_id));
    }

    #[test]
    fn cleanup_purges_timed_out_clients_from_index() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let config = ServerConfig {
            client_timeout_secs: 1,
            ..ServerConfig::default()
        };
        let mut state = ServerState::new(config, [0u8; 32], pty_tx);
        state
            .ensure_session("work", 80, 24, "forward-only", &[])
            .unwrap();
        let client_id = [8u8; 16];
        let mut client = test_client_state([1u8; 32]);
        // Make the client look stale so cleanup reaps it.
        client.last_seen = Instant::now() - Duration::from_secs(60);
        state.insert_client("work", client_id, client);

        let state = Arc::new(Mutex::new(state));
        cleanup_disconnected_clients(&state);

        let locked = state.lock().unwrap();
        assert!(
            !locked.client_index.contains_key(&client_id),
            "timed-out client must be dropped from the index, not leaked"
        );
        assert!(locked.lookup_client(&client_id).is_none());
    }

    #[test]
    fn implicit_empty_session_timeout_is_bounded_for_update_reconnect() {
        let configured = Duration::from_secs(2_592_000);
        let implicit = protocol::generate_implicit_session_name();
        assert_eq!(
            empty_session_timeout(&implicit, configured),
            Duration::from_secs(IMPLICIT_EMPTY_SESSION_GRACE_SECS)
        );
        assert_eq!(
            empty_session_timeout("work", configured),
            configured,
            "named sessions keep the normal long timeout"
        );
        assert_eq!(
            empty_session_timeout(&implicit, Duration::from_secs(1)),
            Duration::from_secs(1),
            "tests/admins can still configure a shorter timeout"
        );
    }

    #[tokio::test]
    async fn forged_plaintext_detach_does_not_remove_client() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
        state
            .ensure_session("work", 80, 24, "forward-only", &[])
            .unwrap();
        let client_id = [21u8; 16];
        let session_key = [22u8; 32];
        state.insert_client("work", client_id, test_client_state(session_key));
        let state = Arc::new(Mutex::new(state));
        let mut packet = protocol::decode(
            &protocol::encode_plain(PacketKind::Detach, client_id, 1, 0, b"").unwrap(),
        )
        .unwrap();
        packet.header.session_key_id = protocol::session_key_id(&session_key);

        assert!(handle_detach(&state, &packet).await.is_err());

        let locked = state.lock().unwrap();
        assert!(
            locked.lookup_client(&client_id).is_some(),
            "plaintext detach with a copied key id must not remove a live client"
        );
    }

    #[tokio::test]
    async fn authenticated_detach_removes_client() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
        state
            .ensure_session("work", 80, 24, "forward-only", &[])
            .unwrap();
        let client_id = [23u8; 16];
        let session_key = [24u8; 32];
        state.insert_client("work", client_id, test_client_state(session_key));
        let state = Arc::new(Mutex::new(state));
        let packet = protocol::decode(
            &protocol::encode_encrypted(
                PacketKind::Detach,
                client_id,
                1,
                0,
                &session_key,
                CLIENT_TO_SERVER,
                b"",
            )
            .unwrap(),
        )
        .unwrap();

        handle_detach(&state, &packet).await.unwrap();

        let locked = state.lock().unwrap();
        assert!(locked.lookup_client(&client_id).is_none());
    }

    #[tokio::test]
    async fn duplicate_stream_open_for_open_stream_resends_ok() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
        let client_id = [11u8; 16];
        let session_key = [12u8; 32];
        let stream_id = 99;
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let mut client = test_client_state(session_key);
        client.endpoint = receiver.local_addr().unwrap();
        client.opened_streams.insert(stream_id);
        state.sessions.insert(
            "test".to_string(),
            Session {
                pty: None,
                parser: vt100::Parser::new(24, 80, 100),
                clients: HashMap::from([(client_id, client)]),
                output_seq: 0,
                recent: VecDeque::new(),
                empty_since: None,
                holder_control: None,
                persistent: false,
                bytes_since_persist: 0,
                last_persisted_seq: 0,
                last_screen_persist_at: Instant::now(),
            },
        );
        state.client_index.insert(client_id, "test".to_string());
        let state = Arc::new(Mutex::new(state));
        let raw = protocol::encode_encrypted(
            PacketKind::StreamOpen,
            client_id,
            2,
            0,
            &session_key,
            CLIENT_TO_SERVER,
            &protocol::to_body(&StreamOpen {
                stream_id,
                target_host: "127.0.0.1".to_string(),
                target_port: 9,
            })
            .unwrap(),
        )
        .unwrap();
        let packet = protocol::decode(&raw).unwrap();

        handle_stream_open(&state, &sender, receiver.local_addr().unwrap(), &packet)
            .await
            .unwrap();

        let mut buf = [0u8; 2048];
        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
            .await
            .unwrap()
            .unwrap();
        let packet = protocol::decode(&buf[..n]).unwrap();
        assert_eq!(packet.header.kind, PacketKind::StreamOpenOk);
        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
        let ok: StreamOpenOk = protocol::from_body(&plain).unwrap();
        assert_eq!(ok.stream_id, stream_id);
    }

    #[tokio::test]
    async fn retired_stream_open_is_rejected_not_reopened() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
        let client_id = [15u8; 16];
        let session_key = [16u8; 32];
        let stream_id = 100;
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let mut client = test_client_state(session_key);
        client.endpoint = receiver.local_addr().unwrap();
        client.stream_retired.insert(stream_id);
        client.stream_retired_order.push_back(stream_id);
        state.sessions.insert(
            "test".to_string(),
            Session {
                pty: None,
                parser: vt100::Parser::new(24, 80, 100),
                clients: HashMap::from([(client_id, client)]),
                output_seq: 0,
                recent: VecDeque::new(),
                empty_since: None,
                holder_control: None,
                persistent: false,
                bytes_since_persist: 0,
                last_persisted_seq: 0,
                last_screen_persist_at: Instant::now(),
            },
        );
        state.client_index.insert(client_id, "test".to_string());
        let state = Arc::new(Mutex::new(state));
        let raw = protocol::encode_encrypted(
            PacketKind::StreamOpen,
            client_id,
            2,
            0,
            &session_key,
            CLIENT_TO_SERVER,
            &protocol::to_body(&StreamOpen {
                stream_id,
                target_host: "127.0.0.1".to_string(),
                target_port: 9,
            })
            .unwrap(),
        )
        .unwrap();
        let packet = protocol::decode(&raw).unwrap();

        handle_stream_open(&state, &sender, receiver.local_addr().unwrap(), &packet)
            .await
            .unwrap();

        let mut buf = [0u8; 2048];
        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
            .await
            .unwrap()
            .unwrap();
        let packet = protocol::decode(&buf[..n]).unwrap();
        assert_eq!(packet.header.kind, PacketKind::StreamOpenReject);
        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
        let reject: StreamOpenReject = protocol::from_body(&plain).unwrap();
        assert_eq!(reject.stream_id, stream_id);
        assert!(reject.reason.contains("closed"));

        let locked = state.lock().unwrap();
        let client = locked.sessions["test"].clients.get(&client_id).unwrap();
        assert!(!client.opened_streams.contains(&stream_id));
    }

    #[tokio::test]
    async fn pending_server_stream_open_is_retransmitted() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
        let client_id = [13u8; 16];
        let session_key = [14u8; 32];
        let stream_id = 123;
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let mut client = test_client_state(session_key);
        client.endpoint = receiver.local_addr().unwrap();
        client.stream_pending_opens.insert(
            stream_id,
            PendingStreamOpen {
                target_host: "127.0.0.1".to_string(),
                target_port: 8080,
                last_sent: Instant::now() - Duration::from_millis(250),
                attempts: 1,
            },
        );
        state.sessions.insert(
            "test".to_string(),
            Session {
                pty: None,
                parser: vt100::Parser::new(24, 80, 100),
                clients: HashMap::from([(client_id, client)]),
                output_seq: 0,
                recent: VecDeque::new(),
                empty_since: None,
                holder_control: None,
                persistent: false,
                bytes_since_persist: 0,
                last_persisted_seq: 0,
                last_screen_persist_at: Instant::now(),
            },
        );
        state.client_index.insert(client_id, "test".to_string());
        let state = Arc::new(Mutex::new(state));

        retransmit_pending(&state, &sender).await.unwrap();

        let mut buf = [0u8; 2048];
        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
            .await
            .unwrap()
            .unwrap();
        let packet = protocol::decode(&buf[..n]).unwrap();
        assert_eq!(packet.header.kind, PacketKind::StreamOpen);
        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
        let open: StreamOpen = protocol::from_body(&plain).unwrap();
        assert_eq!(open.stream_id, stream_id);
        assert_eq!(open.target_port, 8080);
    }

    #[test]
    fn forward_only_session_does_not_allocate_pty() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);

        state
            .ensure_session("forward", 80, 24, "forward-only", &[])
            .unwrap();

        assert!(state.sessions["forward"].pty.is_none());
    }

    #[test]
    fn cleanup_reaps_abandoned_sessions_but_keeps_prewarmed() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let config = ServerConfig {
            client_timeout_secs: 1,
            prewarm_sessions: vec!["default".to_string()],
            ..ServerConfig::default()
        };
        let mut state = ServerState::new(config, [0u8; 32], pty_tx);
        // Both sessions are clientless; use forward-only so no real shell spawns.
        state
            .ensure_session("default", 80, 24, "forward-only", &[])
            .unwrap();
        state
            .ensure_session("abandoned", 80, 24, "forward-only", &[])
            .unwrap();
        // Mark both as empty long enough ago to be past the grace period.
        let stale = Instant::now() - Duration::from_secs(60);
        state.sessions.get_mut("default").unwrap().empty_since = Some(stale);
        state.sessions.get_mut("abandoned").unwrap().empty_since = Some(stale);

        let state = Arc::new(Mutex::new(state));
        cleanup_disconnected_clients(&state);

        let locked = state.lock().unwrap();
        assert!(
            locked.sessions.contains_key("default"),
            "prewarmed session must stay hot"
        );
        assert!(
            !locked.sessions.contains_key("abandoned"),
            "clientless non-prewarmed session must be reaped"
        );
    }

    #[test]
    fn accepted_env_filters_by_server_patterns() {
        let config = ServerConfig {
            accept_env: vec!["LANG".to_string(), "LC_*".to_string()],
            ..ServerConfig::default()
        };

        assert_eq!(
            accepted_env(
                &config,
                &[
                    EnvVar {
                        name: "LANG".to_string(),
                        value: "en_US.UTF-8".to_string()
                    },
                    EnvVar {
                        name: "LC_TIME".to_string(),
                        value: "C".to_string()
                    },
                    EnvVar {
                        name: "SECRET_TOKEN".to_string(),
                        value: "nope".to_string()
                    }
                ]
            )
            .unwrap(),
            vec![
                ("LANG".to_string(), "en_US.UTF-8".to_string()),
                ("LC_TIME".to_string(), "C".to_string())
            ]
        );
    }

    #[test]
    fn accepted_env_rejects_invalid_names() {
        let err = accepted_env(
            &ServerConfig::default(),
            &[EnvVar {
                name: "BAD-NAME".to_string(),
                value: "value".to_string(),
            }],
        )
        .unwrap_err();

        assert!(
            err.to_string()
                .contains("invalid environment variable name")
        );
    }

    #[tokio::test]
    async fn stream_data_waits_for_open_and_credit() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
        let client_id = [7u8; 16];
        let session_key = [9u8; 32];
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let endpoint = receiver.local_addr().unwrap();
        state.sessions.insert(
            "test".to_string(),
            Session {
                pty: None,
                parser: vt100::Parser::new(24, 80, 100),
                clients: HashMap::from([(
                    client_id,
                    ClientState {
                        endpoint,
                        mode: "forward-only".to_string(),
                        session_key,
                        last_acked: 0,
                        replay: ReplayWindow::default(),
                        send_seq: 1,
                        cols: 80,
                        rows: 24,
                        last_seen: Instant::now(),
                        pending: VecDeque::new(),
                        allowed_forwardings: Vec::new(),
                        stream_writers: HashMap::new(),
                        opened_streams: HashSet::new(),
                        stream_send_credit: HashMap::new(),
                        stream_pending_data: HashMap::new(),
                        stream_pending_opens: HashMap::new(),
                        stream_next_send_offset: HashMap::new(),
                        stream_sent_data: HashMap::new(),
                        stream_next_recv_offset: HashMap::new(),
                        stream_recv_pending: HashMap::new(),
                        stream_retired: HashSet::new(),
                        stream_retired_order: VecDeque::new(),
                        epoch: 0,
                        epoch_started: Instant::now(),
                        epoch_packets: 0,
                        previous_session_key: None,
                    },
                )]),
                output_seq: 0,
                recent: VecDeque::new(),
                empty_since: None,
                holder_control: None,
                persistent: false,
                bytes_since_persist: 0,
                last_persisted_seq: 0,
                last_screen_persist_at: Instant::now(),
            },
        );
        // Keep the O(1) conn_id index in sync, matching `insert_client`.
        state.client_index.insert(client_id, "test".to_string());
        let state = Arc::new(Mutex::new(state));

        queue_or_send_stream_data_to_client(&state, &sender, client_id, 42, b"hello".to_vec())
            .await
            .unwrap();
        {
            let locked = state.lock().unwrap();
            let client = locked.sessions["test"].clients.get(&client_id).unwrap();
            assert_eq!(client.stream_pending_data[&42].len(), 1);
        }

        {
            let mut locked = state.lock().unwrap();
            let client = locked
                .sessions
                .get_mut("test")
                .unwrap()
                .clients
                .get_mut(&client_id)
                .unwrap();
            client.opened_streams.insert(42);
            client.stream_send_credit.insert(42, STREAM_INITIAL_WINDOW);
        }
        flush_stream_pending_data_to_client(&state, &sender, client_id, 42)
            .await
            .unwrap();

        let mut buf = [0u8; 2048];
        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
            .await
            .unwrap()
            .unwrap();
        let packet = protocol::decode(&buf[..n]).unwrap();
        assert_eq!(packet.header.kind, PacketKind::StreamData);
        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
        let data: StreamData = protocol::from_body(&plain).unwrap();
        assert_eq!(data.stream_id, 42);
        assert_eq!(data.bytes, b"hello");
    }

    #[tokio::test]
    async fn blocked_stream_data_does_not_block_terminal_frames() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
        let client_id = [8u8; 16];
        let session_key = [10u8; 32];
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let endpoint = receiver.local_addr().unwrap();
        state.sessions.insert(
            "test".to_string(),
            Session {
                pty: None,
                parser: vt100::Parser::new(24, 80, 100),
                clients: HashMap::from([(
                    client_id,
                    ClientState {
                        endpoint,
                        mode: "read-write".to_string(),
                        session_key,
                        last_acked: 0,
                        replay: ReplayWindow::default(),
                        send_seq: 1,
                        cols: 80,
                        rows: 24,
                        last_seen: Instant::now(),
                        pending: VecDeque::new(),
                        allowed_forwardings: Vec::new(),
                        stream_writers: HashMap::new(),
                        opened_streams: HashSet::from([42]),
                        stream_send_credit: HashMap::from([(42, 0)]),
                        stream_pending_data: HashMap::new(),
                        stream_pending_opens: HashMap::new(),
                        stream_next_send_offset: HashMap::new(),
                        stream_sent_data: HashMap::new(),
                        stream_next_recv_offset: HashMap::new(),
                        stream_recv_pending: HashMap::new(),
                        stream_retired: HashSet::new(),
                        stream_retired_order: VecDeque::new(),
                        epoch: 0,
                        epoch_started: Instant::now(),
                        epoch_packets: 0,
                        previous_session_key: None,
                    },
                )]),
                output_seq: 0,
                recent: VecDeque::new(),
                empty_since: None,
                holder_control: None,
                persistent: false,
                bytes_since_persist: 0,
                last_persisted_seq: 0,
                last_screen_persist_at: Instant::now(),
            },
        );
        state.client_index.insert(client_id, "test".to_string());
        let state = Arc::new(Mutex::new(state));

        queue_or_send_stream_data_to_client(
            &state,
            &sender,
            client_id,
            42,
            b"blocked-forward-bytes".to_vec(),
        )
        .await
        .unwrap();
        {
            let locked = state.lock().unwrap();
            let client = locked.sessions["test"].clients.get(&client_id).unwrap();
            assert_eq!(client.stream_pending_data[&42].len(), 1);
            assert_eq!(client.stream_send_credit[&42], 0);
        }

        broadcast_output(
            &state,
            &sender,
            PtyOutput {
                session: "test".to_string(),
                bytes: b"terminal-priority".to_vec(),
                exited: false,
            },
        )
        .await
        .unwrap();

        let mut buf = [0u8; 2048];
        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
            .await
            .unwrap()
            .unwrap();
        let packet = protocol::decode(&buf[..n]).unwrap();
        assert_eq!(packet.header.kind, PacketKind::Frame);
        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
        let frame: Frame = protocol::from_body(&plain).unwrap();
        assert_eq!(frame.bytes, b"terminal-priority");
    }

    #[tokio::test]
    async fn live_terminal_output_is_never_replaced_by_paced_snapshots() {
        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
        let config = ServerConfig {
            output_frame_interval_ms: 1000,
            ..ServerConfig::default()
        };
        let mut state = ServerState::new(config, [0u8; 32], pty_tx);
        state
            .ensure_session("test", 80, 24, "forward-only", &[])
            .unwrap();
        let client_id = [42u8; 16];
        let session_key = [7u8; 32];
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
        let mut client = test_client_state(session_key);
        client.endpoint = receiver.local_addr().unwrap();
        state.insert_client("test", client_id, client);
        let state = Arc::new(Mutex::new(state));

        broadcast_output(
            &state,
            &sender,
            PtyOutput {
                session: "test".to_string(),
                bytes: b"first".to_vec(),
                exited: false,
            },
        )
        .await
        .unwrap();
        broadcast_output(
            &state,
            &sender,
            PtyOutput {
                session: "test".to_string(),
                bytes: b"second".to_vec(),
                exited: false,
            },
        )
        .await
        .unwrap();

        let mut buf = [0u8; 4096];
        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
            .await
            .unwrap()
            .unwrap();
        let packet = protocol::decode(&buf[..n]).unwrap();
        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
        let frame: Frame = protocol::from_body(&plain).unwrap();
        assert!(!frame.snapshot);
        assert_eq!(frame.bytes, b"first");
        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
            .await
            .unwrap()
            .unwrap();
        let packet = protocol::decode(&buf[..n]).unwrap();
        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
        let frame: Frame = protocol::from_body(&plain).unwrap();
        assert!(!frame.snapshot);
        assert_eq!(frame.output_seq, 2);
        assert_eq!(frame.bytes, b"second");
    }

    #[test]
    fn remote_non_loopback_bind_requires_explicit_config() {
        let config = ServerConfig::default();
        remote_bind_allowed(&config, "127.0.0.1").unwrap();
        assert!(remote_bind_allowed(&config, "0.0.0.0").is_err());

        let config = ServerConfig {
            allow_remote_non_loopback_bind: true,
            ..ServerConfig::default()
        };
        remote_bind_allowed(&config, "0.0.0.0").unwrap();
    }

    #[cfg(unix)]
    #[test]
    fn server_udp_send_classifier_treats_network_roaming_errors_as_transient() {
        for code in [
            libc::EADDRNOTAVAIL,
            libc::ECONNREFUSED,
            libc::ECONNRESET,
            libc::EHOSTDOWN,
            libc::EHOSTUNREACH,
            libc::ENETDOWN,
            libc::ENETRESET,
            libc::ENETUNREACH,
            libc::ETIMEDOUT,
        ] {
            assert!(
                is_transient_udp_send_error(&std::io::Error::from_raw_os_error(code)),
                "expected errno {code} to be transient"
            );
        }
    }
}