Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 51 Wiki Activity
Files
41256b66b717e680c4b218a745d34d0b6df30d87
dosh/src/bin/dosh-server.rs
T
DuProcessandClaude Opus 4.8 8b1af51bc6 Persist sessions across server restarts via per-session holders 
Make dosh terminal sessions survive a dosh-server restart
(crash/upgrade/systemctl restart): a reattaching client lands on the
SAME shell with screen state intact, instead of a fresh one.

Design: when persist_sessions is on (new config, default true) each
terminal session's shell runs in a small detached per-session holder
process (the dosh-server binary re-exec'd as `dosh-server hold`). The
holder double-forks + setsid into its own session, opens the PTY, spawns
the shell as ITS child, and serves the PTY master fd over a per-session
Unix socket via SCM_RIGHTS. The server adopts that fd to build a PtyHandle
that DETACHES (does not kill the shell) on drop, so a server exit leaves
holder + shell alive. On startup the server scans the runtime dir under
sessions_dir/run, reconnects to each live holder, receives the master fd
again, restores the persisted vt100 screen, and rebuilds the Session.

Screen/scrollback (server-memory only) is mirrored to disk atomically:
byte-throttled on the output hot path plus a 2s periodic flush, and
restored on re-adoption so reattach repaints. Truly-abandoned persistent
sessions are still reaped per the existing grace logic by asking the
holder to shut down. Anything in the persistent path failing degrades
gracefully to the original in-process shell.

PtyHandle gains an Owned (kill-on-drop, unchanged default) vs Adopted
(detach-on-drop) backing; resize on an adopted master uses TIOCSWINSZ.

No wire-format change, so protocol::VERSION stays 3. New deps: libc
(direct). New config: persist_sessions (default true); existing
integration tests pin it false to keep exercising the non-persistent
path unchanged.

Adds tests/integration_smoke.rs::session_survives_server_restart_same_
shell_and_screen: attaches, sets `cd /tmp; export MARK=...`, paints a
screen marker, KILLs the server, restarts it on the same config, reattaches
and asserts same shell pid + MARK + PWD + restored screen. Plus persist.rs
unit tests (name round-trip, screen save/load, dead-holder scan cleanup,
SCM_RIGHTS fd round-trip). cargo fmt --check and cargo test green.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-14 14:19:09 -04:00

3241 lines
118 KiB
Rust
Raw Blame History

use anyhow::{Context, Result, anyhow};
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::native::{
EnvVar, ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello,
derive_native_session_key, generate_native_ephemeral, host_public_key, host_public_key_line,
load_or_create_host_key, sign_server_hello, verify_native_user_auth_from_config,
};
use dosh::persist;
use dosh::protocol::{
self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody,
NativeUserAuthBody, PacketKind, ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT,
StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
};
use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session};
use std::collections::{HashMap, HashSet, VecDeque};
use std::net::{IpAddr, SocketAddr};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::UnixStream as StdUnixStream;
use std::path::PathBuf;
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::sync::mpsc;
const STREAM_INITIAL_WINDOW: usize = 1024 * 1024;
/// Persist a persistent session's screen to disk after this many bytes of output
/// have accumulated since the last mirror. Keeps the atomic write off the
/// per-packet hot path while bounding how much fresh output a crash can lose.
const SCREEN_PERSIST_BYTES: usize = 4096;
/// Sentinel `target_host` sent to the client on a server-initiated `StreamOpen`
/// that carries an SSH-agent connection (the client splices it into its local
/// `SSH_AUTH_SOCK` rather than dialing TCP). Must match the client's constant.
const AGENT_STREAM_SENTINEL: &str = "@dosh-agent";
/// Monotonic counter for unique agent proxy socket filenames within this process.
static AGENT_SOCK_COUNTER: AtomicU64 = AtomicU64::new(0);
/// How long a pending native handshake (ClientHello awaiting UserAuth) is kept
/// before it is swept by the cleanup task.
const NATIVE_HANDSHAKE_TTL_SECS: u64 = 30;
#[derive(Debug, Parser)]
#[command(name = "dosh-server")]
struct Args {
#[command(subcommand)]
command: Command,
}
#[derive(Debug, Subcommand)]
enum Command {
Serve {
#[arg(long)]
config: Option<std::path::PathBuf>,
},
Auth {
#[arg(long, default_value_t = 1)]
protocol: u8,
#[arg(long)]
nonce: Option<String>,
#[arg(long)]
host_key: bool,
#[arg(long, default_value = "default")]
session: String,
#[arg(long, default_value = "read-write")]
mode: String,
#[arg(long, default_value = "80x24")]
size: String,
#[arg(long, default_value = "dev")]
client_version: String,
#[arg(long)]
udp_host: Option<String>,
},
/// Internal: run as a per-session shell holder so the shell survives a
/// server restart. Spawned by the server itself; not meant to be run by
/// hand. Daemonizes, owns the PTY, and serves the master fd over a socket.
#[command(hide = true)]
Hold {
#[arg(long)]
runtime_dir: PathBuf,
#[arg(long)]
session: String,
#[arg(long)]
shell: String,
#[arg(long)]
cols: u16,
#[arg(long)]
rows: u16,
/// `NAME=VALUE` pairs for the shell environment. Repeatable.
#[arg(long = "env")]
env: Vec<String>,
},
}
fn main() -> Result<()> {
let args = Args::parse();
// The holder is handled BEFORE any tokio runtime exists: it forks/setsids to
// daemonize, and forking out of a live async runtime is best avoided. It also
// never uses tokio, so it has no reason to pay for one.
if let Command::Hold {
runtime_dir,
session,
shell,
cols,
rows,
env,
} = &args.command
{
let env = parse_env_pairs(env)?;
return persist::run_holder(runtime_dir, session, shell, *cols, *rows, &env);
}
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.context("build tokio runtime")?;
runtime.block_on(async move { run_command(args.command).await })
}
async fn run_command(command: Command) -> Result<()> {
match command {
Command::Serve { config } => serve(config).await,
Command::Hold { .. } => unreachable!("Hold is handled before the runtime starts"),
Command::Auth {
protocol,
nonce,
host_key,
session,
mode,
size,
client_version: _,
udp_host,
} => {
anyhow::ensure!(protocol == 1, "unsupported protocol {protocol}");
let config = load_server_config(None)?;
if host_key {
let signing_key = load_or_create_host_key(&config)?;
println!("{}", host_public_key_line(&host_public_key(&signing_key)));
return Ok(());
}
let secret = load_or_create_server_secret(&config)?;
let nonce = parse_nonce(nonce.as_deref().context("--nonce is required")?)?;
let size = parse_size(&size)?;
let user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
let udp_host = udp_host.unwrap_or_else(|| "127.0.0.1".to_string());
let resp =
build_bootstrap(&config, &secret, user, session, mode, size, nonce, udp_host)?;
println!("{}", encode_bootstrap(&resp)?);
Ok(())
}
}
}
async fn serve(config_path: Option<std::path::PathBuf>) -> Result<()> {
let config = load_server_config(config_path)?;
let secret = load_or_create_server_secret(&config)?;
let bind = format!("{}:{}", config.bind, config.port);
let socket = Arc::new(
UdpSocket::bind(&bind)
.await
.with_context(|| format!("bind {bind}"))?,
);
eprintln!("dosh-server listening on {bind}");
let (pty_tx, mut pty_rx) = mpsc::unbounded_channel();
let state = Arc::new(Mutex::new(ServerState::new(
config.clone(),
secret,
pty_tx.clone(),
)));
{
let mut locked = state.lock().expect("server state poisoned");
// Re-adopt any persistent sessions whose holders survived a previous
// server's exit, BEFORE prewarming, so a prewarmed name reattaches to
// its existing shell instead of spawning a duplicate.
if config.persist_sessions {
locked.readopt_persistent_sessions();
}
for session in config.prewarm_sessions.clone() {
locked.ensure_session(&session, 80, 24, "read-write", &[])?;
}
}
let output_state = Arc::clone(&state);
let output_socket = Arc::clone(&socket);
tokio::spawn(async move {
while let Some(output) = pty_rx.recv().await {
if let Err(err) = broadcast_output(&output_state, &output_socket, output).await {
eprintln!("broadcast error: {err:#}");
}
}
});
let retransmit_state = Arc::clone(&state);
let retransmit_socket = Arc::clone(&socket);
tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_millis(100));
loop {
interval.tick().await;
if let Err(err) = retransmit_pending(&retransmit_state, &retransmit_socket).await {
eprintln!("retransmit error: {err:#}");
}
if let Err(err) = maybe_rekey_clients(&retransmit_state, &retransmit_socket).await {
eprintln!("rekey error: {err:#}");
}
}
});
let cleanup_state = Arc::clone(&state);
tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(5));
loop {
interval.tick().await;
cleanup_disconnected_clients(&cleanup_state);
}
});
// Periodically mirror persistent sessions' screens to disk, bounding how much
// recent output a sudden crash can lose regardless of throughput (the
// per-packet path only persists every few KB). Cheap: one atomic write per
// persistent session per tick, and only when its screen changed.
if config.persist_sessions {
let flush_state = Arc::clone(&state);
tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(2));
loop {
interval.tick().await;
flush_persistent_screens(&flush_state);
}
});
}
let mut buf = vec![0u8; 65535];
loop {
let (n, peer) = socket.recv_from(&mut buf).await?;
if let Err(err) = handle_packet(&state, &socket, peer, &buf[..n]).await {
eprintln!("packet from {peer}: {err:#}");
}
}
}
struct ServerState {
config: ServerConfig,
secret: [u8; 32],
pty_tx: mpsc::UnboundedSender<PtyOutput>,
sessions: HashMap<String, Session>,
pending_native: HashMap<[u8; 16], PendingNativeAuth>,
next_server_stream_id: u64,
/// O(1) reverse index from a live client's `conn_id` to its session name, so
/// per-packet handlers don't linear-scan every session's client map. Kept in
/// lockstep with `Session::clients` on every insert and remove.
client_index: HashMap<[u8; 16], String>,
/// Per-source-IP token bucket for native ClientHello flood protection,
/// checked before any X25519/Ed25519 work.
native_auth_limiter: NativeAuthRateLimiter,
}
/// Per-source-IP token bucket throttling native auth ClientHellos.
///
/// Threat model item: a malicious unauthenticated client flooding auth attempts.
/// We charge one token per ClientHello *before* doing any X25519/Ed25519 crypto,
/// so a flood cannot burn server CPU. The bucket refills continuously at
/// `per_minute / 60` tokens per second up to a burst of `per_minute`, matching
/// the `native_auth_rate_limit_per_minute` value advertised in `ServerHello`.
struct NativeAuthRateLimiter {
per_minute: u32,
buckets: HashMap<IpAddr, TokenBucket>,
}
#[derive(Clone, Copy)]
struct TokenBucket {
tokens: f64,
last_refill: Instant,
}
impl NativeAuthRateLimiter {
fn new(per_minute: u32) -> Self {
Self {
per_minute,
buckets: HashMap::new(),
}
}
/// Try to spend one token for `ip` at `now`. Returns `Ok(remaining)` whole
/// tokens left when allowed, or `Err(())` when the bucket is empty. A
/// `per_minute` of 0 disables native auth entirely (no tokens ever).
fn check(&mut self, ip: IpAddr, now: Instant) -> std::result::Result<u32, ()> {
if self.per_minute == 0 {
return Err(());
}
let capacity = self.per_minute as f64;
let refill_per_sec = capacity / 60.0;
let bucket = self.buckets.entry(ip).or_insert(TokenBucket {
tokens: capacity,
last_refill: now,
});
let elapsed = now
.saturating_duration_since(bucket.last_refill)
.as_secs_f64();
bucket.tokens = (bucket.tokens + elapsed * refill_per_sec).min(capacity);
bucket.last_refill = now;
if bucket.tokens < 1.0 {
return Err(());
}
bucket.tokens -= 1.0;
Ok(bucket.tokens as u32)
}
/// Drop buckets that have fully refilled, so the map can't grow unbounded
/// under a spoofed-source-IP flood. Called from the periodic cleanup task.
fn evict_full(&mut self, now: Instant) {
if self.per_minute == 0 {
self.buckets.clear();
return;
}
let capacity = self.per_minute as f64;
let refill_per_sec = capacity / 60.0;
self.buckets.retain(|_, bucket| {
let elapsed = now
.saturating_duration_since(bucket.last_refill)
.as_secs_f64();
(bucket.tokens + elapsed * refill_per_sec) < capacity
});
}
}
struct Session {
pty: Option<PtyHandle>,
parser: vt100::Parser,
clients: HashMap<[u8; 16], ClientState>,
output_seq: u64,
recent: VecDeque<Vec<u8>>,
/// When the session last dropped to zero clients, used to reap abandoned
/// sessions (and their shells) after a grace period. `None` while at least
/// one client is attached.
empty_since: Option<Instant>,
/// Open control connection to this session's holder process, when the shell
/// runs persistently out-of-process. Dropping it (server exit) leaves the
/// holder running; sending a shutdown byte over it reaps the holder. `None`
/// for non-persistent (in-process) sessions.
holder_control: Option<StdUnixStream>,
/// Whether this session's shell lives in a holder process (persistent).
persistent: bool,
/// Bytes of session output since the screen was last mirrored to disk, used
/// to throttle the (atomic) screen-persistence writes.
bytes_since_persist: usize,
/// `output_seq` at the last successful screen mirror, so the periodic flush
/// skips sessions whose screen has not changed.
last_persisted_seq: u64,
}
#[derive(Clone)]
struct ClientState {
endpoint: SocketAddr,
mode: String,
session_key: [u8; 32],
last_acked: u64,
replay: ReplayWindow,
send_seq: u64,
cols: u16,
rows: u16,
last_seen: Instant,
pending: VecDeque<PendingFrame>,
allowed_forwardings: Vec<ForwardingRequest>,
stream_writers: HashMap<u64, mpsc::Sender<Vec<u8>>>,
opened_streams: HashSet<u64>,
stream_send_credit: HashMap<u64, usize>,
stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
/// Current transport key epoch (0 = original handshake key). Bumped on rekey.
epoch: u64,
/// When the current epoch began, for the wall-clock rekey trigger.
epoch_started: Instant,
/// Packets seen/sent in the current epoch, for the packet-count trigger.
epoch_packets: u64,
/// The previous epoch's key, retained briefly so in-flight pre-rekey packets
/// still decrypt (matched via `session_key_id`) instead of dropping fatally.
previous_session_key: Option<[u8; 32]>,
}
#[derive(Clone)]
struct PendingFrame {
output_seq: u64,
packet: Vec<u8>,
last_sent: Instant,
attempts: u8,
}
struct PendingNativeAuth {
client: dosh::native::NativeClientHello,
server: NativeServerHello,
session_key: [u8; 32],
peer: SocketAddr,
created_at: Instant,
}
impl ServerState {
fn new(
config: ServerConfig,
secret: [u8; 32],
pty_tx: mpsc::UnboundedSender<PtyOutput>,
) -> Self {
let per_minute = config.native_auth_rate_limit_per_minute;
Self {
config,
secret,
pty_tx,
sessions: HashMap::new(),
pending_native: HashMap::new(),
next_server_stream_id: 1u64 << 63,
client_index: HashMap::new(),
native_auth_limiter: NativeAuthRateLimiter::new(per_minute),
}
}
fn ensure_session(
&mut self,
name: &str,
cols: u16,
rows: u16,
mode: &str,
requested_env: &[EnvVar],
) -> Result<()> {
let env = accepted_env(&self.config, requested_env)?;
// Whether this existing session needs a (re)spawned PTY. Computed without
// holding a mutable borrow of `self.sessions` so we can call the
// `&self`-borrowing spawn helper.
let needs_pty = self
.sessions
.get(name)
.map(|session| session.pty.is_none() && mode_uses_pty(mode));
if let Some(needs_pty) = needs_pty {
if needs_pty {
let (pty, control, persistent) = self.spawn_session_pty(name, cols, rows, &env)?;
if let Some(session) = self.sessions.get_mut(name) {
session.pty = Some(pty);
session.holder_control = control;
session.persistent = persistent;
}
}
return Ok(());
}
let (pty, control, persistent) = if mode_uses_pty(mode) {
let (pty, control, persistent) = self.spawn_session_pty(name, cols, rows, &env)?;
(Some(pty), control, persistent)
} else {
(None, None, false)
};
self.sessions.insert(
name.to_string(),
Session {
pty,
parser: vt100::Parser::new(rows.max(1), cols.max(1), self.config.scrollback),
clients: HashMap::new(),
output_seq: 0,
recent: VecDeque::with_capacity(self.config.scrollback),
empty_since: None,
holder_control: control,
persistent,
bytes_since_persist: 0,
last_persisted_seq: 0,
},
);
Ok(())
}
/// Spawn (or, in the persistent case, spawn-and-adopt) a PTY for a session.
///
/// With `persist_sessions` on, the shell is launched in a detached holder
/// process whose PTY master fd is passed back via SCM_RIGHTS, so the shell
/// outlives this server. If anything in that path fails we degrade
/// gracefully to the original in-process spawn (a shell that dies with the
/// server), so persistence never makes a session un-attachable.
fn spawn_session_pty(
&self,
name: &str,
cols: u16,
rows: u16,
env: &[(String, String)],
) -> Result<(PtyHandle, Option<StdUnixStream>, bool)> {
let cols = cols.max(1);
let rows = rows.max(1);
if self.config.persist_sessions {
match self.spawn_persistent_pty(name, cols, rows, env) {
Ok(pair) => return Ok((pair.0, Some(pair.1), true)),
Err(err) => {
eprintln!(
"persistence unavailable for session {name}, using in-process shell: {err:#}"
);
}
}
}
let pty = spawn_pty_session(
name.to_string(),
&self.config.shell,
cols,
rows,
env,
self.pty_tx.clone(),
)?;
Ok((pty, None, false))
}
/// Spawn a fresh holder for `name` and adopt its PTY master fd.
fn spawn_persistent_pty(
&self,
name: &str,
cols: u16,
rows: u16,
env: &[(String, String)],
) -> Result<(PtyHandle, StdUnixStream)> {
let sessions_dir = self.sessions_dir();
persist::spawn_holder(&sessions_dir, name, &self.config.shell, cols, rows, env)?;
self.adopt_persistent_pty(name)
}
/// Connect to an existing holder for `name` and build an adopted PtyHandle
/// from the master fd it hands back. Used both right after spawning a holder
/// and when re-adopting a holder left behind by a previous server.
fn adopt_persistent_pty(&self, name: &str) -> Result<(PtyHandle, StdUnixStream)> {
let sessions_dir = self.sessions_dir();
let (fd, control) = persist::adopt_holder(&sessions_dir, name)?;
let pty = match adopt_pty_from_fd(name.to_string(), fd, self.pty_tx.clone()) {
Ok(pty) => pty,
Err(err) => {
// We own `fd`; close it so the failed adopt doesn't leak it.
drop(persist::fd_into_file(fd));
return Err(err);
}
};
Ok((pty, control))
}
fn sessions_dir(&self) -> PathBuf {
expand_tilde(&self.config.sessions_dir)
}
/// On startup, re-adopt persistent sessions whose holder processes outlived a
/// previous server. For each live holder we reconnect, receive the PTY master
/// fd, restore the persisted screen/scrollback, and rebuild the in-memory
/// `Session` so a reattaching client lands on the same shell with its screen
/// intact. A holder we cannot adopt is left for the cleanup pass / fresh
/// re-create, never crashing startup.
fn readopt_persistent_sessions(&mut self) {
let sessions_dir = self.sessions_dir();
let holders = persist::scan_existing_holders(&sessions_dir);
for meta in holders {
let name = meta.session.clone();
if self.sessions.contains_key(&name) {
continue;
}
let (pty, control) = match self.adopt_persistent_pty(&name) {
Ok(pair) => pair,
Err(err) => {
eprintln!("could not re-adopt session {name}: {err:#}");
persist::remove_runtime_dir(&sessions_dir, &name);
continue;
}
};
// Restore screen + scrollback geometry/seq if we saved one. The
// restored snapshot is replayed into a fresh parser so a reattaching
// client (and `screen_snapshot`) sees the pre-restart screen.
let saved = persist::load_screen(&sessions_dir, &name);
let (cols, rows) = saved
.as_ref()
.map(|s| (s.cols.max(1), s.rows.max(1)))
.unwrap_or((80, 24));
let mut parser = vt100::Parser::new(rows, cols, self.config.scrollback);
let mut output_seq = 0;
if let Some(saved) = saved {
parser.process(&saved.snapshot);
output_seq = saved.output_seq;
}
self.sessions.insert(
name.clone(),
Session {
pty: Some(pty),
parser,
clients: HashMap::new(),
output_seq,
recent: VecDeque::with_capacity(self.config.scrollback),
empty_since: Some(Instant::now()),
holder_control: Some(control),
persistent: true,
bytes_since_persist: 0,
last_persisted_seq: output_seq,
},
);
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);
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) {
if let Some(session) = self.sessions.get_mut(&session_name) {
return session.clients.remove(client_id).is_some();
}
}
false
}
/// O(1) resolve of a live client's session key and session name via the
/// index, replacing the former O(sessions) linear scan.
fn lookup_client(&self, client_id: &[u8; 16]) -> Option<([u8; 32], String)> {
let session_name = self.client_index.get(client_id)?;
let session = self.sessions.get(session_name)?;
let client = session.clients.get(client_id)?;
Some((client.session_key, session_name.clone()))
}
/// Resolve the decryption key for an inbound packet, honoring rekey grace.
///
/// During the brief window after a rekey, in-flight packets sealed under the
/// previous epoch's key must still decrypt instead of being dropped as fatal
/// (spec: "stale encrypted packets ... ignored, not fatal"). We pick whichever
/// of the current or retained-previous key matches the packet's
/// `session_key_id`, falling back to the current key when neither matches (so
/// a genuinely stale/forged packet still fails the AEAD check, not silently).
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()));
}
Some((client.session_key, session_name.clone()))
}
/// O(1) mutable access to a live client via the conn_id index.
fn client_mut(&mut self, client_id: &[u8; 16]) -> Option<&mut ClientState> {
let session_name = self.client_index.get(client_id)?;
self.sessions
.get_mut(session_name)?
.clients
.get_mut(client_id)
}
}
fn mode_uses_pty(mode: &str) -> bool {
mode != "forward-only"
}
/// Resize the PTY and terminal parser for a session.
///
/// `cols`/`rows` come straight off the wire from the client, where a terminal
/// that has not been sized yet (or a buggy/hostile client) can report `0`.
/// `vt100::Parser::set_size(0, _)` panics with a subtract overflow and would
/// take down the whole single-threaded daemon, so clamp to at least `1` here,
/// the same way `Parser::new` is constructed in `ensure_session`.
fn apply_terminal_size(
pty: Option<&PtyHandle>,
parser: &mut vt100::Parser,
cols: u16,
rows: u16,
) -> Result<()> {
let cols = cols.max(1);
let rows = rows.max(1);
pty.context("terminal session has no pty")?
.resize(cols, rows)?;
parser.set_size(rows, cols);
Ok(())
}
fn mode_allows_terminal_updates(mode: &str) -> bool {
mode != "view-only" && mode != "forward-only"
}
fn accepted_env(config: &ServerConfig, requested: &[EnvVar]) -> Result<Vec<(String, String)>> {
let mut env = Vec::new();
for entry in requested {
anyhow::ensure!(
valid_env_name(&entry.name),
"invalid environment variable name {:?}",
entry.name
);
anyhow::ensure!(
!entry.value.as_bytes().contains(&0),
"environment variable {:?} contains NUL",
entry.name
);
if config
.accept_env
.iter()
.any(|pattern| glob_matches(pattern, &entry.name))
{
env.push((entry.name.clone(), entry.value.clone()));
}
}
Ok(env)
}
fn valid_env_name(name: &str) -> bool {
let mut chars = name.chars();
let Some(first) = chars.next() else {
return false;
};
if !(first == '_' || first.is_ascii_alphabetic()) {
return false;
}
chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
}
fn glob_matches(pattern: &str, value: &str) -> bool {
let pattern = pattern.as_bytes();
let value = value.as_bytes();
let (mut p, mut v) = (0, 0);
let mut star = None;
let mut star_value = 0;
while v < value.len() {
if p < pattern.len() && (pattern[p] == b'?' || pattern[p] == value[v]) {
p += 1;
v += 1;
} else if p < pattern.len() && pattern[p] == b'*' {
star = Some(p);
p += 1;
star_value = v;
} else if let Some(star_pos) = star {
p = star_pos + 1;
star_value += 1;
v = star_value;
} else {
return false;
}
}
while p < pattern.len() && pattern[p] == b'*' {
p += 1;
}
p == pattern.len()
}
async fn handle_packet(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
raw: &[u8],
) -> Result<()> {
// A peer running a different wire protocol version cannot be decoded at all,
// because the header VERSION byte is part of the AEAD AAD and the framing.
// Rather than drop it (which the peer sees as a silent timeout), answer with
// a clear, named version-mismatch reject so the user is told to upgrade.
if let Some(foreign) = protocol::peek_foreign_wire_version(raw) {
let _ = foreign;
return send_reject(socket, peer, protocol::VERSION_MISMATCH_REASON).await;
}
let packet = protocol::decode(raw)?;
match packet.header.kind {
PacketKind::NativeClientHello => {
handle_native_client_hello(state, socket, peer, packet.body).await
}
PacketKind::NativeUserAuth => handle_native_user_auth(state, socket, peer, &packet).await,
PacketKind::BootstrapAttachRequest => {
handle_bootstrap_attach(state, socket, peer, packet.body).await
}
PacketKind::TicketAttachRequest => {
handle_ticket_attach(state, socket, peer, packet.body).await
}
PacketKind::ResumeRequest => handle_resume(state, socket, peer, &packet).await,
PacketKind::Input
| PacketKind::Resize
| PacketKind::Ping
| PacketKind::Ack
| PacketKind::StreamOpen
| PacketKind::StreamOpenOk
| PacketKind::StreamOpenReject
| PacketKind::StreamData
| PacketKind::StreamClose
| PacketKind::StreamEof
| PacketKind::StreamWindowAdjust
| PacketKind::RekeyAck
if find_client_key(state, &packet.header.conn_id).is_err() =>
{
send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client").await
}
PacketKind::Input => handle_input(state, peer, &packet).await,
PacketKind::Resize => handle_resize(state, peer, &packet).await,
PacketKind::Ping => handle_ping(state, socket, peer, &packet).await,
PacketKind::Ack => handle_ack(state, peer, &packet).await,
PacketKind::RekeyAck => handle_rekey_ack(state, peer, &packet).await,
PacketKind::Detach => handle_detach(state, &packet).await,
PacketKind::StreamOpen => handle_stream_open(state, socket, peer, &packet).await,
PacketKind::StreamOpenOk => handle_stream_open_ok(state, socket, peer, &packet).await,
PacketKind::StreamOpenReject => handle_stream_open_reject(state, peer, &packet).await,
PacketKind::StreamData => handle_stream_data(state, socket, peer, &packet).await,
PacketKind::StreamWindowAdjust => {
handle_stream_window_adjust(state, socket, peer, &packet).await
}
PacketKind::StreamClose => handle_stream_close(state, peer, &packet).await,
_ => Ok(()),
}
}
async fn handle_native_client_hello(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
body: Vec<u8>,
) -> Result<()> {
let req: NativeClientHelloBody = protocol::from_body(&body)?;
// Defense in depth against a peer that shares our wire VERSION but advertises
// a different native handshake protocol_version (the spec's negotiation
// point). Reject with a named, actionable reason before any crypto.
if let Err(err) =
dosh::native::check_native_protocol_version(req.hello.protocol_version, "client")
{
return send_reject(socket, peer, &err.to_string()).await;
}
// Rate-limit native auth *before* touching the host key or any X25519/Ed25519
// work, so a flood of ClientHellos from one source cannot burn server CPU.
// A separate, non-crypto reject lets a legitimate client back off cleanly.
let rate_limit_remaining = {
let mut locked = state.lock().expect("server state poisoned");
locked.native_auth_limiter.check(peer.ip(), Instant::now())
};
let rate_limit_remaining = match rate_limit_remaining {
Ok(remaining) => remaining,
Err(()) => {
return send_reject(socket, peer, "native auth rate limit exceeded").await;
}
};
let built: Result<([u8; 16], NativeServerHello)> = {
let mut locked = state.lock().expect("server state poisoned");
if !locked.config.native_auth {
Err(anyhow!("native auth disabled"))
} else if !req
.hello
.supported_aead
.iter()
.any(|algorithm| algorithm == "chacha20poly1305")
{
Err(anyhow!("native auth requires chacha20poly1305"))
} else if !req
.hello
.supported_user_key_algorithms
.iter()
.any(|algorithm| algorithm == "ssh-ed25519")
{
Err(anyhow!("native auth requires ssh-ed25519"))
} 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)?;
socket.send_to(&out, peer).await?;
Ok(())
}
async fn handle_native_user_auth(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let pending = {
let mut locked = state.lock().expect("server state poisoned");
locked.pending_native.remove(&packet.header.conn_id)
};
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) {
return send_reject_to_client(socket, peer, packet.header.conn_id, "native auth expired")
.await;
}
let body = protocol::decrypt_body(packet, &pending.session_key, CLIENT_TO_SERVER)?;
let req: NativeUserAuthBody = protocol::from_body(&body)?;
if pending.client.requested_mode == "doctor" {
let check = {
let locked = state.lock().expect("server state poisoned");
if let Err(err) = verify_native_user_auth_from_config(
&locked.config,
&pending.client,
&pending.server,
&req.auth,
Some(peer.ip()),
)
.context("verify native user auth")
{
return send_reject_to_client(
socket,
peer,
packet.header.conn_id,
&format!("{err:#}"),
)
.await;
}
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,
policy_flags: Vec::new(),
server_version: env!("CARGO_PKG_VERSION").to_string(),
}
};
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,
)?;
socket.send_to(&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(),
});
}
let attached: Result<(
[u8; 16],
[u8; 16],
Vec<u8>,
[u8; 32],
String,
String,
u64,
Vec<u8>,
Vec<ForwardingRequest>,
)> = {
let mut locked = state.lock().expect("server state poisoned");
verify_native_user_auth_from_config(
&locked.config,
&pending.client,
&pending.server,
&req.auth,
Some(peer.ip()),
)
.context("verify native user auth")
.and_then(|_| {
let session_name = pending.client.requested_session.clone();
let mode = pending.client.requested_mode.clone();
let cols = pending.client.terminal_size.0;
let rows = pending.client.terminal_size.1;
if !locked.sessions.contains_key(&session_name) && !locked.config.create_on_attach {
return Err(anyhow!("session does not exist"));
}
locked.ensure_session(&session_name, cols, rows, &mode, &session_env)?;
let session_key = pending.session_key;
let key_id = protocol::session_key_id(&session_key);
let attach_ticket_psk = crypto::random_32();
let issued_at = now_secs()?;
let attach_ticket = build_attach_ticket(
&locked.secret,
crypto::sha256(&locked.secret),
1,
pending.client.requested_user.clone(),
session_name.clone(),
mode.clone(),
issued_at,
issued_at + locked.config.attach_ticket_ttl_secs,
&attach_ticket_psk,
)?;
let session = locked
.sessions
.get_mut(&session_name)
.expect("session exists");
if mode_allows_terminal_updates(&mode) {
apply_terminal_size(session.pty.as_ref(), &mut session.parser, cols, rows)?;
}
let client_id = crypto::random_16();
let snapshot = screen_snapshot(session.parser.screen());
let output_seq = session.output_seq;
locked.insert_client(
&session_name,
client_id,
ClientState {
endpoint: peer,
mode: mode.clone(),
session_key,
last_acked: output_seq,
replay: ReplayWindow::default(),
send_seq: 1,
cols,
rows,
last_seen: Instant::now(),
pending: VecDeque::new(),
allowed_forwardings: req.auth.requested_forwardings.clone(),
stream_writers: HashMap::new(),
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
previous_session_key: None,
},
);
Ok((
client_id,
key_id,
attach_ticket,
attach_ticket_psk,
session_name,
mode,
output_seq,
snapshot,
req.auth.requested_forwardings.clone(),
))
})
};
let (
client_id,
key_id,
attach_ticket,
attach_ticket_psk,
session_name,
mode,
output_seq,
snapshot,
requested_forwardings,
) = 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;
}
};
if let Err(err) = start_remote_forwards(
state,
socket,
client_id,
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);
}
let ok = NativeAuthOk {
client_id,
session: session_name,
mode,
session_key: pending.session_key,
session_key_id: key_id,
attach_ticket,
attach_ticket_psk,
initial_seq: output_seq,
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,
)?;
socket.send_to(&out, peer).await?;
Ok(())
}
async fn handle_bootstrap_attach(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
body: Vec<u8>,
) -> Result<()> {
let req: BootstrapAttachRequest = protocol::from_body(&body)?;
let (client_id, key, key_id, session_name, mode, output_seq, snapshot) = {
let mut locked = state.lock().expect("server state poisoned");
if !verify_bootstrap(&req.bootstrap, &locked.secret)? {
return send_reject(socket, peer, "invalid or expired bootstrap").await;
}
if !locked.sessions.contains_key(&req.bootstrap.session) && !locked.config.create_on_attach
{
return send_reject(socket, peer, "session does not exist").await;
}
locked.ensure_session(
&req.bootstrap.session,
req.cols,
req.rows,
&req.bootstrap.mode,
&req.requested_env,
)?;
let session = locked
.sessions
.get_mut(&req.bootstrap.session)
.expect("session exists");
if mode_allows_terminal_updates(&req.bootstrap.mode) {
apply_terminal_size(
session.pty.as_ref(),
&mut session.parser,
req.cols,
req.rows,
)?;
}
let client_id = crypto::random_16();
let snapshot = screen_snapshot(session.parser.screen());
let output_seq = session.output_seq;
let bootstrap_session = req.bootstrap.session.clone();
locked.insert_client(
&bootstrap_session,
client_id,
ClientState {
endpoint: peer,
mode: req.bootstrap.mode.clone(),
session_key: req.bootstrap.session_key,
last_acked: output_seq,
replay: ReplayWindow::default(),
send_seq: 1,
cols: req.cols,
rows: req.rows,
last_seen: Instant::now(),
pending: VecDeque::new(),
allowed_forwardings: Vec::new(),
stream_writers: HashMap::new(),
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
previous_session_key: None,
},
);
(
client_id,
req.bootstrap.session_key,
req.bootstrap.session_key_id,
req.bootstrap.session.clone(),
req.bootstrap.mode.clone(),
output_seq,
snapshot,
)
};
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,
)?;
socket.send_to(&out, peer).await?;
Ok(())
}
async fn handle_ticket_attach(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
body: Vec<u8>,
) -> Result<()> {
let env: TicketAttachEnvelope = protocol::from_body(&body)?;
let (ticket, request_plain) = {
let locked = state.lock().expect("server state poisoned");
if !locked.config.allow_attach_tickets {
return send_reject(socket, peer, "attach tickets disabled").await;
}
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,
)?;
(ticket, request_plain)
};
let req: TicketAttachBody = protocol::from_body(&request_plain)?;
if req.session != ticket.session || req.mode != ticket.mode {
return send_reject(socket, peer, "ticket scope mismatch").await;
}
let session_key = crypto::random_32();
let session_key_id = protocol::session_key_id(&session_key);
let (client_id, output_seq, snapshot) = {
let mut locked = state.lock().expect("server state poisoned");
if !locked.sessions.contains_key(&req.session) && !locked.config.create_on_attach {
return send_reject(socket, peer, "session does not exist").await;
}
locked.ensure_session(
&req.session,
req.cols,
req.rows,
&req.mode,
&req.requested_env,
)?;
let session = locked
.sessions
.get_mut(&req.session)
.expect("session exists");
if mode_allows_terminal_updates(&req.mode) {
apply_terminal_size(
session.pty.as_ref(),
&mut session.parser,
req.cols,
req.rows,
)?;
}
let client_id = crypto::random_16();
let snapshot = screen_snapshot(session.parser.screen());
let output_seq = session.output_seq;
let ticket_session = req.session.clone();
locked.insert_client(
&ticket_session,
client_id,
ClientState {
endpoint: peer,
mode: req.mode.clone(),
session_key,
last_acked: output_seq,
replay: ReplayWindow::default(),
send_seq: 1,
cols: req.cols,
rows: req.rows,
last_seen: Instant::now(),
pending: VecDeque::new(),
allowed_forwardings: Vec::new(),
stream_writers: HashMap::new(),
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
previous_session_key: None,
},
);
(client_id, output_seq, snapshot)
};
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)?;
socket.send_to(&out, peer).await?;
Ok(())
}
async fn send_reject(socket: &UdpSocket, peer: SocketAddr, reason: &str) -> Result<()> {
send_reject_to_client(socket, peer, [0u8; 16], reason).await
}
async fn send_reject_to_client(
socket: &UdpSocket,
peer: SocketAddr,
client_id: [u8; 16],
reason: &str,
) -> Result<()> {
let body = protocol::to_body(&AttachReject {
reason: reason.to_string(),
})?;
let out = protocol::encode_plain(PacketKind::AttachReject, client_id, 0, 0, &body)?;
socket.send_to(&out, peer).await?;
Ok(())
}
async fn handle_resume(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = match find_client_decrypt_key(state, &packet.header) {
Ok(found) => found,
Err(_) => return send_reject(socket, peer, "unknown client").await,
};
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let req: ResumeRequest = protocol::from_body(&body)?;
let (send_seq, output_seq, snapshot) = {
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
.get_mut(&req.session)
.ok_or_else(|| anyhow!("unknown session"))?;
let client = session
.clients
.get_mut(&packet.header.conn_id)
.ok_or_else(|| anyhow!("unknown client"))?;
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
client.endpoint = peer;
client.last_acked = req.last_rendered_seq;
client.cols = req.cols;
client.rows = req.rows;
client.last_seen = Instant::now();
client.send_seq += 1;
if mode_allows_terminal_updates(&client.mode) {
apply_terminal_size(
session.pty.as_ref(),
&mut session.parser,
req.cols,
req.rows,
)?;
}
let snapshot = screen_snapshot(session.parser.screen());
(client.send_seq, session.output_seq, snapshot)
};
let frame = Frame {
session: session_name,
output_seq,
bytes: snapshot,
snapshot: true,
closed: false,
};
let body = protocol::to_body(&frame)?;
let out = protocol::encode_encrypted(
PacketKind::ResumeOk,
packet.header.conn_id,
send_seq,
packet.header.seq,
&key,
SERVER_TO_CLIENT,
&body,
)?;
socket.send_to(&out, peer).await?;
Ok(())
}
async fn handle_input(
state: &Arc<Mutex<ServerState>>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let input: Input = protocol::from_body(&body)?;
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
.get_mut(&session_name)
.ok_or_else(|| anyhow!("unknown session"))?;
let client = session
.clients
.get_mut(&packet.header.conn_id)
.ok_or_else(|| anyhow!("unknown client"))?;
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
if client.endpoint != peer {
client.endpoint = peer;
}
client.last_seen = Instant::now();
if !mode_allows_terminal_updates(&client.mode) {
return Ok(());
}
session
.pty
.as_ref()
.context("terminal session has no pty")?
.write_all(&input.bytes)?;
Ok(())
}
async fn handle_resize(
state: &Arc<Mutex<ServerState>>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let resize: Resize = protocol::from_body(&body)?;
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
.get_mut(&session_name)
.ok_or_else(|| anyhow!("unknown session"))?;
let client = session
.clients
.get_mut(&packet.header.conn_id)
.ok_or_else(|| anyhow!("unknown client"))?;
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
// Connection migration happens for any authenticated, fresh packet (spec §11),
// independent of the session mode; a view-only client roams too.
client.endpoint = peer;
client.last_seen = Instant::now();
if mode_allows_terminal_updates(&client.mode) {
client.cols = resize.cols;
client.rows = resize.rows;
apply_terminal_size(
session.pty.as_ref(),
&mut session.parser,
resize.cols,
resize.rows,
)?;
}
Ok(())
}
async fn handle_ping(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, _) = find_client_decrypt_key(state, &packet.header)?;
// Authenticate the ping (verify its AEAD tag) BEFORE acting on it. Pings carry
// an empty plaintext but still a tag, so this rejects a spoofed ping that
// could otherwise replay-poison or, worse, hijack `endpoint` migration.
let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let seq = {
let mut locked = state.lock().expect("server state poisoned");
let client = locked
.client_mut(&packet.header.conn_id)
.ok_or_else(|| anyhow!("unknown client"))?;
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
// Connection migration: any authenticated, fresh packet from a new source
// address roams the session there (spec §11). Gated on replay-accept so a
// captured old packet can't redirect the stream.
client.endpoint = peer;
client.last_seen = Instant::now();
client.send_seq += 1;
client.send_seq
};
let out = protocol::encode_encrypted(
PacketKind::Pong,
packet.header.conn_id,
seq,
packet.header.seq,
&key,
SERVER_TO_CLIENT,
b"",
)?;
socket.send_to(&out, peer).await?;
Ok(())
}
async fn handle_detach(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
let mut locked = state.lock().expect("server state poisoned");
locked.remove_client_everywhere(&packet.header.conn_id);
Ok(())
}
fn remove_client(state: &Arc<Mutex<ServerState>>, client_id: [u8; 16]) {
let mut locked = state.lock().expect("server state poisoned");
locked.remove_client_everywhere(&client_id);
}
async fn handle_ack(
state: &Arc<Mutex<ServerState>>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, _) = find_client_decrypt_key(state, &packet.header)?;
let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let mut locked = state.lock().expect("server state poisoned");
if let Some(client) = locked.client_mut(&packet.header.conn_id) {
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
// Connection migration on any authenticated, fresh packet (spec §11).
client.endpoint = peer;
client.last_seen = Instant::now();
client.last_acked = packet.header.ack;
while client
.pending
.front()
.is_some_and(|pending| pending.output_seq <= packet.header.ack)
{
client.pending.pop_front();
}
}
Ok(())
}
/// The client confirms it has switched to the new epoch key. Sealed under the
/// new (now-current) key, so decrypting it authenticates the confirmation. We
/// retire the previous-epoch key immediately, ending the grace early.
async fn handle_rekey_ack(
state: &Arc<Mutex<ServerState>>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, _) = find_client_decrypt_key(state, &packet.header)?;
let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let mut locked = state.lock().expect("server state poisoned");
if let Some(client) = locked.client_mut(&packet.header.conn_id) {
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
client.endpoint = peer;
client.last_seen = Instant::now();
// Only retire the old key if the ack was sealed under the *current* key,
// confirming the client really is on the new epoch.
if protocol::session_key_id(&client.session_key) == packet.header.session_key_id {
client.previous_session_key = None;
}
}
Ok(())
}
async fn handle_stream_open(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let open: StreamOpen = protocol::from_body(&body)?;
let allowed = {
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();
stream_open_allowed(&config, client, &open)
};
if let Err(err) = allowed {
return send_stream_open_reject(
state,
socket,
packet.header.conn_id,
open.stream_id,
err.to_string(),
)
.await;
}
let target = format!("{}:{}", open.target_host, open.target_port);
let stream = match TcpStream::connect((open.target_host.as_str(), open.target_port)).await {
Ok(stream) => stream,
Err(err) => {
return send_stream_open_reject(
state,
socket,
packet.header.conn_id,
open.stream_id,
format!("connect {target}: {err}"),
)
.await;
}
};
let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
{
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"))?;
client.stream_writers.insert(open.stream_id, writer_tx);
client.opened_streams.insert(open.stream_id);
client
.stream_send_credit
.insert(open.stream_id, STREAM_INITIAL_WINDOW);
}
tokio::spawn(async move {
while let Some(bytes) = writer_rx.recv().await {
if writer.write_all(&bytes).await.is_err() {
break;
}
}
});
send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;
let state = Arc::clone(state);
let socket = Arc::clone(socket);
let client_id = packet.header.conn_id;
let stream_id = open.stream_id;
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if let Err(err) = send_stream_data_to_client(
&state,
&socket,
client_id,
stream_id,
buf[..n].to_vec(),
)
.await
{
eprintln!("stream send error: {err:#}");
break;
}
}
Err(_) => break,
}
}
let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
});
Ok(())
}
async fn handle_stream_data(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let data: StreamData = protocol::from_body(&body)?;
let writer = {
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();
client.stream_writers.get(&data.stream_id).cloned()
};
let len = data.bytes.len();
if let Some(writer) = writer {
let _ = writer.send(data.bytes).await;
}
// Always return flow-control credit, even when the stream's writer is already
// gone (closed/unknown). 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, data.stream_id, len)
.await?;
Ok(())
}
async fn handle_stream_open_ok(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let ok: StreamOpenOk = protocol::from_body(&body)?;
{
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
.get_mut(&session_name)
.ok_or_else(|| anyhow!("unknown session"))?;
let client = session
.clients
.get_mut(&packet.header.conn_id)
.ok_or_else(|| anyhow!("unknown client"))?;
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
client.endpoint = peer;
client.last_seen = Instant::now();
client.opened_streams.insert(ok.stream_id);
client
.stream_send_credit
.entry(ok.stream_id)
.or_insert(STREAM_INITIAL_WINDOW);
}
flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, ok.stream_id).await
}
async fn handle_stream_open_reject(
state: &Arc<Mutex<ServerState>>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let reject: StreamOpenReject = protocol::from_body(&body)?;
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
.get_mut(&session_name)
.ok_or_else(|| anyhow!("unknown session"))?;
let client = session
.clients
.get_mut(&packet.header.conn_id)
.ok_or_else(|| anyhow!("unknown client"))?;
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
client.endpoint = peer;
client.last_seen = Instant::now();
client.stream_writers.remove(&reject.stream_id);
client.opened_streams.remove(&reject.stream_id);
client.stream_send_credit.remove(&reject.stream_id);
client.stream_pending_data.remove(&reject.stream_id);
Ok(())
}
async fn handle_stream_window_adjust(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let adjust: StreamWindowAdjust = protocol::from_body(&body)?;
{
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
.get_mut(&session_name)
.ok_or_else(|| anyhow!("unknown session"))?;
let client = session
.clients
.get_mut(&packet.header.conn_id)
.ok_or_else(|| anyhow!("unknown client"))?;
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
client.endpoint = peer;
client.last_seen = Instant::now();
add_stream_credit(
&mut client.stream_send_credit,
adjust.stream_id,
adjust.bytes,
);
}
flush_stream_pending_data_to_client(state, socket, packet.header.conn_id, adjust.stream_id)
.await
}
async fn handle_stream_close(
state: &Arc<Mutex<ServerState>>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let close: StreamClose = protocol::from_body(&body)?;
let mut locked = state.lock().expect("server state poisoned");
let Some(session) = locked.sessions.get_mut(&session_name) else {
return Ok(());
};
let Some(client) = session.clients.get_mut(&packet.header.conn_id) else {
return Ok(());
};
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
// Connection migration on any authenticated, fresh packet (spec §11).
client.endpoint = peer;
client.last_seen = Instant::now();
client.stream_writers.remove(&close.stream_id);
client.opened_streams.remove(&close.stream_id);
client.stream_send_credit.remove(&close.stream_id);
client.stream_pending_data.remove(&close.stream_id);
Ok(())
}
async fn start_remote_forwards(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
requested_forwardings: Vec<ForwardingRequest>,
session_name: String,
) -> Result<()> {
let remote_forwards = requested_forwardings
.into_iter()
.filter(|forward| forward.kind == ForwardingKind::Remote)
.collect::<Vec<_>>();
if remote_forwards.is_empty() {
return Ok(());
}
let config = {
let locked = state.lock().expect("server state poisoned");
locked.config.clone()
};
anyhow::ensure!(
config.allow_remote_forwarding,
"remote forwarding disabled on server"
);
for forward in remote_forwards {
let bind_host = forward
.bind_host
.clone()
.unwrap_or_else(|| "127.0.0.1".to_string());
remote_bind_allowed(&config, &bind_host)?;
let listen_port = forward.listen_port;
let target_host = forward
.target_host
.clone()
.ok_or_else(|| anyhow!("remote forward target host is required"))?;
let target_port = forward
.target_port
.ok_or_else(|| anyhow!("remote forward target port is required"))?;
let bind = format!("{bind_host}:{listen_port}");
let listener = TcpListener::bind(&bind)
.await
.with_context(|| format!("bind remote forward {bind}"))?;
let state = Arc::clone(state);
let socket = Arc::clone(socket);
let session_name = session_name.clone();
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
break;
};
let stream_id = allocate_server_stream_id(&state);
let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
{
let mut locked = state.lock().expect("server state poisoned");
if let Some(session) = locked.sessions.get_mut(&session_name)
&& let Some(client) = session.clients.get_mut(&client_id)
{
client.stream_writers.insert(stream_id, writer_tx);
client
.stream_send_credit
.insert(stream_id, STREAM_INITIAL_WINDOW);
} else {
break;
}
}
tokio::spawn(async move {
while let Some(bytes) = writer_rx.recv().await {
if writer.write_all(&bytes).await.is_err() {
break;
}
}
});
if let Err(err) = send_stream_open_to_client(
&state,
&socket,
client_id,
stream_id,
target_host.clone(),
target_port,
)
.await
{
eprintln!("remote forward open error: {err:#}");
break;
}
let state = Arc::clone(&state);
let socket = Arc::clone(&socket);
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if let Err(err) = send_stream_data_to_client(
&state,
&socket,
client_id,
stream_id,
buf[..n].to_vec(),
)
.await
{
eprintln!("remote forward data error: {err:#}");
break;
}
}
Err(_) => break,
}
}
let _ =
send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
});
}
});
}
Ok(())
}
fn allocate_server_stream_id(state: &Arc<Mutex<ServerState>>) -> u64 {
let mut locked = state.lock().expect("server state poisoned");
let stream_id = locked.next_server_stream_id;
locked.next_server_stream_id = locked.next_server_stream_id.wrapping_add(1).max(1u64 << 63);
stream_id
}
/// Whether the client requested SSH-agent forwarding during auth.
fn agent_forwarding_requested(forwardings: &[ForwardingRequest]) -> bool {
forwardings.iter().any(|f| f.kind == ForwardingKind::Agent)
}
/// Bind the per-session SSH-agent proxy unix socket, returning its `UnixListener`
/// and path. SECURITY: the socket lives in a process-private directory created
/// 0700 and the socket itself is chmod 0600, so only this user can connect — the
/// forwarded agent is never world-reachable. Called only when the client opted in
/// AND `allow_agent_forwarding` is set; the path is then exported as the remote
/// session's `SSH_AUTH_SOCK`.
fn bind_agent_proxy_socket() -> Result<(UnixListener, PathBuf)> {
let dir = std::env::temp_dir().join(format!("dosh-agent-{}", std::process::id()));
std::fs::create_dir_all(&dir).with_context(|| format!("create agent dir {}", dir.display()))?;
std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700))
.with_context(|| format!("chmod agent dir {}", dir.display()))?;
let n = AGENT_SOCK_COUNTER.fetch_add(1, Ordering::Relaxed);
let path = dir.join(format!("agent.{}.{n}.sock", std::process::id()));
// Stale socket from a crashed prior run with the same name: remove first.
let _ = std::fs::remove_file(&path);
let listener = UnixListener::bind(&path)
.with_context(|| format!("bind agent socket {}", path.display()))?;
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))
.with_context(|| format!("chmod agent socket {}", path.display()))?;
Ok((listener, path))
}
/// Run the accept loop for a forwarded SSH-agent: each connection from a remote
/// process to the proxy socket is tunneled to the client (which splices it into
/// the user's real ssh-agent) over a Dosh stream using the agent sentinel target.
/// The socket file is removed when the loop ends (client gone / session closed).
fn spawn_agent_forward(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
listener: UnixListener,
socket_path: PathBuf,
) {
let state = Arc::clone(state);
let socket = Arc::clone(socket);
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
break;
};
// Stop accepting once the client has gone away.
{
let locked = state.lock().expect("server state poisoned");
if locked.lookup_client(&client_id).is_none() {
break;
}
}
let stream_id = allocate_server_stream_id(&state);
let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
{
let mut locked = state.lock().expect("server state poisoned");
if let Some(client) = locked.client_mut(&client_id) {
client.stream_writers.insert(stream_id, writer_tx);
client
.stream_send_credit
.insert(stream_id, STREAM_INITIAL_WINDOW);
} else {
break;
}
}
tokio::spawn(async move {
while let Some(bytes) = writer_rx.recv().await {
if writer.write_all(&bytes).await.is_err() {
break;
}
}
});
if let Err(err) = send_stream_open_to_client(
&state,
&socket,
client_id,
stream_id,
AGENT_STREAM_SENTINEL.to_string(),
0,
)
.await
{
eprintln!("agent forward open error: {err:#}");
break;
}
let state = Arc::clone(&state);
let socket = Arc::clone(&socket);
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if let Err(err) = send_stream_data_to_client(
&state,
&socket,
client_id,
stream_id,
buf[..n].to_vec(),
)
.await
{
eprintln!("agent forward data error: {err:#}");
break;
}
}
Err(_) => break,
}
}
let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
});
}
let _ = std::fs::remove_file(&socket_path);
});
}
fn is_loopback_bind(host: &str) -> bool {
matches!(host, "127.0.0.1" | "localhost" | "::1")
}
fn remote_bind_allowed(config: &ServerConfig, host: &str) -> Result<()> {
if !is_loopback_bind(host) {
anyhow::ensure!(
config.allow_remote_non_loopback_bind,
"remote forwarding non-loopback bind disabled"
);
}
Ok(())
}
fn stream_open_allowed(
config: &ServerConfig,
client: &ClientState,
open: &StreamOpen,
) -> Result<()> {
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(())
}
async fn send_stream_open_ok(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
) -> Result<()> {
let body = protocol::to_body(&StreamOpenOk { stream_id })?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenOk, body).await
}
async fn send_stream_open_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
target_host: String,
target_port: u16,
) -> Result<()> {
let body = protocol::to_body(&StreamOpen {
stream_id,
target_host,
target_port,
})?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpen, body).await
}
async fn send_stream_open_reject(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
reason: String,
) -> Result<()> {
let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenReject, body).await
}
async fn send_stream_data_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
bytes: Vec<u8>,
) -> Result<()> {
queue_or_send_stream_data_to_client(state, socket, client_id, stream_id, bytes).await
}
async fn queue_or_send_stream_data_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
bytes: Vec<u8>,
) -> Result<()> {
let send = {
let mut locked = state.lock().expect("server state poisoned");
let mut send = None;
if let Some(client) = locked.client_mut(&client_id) {
if !client.opened_streams.contains(&stream_id)
|| client
.stream_send_credit
.get(&stream_id)
.copied()
.unwrap_or(0)
< bytes.len()
|| client
.stream_pending_data
.get(&stream_id)
.is_some_and(|pending| !pending.is_empty())
{
client
.stream_pending_data
.entry(stream_id)
.or_default()
.push_back(bytes);
return Ok(());
}
*client.stream_send_credit.entry(stream_id).or_default() -= bytes.len();
client.send_seq += 1;
let body = protocol::to_body(&StreamData { stream_id, bytes })?;
let packet = protocol::encode_encrypted(
PacketKind::StreamData,
client_id,
client.send_seq,
client.last_acked,
&client.session_key,
SERVER_TO_CLIENT,
&body,
)?;
send = Some((client.endpoint, packet));
}
send
};
if let Some((endpoint, packet)) = send {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
async fn flush_stream_pending_data_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
) -> Result<()> {
loop {
let send = {
let mut locked = state.lock().expect("server state poisoned");
let mut send = None;
if let Some(client) = locked.client_mut(&client_id) {
let Some(pending) = client.stream_pending_data.get_mut(&stream_id) else {
return Ok(());
};
let Some(bytes) = pending.front() else {
client.stream_pending_data.remove(&stream_id);
return Ok(());
};
let credit = client
.stream_send_credit
.get(&stream_id)
.copied()
.unwrap_or(0);
if credit < bytes.len() {
return Ok(());
}
let bytes = pending.pop_front().expect("pending front exists");
if pending.is_empty() {
client.stream_pending_data.remove(&stream_id);
}
*client.stream_send_credit.entry(stream_id).or_default() -= bytes.len();
client.send_seq += 1;
let body = protocol::to_body(&StreamData { stream_id, bytes })?;
let packet = protocol::encode_encrypted(
PacketKind::StreamData,
client_id,
client.send_seq,
client.last_acked,
&client.session_key,
SERVER_TO_CLIENT,
&body,
)?;
send = Some((client.endpoint, packet));
}
send
};
let Some((endpoint, packet)) = send else {
return Ok(());
};
socket.send_to(&packet, endpoint).await?;
}
}
fn add_stream_credit(stream_send_credit: &mut HashMap<u64, usize>, stream_id: u64, bytes: u32) {
let credit = stream_send_credit.entry(stream_id).or_default();
*credit = credit.saturating_add(bytes as usize);
}
async fn send_stream_window_adjust_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
bytes: usize,
) -> Result<()> {
let body = protocol::to_body(&StreamWindowAdjust {
stream_id,
bytes: bytes.min(u32::MAX as usize) as u32,
})?;
send_stream_packet_to_client(
state,
socket,
client_id,
PacketKind::StreamWindowAdjust,
body,
)
.await
}
async fn send_stream_close_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
) -> Result<()> {
{
let mut locked = state.lock().expect("server state poisoned");
if let Some(client) = locked.client_mut(&client_id) {
client.stream_writers.remove(&stream_id);
client.opened_streams.remove(&stream_id);
client.stream_send_credit.remove(&stream_id);
client.stream_pending_data.remove(&stream_id);
}
}
let body = protocol::to_body(&StreamClose { stream_id })?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamClose, body).await
}
async fn send_stream_packet_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
kind: PacketKind,
body: Vec<u8>,
) -> Result<()> {
let send = {
let mut locked = state.lock().expect("server state poisoned");
let mut found = None;
if let Some(client) = locked.client_mut(&client_id) {
client.send_seq += 1;
let packet = protocol::encode_encrypted(
kind,
client_id,
client.send_seq,
client.last_acked,
&client.session_key,
SERVER_TO_CLIENT,
&body,
)?;
found = Some((client.endpoint, packet));
}
found
};
if let Some((endpoint, packet)) = send {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
async fn broadcast_output(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
output: PtyOutput,
) -> Result<()> {
if output.exited {
return broadcast_session_exit(state, socket, output.session).await;
}
let mut persist_screen: Option<(String, u16, u16, u64, Vec<u8>)> = None;
let sends = {
let mut locked = state.lock().expect("server state poisoned");
let scrollback = locked.config.scrollback;
let retransmit_window = locked.config.retransmit_window;
let session = locked
.sessions
.get_mut(&output.session)
.ok_or_else(|| anyhow!("unknown session"))?;
session.parser.process(&output.bytes);
session.output_seq += 1;
let output_seq = session.output_seq;
session.recent.push_back(output.bytes.clone());
while session.recent.len() > scrollback {
session.recent.pop_front();
}
// For a persistent session, mirror the screen to disk periodically so a
// post-restart reattach repaints. Throttled by accumulated bytes to keep
// the (atomic) write off the per-packet hot path. The actual file write
// happens after the lock is dropped.
if session.persistent {
session.bytes_since_persist = session
.bytes_since_persist
.saturating_add(output.bytes.len());
if session.bytes_since_persist >= SCREEN_PERSIST_BYTES {
session.bytes_since_persist = 0;
session.last_persisted_seq = output_seq;
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);
// Only materialize a screen snapshot when this client has fallen too
// far behind; the common case sends the raw output bytes and must not
// clone the whole vt100 grid per client per packet.
let (bytes, snapshot) = if client.pending.len() >= retransmit_window {
client.pending.clear();
(screen_snapshot(session.parser.screen()), true)
} else {
(output.bytes.clone(), false)
};
let frame = Frame {
session: output.session.clone(),
output_seq,
bytes,
snapshot,
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,
)?;
while client.pending.len() >= retransmit_window {
client.pending.pop_front();
}
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:#}");
}
}
for (endpoint, packet) in sends {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
async fn broadcast_session_exit(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
session_name: String,
) -> Result<()> {
let (sends, persistent) = {
let mut locked = state.lock().expect("server state poisoned");
let Some(mut session) = locked.sessions.remove(&session_name) else {
return Ok(());
};
let persistent = session.persistent;
session.output_seq += 1;
let output_seq = session.output_seq;
let mut sends = Vec::new();
for (client_id, mut client) in session.clients.drain() {
locked.client_index.remove(&client_id);
client.send_seq += 1;
let frame = Frame {
session: session_name.clone(),
output_seq,
bytes: b"\r\n[dosh session exited]\r\n".to_vec(),
snapshot: false,
closed: true,
};
let body = protocol::to_body(&frame)?;
let packet = protocol::encode_encrypted(
PacketKind::Frame,
client_id,
client.send_seq,
client.last_acked,
&client.session_key,
SERVER_TO_CLIENT,
&body,
)?;
sends.push((client.endpoint, packet));
}
(sends, persistent)
};
if persistent {
// The shell exited; its holder cleans up its own runtime dir on the same
// event, but remove ours too so a re-adopt scan never sees a dead entry.
let sessions_dir = {
let locked = state.lock().expect("server state poisoned");
locked.sessions_dir()
};
persist::remove_runtime_dir(&sessions_dir, &session_name);
}
for (endpoint, packet) in sends {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
fn screen_snapshot(screen: &vt100::Screen) -> Vec<u8> {
let mut bytes = Vec::new();
if screen.alternate_screen() {
bytes.extend_from_slice(b"\x1b[?1049h");
} else {
bytes.extend_from_slice(b"\x1b[?1049l");
}
bytes.extend(screen.state_formatted());
bytes
}
async fn retransmit_pending(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
) -> Result<()> {
let sends = {
let mut locked = state.lock().expect("server state poisoned");
let now = Instant::now();
let mut sends = Vec::new();
for session in locked.sessions.values_mut() {
for client in session.clients.values_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.attempts < 8
{
pending.last_sent = now;
pending.attempts += 1;
sends.push((client.endpoint, pending.packet.clone()));
}
}
}
}
sends
};
for (endpoint, packet) in sends {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
/// How long the previous epoch's key is retained after a rekey so in-flight
/// pre-rekey packets still decrypt instead of being dropped as fatal.
const REKEY_PREVIOUS_KEY_GRACE_SECS: u64 = 5;
/// Rotate transport traffic keys for any client past its rekey threshold (spec
/// §11). Triggers on either a packet count or a wall-clock interval, configured
/// via `rekey_after_packets` / `rekey_after_secs`. The new key is derived from
/// fresh server material independent of the handshake keys
/// (see [`dosh::native::derive_rekey_session_key`]); the server installs it
/// immediately, retains the old key for a short grace, and ships the fresh
/// material to the client in a `Rekey` packet sealed under the *current* key.
async fn maybe_rekey_clients(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
) -> Result<()> {
let sends = {
let mut locked = state.lock().expect("server state poisoned");
let after_packets = locked.config.rekey_after_packets;
let after_secs = locked.config.rekey_after_secs;
if after_packets == 0 && after_secs == 0 {
return Ok(());
}
let now = Instant::now();
let grace = Duration::from_secs(REKEY_PREVIOUS_KEY_GRACE_SECS);
let mut sends = Vec::new();
for session in locked.sessions.values_mut() {
for (client_id, client) in session.clients.iter_mut() {
// Expire the retained previous-epoch key once its grace is over.
if client.previous_session_key.is_some()
&& now.duration_since(client.epoch_started) >= grace
{
client.previous_session_key = None;
}
let by_packets = after_packets != 0 && client.epoch_packets >= after_packets;
let by_time = after_secs != 0
&& now.duration_since(client.epoch_started).as_secs() >= after_secs;
if !(by_packets || by_time) {
continue;
}
// Don't start a new rekey while the previous epoch's grace is
// still active, so we never juggle three live keys at once.
if client.previous_session_key.is_some() {
continue;
}
let previous_key = client.session_key;
let previous_key_id = protocol::session_key_id(&previous_key);
let new_epoch = client.epoch.saturating_add(1);
let rekey_material = crypto::random_32();
let new_key = match dosh::native::derive_rekey_session_key(
&previous_key,
&rekey_material,
&previous_key_id,
new_epoch,
) {
Ok(key) => key,
Err(err) => {
eprintln!("rekey derive error: {err:#}");
continue;
}
};
let new_key_id = protocol::session_key_id(&new_key);
let rekey = protocol::Rekey {
epoch: new_epoch,
rekey_material,
new_session_key_id: new_key_id,
};
let body = protocol::to_body(&rekey)?;
client.send_seq += 1;
// The Rekey itself is sealed under the CURRENT key so the client
// can decrypt it before switching.
let packet = protocol::encode_encrypted(
PacketKind::Rekey,
*client_id,
client.send_seq,
client.last_acked,
&previous_key,
SERVER_TO_CLIENT,
&body,
)?;
// Install the new key, keep the old one for the grace window.
client.previous_session_key = Some(previous_key);
client.session_key = new_key;
client.epoch = new_epoch;
client.epoch_started = now;
client.epoch_packets = 0;
sends.push((client.endpoint, packet));
}
}
sends
};
for (endpoint, packet) in sends {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
/// Mirror every persistent session's current screen to disk if it changed since
/// the last mirror. Runs on a slow timer so a crash loses at most a couple of
/// seconds of screen state even for a low-throughput interactive session that
/// never crosses the per-packet byte threshold. Screen bytes are captured under
/// the lock; the (atomic) file writes happen after it is released.
fn flush_persistent_screens(state: &Arc<Mutex<ServerState>>) {
let (sessions_dir, to_write) = {
let mut locked = state.lock().expect("server state poisoned");
if !locked.config.persist_sessions {
return;
}
let sessions_dir = locked.sessions_dir();
let mut to_write: Vec<(String, u16, u16, u64, Vec<u8>)> = Vec::new();
for (name, session) in locked.sessions.iter_mut() {
if !session.persistent || session.output_seq == session.last_persisted_seq {
continue;
}
session.last_persisted_seq = session.output_seq;
session.bytes_since_persist = 0;
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:#}");
}
}
}
fn cleanup_disconnected_clients(state: &Arc<Mutex<ServerState>>) {
// Sessions removed here are dropped after the lock is released so their
// shells are killed (PtyHandle::drop) without holding up the event loop.
// For a persistent session, dropping the handle only DETACHES (the shell
// lives in the holder), so we also tell the holder to shut down — otherwise
// a truly-abandoned shell would linger forever.
let sessions_dir = {
let locked = state.lock().expect("server state poisoned");
locked.sessions_dir()
};
let reaped: Vec<(String, Session)> = {
let mut locked = state.lock().expect("server state poisoned");
let timeout = Duration::from_secs(locked.config.client_timeout_secs.max(1));
let now = Instant::now();
let prewarm: HashSet<String> = locked.config.prewarm_sessions.iter().cloned().collect();
// Collect timed-out clients first so we can purge them from the O(1)
// conn_id index too, keeping it in lockstep with `Session::clients`.
let mut timed_out: Vec<[u8; 16]> = Vec::new();
for session in locked.sessions.values_mut() {
session.clients.retain(|client_id, client| {
let keep = now.duration_since(client.last_seen) <= timeout;
if !keep {
timed_out.push(*client_id);
}
keep
});
if session.clients.is_empty() {
session.empty_since.get_or_insert(now);
} else {
session.empty_since = None;
}
}
for client_id in timed_out {
locked.client_index.remove(&client_id);
}
// Evict half-finished native handshakes. Each unauthenticated ClientHello
// inserts a PendingNativeAuth that is otherwise only removed by a matching
// UserAuth, so without this a flood of hellos grows the map without bound.
let handshake_ttl = Duration::from_secs(NATIVE_HANDSHAKE_TTL_SECS);
locked
.pending_native
.retain(|_, pending| now.duration_since(pending.created_at) < handshake_ttl);
// Drop fully-refilled rate-limit buckets so a spoofed-source flood cannot
// grow the map without bound.
locked.native_auth_limiter.evict_full(now);
// Reap sessions that have been clientless past the grace period. Prewarmed
// sessions stay hot on purpose, even with no clients attached.
let to_reap: Vec<String> = locked
.sessions
.iter()
.filter(|(name, session)| {
!prewarm.contains(name.as_str())
&& session
.empty_since
.is_some_and(|since| now.duration_since(since) >= 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 find_client_key(
state: &Arc<Mutex<ServerState>>,
client_id: &[u8; 16],
) -> Result<([u8; 32], String)> {
let locked = state.lock().expect("server state poisoned");
locked
.lookup_client(client_id)
.ok_or_else(|| anyhow!("unknown client"))
}
/// Like [`find_client_key`] but selects the 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.
fn find_client_decrypt_key(
state: &Arc<Mutex<ServerState>>,
header: &protocol::Header,
) -> Result<([u8; 32], String)> {
let locked = state.lock().expect("server state poisoned");
locked
.lookup_client_key_for(&header.conn_id, &header.session_key_id)
.ok_or_else(|| anyhow!("unknown client"))
}
fn parse_nonce(raw: &str) -> Result<[u8; 12]> {
let bytes = base64::Engine::decode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, raw)
.context("decode nonce")?;
anyhow::ensure!(bytes.len() == 12, "nonce must decode to 12 bytes");
let mut out = [0u8; 12];
out.copy_from_slice(&bytes);
Ok(out)
}
fn parse_size(raw: &str) -> Result<(u16, u16)> {
let (cols, rows) = raw.split_once('x').context("size must be COLSxROWS")?;
Ok((cols.parse()?, rows.parse()?))
}
/// Parse `--env NAME=VALUE` arguments passed to the holder back into pairs. The
/// value may itself contain `=`, so split only on the first.
fn parse_env_pairs(raw: &[String]) -> Result<Vec<(String, String)>> {
raw.iter()
.map(|entry| {
let (name, value) = entry
.split_once('=')
.with_context(|| format!("env entry {entry:?} must be NAME=VALUE"))?;
Ok((name.to_string(), value.to_string()))
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn 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());
}
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(),
epoch: 0,
epoch_started: Instant::now(),
epoch_packets: 0,
previous_session_key: None,
}
}
#[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());
// Removing purges both the session map and the index.
assert!(state.remove_client_everywhere(&client_id));
assert!(!state.client_index.contains_key(&client_id));
assert!(state.lookup_client(&client_id).is_none());
assert!(state.client_mut(&client_id).is_none());
assert!(!state.sessions["work"].clients.contains_key(&client_id));
// Removing an absent client is a no-op.
assert!(!state.remove_client_everywhere(&client_id));
}
#[test]
fn cleanup_purges_timed_out_clients_from_index() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
let config = ServerConfig {
client_timeout_secs: 1,
..ServerConfig::default()
};
let mut state = ServerState::new(config, [0u8; 32], pty_tx);
state
.ensure_session("work", 80, 24, "forward-only", &[])
.unwrap();
let client_id = [8u8; 16];
let mut client = test_client_state([1u8; 32]);
// Make the client look stale so cleanup reaps it.
client.last_seen = Instant::now() - Duration::from_secs(60);
state.insert_client("work", client_id, client);
let state = Arc::new(Mutex::new(state));
cleanup_disconnected_clients(&state);
let locked = state.lock().unwrap();
assert!(
!locked.client_index.contains_key(&client_id),
"timed-out client must be dropped from the index, not leaked"
);
assert!(locked.lookup_client(&client_id).is_none());
}
#[test]
fn forward_only_session_does_not_allocate_pty() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
state
.ensure_session("forward", 80, 24, "forward-only", &[])
.unwrap();
assert!(state.sessions["forward"].pty.is_none());
}
#[test]
fn cleanup_reaps_abandoned_sessions_but_keeps_prewarmed() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
let config = ServerConfig {
client_timeout_secs: 1,
prewarm_sessions: vec!["default".to_string()],
..ServerConfig::default()
};
let mut state = ServerState::new(config, [0u8; 32], pty_tx);
// Both sessions are clientless; use forward-only so no real shell spawns.
state
.ensure_session("default", 80, 24, "forward-only", &[])
.unwrap();
state
.ensure_session("abandoned", 80, 24, "forward-only", &[])
.unwrap();
// Mark both as empty long enough ago to be past the grace period.
let stale = Instant::now() - Duration::from_secs(60);
state.sessions.get_mut("default").unwrap().empty_since = Some(stale);
state.sessions.get_mut("abandoned").unwrap().empty_since = Some(stale);
let state = Arc::new(Mutex::new(state));
cleanup_disconnected_clients(&state);
let locked = state.lock().unwrap();
assert!(
locked.sessions.contains_key("default"),
"prewarmed session must stay hot"
);
assert!(
!locked.sessions.contains_key("abandoned"),
"clientless non-prewarmed session must be reaped"
);
}
#[test]
fn accepted_env_filters_by_server_patterns() {
let config = ServerConfig {
accept_env: vec!["LANG".to_string(), "LC_*".to_string()],
..ServerConfig::default()
};
assert_eq!(
accepted_env(
&config,
&[
EnvVar {
name: "LANG".to_string(),
value: "en_US.UTF-8".to_string()
},
EnvVar {
name: "LC_TIME".to_string(),
value: "C".to_string()
},
EnvVar {
name: "SECRET_TOKEN".to_string(),
value: "nope".to_string()
}
]
)
.unwrap(),
vec![
("LANG".to_string(), "en_US.UTF-8".to_string()),
("LC_TIME".to_string(), "C".to_string())
]
);
}
#[test]
fn accepted_env_rejects_invalid_names() {
let err = accepted_env(
&ServerConfig::default(),
&[EnvVar {
name: "BAD-NAME".to_string(),
value: "value".to_string(),
}],
)
.unwrap_err();
assert!(
err.to_string()
.contains("invalid environment variable name")
);
}
#[tokio::test]
async fn stream_data_waits_for_open_and_credit() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
let client_id = [7u8; 16];
let session_key = [9u8; 32];
let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
let endpoint = receiver.local_addr().unwrap();
state.sessions.insert(
"test".to_string(),
Session {
pty: None,
parser: vt100::Parser::new(24, 80, 100),
clients: HashMap::from([(
client_id,
ClientState {
endpoint,
mode: "forward-only".to_string(),
session_key,
last_acked: 0,
replay: ReplayWindow::default(),
send_seq: 1,
cols: 80,
rows: 24,
last_seen: Instant::now(),
pending: VecDeque::new(),
allowed_forwardings: Vec::new(),
stream_writers: HashMap::new(),
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
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,
},
);
// 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");
}
#[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();
}
}
Reference in New Issue View Git Blame Copy Permalink