diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index f93443d..cab42b9 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -1,6233 +1,8 @@ -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, StreamEof, 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_error, 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; -const STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS: u32 = 8; - -/// 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; - -fn hex_id(id: [u8; 16]) -> String { - let mut out = String::with_capacity(32); - for byte in id { - use std::fmt::Write as _; - let _ = write!(&mut out, "{byte:02x}"); - } - out -} - -#[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, - }, - Auth { - #[arg(long, default_value_t = 1)] - protocol: u8, - #[arg(long)] - nonce: Option, - #[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, - }, - /// 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, - }, -} - -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) -> 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}"); - dosh::trace::event( - "server.start", - &[ - ("version", dosh::build_info::VERSION.to_string()), - ("bind", bind.clone()), - ], - ); - - 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) = match socket.recv_from(&mut buf).await { - Ok(value) => value, - Err(err) if is_transient_udp_error(&err) => continue, - Err(err) => return Err(err.into()), - }; - 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, - sessions: HashMap, - 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, -} - -#[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 { - 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, - parser: vt100::Parser, - restored_screen: Option, - clients: HashMap<[u8; 16], ClientState>, - output_seq: u64, - recent: VecDeque>, - /// 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, - /// 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, - /// 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 RestoredScreenSnapshot { - cols: u16, - rows: u16, - snapshot: Vec, -} - -#[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, - allowed_forwardings: Vec, - stream_writers: HashMap>>, - opened_streams: HashSet, - stream_send_credit: HashMap, - stream_pending_data: HashMap>>, - stream_pending_opens: HashMap, - stream_next_send_offset: HashMap, - stream_sent_data: HashMap>, - stream_next_recv_offset: HashMap, - stream_recv_pending: HashMap>>, - stream_local_eof_sent: HashSet, - stream_remote_eof_received: HashSet, - stream_eof_retransmit: HashMap, - stream_close_retransmit: HashMap, - stream_window_adjust_retransmit: HashMap, - stream_retired: HashSet, - stream_retired_order: VecDeque, - stream_retransmit_srtt: Option, - /// 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, - last_sent: Instant, - attempts: u8, -} - -#[derive(Clone)] -struct PendingStreamChunk { - offset: u64, - bytes: Vec, - last_sent: Instant, - attempts: u32, -} - -#[derive(Clone)] -struct PendingStreamOpen { - target_host: String, - target_port: u16, - last_sent: Instant, - attempts: u32, -} - -#[derive(Clone)] -struct PendingStreamControl { - last_sent: Instant, - attempts: u32, -} - -#[derive(Clone)] -struct PendingWindowAdjust { - received_offset: u64, - bytes: usize, - 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, - ) -> 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), - restored_screen: None, - 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, 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; - let mut restored_screen = None; - if let Some(saved) = saved { - parser.process(&saved.snapshot); - output_seq = saved.output_seq; - restored_screen = Some(RestoredScreenSnapshot { - cols, - rows, - snapshot: saved.snapshot, - }); - } - self.sessions.insert( - name.clone(), - Session { - pty: Some(pty), - parser, - restored_screen, - 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) - { - let removed = session.clients.remove(client_id).is_some(); - if removed && session.clients.is_empty() { - session.empty_since.get_or_insert_with(Instant::now); - } - return removed; - } - 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 attach_snapshot(session: &Session, cols: u16, rows: u16) -> Vec { - let cols = cols.max(1); - let rows = rows.max(1); - if let Some(restored) = &session.restored_screen - && restored.cols == cols - && restored.rows == rows - { - return restored.snapshot.clone(); - } - screen_snapshot(session.parser.screen()) -} - -fn restored_screen_is_pending_reattach(session: &Session) -> bool { - session.restored_screen.is_some() -} - -fn mode_allows_terminal_updates(mode: &str) -> bool { - mode != "view-only" && mode != "forward-only" -} - -fn accepted_env(config: &ServerConfig, requested: &[EnvVar]) -> Result> { - 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>, - socket: &Arc, - 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::StreamEof => handle_stream_eof(state, peer, &packet).await, - PacketKind::StreamClose => handle_stream_close(state, peer, &packet).await, - _ => Ok(()), - } -} - -async fn handle_native_client_hello( - state: &Arc>, - socket: &Arc, - peer: SocketAddr, - body: Vec, -) -> 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>, - socket: &Arc, - 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, - attach_ticket_psk: [u8; 32], - session_name: String, - mode: String, - output_seq: u64, - snapshot: Vec, - requested_forwardings: Vec, - } - - let attached: Result = { - 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 = attach_snapshot(session, cols, rows); - 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_local_eof_sent: HashSet::new(), - stream_remote_eof_received: HashSet::new(), - stream_eof_retransmit: HashMap::new(), - stream_close_retransmit: HashMap::new(), - stream_window_adjust_retransmit: HashMap::new(), - stream_retired: HashSet::new(), - stream_retired_order: VecDeque::new(), - stream_retransmit_srtt: None, - 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>, - socket: &Arc, - peer: SocketAddr, - body: Vec, -) -> Result<()> { - let req: BootstrapAttachRequest = protocol::from_body(&body)?; - type BootstrapAttachResult = ([u8; 16], [u8; 32], [u8; 16], String, String, u64, Vec); - let attached: Result = { - 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 = attach_snapshot(session, req.cols, req.rows); - 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_local_eof_sent: HashSet::new(), - stream_remote_eof_received: HashSet::new(), - stream_eof_retransmit: HashMap::new(), - stream_close_retransmit: HashMap::new(), - stream_window_adjust_retransmit: HashMap::new(), - stream_retired: HashSet::new(), - stream_retired_order: VecDeque::new(), - stream_retransmit_srtt: None, - 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>, - socket: &Arc, - peer: SocketAddr, - body: Vec, -) -> Result<()> { - dosh::trace::event("server.ticket_attach_start", &[("peer", peer.to_string())]); - 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) => { - let reason = err.to_string(); - dosh::trace::event("server.ticket_attach_reject", &[("reason", reason.clone())]); - return send_reject(socket, peer, &reason).await; - } - }; - let req: TicketAttachBody = protocol::from_body(&request_plain)?; - if req.session != ticket.session || req.mode != ticket.mode { - dosh::trace::event( - "server.ticket_attach_reject", - &[ - ("session", req.session.clone()), - ("mode", req.mode.clone()), - ("reason", "ticket scope mismatch".to_string()), - ], - ); - 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 = attach_snapshot(session, req.cols, req.rows); - 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_local_eof_sent: HashSet::new(), - stream_remote_eof_received: HashSet::new(), - stream_eof_retransmit: HashMap::new(), - stream_close_retransmit: HashMap::new(), - stream_window_adjust_retransmit: HashMap::new(), - stream_retired: HashSet::new(), - stream_retired_order: VecDeque::new(), - stream_retransmit_srtt: None, - 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) => { - let reason = err.to_string(); - dosh::trace::event( - "server.ticket_attach_reject", - &[ - ("session", req.session.clone()), - ("mode", req.mode.clone()), - ("reason", reason.clone()), - ], - ); - return send_reject(socket, peer, &reason).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?; - dosh::trace::event( - "server.ticket_attach_ok", - &[ - ("session", ok.session.clone()), - ("mode", ok.mode.clone()), - ("client", hex_id(client_id)), - ("output_seq", output_seq.to_string()), - ("bytes", ok.snapshot.len().to_string()), - ], - ); - 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>, - socket: &Arc, - peer: SocketAddr, - packet: &protocol::Packet, -) -> Result<()> { - dosh::trace::event( - "server.resume_start", - &[ - ("peer", peer.to_string()), - ("client", hex_id(packet.header.conn_id)), - ("seq", packet.header.seq.to_string()), - ("ack", packet.header.ack.to_string()), - ], - ); - let (key, session_name) = match find_client_decrypt_key(state, &packet.header) { - Ok(found) => found, - Err(_) => { - dosh::trace::event( - "server.resume_unknown_client", - &[("client", hex_id(packet.header.conn_id))], - ); - 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 (send_seq, allow_terminal_updates) = { - let client = session - .clients - .get_mut(&packet.header.conn_id) - .ok_or_else(|| anyhow!("unknown client"))?; - if !client.replay.accept(packet.header.seq) { - dosh::trace::event( - "server.resume_replay_drop", - &[ - ("session", req.session.clone()), - ("seq", packet.header.seq.to_string()), - ], - ); - return Ok(()); - } - if client.endpoint != peer { - dosh::trace::event( - "server.resume_roam", - &[ - ("session", req.session.clone()), - ("from", client.endpoint.to_string()), - ("to", peer.to_string()), - ], - ); - } - 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; - (client.send_seq, mode_allows_terminal_updates(&client.mode)) - }; - if allow_terminal_updates { - apply_terminal_size( - session.pty.as_ref(), - &mut session.parser, - req.cols, - req.rows, - )?; - } - let snapshot = attach_snapshot(session, req.cols, req.rows); - (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?; - dosh::trace::event( - "server.resume_ok", - &[ - ("session", frame.session), - ("output_seq", frame.output_seq.to_string()), - ("bytes", frame.bytes.len().to_string()), - ], - ); - Ok(()) -} - -async fn handle_input( - state: &Arc>, - 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 input_summary = dosh::trace::bytes_summary(&input.bytes); - 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) { - dosh::trace::event( - "server.input_replay_drop", - &[ - ("session", session_name.clone()), - ("seq", packet.header.seq.to_string()), - ("summary", input_summary), - ], - ); - return Ok(()); - } - if client.endpoint != peer { - dosh::trace::event( - "server.input_roam", - &[ - ("session", session_name.clone()), - ("from", client.endpoint.to_string()), - ("to", peer.to_string()), - ], - ); - client.endpoint = peer; - } - client.last_seen = Instant::now(); - if !mode_allows_terminal_updates(&client.mode) { - dosh::trace::event( - "server.input_rejected_mode", - &[ - ("session", session_name), - ("mode", client.mode.clone()), - ("summary", input_summary), - ], - ); - return Ok(()); - } - session - .pty - .as_ref() - .context("terminal session has no pty")? - .write_all(&input.bytes)?; - dosh::trace::event( - "server.pty_write", - &[ - ("session", session_name), - ("seq", packet.header.seq.to_string()), - ("summary", input_summary), - ], - ); - Ok(()) -} - -async fn handle_resize( - state: &Arc>, - 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>, - socket: &Arc, - 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>, 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>, 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 { - 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>, - 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>, - 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>, - socket: &Arc, - 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::>(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::>(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::>(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; - } - } - let _ = writer.shutdown().await; - }); - - 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]; - let mut graceful_eof = false; - loop { - match reader.read(&mut buf).await { - Ok(0) => { - graceful_eof = true; - 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 _ = if graceful_eof { - send_stream_eof_to_client(&state, &socket, client_id, stream_id).await - } else { - send_stream_close_to_client(&state, &socket, client_id, stream_id).await - }; - }); - - Ok(()) -} - -fn register_opened_stream( - state: &Arc>, - session_name: &str, - client_id: [u8; 16], - stream_id: u64, - writer_tx: mpsc::Sender>, -) -> 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>, - socket: &Arc, - 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>, - socket: &Arc, - 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(); - let Some(pending_open) = client.stream_pending_opens.remove(&ok.stream_id) else { - return Ok(()); - }; - if pending_open.attempts == 1 { - dosh::transport::observe_retransmit_rtt( - &mut client.stream_retransmit_srtt, - pending_open.last_sent.elapsed(), - ); - } - 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?; - let should_send_eof = { - let mut locked = state.lock().expect("server state poisoned"); - locked - .client_mut(&packet.header.conn_id) - .is_some_and(|client| client.stream_local_eof_sent.contains(&ok.stream_id)) - }; - if should_send_eof { - let body = protocol::to_body(&StreamEof { - stream_id: ok.stream_id, - })?; - { - let mut locked = state.lock().expect("server state poisoned"); - if let Some(client) = locked.client_mut(&packet.header.conn_id) { - client.stream_eof_retransmit.insert( - ok.stream_id, - PendingStreamControl { - last_sent: Instant::now(), - attempts: 1, - }, - ); - } - } - send_stream_packet_to_client( - state, - socket, - packet.header.conn_id, - PacketKind::StreamEof, - body, - ) - .await?; - } - Ok(()) -} - -async fn handle_stream_open_reject( - state: &Arc>, - 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>, - socket: &Arc, - 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_eof( - state: &Arc>, - 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 eof: StreamEof = 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(()); - } - client.endpoint = peer; - client.last_seen = Instant::now(); - if !client.opened_streams.contains(&eof.stream_id) { - return Ok(()); - } - client.stream_remote_eof_received.insert(eof.stream_id); - client.stream_writers.remove(&eof.stream_id); - if client_stream_eof_complete(client, eof.stream_id) { - retire_client_stream(client, eof.stream_id); - } - Ok(()) -} - -async fn handle_stream_close( - state: &Arc>, - 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(); - client.stream_close_retransmit.remove(&close.stream_id); - retire_client_stream(client, close.stream_id); - Ok(()) -} - -async fn start_remote_forwards( - state: &Arc>, - socket: &Arc, - client_id: [u8; 16], - requested_forwardings: Vec, - session_name: String, -) -> Result<()> { - let remote_forwards = requested_forwardings - .into_iter() - .filter(|forward| forward.kind == ForwardingKind::Remote) - .collect::>(); - 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::>(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; - } - } - let _ = writer.shutdown().await; - }); - 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]; - let mut graceful_eof = false; - loop { - match reader.read(&mut buf).await { - Ok(0) => { - graceful_eof = true; - 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 _ = if graceful_eof { - send_stream_eof_to_client(&state, &socket, client_id, stream_id).await - } else { - send_stream_close_to_client(&state, &socket, client_id, stream_id).await - }; - }); - } - }); - } - Ok(()) -} - -fn allocate_server_stream_id(state: &Arc>) -> 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>, - socket: &Arc, - 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::>(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; - } - } - let _ = writer.shutdown().await; - }); - 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]; - let mut graceful_eof = false; - loop { - match reader.read(&mut buf).await { - Ok(0) => { - graceful_eof = true; - 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 _ = if graceful_eof { - send_stream_eof_to_client(&state, &socket, client_id, stream_id).await - } else { - 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, - file: fs::File, - hasher: Sha256, - written: u64, - expected_size: u64, - mode: Option, - modified_secs: Option, - atomic: bool, - original_len: Option, -} - -async fn run_file_stream_service( - state: Arc>, - socket: Arc, - client_id: [u8; 16], - stream_id: u64, - mut rx: mpsc::Receiver>, -) -> 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 = 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>, - socket: &Arc, - client_id: [u8; 16], - stream_id: u64, - home: &Path, - upload: &mut Option, - 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::Readlink { path } => { - let path = clean_remote_path(&path, home)?; - let target = - fs::read_link(&path).with_context(|| format!("readlink {}", path.display()))?; - let target = target - .to_str() - .ok_or_else(|| anyhow!("symlink target is not valid UTF-8"))? - .to_string(); - send_file_response_to_client( - state, - socket, - client_id, - stream_id, - FileResponse::LinkTarget { target }, - ) - .await - } - FileRequest::Symlink { - path, - target, - overwrite, - } => { - let path = clean_remote_path(&path, home)?; - if let Some(parent) = path.parent() { - fs::create_dir_all(parent) - .with_context(|| format!("create {}", parent.display()))?; - } - if let Ok(metadata) = fs::symlink_metadata(&path) { - anyhow::ensure!(overwrite, "destination exists: {}", path.display()); - anyhow::ensure!( - !metadata.is_dir(), - "destination is a directory: {}", - path.display() - ); - fs::remove_file(&path).with_context(|| format!("remove {}", path.display()))?; - } - create_symlink(&target, &path)?; - 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, original_len, prefix_sha256) = - if resume && final_path.exists() { - let original_len = fs::metadata(&final_path) - .with_context(|| format!("stat {}", final_path.display()))? - .len(); - let written = original_len.min(size); - if written > 0 { - let mut prefix = fs::File::open(&final_path)?; - hash_exact_prefix(&mut prefix, written, &mut hasher)?; - } - let prefix_sha256 = (written > 0).then(|| hasher.clone().finalize().into()); - let mut file = fs::OpenOptions::new() - .read(true) - .append(true) - .open(&final_path) - .with_context(|| format!("open {}", final_path.display()))?; - file.seek(SeekFrom::Start(written))?; - ( - file, - None, - written, - false, - Some(original_len), - prefix_sha256, - ) - } 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, None, None) - }; - *upload = Some(UploadState { - final_path, - temp_path, - file, - hasher, - written, - expected_size: size, - mode, - modified_secs, - atomic, - original_len, - }); - send_file_response_to_client( - state, - socket, - client_id, - stream_id, - FileResponse::Resume { - offset: written, - prefix_sha256, - }, - ) - .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"); - }; - 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); - } else if let Some(original_len) = active.original_len { - let _ = active.file.set_len(original_len); - } - bail!("upload checksum mismatch"); - } - active - .file - .set_len(active.expected_size) - .with_context(|| format!("truncate {}", active.final_path.display()))?; - active.file.flush()?; - active.file.sync_all()?; - 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(mode) = active.mode { - set_mode(&active.final_path, mode)?; - } - 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 - ); - anyhow::ensure!( - offset <= meta.len, - "download offset {offset} is beyond file size {}", - meta.len - ); - let mut file = - fs::File::open(&path).with_context(|| format!("open {}", path.display()))?; - file.seek(SeekFrom::Start(offset))?; - let mut hasher = Sha256::new(); - if offset > 0 { - let mut prefix = fs::File::open(&path)?; - hash_exact_prefix(&mut prefix, offset, &mut hasher)?; - } - let prefix_sha256 = (offset > 0).then(|| hasher.clone().finalize().into()); - send_file_response_to_client( - state, - socket, - client_id, - stream_id, - FileResponse::Start { - meta: meta.clone(), - offset, - prefix_sha256, - }, - ) - .await?; - 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>, - socket: &Arc, - 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>, - socket: Arc, - client_id: [u8; 16], - stream_id: u64, - mut rx: mpsc::Receiver>, -) -> 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>, - socket: &Arc, - 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 { - 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())) -} - #[cfg(unix)] -fn create_symlink(target: &str, path: &Path) -> Result<()> { - std::os::unix::fs::symlink(target, path) - .with_context(|| format!("symlink {} -> {}", path.display(), target)) -} +include!("dosh-server/unix.rs"); #[cfg(not(unix))] -fn create_symlink(_target: &str, _path: &Path) -> Result<()> { - bail!("symlink creation is not supported on this server platform") -} - -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>, - socket: &Arc, - 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>, - socket: &Arc, - 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>, - socket: &Arc, - 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>, - socket: &Arc, - client_id: [u8; 16], - stream_id: u64, - bytes: Vec, -) -> 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>, - socket: &Arc, - client_id: [u8; 16], - stream_id: u64, - bytes: Vec, -) -> Result<()> { - let sends = { - let mut locked = state.lock().expect("server state poisoned"); - let mut sends = Vec::new(); - if let Some(client) = locked.client_mut(&client_id) { - for chunk in dosh::transport::split_stream_data_bytes(bytes, STREAM_INITIAL_WINDOW) { - if !client.opened_streams.contains(&stream_id) - || client - .stream_send_credit - .get(&stream_id) - .copied() - .unwrap_or(0) - < chunk.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(chunk); - continue; - } - *client.stream_send_credit.entry(stream_id).or_default() -= chunk.len(); - let offset = *client.stream_next_send_offset.entry(stream_id).or_default(); - client - .stream_next_send_offset - .insert(stream_id, offset.saturating_add(chunk.len() as u64)); - client.send_seq += 1; - let body = protocol::to_body(&StreamData { - stream_id, - offset, - bytes: chunk.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: chunk, - last_sent: Instant::now(), - attempts: 1, - }, - ); - sends.push((client.endpoint, packet)); - } - } - sends - }; - for (endpoint, packet) in sends { - let _ = send_udp(socket, &packet, endpoint).await?; - } - Ok(()) -} - -async fn flush_stream_pending_data_to_client( - state: &Arc>, - socket: &Arc, - 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(front_len) = pending.front().map(Vec::len) else { - client.stream_pending_data.remove(&stream_id); - return Ok(()); - }; - if front_len > dosh::transport::stream_data_chunk_limit(STREAM_INITIAL_WINDOW) { - let Some(bytes) = pending.pop_front() else { - client.stream_pending_data.remove(&stream_id); - return Ok(()); - }; - for chunk in - dosh::transport::split_stream_data_bytes(bytes, STREAM_INITIAL_WINDOW) - .into_iter() - .rev() - { - pending.push_front(chunk); - } - continue; - } - let credit = client - .stream_send_credit - .get(&stream_id) - .copied() - .unwrap_or(0); - if credit < front_len { - return Ok(()); - } - let Some(bytes) = pending.pop_front() else { - client.stream_pending_data.remove(&stream_id); - return Ok(()); - }; - 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>, usize, u64) { - dosh::transport::accept_stream_data_chunks( - &mut client.stream_next_recv_offset, - &mut client.stream_recv_pending, - data.stream_id, - data.offset, - data.bytes, - STREAM_INITIAL_WINDOW, - ) -} - -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); - client.stream_local_eof_sent.remove(&stream_id); - client.stream_remote_eof_received.remove(&stream_id); - client.stream_eof_retransmit.remove(&stream_id); - client.stream_window_adjust_retransmit.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) - || client.stream_local_eof_sent.contains(&stream_id) - || client.stream_remote_eof_received.contains(&stream_id) - || client.stream_eof_retransmit.contains_key(&stream_id) - || client - .stream_window_adjust_retransmit - .contains_key(&stream_id) -} - -fn client_stream_eof_complete(client: &ClientState, stream_id: u64) -> bool { - client.stream_local_eof_sent.contains(&stream_id) - && client.stream_remote_eof_received.contains(&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 mut acked_bytes = 0usize; - let mut samples = Vec::new(); - let remove_stream = { - let Some(sent) = client.stream_sent_data.get_mut(&stream_id) else { - return; - }; - let acked_offsets: Vec = sent - .iter() - .filter_map(|(offset, chunk)| { - let end = chunk.offset.saturating_add(chunk.bytes.len() as u64); - (end <= received_offset).then_some(*offset) - }) - .collect(); - for offset in acked_offsets { - if let Some(chunk) = sent.remove(&offset) { - acked_bytes = acked_bytes.saturating_add(chunk.bytes.len()); - if chunk.attempts == 1 { - samples.push(chunk.last_sent.elapsed()); - } - } - } - sent.is_empty() - }; - if remove_stream { - client.stream_sent_data.remove(&stream_id); - } - for sample in samples { - dosh::transport::observe_retransmit_rtt(&mut client.stream_retransmit_srtt, sample); - } - add_stream_credit(&mut client.stream_send_credit, stream_id, acked_bytes); -} - -fn add_stream_credit(stream_send_credit: &mut HashMap, 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>, - socket: &Arc, - client_id: [u8; 16], - stream_id: u64, - bytes: usize, - received_offset: u64, -) -> Result<()> { - if bytes > 0 { - let mut locked = state.lock().expect("server state poisoned"); - if let Some(client) = locked.client_mut(&client_id) { - client.stream_window_adjust_retransmit.insert( - stream_id, - PendingWindowAdjust { - received_offset, - bytes, - last_sent: Instant::now(), - attempts: 1, - }, - ); - } - } - 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>, - socket: &Arc, - 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); - client.stream_close_retransmit.insert( - stream_id, - PendingStreamControl { - last_sent: Instant::now(), - attempts: 1, - }, - ); - } - } - 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_eof_to_client( - state: &Arc>, - socket: &Arc, - 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) { - client.stream_local_eof_sent.insert(stream_id); - client.stream_eof_retransmit.insert( - stream_id, - PendingStreamControl { - last_sent: Instant::now(), - attempts: 1, - }, - ); - } - } - let body = protocol::to_body(&StreamEof { stream_id })?; - send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamEof, body).await -} - -async fn send_stream_packet_to_client( - state: &Arc>, - socket: &Arc, - client_id: [u8; 16], - kind: PacketKind, - body: Vec, -) -> 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>, - socket: &Arc, - 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)> = 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"))?; - let has_clients = !session.clients.is_empty(); - session.parser.process(&output.bytes); - if has_clients { - session.restored_screen = None; - } - 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 && !restored_screen_is_pending_reattach(session) { - session.bytes_since_persist = session - .bytes_since_persist - .saturating_add(output.bytes.len()); - let now = Instant::now(); - if should_persist_screen_immediately(&output.bytes) - || 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>, - socket: &Arc, - 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 { - 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 -} - -fn should_persist_screen_immediately(bytes: &[u8]) -> bool { - bytes.len() <= 1024 && bytes.contains(&b'\n') -} - -async fn retransmit_pending( - state: &Arc>, - socket: &Arc, -) -> 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 = - client.stream_pending_opens.keys().copied().collect(); - let mut retransmit_opens = Vec::new(); - let stream_retransmit_after = dosh::transport::adaptive_retransmit_after( - dosh::transport::DEFAULT_RETRANSMIT_AFTER, - client.stream_retransmit_srtt, - ); - 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) < stream_retransmit_after { - 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_eof_ids: Vec = client - .stream_local_eof_sent - .iter() - .copied() - .filter(|stream_id| client.opened_streams.contains(stream_id)) - .collect(); - for stream_id in stream_eof_ids { - let pending = client.stream_eof_retransmit.entry(stream_id).or_insert( - PendingStreamControl { - last_sent: now - stream_retransmit_after, - attempts: 0, - }, - ); - if now.duration_since(pending.last_sent) < stream_retransmit_after { - continue; - } - pending.last_sent = now; - pending.attempts = pending.attempts.saturating_add(1); - client.send_seq += 1; - let body = protocol::to_body(&StreamEof { stream_id })?; - let packet = protocol::encode_encrypted( - PacketKind::StreamEof, - *client_id, - client.send_seq, - client.last_acked, - &client.session_key, - SERVER_TO_CLIENT, - &body, - )?; - sends.push((client.endpoint, packet)); - } - let stream_close_ids: Vec = - client.stream_close_retransmit.keys().copied().collect(); - for stream_id in stream_close_ids { - let Some(pending) = client.stream_close_retransmit.get_mut(&stream_id) else { - continue; - }; - if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { - client.stream_close_retransmit.remove(&stream_id); - continue; - } - if now.duration_since(pending.last_sent) < stream_retransmit_after { - continue; - } - pending.last_sent = now; - pending.attempts = pending.attempts.saturating_add(1); - client.send_seq += 1; - let body = protocol::to_body(&StreamClose { stream_id })?; - let packet = protocol::encode_encrypted( - PacketKind::StreamClose, - *client_id, - client.send_seq, - client.last_acked, - &client.session_key, - SERVER_TO_CLIENT, - &body, - )?; - sends.push((client.endpoint, packet)); - } - let stream_adjust_ids: Vec = client - .stream_window_adjust_retransmit - .keys() - .copied() - .collect(); - for stream_id in stream_adjust_ids { - let Some(pending) = client.stream_window_adjust_retransmit.get_mut(&stream_id) - else { - continue; - }; - if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { - client.stream_window_adjust_retransmit.remove(&stream_id); - continue; - } - if now.duration_since(pending.last_sent) < stream_retransmit_after { - continue; - } - pending.last_sent = now; - pending.attempts = pending.attempts.saturating_add(1); - client.send_seq += 1; - let body = protocol::to_body(&StreamWindowAdjust { - stream_id, - received_offset: pending.received_offset, - bytes: pending.bytes.min(u32::MAX as usize) as u32, - })?; - let packet = protocol::encode_encrypted( - PacketKind::StreamWindowAdjust, - *client_id, - client.send_seq, - client.last_acked, - &client.session_key, - SERVER_TO_CLIENT, - &body, - )?; - sends.push((client.endpoint, packet)); - } - let stream_ids: Vec = 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) < stream_retransmit_after { - 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>, - socket: &Arc, -) -> 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>) { - 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)> = Vec::new(); - for (name, session) in locked.sessions.iter_mut() { - if !session.persistent - || restored_screen_is_pending_reattach(session) - || 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>, 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>) { - // 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 = 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 = 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>, - 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> { - 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); - } - - #[test] - fn attach_snapshot_prefers_restored_screen_for_same_size_reattach() { - let restored = b"\x1b[?1049lRESTORED_SCREEN_MARKER".to_vec(); - let mut session = Session { - pty: None, - parser: vt100::Parser::new(24, 80, 100), - restored_screen: Some(RestoredScreenSnapshot { - cols: 80, - rows: 24, - snapshot: restored.clone(), - }), - clients: HashMap::new(), - output_seq: 7, - recent: VecDeque::new(), - empty_since: Some(Instant::now()), - holder_control: None, - persistent: true, - bytes_since_persist: 0, - last_persisted_seq: 7, - last_screen_persist_at: Instant::now(), - }; - - session.parser.set_size(24, 80); - assert_eq!(attach_snapshot(&session, 80, 24), restored); - assert_ne!(attach_snapshot(&session, 100, 30), restored); - } - - #[test] - fn first_reattach_prefers_restored_screen_even_after_pre_attach_pty_noise() { - let restored = b"\x1b[?1049lRESTORED_AFTER_RESTART".to_vec(); - let mut session = Session { - pty: None, - parser: vt100::Parser::new(24, 80, 100), - restored_screen: Some(RestoredScreenSnapshot { - cols: 80, - rows: 24, - snapshot: restored.clone(), - }), - clients: HashMap::new(), - output_seq: 8, - recent: VecDeque::new(), - empty_since: Some(Instant::now()), - holder_control: None, - persistent: true, - bytes_since_persist: 0, - last_persisted_seq: 7, - last_screen_persist_at: Instant::now(), - }; - session.parser.process(b"$ "); - - assert_eq!(attach_snapshot(&session, 80, 24), restored); - - session - .clients - .insert([1u8; 16], test_client_state([2u8; 32])); - assert_eq!(attach_snapshot(&session, 80, 24), restored); - - session.restored_screen = None; - assert_ne!(attach_snapshot(&session, 80, 24), restored); - } - - #[tokio::test] - async fn no_client_pty_output_before_first_reattach_keeps_restored_screen() { - let restored = b"\x1b[?1049lRESTORED_AFTER_RESTART".to_vec(); - let (pty_tx, _rx) = mpsc::unbounded_channel(); - let state = Arc::new(Mutex::new(ServerState::new( - ServerConfig { - persist_sessions: true, - ..ServerConfig::default() - }, - [0u8; 32], - pty_tx, - ))); - { - let mut locked = state.lock().unwrap(); - locked.sessions.insert( - "work".to_string(), - Session { - pty: None, - parser: vt100::Parser::new(24, 80, 100), - restored_screen: Some(RestoredScreenSnapshot { - cols: 80, - rows: 24, - snapshot: restored.clone(), - }), - clients: HashMap::new(), - output_seq: 7, - recent: VecDeque::new(), - empty_since: Some(Instant::now()), - holder_control: None, - persistent: true, - bytes_since_persist: 0, - last_persisted_seq: 7, - last_screen_persist_at: Instant::now(), - }, - ); - } - - let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); - broadcast_output( - &state, - &socket, - PtyOutput { - session: "work".to_string(), - bytes: b"$ ".to_vec(), - exited: false, - }, - ) - .await - .unwrap(); - - let locked = state.lock().unwrap(); - let session = locked.sessions.get("work").unwrap(); - assert!(session.restored_screen.is_some()); - assert_eq!(attach_snapshot(session, 80, 24), restored); - assert_eq!( - session.last_persisted_seq, 7, - "pre-reattach PTY output must not overwrite the persisted restored screen" - ); - drop(locked); - - flush_persistent_screens(&state); - - let locked = state.lock().unwrap(); - let session = locked.sessions.get("work").unwrap(); - assert_eq!( - session.last_persisted_seq, 7, - "periodic flush must not replace restored screen before client-visible output" - ); - } - - #[test] - fn line_oriented_shell_output_persists_screen_immediately() { - assert!(should_persist_screen_immediately( - b"IMPLICIT_SCREEN_MARKER\r\n$ " - )); - assert!(!should_persist_screen_immediately(b"\x1b[2K\x1b[H")); - assert!(!should_persist_screen_immediately(&vec![b'x'; 2048])); - } - - #[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_local_eof_sent: HashSet::new(), - stream_remote_eof_received: HashSet::new(), - stream_eof_retransmit: HashMap::new(), - stream_close_retransmit: HashMap::new(), - stream_window_adjust_retransmit: HashMap::new(), - stream_retired: HashSet::new(), - stream_retired_order: VecDeque::new(), - stream_retransmit_srtt: None, - 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])])); - client.stream_window_adjust_retransmit.insert( - stream_id, - PendingWindowAdjust { - received_offset: 77, - bytes: 5, - last_sent: Instant::now(), - attempts: 1, - }, - ); - client.stream_window_adjust_retransmit.insert( - other_stream_id, - PendingWindowAdjust { - received_offset: 88, - bytes: 6, - last_sent: Instant::now(), - attempts: 1, - }, - ); - - 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 - .stream_window_adjust_retransmit - .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)); - assert!( - client - .stream_window_adjust_retransmit - .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)); - assert!( - state.sessions["work"].empty_since.is_none(), - "removing one of multiple clients must not mark the session empty" - ); - assert!(state.remove_client_everywhere(&[7u8; 16])); - assert!( - state.sessions["work"].empty_since.is_some(), - "removing the final client should start empty-session cleanup grace immediately" - ); - // 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), - restored_screen: None, - 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), - restored_screen: None, - 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), - restored_screen: None, - 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); - } - - #[tokio::test] - async fn server_stream_retransmit_uses_observed_rtt() { - 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 = 124; - 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(); - dosh::transport::observe_retransmit_rtt( - &mut client.stream_retransmit_srtt, - Duration::from_millis(25), - ); - client.stream_pending_opens.insert( - stream_id, - PendingStreamOpen { - target_host: "127.0.0.1".to_string(), - target_port: 8081, - last_sent: Instant::now() - Duration::from_millis(50), - attempts: 1, - }, - ); - state.sessions.insert( - "test".to_string(), - Session { - pty: None, - parser: vt100::Parser::new(24, 80, 100), - restored_screen: None, - 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, 8081); - } - - #[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), - restored_screen: None, - 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_local_eof_sent: HashSet::new(), - stream_remote_eof_received: HashSet::new(), - stream_eof_retransmit: HashMap::new(), - stream_close_retransmit: HashMap::new(), - stream_window_adjust_retransmit: HashMap::new(), - stream_retired: HashSet::new(), - stream_retired_order: VecDeque::new(), - stream_retransmit_srtt: None, - 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 stream_eof_to_client_does_not_retire_stream() { - 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 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), - restored_screen: None, - 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::from([42]), - stream_send_credit: HashMap::from([(42, STREAM_INITIAL_WINDOW)]), - 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_local_eof_sent: HashSet::new(), - stream_remote_eof_received: HashSet::new(), - stream_eof_retransmit: HashMap::new(), - stream_close_retransmit: HashMap::new(), - stream_window_adjust_retransmit: HashMap::new(), - stream_retired: HashSet::new(), - stream_retired_order: VecDeque::new(), - stream_retransmit_srtt: None, - 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)); - - send_stream_eof_to_client(&state, &sender, client_id, 42) - .await - .unwrap(); - - { - let locked = state.lock().unwrap(); - let client = locked.sessions["test"].clients.get(&client_id).unwrap(); - assert!(client.opened_streams.contains(&42)); - assert!(client.stream_local_eof_sent.contains(&42)); - assert!(!client.stream_retired.contains(&42)); - } - - 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::StreamEof); - let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); - let eof: StreamEof = protocol::from_body(&plain).unwrap(); - assert_eq!(eof.stream_id, 42); - } - - #[tokio::test] - async fn pending_server_stream_eof_is_retransmitted() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); - let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); - let client_id = [17u8; 16]; - let session_key = [18u8; 32]; - let stream_id = 125; - 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); - client.stream_local_eof_sent.insert(stream_id); - client.stream_local_eof_sent.insert(999); - client.stream_eof_retransmit.insert( - stream_id, - PendingStreamControl { - 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), - restored_screen: None, - 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::StreamEof); - let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); - let eof: StreamEof = protocol::from_body(&plain).unwrap(); - assert_eq!(eof.stream_id, stream_id); - - let locked = state.lock().unwrap(); - let client = locked.sessions["test"].clients.get(&client_id).unwrap(); - assert_eq!(client.stream_eof_retransmit[&stream_id].attempts, 2); - assert!(!client.stream_eof_retransmit.contains_key(&999)); - } - - #[tokio::test] - async fn pending_server_stream_close_is_retransmitted() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); - let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); - let client_id = [19u8; 16]; - let session_key = [20u8; 32]; - let stream_id = 126; - 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_close_retransmit.insert( - stream_id, - PendingStreamControl { - 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), - restored_screen: None, - 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::StreamClose); - let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); - let close: StreamClose = protocol::from_body(&plain).unwrap(); - assert_eq!(close.stream_id, stream_id); - - let locked = state.lock().unwrap(); - let client = locked.sessions["test"].clients.get(&client_id).unwrap(); - assert_eq!(client.stream_close_retransmit[&stream_id].attempts, 2); - } - - #[tokio::test] - async fn pending_server_stream_close_expires_after_attempt_cap() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); - let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); - let client_id = [21u8; 16]; - let session_key = [22u8; 32]; - let stream_id = 127; - 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_close_retransmit.insert( - stream_id, - PendingStreamControl { - last_sent: Instant::now() - Duration::from_millis(250), - attempts: STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS, - }, - ); - state.sessions.insert( - "test".to_string(), - Session { - pty: None, - parser: vt100::Parser::new(24, 80, 100), - restored_screen: None, - 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 locked = state.lock().unwrap(); - let client = locked.sessions["test"].clients.get(&client_id).unwrap(); - assert!(!client.stream_close_retransmit.contains_key(&stream_id)); - } - - #[tokio::test] - async fn pending_server_stream_window_adjust_is_retransmitted() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); - let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); - let client_id = [23u8; 16]; - let session_key = [24u8; 32]; - let stream_id = 128; - 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_window_adjust_retransmit.insert( - stream_id, - PendingWindowAdjust { - received_offset: 4567, - bytes: 123, - 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), - restored_screen: None, - 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::StreamWindowAdjust); - let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); - let adjust: StreamWindowAdjust = protocol::from_body(&plain).unwrap(); - assert_eq!(adjust.stream_id, stream_id); - assert_eq!(adjust.received_offset, 4567); - assert_eq!(adjust.bytes, 123); - - let locked = state.lock().unwrap(); - let client = locked.sessions["test"].clients.get(&client_id).unwrap(); - assert_eq!( - client.stream_window_adjust_retransmit[&stream_id].attempts, - 2 - ); - } - - #[tokio::test] - async fn pending_server_stream_window_adjust_expires_after_attempt_cap() { - let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); - let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); - let client_id = [25u8; 16]; - let session_key = [26u8; 32]; - let stream_id = 129; - 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_window_adjust_retransmit.insert( - stream_id, - PendingWindowAdjust { - received_offset: 4567, - bytes: 123, - last_sent: Instant::now() - Duration::from_millis(250), - attempts: STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS, - }, - ); - state.sessions.insert( - "test".to_string(), - Session { - pty: None, - parser: vt100::Parser::new(24, 80, 100), - restored_screen: None, - 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 locked = state.lock().unwrap(); - let client = locked.sessions["test"].clients.get(&client_id).unwrap(); - assert!( - !client - .stream_window_adjust_retransmit - .contains_key(&stream_id) - ); - } - - #[tokio::test] - async fn server_stream_send_splits_large_writes() { - 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 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), - restored_screen: None, - 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::from([42]), - stream_send_credit: HashMap::from([(42, STREAM_INITIAL_WINDOW)]), - 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_local_eof_sent: HashSet::new(), - stream_remote_eof_received: HashSet::new(), - stream_eof_retransmit: HashMap::new(), - stream_close_retransmit: HashMap::new(), - stream_window_adjust_retransmit: HashMap::new(), - stream_retired: HashSet::new(), - stream_retired_order: VecDeque::new(), - stream_retransmit_srtt: None, - 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)); - let bytes = vec![3; dosh::transport::MAX_STREAM_DATA_BYTES * 2 + 17]; - - queue_or_send_stream_data_to_client(&state, &sender, client_id, 42, bytes) - .await - .unwrap(); - - { - let locked = state.lock().unwrap(); - let client = locked.sessions["test"].clients.get(&client_id).unwrap(); - assert!(client.stream_pending_data.is_empty()); - assert_eq!(client.stream_sent_data[&42].len(), 3); - assert_eq!( - client.stream_next_send_offset[&42], - (dosh::transport::MAX_STREAM_DATA_BYTES * 2 + 17) as u64 - ); - } - - let mut seen = Vec::new(); - let mut buf = vec![0u8; u16::MAX as usize]; - for _ in 0..3 { - 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(); - seen.push((data.offset, data.bytes.len())); - } - assert_eq!( - seen, - vec![ - (0, dosh::transport::MAX_STREAM_DATA_BYTES), - ( - dosh::transport::MAX_STREAM_DATA_BYTES as u64, - dosh::transport::MAX_STREAM_DATA_BYTES - ), - ((dosh::transport::MAX_STREAM_DATA_BYTES * 2) as u64, 17) - ] - ); - } - - #[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), - restored_screen: None, - 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_local_eof_sent: HashSet::new(), - stream_remote_eof_received: HashSet::new(), - stream_eof_retransmit: HashMap::new(), - stream_close_retransmit: HashMap::new(), - stream_window_adjust_retransmit: HashMap::new(), - stream_retired: HashSet::new(), - stream_retired_order: VecDeque::new(), - stream_retransmit_srtt: None, - 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" - ); - } - } +fn main() { + eprintln!("dosh-server is only supported on Unix-like hosts"); + std::process::exit(1); } diff --git a/src/bin/dosh-server/unix.rs b/src/bin/dosh-server/unix.rs new file mode 100644 index 0000000..f93443d --- /dev/null +++ b/src/bin/dosh-server/unix.rs @@ -0,0 +1,6233 @@ +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, StreamEof, 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_error, 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; +const STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS: u32 = 8; + +/// 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; + +fn hex_id(id: [u8; 16]) -> String { + let mut out = String::with_capacity(32); + for byte in id { + use std::fmt::Write as _; + let _ = write!(&mut out, "{byte:02x}"); + } + out +} + +#[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, + }, + Auth { + #[arg(long, default_value_t = 1)] + protocol: u8, + #[arg(long)] + nonce: Option, + #[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, + }, + /// 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, + }, +} + +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) -> 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}"); + dosh::trace::event( + "server.start", + &[ + ("version", dosh::build_info::VERSION.to_string()), + ("bind", bind.clone()), + ], + ); + + 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) = match socket.recv_from(&mut buf).await { + Ok(value) => value, + Err(err) if is_transient_udp_error(&err) => continue, + Err(err) => return Err(err.into()), + }; + 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, + sessions: HashMap, + 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, +} + +#[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 { + 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, + parser: vt100::Parser, + restored_screen: Option, + clients: HashMap<[u8; 16], ClientState>, + output_seq: u64, + recent: VecDeque>, + /// 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, + /// 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, + /// 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 RestoredScreenSnapshot { + cols: u16, + rows: u16, + snapshot: Vec, +} + +#[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, + allowed_forwardings: Vec, + stream_writers: HashMap>>, + opened_streams: HashSet, + stream_send_credit: HashMap, + stream_pending_data: HashMap>>, + stream_pending_opens: HashMap, + stream_next_send_offset: HashMap, + stream_sent_data: HashMap>, + stream_next_recv_offset: HashMap, + stream_recv_pending: HashMap>>, + stream_local_eof_sent: HashSet, + stream_remote_eof_received: HashSet, + stream_eof_retransmit: HashMap, + stream_close_retransmit: HashMap, + stream_window_adjust_retransmit: HashMap, + stream_retired: HashSet, + stream_retired_order: VecDeque, + stream_retransmit_srtt: Option, + /// 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, + last_sent: Instant, + attempts: u8, +} + +#[derive(Clone)] +struct PendingStreamChunk { + offset: u64, + bytes: Vec, + last_sent: Instant, + attempts: u32, +} + +#[derive(Clone)] +struct PendingStreamOpen { + target_host: String, + target_port: u16, + last_sent: Instant, + attempts: u32, +} + +#[derive(Clone)] +struct PendingStreamControl { + last_sent: Instant, + attempts: u32, +} + +#[derive(Clone)] +struct PendingWindowAdjust { + received_offset: u64, + bytes: usize, + 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, + ) -> 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), + restored_screen: None, + 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, 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; + let mut restored_screen = None; + if let Some(saved) = saved { + parser.process(&saved.snapshot); + output_seq = saved.output_seq; + restored_screen = Some(RestoredScreenSnapshot { + cols, + rows, + snapshot: saved.snapshot, + }); + } + self.sessions.insert( + name.clone(), + Session { + pty: Some(pty), + parser, + restored_screen, + 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) + { + let removed = session.clients.remove(client_id).is_some(); + if removed && session.clients.is_empty() { + session.empty_since.get_or_insert_with(Instant::now); + } + return removed; + } + 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 attach_snapshot(session: &Session, cols: u16, rows: u16) -> Vec { + let cols = cols.max(1); + let rows = rows.max(1); + if let Some(restored) = &session.restored_screen + && restored.cols == cols + && restored.rows == rows + { + return restored.snapshot.clone(); + } + screen_snapshot(session.parser.screen()) +} + +fn restored_screen_is_pending_reattach(session: &Session) -> bool { + session.restored_screen.is_some() +} + +fn mode_allows_terminal_updates(mode: &str) -> bool { + mode != "view-only" && mode != "forward-only" +} + +fn accepted_env(config: &ServerConfig, requested: &[EnvVar]) -> Result> { + 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>, + socket: &Arc, + 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::StreamEof => handle_stream_eof(state, peer, &packet).await, + PacketKind::StreamClose => handle_stream_close(state, peer, &packet).await, + _ => Ok(()), + } +} + +async fn handle_native_client_hello( + state: &Arc>, + socket: &Arc, + peer: SocketAddr, + body: Vec, +) -> 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>, + socket: &Arc, + 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, + attach_ticket_psk: [u8; 32], + session_name: String, + mode: String, + output_seq: u64, + snapshot: Vec, + requested_forwardings: Vec, + } + + let attached: Result = { + 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 = attach_snapshot(session, cols, rows); + 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_local_eof_sent: HashSet::new(), + stream_remote_eof_received: HashSet::new(), + stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), + stream_window_adjust_retransmit: HashMap::new(), + stream_retired: HashSet::new(), + stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, + 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>, + socket: &Arc, + peer: SocketAddr, + body: Vec, +) -> Result<()> { + let req: BootstrapAttachRequest = protocol::from_body(&body)?; + type BootstrapAttachResult = ([u8; 16], [u8; 32], [u8; 16], String, String, u64, Vec); + let attached: Result = { + 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 = attach_snapshot(session, req.cols, req.rows); + 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_local_eof_sent: HashSet::new(), + stream_remote_eof_received: HashSet::new(), + stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), + stream_window_adjust_retransmit: HashMap::new(), + stream_retired: HashSet::new(), + stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, + 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>, + socket: &Arc, + peer: SocketAddr, + body: Vec, +) -> Result<()> { + dosh::trace::event("server.ticket_attach_start", &[("peer", peer.to_string())]); + 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) => { + let reason = err.to_string(); + dosh::trace::event("server.ticket_attach_reject", &[("reason", reason.clone())]); + return send_reject(socket, peer, &reason).await; + } + }; + let req: TicketAttachBody = protocol::from_body(&request_plain)?; + if req.session != ticket.session || req.mode != ticket.mode { + dosh::trace::event( + "server.ticket_attach_reject", + &[ + ("session", req.session.clone()), + ("mode", req.mode.clone()), + ("reason", "ticket scope mismatch".to_string()), + ], + ); + 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 = attach_snapshot(session, req.cols, req.rows); + 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_local_eof_sent: HashSet::new(), + stream_remote_eof_received: HashSet::new(), + stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), + stream_window_adjust_retransmit: HashMap::new(), + stream_retired: HashSet::new(), + stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, + 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) => { + let reason = err.to_string(); + dosh::trace::event( + "server.ticket_attach_reject", + &[ + ("session", req.session.clone()), + ("mode", req.mode.clone()), + ("reason", reason.clone()), + ], + ); + return send_reject(socket, peer, &reason).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?; + dosh::trace::event( + "server.ticket_attach_ok", + &[ + ("session", ok.session.clone()), + ("mode", ok.mode.clone()), + ("client", hex_id(client_id)), + ("output_seq", output_seq.to_string()), + ("bytes", ok.snapshot.len().to_string()), + ], + ); + 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>, + socket: &Arc, + peer: SocketAddr, + packet: &protocol::Packet, +) -> Result<()> { + dosh::trace::event( + "server.resume_start", + &[ + ("peer", peer.to_string()), + ("client", hex_id(packet.header.conn_id)), + ("seq", packet.header.seq.to_string()), + ("ack", packet.header.ack.to_string()), + ], + ); + let (key, session_name) = match find_client_decrypt_key(state, &packet.header) { + Ok(found) => found, + Err(_) => { + dosh::trace::event( + "server.resume_unknown_client", + &[("client", hex_id(packet.header.conn_id))], + ); + 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 (send_seq, allow_terminal_updates) = { + let client = session + .clients + .get_mut(&packet.header.conn_id) + .ok_or_else(|| anyhow!("unknown client"))?; + if !client.replay.accept(packet.header.seq) { + dosh::trace::event( + "server.resume_replay_drop", + &[ + ("session", req.session.clone()), + ("seq", packet.header.seq.to_string()), + ], + ); + return Ok(()); + } + if client.endpoint != peer { + dosh::trace::event( + "server.resume_roam", + &[ + ("session", req.session.clone()), + ("from", client.endpoint.to_string()), + ("to", peer.to_string()), + ], + ); + } + 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; + (client.send_seq, mode_allows_terminal_updates(&client.mode)) + }; + if allow_terminal_updates { + apply_terminal_size( + session.pty.as_ref(), + &mut session.parser, + req.cols, + req.rows, + )?; + } + let snapshot = attach_snapshot(session, req.cols, req.rows); + (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?; + dosh::trace::event( + "server.resume_ok", + &[ + ("session", frame.session), + ("output_seq", frame.output_seq.to_string()), + ("bytes", frame.bytes.len().to_string()), + ], + ); + Ok(()) +} + +async fn handle_input( + state: &Arc>, + 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 input_summary = dosh::trace::bytes_summary(&input.bytes); + 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) { + dosh::trace::event( + "server.input_replay_drop", + &[ + ("session", session_name.clone()), + ("seq", packet.header.seq.to_string()), + ("summary", input_summary), + ], + ); + return Ok(()); + } + if client.endpoint != peer { + dosh::trace::event( + "server.input_roam", + &[ + ("session", session_name.clone()), + ("from", client.endpoint.to_string()), + ("to", peer.to_string()), + ], + ); + client.endpoint = peer; + } + client.last_seen = Instant::now(); + if !mode_allows_terminal_updates(&client.mode) { + dosh::trace::event( + "server.input_rejected_mode", + &[ + ("session", session_name), + ("mode", client.mode.clone()), + ("summary", input_summary), + ], + ); + return Ok(()); + } + session + .pty + .as_ref() + .context("terminal session has no pty")? + .write_all(&input.bytes)?; + dosh::trace::event( + "server.pty_write", + &[ + ("session", session_name), + ("seq", packet.header.seq.to_string()), + ("summary", input_summary), + ], + ); + Ok(()) +} + +async fn handle_resize( + state: &Arc>, + 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>, + socket: &Arc, + 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>, 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>, 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 { + 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>, + 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>, + 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>, + socket: &Arc, + 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::>(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::>(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::>(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; + } + } + let _ = writer.shutdown().await; + }); + + 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]; + let mut graceful_eof = false; + loop { + match reader.read(&mut buf).await { + Ok(0) => { + graceful_eof = true; + 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 _ = if graceful_eof { + send_stream_eof_to_client(&state, &socket, client_id, stream_id).await + } else { + send_stream_close_to_client(&state, &socket, client_id, stream_id).await + }; + }); + + Ok(()) +} + +fn register_opened_stream( + state: &Arc>, + session_name: &str, + client_id: [u8; 16], + stream_id: u64, + writer_tx: mpsc::Sender>, +) -> 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>, + socket: &Arc, + 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>, + socket: &Arc, + 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(); + let Some(pending_open) = client.stream_pending_opens.remove(&ok.stream_id) else { + return Ok(()); + }; + if pending_open.attempts == 1 { + dosh::transport::observe_retransmit_rtt( + &mut client.stream_retransmit_srtt, + pending_open.last_sent.elapsed(), + ); + } + 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?; + let should_send_eof = { + let mut locked = state.lock().expect("server state poisoned"); + locked + .client_mut(&packet.header.conn_id) + .is_some_and(|client| client.stream_local_eof_sent.contains(&ok.stream_id)) + }; + if should_send_eof { + let body = protocol::to_body(&StreamEof { + stream_id: ok.stream_id, + })?; + { + let mut locked = state.lock().expect("server state poisoned"); + if let Some(client) = locked.client_mut(&packet.header.conn_id) { + client.stream_eof_retransmit.insert( + ok.stream_id, + PendingStreamControl { + last_sent: Instant::now(), + attempts: 1, + }, + ); + } + } + send_stream_packet_to_client( + state, + socket, + packet.header.conn_id, + PacketKind::StreamEof, + body, + ) + .await?; + } + Ok(()) +} + +async fn handle_stream_open_reject( + state: &Arc>, + 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>, + socket: &Arc, + 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_eof( + state: &Arc>, + 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 eof: StreamEof = 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(()); + } + client.endpoint = peer; + client.last_seen = Instant::now(); + if !client.opened_streams.contains(&eof.stream_id) { + return Ok(()); + } + client.stream_remote_eof_received.insert(eof.stream_id); + client.stream_writers.remove(&eof.stream_id); + if client_stream_eof_complete(client, eof.stream_id) { + retire_client_stream(client, eof.stream_id); + } + Ok(()) +} + +async fn handle_stream_close( + state: &Arc>, + 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(); + client.stream_close_retransmit.remove(&close.stream_id); + retire_client_stream(client, close.stream_id); + Ok(()) +} + +async fn start_remote_forwards( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + requested_forwardings: Vec, + session_name: String, +) -> Result<()> { + let remote_forwards = requested_forwardings + .into_iter() + .filter(|forward| forward.kind == ForwardingKind::Remote) + .collect::>(); + 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::>(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; + } + } + let _ = writer.shutdown().await; + }); + 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]; + let mut graceful_eof = false; + loop { + match reader.read(&mut buf).await { + Ok(0) => { + graceful_eof = true; + 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 _ = if graceful_eof { + send_stream_eof_to_client(&state, &socket, client_id, stream_id).await + } else { + send_stream_close_to_client(&state, &socket, client_id, stream_id).await + }; + }); + } + }); + } + Ok(()) +} + +fn allocate_server_stream_id(state: &Arc>) -> 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>, + socket: &Arc, + 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::>(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; + } + } + let _ = writer.shutdown().await; + }); + 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]; + let mut graceful_eof = false; + loop { + match reader.read(&mut buf).await { + Ok(0) => { + graceful_eof = true; + 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 _ = if graceful_eof { + send_stream_eof_to_client(&state, &socket, client_id, stream_id).await + } else { + 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, + file: fs::File, + hasher: Sha256, + written: u64, + expected_size: u64, + mode: Option, + modified_secs: Option, + atomic: bool, + original_len: Option, +} + +async fn run_file_stream_service( + state: Arc>, + socket: Arc, + client_id: [u8; 16], + stream_id: u64, + mut rx: mpsc::Receiver>, +) -> 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 = 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>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + home: &Path, + upload: &mut Option, + 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::Readlink { path } => { + let path = clean_remote_path(&path, home)?; + let target = + fs::read_link(&path).with_context(|| format!("readlink {}", path.display()))?; + let target = target + .to_str() + .ok_or_else(|| anyhow!("symlink target is not valid UTF-8"))? + .to_string(); + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::LinkTarget { target }, + ) + .await + } + FileRequest::Symlink { + path, + target, + overwrite, + } => { + let path = clean_remote_path(&path, home)?; + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("create {}", parent.display()))?; + } + if let Ok(metadata) = fs::symlink_metadata(&path) { + anyhow::ensure!(overwrite, "destination exists: {}", path.display()); + anyhow::ensure!( + !metadata.is_dir(), + "destination is a directory: {}", + path.display() + ); + fs::remove_file(&path).with_context(|| format!("remove {}", path.display()))?; + } + create_symlink(&target, &path)?; + 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, original_len, prefix_sha256) = + if resume && final_path.exists() { + let original_len = fs::metadata(&final_path) + .with_context(|| format!("stat {}", final_path.display()))? + .len(); + let written = original_len.min(size); + if written > 0 { + let mut prefix = fs::File::open(&final_path)?; + hash_exact_prefix(&mut prefix, written, &mut hasher)?; + } + let prefix_sha256 = (written > 0).then(|| hasher.clone().finalize().into()); + let mut file = fs::OpenOptions::new() + .read(true) + .append(true) + .open(&final_path) + .with_context(|| format!("open {}", final_path.display()))?; + file.seek(SeekFrom::Start(written))?; + ( + file, + None, + written, + false, + Some(original_len), + prefix_sha256, + ) + } 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, None, None) + }; + *upload = Some(UploadState { + final_path, + temp_path, + file, + hasher, + written, + expected_size: size, + mode, + modified_secs, + atomic, + original_len, + }); + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::Resume { + offset: written, + prefix_sha256, + }, + ) + .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"); + }; + 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); + } else if let Some(original_len) = active.original_len { + let _ = active.file.set_len(original_len); + } + bail!("upload checksum mismatch"); + } + active + .file + .set_len(active.expected_size) + .with_context(|| format!("truncate {}", active.final_path.display()))?; + active.file.flush()?; + active.file.sync_all()?; + 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(mode) = active.mode { + set_mode(&active.final_path, mode)?; + } + 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 + ); + anyhow::ensure!( + offset <= meta.len, + "download offset {offset} is beyond file size {}", + meta.len + ); + let mut file = + fs::File::open(&path).with_context(|| format!("open {}", path.display()))?; + file.seek(SeekFrom::Start(offset))?; + let mut hasher = Sha256::new(); + if offset > 0 { + let mut prefix = fs::File::open(&path)?; + hash_exact_prefix(&mut prefix, offset, &mut hasher)?; + } + let prefix_sha256 = (offset > 0).then(|| hasher.clone().finalize().into()); + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::Start { + meta: meta.clone(), + offset, + prefix_sha256, + }, + ) + .await?; + 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>, + socket: &Arc, + 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>, + socket: Arc, + client_id: [u8; 16], + stream_id: u64, + mut rx: mpsc::Receiver>, +) -> 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>, + socket: &Arc, + 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 { + 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())) +} + +#[cfg(unix)] +fn create_symlink(target: &str, path: &Path) -> Result<()> { + std::os::unix::fs::symlink(target, path) + .with_context(|| format!("symlink {} -> {}", path.display(), target)) +} + +#[cfg(not(unix))] +fn create_symlink(_target: &str, _path: &Path) -> Result<()> { + bail!("symlink creation is not supported on this server platform") +} + +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>, + socket: &Arc, + 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>, + socket: &Arc, + 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>, + socket: &Arc, + 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>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + bytes: Vec, +) -> 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>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + bytes: Vec, +) -> Result<()> { + let sends = { + let mut locked = state.lock().expect("server state poisoned"); + let mut sends = Vec::new(); + if let Some(client) = locked.client_mut(&client_id) { + for chunk in dosh::transport::split_stream_data_bytes(bytes, STREAM_INITIAL_WINDOW) { + if !client.opened_streams.contains(&stream_id) + || client + .stream_send_credit + .get(&stream_id) + .copied() + .unwrap_or(0) + < chunk.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(chunk); + continue; + } + *client.stream_send_credit.entry(stream_id).or_default() -= chunk.len(); + let offset = *client.stream_next_send_offset.entry(stream_id).or_default(); + client + .stream_next_send_offset + .insert(stream_id, offset.saturating_add(chunk.len() as u64)); + client.send_seq += 1; + let body = protocol::to_body(&StreamData { + stream_id, + offset, + bytes: chunk.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: chunk, + last_sent: Instant::now(), + attempts: 1, + }, + ); + sends.push((client.endpoint, packet)); + } + } + sends + }; + for (endpoint, packet) in sends { + let _ = send_udp(socket, &packet, endpoint).await?; + } + Ok(()) +} + +async fn flush_stream_pending_data_to_client( + state: &Arc>, + socket: &Arc, + 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(front_len) = pending.front().map(Vec::len) else { + client.stream_pending_data.remove(&stream_id); + return Ok(()); + }; + if front_len > dosh::transport::stream_data_chunk_limit(STREAM_INITIAL_WINDOW) { + let Some(bytes) = pending.pop_front() else { + client.stream_pending_data.remove(&stream_id); + return Ok(()); + }; + for chunk in + dosh::transport::split_stream_data_bytes(bytes, STREAM_INITIAL_WINDOW) + .into_iter() + .rev() + { + pending.push_front(chunk); + } + continue; + } + let credit = client + .stream_send_credit + .get(&stream_id) + .copied() + .unwrap_or(0); + if credit < front_len { + return Ok(()); + } + let Some(bytes) = pending.pop_front() else { + client.stream_pending_data.remove(&stream_id); + return Ok(()); + }; + 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>, usize, u64) { + dosh::transport::accept_stream_data_chunks( + &mut client.stream_next_recv_offset, + &mut client.stream_recv_pending, + data.stream_id, + data.offset, + data.bytes, + STREAM_INITIAL_WINDOW, + ) +} + +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); + client.stream_local_eof_sent.remove(&stream_id); + client.stream_remote_eof_received.remove(&stream_id); + client.stream_eof_retransmit.remove(&stream_id); + client.stream_window_adjust_retransmit.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) + || client.stream_local_eof_sent.contains(&stream_id) + || client.stream_remote_eof_received.contains(&stream_id) + || client.stream_eof_retransmit.contains_key(&stream_id) + || client + .stream_window_adjust_retransmit + .contains_key(&stream_id) +} + +fn client_stream_eof_complete(client: &ClientState, stream_id: u64) -> bool { + client.stream_local_eof_sent.contains(&stream_id) + && client.stream_remote_eof_received.contains(&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 mut acked_bytes = 0usize; + let mut samples = Vec::new(); + let remove_stream = { + let Some(sent) = client.stream_sent_data.get_mut(&stream_id) else { + return; + }; + let acked_offsets: Vec = sent + .iter() + .filter_map(|(offset, chunk)| { + let end = chunk.offset.saturating_add(chunk.bytes.len() as u64); + (end <= received_offset).then_some(*offset) + }) + .collect(); + for offset in acked_offsets { + if let Some(chunk) = sent.remove(&offset) { + acked_bytes = acked_bytes.saturating_add(chunk.bytes.len()); + if chunk.attempts == 1 { + samples.push(chunk.last_sent.elapsed()); + } + } + } + sent.is_empty() + }; + if remove_stream { + client.stream_sent_data.remove(&stream_id); + } + for sample in samples { + dosh::transport::observe_retransmit_rtt(&mut client.stream_retransmit_srtt, sample); + } + add_stream_credit(&mut client.stream_send_credit, stream_id, acked_bytes); +} + +fn add_stream_credit(stream_send_credit: &mut HashMap, 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>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + bytes: usize, + received_offset: u64, +) -> Result<()> { + if bytes > 0 { + let mut locked = state.lock().expect("server state poisoned"); + if let Some(client) = locked.client_mut(&client_id) { + client.stream_window_adjust_retransmit.insert( + stream_id, + PendingWindowAdjust { + received_offset, + bytes, + last_sent: Instant::now(), + attempts: 1, + }, + ); + } + } + 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>, + socket: &Arc, + 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); + client.stream_close_retransmit.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now(), + attempts: 1, + }, + ); + } + } + 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_eof_to_client( + state: &Arc>, + socket: &Arc, + 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) { + client.stream_local_eof_sent.insert(stream_id); + client.stream_eof_retransmit.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now(), + attempts: 1, + }, + ); + } + } + let body = protocol::to_body(&StreamEof { stream_id })?; + send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamEof, body).await +} + +async fn send_stream_packet_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + kind: PacketKind, + body: Vec, +) -> 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>, + socket: &Arc, + 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)> = 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"))?; + let has_clients = !session.clients.is_empty(); + session.parser.process(&output.bytes); + if has_clients { + session.restored_screen = None; + } + 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 && !restored_screen_is_pending_reattach(session) { + session.bytes_since_persist = session + .bytes_since_persist + .saturating_add(output.bytes.len()); + let now = Instant::now(); + if should_persist_screen_immediately(&output.bytes) + || 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>, + socket: &Arc, + 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 { + 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 +} + +fn should_persist_screen_immediately(bytes: &[u8]) -> bool { + bytes.len() <= 1024 && bytes.contains(&b'\n') +} + +async fn retransmit_pending( + state: &Arc>, + socket: &Arc, +) -> 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 = + client.stream_pending_opens.keys().copied().collect(); + let mut retransmit_opens = Vec::new(); + let stream_retransmit_after = dosh::transport::adaptive_retransmit_after( + dosh::transport::DEFAULT_RETRANSMIT_AFTER, + client.stream_retransmit_srtt, + ); + 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) < stream_retransmit_after { + 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_eof_ids: Vec = client + .stream_local_eof_sent + .iter() + .copied() + .filter(|stream_id| client.opened_streams.contains(stream_id)) + .collect(); + for stream_id in stream_eof_ids { + let pending = client.stream_eof_retransmit.entry(stream_id).or_insert( + PendingStreamControl { + last_sent: now - stream_retransmit_after, + attempts: 0, + }, + ); + if now.duration_since(pending.last_sent) < stream_retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + client.send_seq += 1; + let body = protocol::to_body(&StreamEof { stream_id })?; + let packet = protocol::encode_encrypted( + PacketKind::StreamEof, + *client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + sends.push((client.endpoint, packet)); + } + let stream_close_ids: Vec = + client.stream_close_retransmit.keys().copied().collect(); + for stream_id in stream_close_ids { + let Some(pending) = client.stream_close_retransmit.get_mut(&stream_id) else { + continue; + }; + if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { + client.stream_close_retransmit.remove(&stream_id); + continue; + } + if now.duration_since(pending.last_sent) < stream_retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + client.send_seq += 1; + let body = protocol::to_body(&StreamClose { stream_id })?; + let packet = protocol::encode_encrypted( + PacketKind::StreamClose, + *client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + sends.push((client.endpoint, packet)); + } + let stream_adjust_ids: Vec = client + .stream_window_adjust_retransmit + .keys() + .copied() + .collect(); + for stream_id in stream_adjust_ids { + let Some(pending) = client.stream_window_adjust_retransmit.get_mut(&stream_id) + else { + continue; + }; + if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS { + client.stream_window_adjust_retransmit.remove(&stream_id); + continue; + } + if now.duration_since(pending.last_sent) < stream_retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + client.send_seq += 1; + let body = protocol::to_body(&StreamWindowAdjust { + stream_id, + received_offset: pending.received_offset, + bytes: pending.bytes.min(u32::MAX as usize) as u32, + })?; + let packet = protocol::encode_encrypted( + PacketKind::StreamWindowAdjust, + *client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + sends.push((client.endpoint, packet)); + } + let stream_ids: Vec = 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) < stream_retransmit_after { + 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>, + socket: &Arc, +) -> 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>) { + 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)> = Vec::new(); + for (name, session) in locked.sessions.iter_mut() { + if !session.persistent + || restored_screen_is_pending_reattach(session) + || 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>, 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>) { + // 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 = 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 = 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>, + 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> { + 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); + } + + #[test] + fn attach_snapshot_prefers_restored_screen_for_same_size_reattach() { + let restored = b"\x1b[?1049lRESTORED_SCREEN_MARKER".to_vec(); + let mut session = Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: Some(RestoredScreenSnapshot { + cols: 80, + rows: 24, + snapshot: restored.clone(), + }), + clients: HashMap::new(), + output_seq: 7, + recent: VecDeque::new(), + empty_since: Some(Instant::now()), + holder_control: None, + persistent: true, + bytes_since_persist: 0, + last_persisted_seq: 7, + last_screen_persist_at: Instant::now(), + }; + + session.parser.set_size(24, 80); + assert_eq!(attach_snapshot(&session, 80, 24), restored); + assert_ne!(attach_snapshot(&session, 100, 30), restored); + } + + #[test] + fn first_reattach_prefers_restored_screen_even_after_pre_attach_pty_noise() { + let restored = b"\x1b[?1049lRESTORED_AFTER_RESTART".to_vec(); + let mut session = Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: Some(RestoredScreenSnapshot { + cols: 80, + rows: 24, + snapshot: restored.clone(), + }), + clients: HashMap::new(), + output_seq: 8, + recent: VecDeque::new(), + empty_since: Some(Instant::now()), + holder_control: None, + persistent: true, + bytes_since_persist: 0, + last_persisted_seq: 7, + last_screen_persist_at: Instant::now(), + }; + session.parser.process(b"$ "); + + assert_eq!(attach_snapshot(&session, 80, 24), restored); + + session + .clients + .insert([1u8; 16], test_client_state([2u8; 32])); + assert_eq!(attach_snapshot(&session, 80, 24), restored); + + session.restored_screen = None; + assert_ne!(attach_snapshot(&session, 80, 24), restored); + } + + #[tokio::test] + async fn no_client_pty_output_before_first_reattach_keeps_restored_screen() { + let restored = b"\x1b[?1049lRESTORED_AFTER_RESTART".to_vec(); + let (pty_tx, _rx) = mpsc::unbounded_channel(); + let state = Arc::new(Mutex::new(ServerState::new( + ServerConfig { + persist_sessions: true, + ..ServerConfig::default() + }, + [0u8; 32], + pty_tx, + ))); + { + let mut locked = state.lock().unwrap(); + locked.sessions.insert( + "work".to_string(), + Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: Some(RestoredScreenSnapshot { + cols: 80, + rows: 24, + snapshot: restored.clone(), + }), + clients: HashMap::new(), + output_seq: 7, + recent: VecDeque::new(), + empty_since: Some(Instant::now()), + holder_control: None, + persistent: true, + bytes_since_persist: 0, + last_persisted_seq: 7, + last_screen_persist_at: Instant::now(), + }, + ); + } + + let socket = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap()); + broadcast_output( + &state, + &socket, + PtyOutput { + session: "work".to_string(), + bytes: b"$ ".to_vec(), + exited: false, + }, + ) + .await + .unwrap(); + + let locked = state.lock().unwrap(); + let session = locked.sessions.get("work").unwrap(); + assert!(session.restored_screen.is_some()); + assert_eq!(attach_snapshot(session, 80, 24), restored); + assert_eq!( + session.last_persisted_seq, 7, + "pre-reattach PTY output must not overwrite the persisted restored screen" + ); + drop(locked); + + flush_persistent_screens(&state); + + let locked = state.lock().unwrap(); + let session = locked.sessions.get("work").unwrap(); + assert_eq!( + session.last_persisted_seq, 7, + "periodic flush must not replace restored screen before client-visible output" + ); + } + + #[test] + fn line_oriented_shell_output_persists_screen_immediately() { + assert!(should_persist_screen_immediately( + b"IMPLICIT_SCREEN_MARKER\r\n$ " + )); + assert!(!should_persist_screen_immediately(b"\x1b[2K\x1b[H")); + assert!(!should_persist_screen_immediately(&vec![b'x'; 2048])); + } + + #[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_local_eof_sent: HashSet::new(), + stream_remote_eof_received: HashSet::new(), + stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), + stream_window_adjust_retransmit: HashMap::new(), + stream_retired: HashSet::new(), + stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, + 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])])); + client.stream_window_adjust_retransmit.insert( + stream_id, + PendingWindowAdjust { + received_offset: 77, + bytes: 5, + last_sent: Instant::now(), + attempts: 1, + }, + ); + client.stream_window_adjust_retransmit.insert( + other_stream_id, + PendingWindowAdjust { + received_offset: 88, + bytes: 6, + last_sent: Instant::now(), + attempts: 1, + }, + ); + + 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 + .stream_window_adjust_retransmit + .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)); + assert!( + client + .stream_window_adjust_retransmit + .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)); + assert!( + state.sessions["work"].empty_since.is_none(), + "removing one of multiple clients must not mark the session empty" + ); + assert!(state.remove_client_everywhere(&[7u8; 16])); + assert!( + state.sessions["work"].empty_since.is_some(), + "removing the final client should start empty-session cleanup grace immediately" + ); + // 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), + restored_screen: None, + 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), + restored_screen: None, + 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), + restored_screen: None, + 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); + } + + #[tokio::test] + async fn server_stream_retransmit_uses_observed_rtt() { + 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 = 124; + 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(); + dosh::transport::observe_retransmit_rtt( + &mut client.stream_retransmit_srtt, + Duration::from_millis(25), + ); + client.stream_pending_opens.insert( + stream_id, + PendingStreamOpen { + target_host: "127.0.0.1".to_string(), + target_port: 8081, + last_sent: Instant::now() - Duration::from_millis(50), + attempts: 1, + }, + ); + state.sessions.insert( + "test".to_string(), + Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, + 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, 8081); + } + + #[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), + restored_screen: None, + 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_local_eof_sent: HashSet::new(), + stream_remote_eof_received: HashSet::new(), + stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), + stream_window_adjust_retransmit: HashMap::new(), + stream_retired: HashSet::new(), + stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, + 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 stream_eof_to_client_does_not_retire_stream() { + 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 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), + restored_screen: None, + 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::from([42]), + stream_send_credit: HashMap::from([(42, STREAM_INITIAL_WINDOW)]), + 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_local_eof_sent: HashSet::new(), + stream_remote_eof_received: HashSet::new(), + stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), + stream_window_adjust_retransmit: HashMap::new(), + stream_retired: HashSet::new(), + stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, + 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)); + + send_stream_eof_to_client(&state, &sender, client_id, 42) + .await + .unwrap(); + + { + let locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert!(client.opened_streams.contains(&42)); + assert!(client.stream_local_eof_sent.contains(&42)); + assert!(!client.stream_retired.contains(&42)); + } + + 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::StreamEof); + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); + let eof: StreamEof = protocol::from_body(&plain).unwrap(); + assert_eq!(eof.stream_id, 42); + } + + #[tokio::test] + async fn pending_server_stream_eof_is_retransmitted() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [17u8; 16]; + let session_key = [18u8; 32]; + let stream_id = 125; + 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); + client.stream_local_eof_sent.insert(stream_id); + client.stream_local_eof_sent.insert(999); + client.stream_eof_retransmit.insert( + stream_id, + PendingStreamControl { + 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), + restored_screen: None, + 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::StreamEof); + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); + let eof: StreamEof = protocol::from_body(&plain).unwrap(); + assert_eq!(eof.stream_id, stream_id); + + let locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert_eq!(client.stream_eof_retransmit[&stream_id].attempts, 2); + assert!(!client.stream_eof_retransmit.contains_key(&999)); + } + + #[tokio::test] + async fn pending_server_stream_close_is_retransmitted() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [19u8; 16]; + let session_key = [20u8; 32]; + let stream_id = 126; + 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_close_retransmit.insert( + stream_id, + PendingStreamControl { + 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), + restored_screen: None, + 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::StreamClose); + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); + let close: StreamClose = protocol::from_body(&plain).unwrap(); + assert_eq!(close.stream_id, stream_id); + + let locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert_eq!(client.stream_close_retransmit[&stream_id].attempts, 2); + } + + #[tokio::test] + async fn pending_server_stream_close_expires_after_attempt_cap() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [21u8; 16]; + let session_key = [22u8; 32]; + let stream_id = 127; + 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_close_retransmit.insert( + stream_id, + PendingStreamControl { + last_sent: Instant::now() - Duration::from_millis(250), + attempts: STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS, + }, + ); + state.sessions.insert( + "test".to_string(), + Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, + 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 locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert!(!client.stream_close_retransmit.contains_key(&stream_id)); + } + + #[tokio::test] + async fn pending_server_stream_window_adjust_is_retransmitted() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [23u8; 16]; + let session_key = [24u8; 32]; + let stream_id = 128; + 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_window_adjust_retransmit.insert( + stream_id, + PendingWindowAdjust { + received_offset: 4567, + bytes: 123, + 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), + restored_screen: None, + 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::StreamWindowAdjust); + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap(); + let adjust: StreamWindowAdjust = protocol::from_body(&plain).unwrap(); + assert_eq!(adjust.stream_id, stream_id); + assert_eq!(adjust.received_offset, 4567); + assert_eq!(adjust.bytes, 123); + + let locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert_eq!( + client.stream_window_adjust_retransmit[&stream_id].attempts, + 2 + ); + } + + #[tokio::test] + async fn pending_server_stream_window_adjust_expires_after_attempt_cap() { + let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); + let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx); + let client_id = [25u8; 16]; + let session_key = [26u8; 32]; + let stream_id = 129; + 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_window_adjust_retransmit.insert( + stream_id, + PendingWindowAdjust { + received_offset: 4567, + bytes: 123, + last_sent: Instant::now() - Duration::from_millis(250), + attempts: STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS, + }, + ); + state.sessions.insert( + "test".to_string(), + Session { + pty: None, + parser: vt100::Parser::new(24, 80, 100), + restored_screen: None, + 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 locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert!( + !client + .stream_window_adjust_retransmit + .contains_key(&stream_id) + ); + } + + #[tokio::test] + async fn server_stream_send_splits_large_writes() { + 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 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), + restored_screen: None, + 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::from([42]), + stream_send_credit: HashMap::from([(42, STREAM_INITIAL_WINDOW)]), + 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_local_eof_sent: HashSet::new(), + stream_remote_eof_received: HashSet::new(), + stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), + stream_window_adjust_retransmit: HashMap::new(), + stream_retired: HashSet::new(), + stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, + 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)); + let bytes = vec![3; dosh::transport::MAX_STREAM_DATA_BYTES * 2 + 17]; + + queue_or_send_stream_data_to_client(&state, &sender, client_id, 42, bytes) + .await + .unwrap(); + + { + let locked = state.lock().unwrap(); + let client = locked.sessions["test"].clients.get(&client_id).unwrap(); + assert!(client.stream_pending_data.is_empty()); + assert_eq!(client.stream_sent_data[&42].len(), 3); + assert_eq!( + client.stream_next_send_offset[&42], + (dosh::transport::MAX_STREAM_DATA_BYTES * 2 + 17) as u64 + ); + } + + let mut seen = Vec::new(); + let mut buf = vec![0u8; u16::MAX as usize]; + for _ in 0..3 { + 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(); + seen.push((data.offset, data.bytes.len())); + } + assert_eq!( + seen, + vec![ + (0, dosh::transport::MAX_STREAM_DATA_BYTES), + ( + dosh::transport::MAX_STREAM_DATA_BYTES as u64, + dosh::transport::MAX_STREAM_DATA_BYTES + ), + ((dosh::transport::MAX_STREAM_DATA_BYTES * 2) as u64, 17) + ] + ); + } + + #[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), + restored_screen: None, + 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_local_eof_sent: HashSet::new(), + stream_remote_eof_received: HashSet::new(), + stream_eof_retransmit: HashMap::new(), + stream_close_retransmit: HashMap::new(), + stream_window_adjust_retransmit: HashMap::new(), + stream_retired: HashSet::new(), + stream_retired_order: VecDeque::new(), + stream_retransmit_srtt: None, + 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" + ); + } + } +} diff --git a/tests/hostile_network.rs b/tests/hostile_network.rs index 3799dcc..7bb8565 100644 --- a/tests/hostile_network.rs +++ b/tests/hostile_network.rs @@ -18,6 +18,7 @@ //! Determinism: the relay's drop/reorder/dup behavior is driven by a fixed-seed //! PRNG and by explicit one-shot toggles, never by wall-clock timing, so the //! tests are reproducible and fast. +#![cfg(unix)] #![allow( clippy::collapsible_if, clippy::explicit_counter_loop, diff --git a/tests/integration_smoke.rs b/tests/integration_smoke.rs index 157e3bd..31154c9 100644 --- a/tests/integration_smoke.rs +++ b/tests/integration_smoke.rs @@ -1,3 +1,4 @@ +#![cfg(unix)] #![allow(clippy::collapsible_if)] use std::fs;