From 8b1af51bc63dc49568dff93195585157e4f01449 Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Sun, 14 Jun 2026 14:19:09 -0400 Subject: [PATCH] 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 --- Cargo.lock | 1 + Cargo.toml | 1 + README.md | 3 + docs/NATIVE_V1_SPEC.md | 9 +- src/bin/dosh-server.rs | 388 ++++++++++++++++++++-- src/config.rs | 9 + src/lib.rs | 1 + src/persist.rs | 658 +++++++++++++++++++++++++++++++++++++ src/pty.rs | 149 +++++++-- tests/hostile_network.rs | 1 + tests/integration_smoke.rs | 183 +++++++++++ 11 files changed, 1347 insertions(+), 56 deletions(-) create mode 100644 src/persist.rs diff --git a/Cargo.lock b/Cargo.lock index b00bafe..a5a3732 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -448,6 +448,7 @@ dependencies = [ "ed25519-dalek", "hkdf", "hmac", + "libc", "portable-pty", "rand", "rpassword", diff --git a/Cargo.toml b/Cargo.toml index 90f80f5..60fcdd4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,7 @@ dirs = "5.0" ed25519-dalek = "2.1" hkdf = "0.12" hmac = "0.12" +libc = "0.2" portable-pty = "0.8" rand = "0.8" rpassword = "7.5.4" diff --git a/README.md b/README.md index 9499182..0c8b554 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,9 @@ dosh-server client table per session encrypted UDP protocol tiny SSH-invoked dosh-auth helper mode + persistent sessions (persist_sessions, default on): each shell runs in a + detached per-session holder process; the server adopts its PTY master fd via + SCM_RIGHTS and re-adopts it after a restart, so sessions survive crash/upgrade dosh-client terminal raw mode diff --git a/docs/NATIVE_V1_SPEC.md b/docs/NATIVE_V1_SPEC.md index 7accacb..dcb8e47 100644 --- a/docs/NATIVE_V1_SPEC.md +++ b/docs/NATIVE_V1_SPEC.md @@ -123,8 +123,13 @@ Native v1 must be boring under real use: - A closed laptop must not kill the remote session. - A client crash must not kill the remote session. -- Server restart may drop live PTYs in v1 unless session persistence is implemented, - but the client must fail clearly and reconnect cleanly afterward. +- A server restart must not kill the remote session when `persist_sessions` is on + (the default): each session's shell runs in a detached per-session *holder* + process whose PTY master fd the server passes back to itself via SCM_RIGHTS, so + the shell + scrollback survive a server crash/upgrade/`systemctl restart` and a + reattaching client lands on the same shell with its screen restored. With + `persist_sessions = false` the old behavior applies: live PTYs drop on restart, + but the client must still fail clearly and reconnect cleanly afterward. - Decrypt failures from stale packets must be ignored or trigger reconnect, never terminate the terminal by themselves. - Terminal cleanup must restore cursor, mouse mode, bracketed paste, alternate diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 4fa7258..4e392d1 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -4,13 +4,14 @@ 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, load_server_config}; +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, @@ -18,10 +19,11 @@ use dosh::protocol::{ StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; -use dosh::pty::{PtyHandle, PtyOutput, spawn_pty_session}; +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}; @@ -32,6 +34,11 @@ 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. @@ -74,13 +81,55 @@ enum Command { #[arg(long)] udp_host: Option, }, + /// Internal: run as a per-session shell holder so the shell survives a + /// server restart. Spawned by the server itself; not meant to be run by + /// hand. Daemonizes, owns the PTY, and serves the master fd over a socket. + #[command(hide = true)] + Hold { + #[arg(long)] + runtime_dir: PathBuf, + #[arg(long)] + session: String, + #[arg(long)] + shell: String, + #[arg(long)] + cols: u16, + #[arg(long)] + rows: u16, + /// `NAME=VALUE` pairs for the shell environment. Repeatable. + #[arg(long = "env")] + env: Vec, + }, } -#[tokio::main(flavor = "current_thread")] -async fn main() -> Result<()> { +fn main() -> Result<()> { let args = Args::parse(); - match args.command { + // 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, @@ -130,6 +179,12 @@ async fn serve(config_path: Option) -> Result<()> { ))); { 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", &[])?; } @@ -169,6 +224,21 @@ async fn serve(config_path: Option) -> Result<()> { } }); + // 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?; @@ -273,6 +343,19 @@ struct Session { /// sessions (and their shells) after a grace period. `None` while at least /// one client is attached. empty_since: Option, + /// Open control connection to this session's holder process, when the shell + /// runs persistently out-of-process. Dropping it (server exit) leaves the + /// holder running; sending a shutdown byte over it reaps the holder. `None` + /// for non-persistent (in-process) sessions. + holder_control: Option, + /// Whether this session's shell lives in a holder process (persistent). + persistent: bool, + /// Bytes of session output since the screen was last mirrored to disk, used + /// to throttle the (atomic) screen-persistence writes. + bytes_since_persist: usize, + /// `output_seq` at the last successful screen mirror, so the periodic flush + /// skips sessions whose screen has not changed. + last_persisted_seq: u64, } #[derive(Clone)] @@ -347,30 +430,29 @@ impl ServerState { requested_env: &[EnvVar], ) -> Result<()> { let env = accepted_env(&self.config, requested_env)?; - if let Some(session) = self.sessions.get_mut(name) { - if session.pty.is_none() && mode_uses_pty(mode) { - session.pty = Some(spawn_pty_session( - name.to_string(), - &self.config.shell, - cols.max(1), - rows.max(1), - &env, - self.pty_tx.clone(), - )?); + // 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 = if mode_uses_pty(mode) { - Some(spawn_pty_session( - name.to_string(), - &self.config.shell, - cols.max(1), - rows.max(1), - &env, - self.pty_tx.clone(), - )?) + 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, None, false) }; self.sessions.insert( name.to_string(), @@ -381,11 +463,144 @@ impl ServerState { 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, 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) { @@ -2240,6 +2455,7 @@ async fn broadcast_output( return broadcast_session_exit(state, socket, output.session).await; } + let mut persist_screen: Option<(String, u16, u16, u64, Vec)> = None; let sends = { let mut locked = state.lock().expect("server state poisoned"); let scrollback = locked.config.scrollback; @@ -2255,6 +2471,27 @@ async fn broadcast_output( 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; @@ -2299,6 +2536,17 @@ async fn broadcast_output( } 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?; } @@ -2310,11 +2558,12 @@ async fn broadcast_session_exit( socket: &Arc, session_name: String, ) -> Result<()> { - let sends = { + 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(); @@ -2340,8 +2589,17 @@ async fn broadcast_session_exit( )?; sends.push((client.endpoint, packet)); } - sends + (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?; } @@ -2487,10 +2745,56 @@ async fn maybe_rekey_clients( Ok(()) } +/// Mirror every persistent session's current screen to disk if it changed since +/// the last mirror. Runs on a slow timer so a crash loses at most a couple of +/// seconds of screen state even for a low-throughput interactive session that +/// never crosses the per-packet byte threshold. Screen bytes are captured under +/// the lock; the (atomic) file writes happen after it is released. +fn flush_persistent_screens(state: &Arc>) { + let (sessions_dir, to_write) = { + let mut locked = state.lock().expect("server state poisoned"); + if !locked.config.persist_sessions { + return; + } + let sessions_dir = locked.sessions_dir(); + let mut to_write: Vec<(String, u16, u16, u64, Vec)> = Vec::new(); + for (name, session) in locked.sessions.iter_mut() { + if !session.persistent || 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>) { // Sessions removed here are dropped after the lock is released so their // shells are killed (PtyHandle::drop) without holding up the event loop. - let reaped: Vec = { + // 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(); @@ -2540,10 +2844,17 @@ fn cleanup_disconnected_clients(state: &Arc>) { .collect(); to_reap .into_iter() - .filter_map(|name| locked.sessions.remove(&name)) + .filter_map(|name| locked.sessions.remove(&name).map(|session| (name, session))) .collect() }; - drop(reaped); + 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( @@ -2583,6 +2894,19 @@ fn parse_size(raw: &str) -> Result<(u16, u16)> { Ok((cols.parse()?, rows.parse()?)) } +/// Parse `--env NAME=VALUE` arguments passed to the holder back into pairs. The +/// value may itself contain `=`, so split only on the first. +fn parse_env_pairs(raw: &[String]) -> Result> { + raw.iter() + .map(|entry| { + let (name, value) = entry + .split_once('=') + .with_context(|| format!("env entry {entry:?} must be NAME=VALUE"))?; + Ok((name.to_string(), value.to_string())) + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; @@ -2853,6 +3177,10 @@ mod tests { 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`. diff --git a/src/config.rs b/src/config.rs index d48d15b..c4bbcaa 100644 --- a/src/config.rs +++ b/src/config.rs @@ -46,6 +46,14 @@ pub struct ServerConfig { pub allow_agent_forwarding: bool, #[serde(default = "default_accept_env")] pub accept_env: Vec, + /// Run each terminal session's shell under a small per-session *holder* + /// process so it survives a `dosh-server` restart (crash/upgrade/ + /// `systemctl restart`). On startup the server re-adopts live holders and + /// reattaching clients land on the same shell with screen state intact. When + /// `false`, sessions behave exactly as before: the shell is a child of the + /// server and dies with it. + #[serde(default = "default_true")] + pub persist_sessions: bool, } impl Default for ServerConfig { @@ -76,6 +84,7 @@ impl Default for ServerConfig { allow_remote_non_loopback_bind: false, allow_agent_forwarding: false, accept_env: default_accept_env(), + persist_sessions: true, } } } diff --git a/src/lib.rs b/src/lib.rs index abe852d..51ab7aa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ pub mod auth; pub mod config; pub mod crypto; pub mod native; +pub mod persist; pub mod protocol; pub mod pty; pub mod ssh_agent; diff --git a/src/persist.rs b/src/persist.rs new file mode 100644 index 0000000..171da1a --- /dev/null +++ b/src/persist.rs @@ -0,0 +1,658 @@ +//! Session persistence across server restarts. +//! +//! A persistent dosh session's shell does not run as a child of `dosh-server`. +//! Instead the server spawns a tiny per-session *holder* process (the same +//! `dosh-server` binary re-exec'd as `dosh-server hold ...`). The holder calls +//! `setsid()` to leave the server's process group/session, opens a PTY, spawns +//! the shell as ITS own child, and listens on a Unix socket in a per-session +//! runtime directory. The server connects to that socket and the holder hands it +//! the PTY master fd over `SCM_RIGHTS`. +//! +//! Because the shell belongs to the holder (which is in its own session and is +//! not waited on by the server), killing or restarting `dosh-server` leaves the +//! holder + shell alive. On startup the server scans the runtime directory, +//! reconnects to each live holder, receives the master fd again, and rebuilds the +//! in-memory session so clients reattach to the very same shell. +//! +//! Screen / scrollback state lives only in server memory, so it is also mirrored +//! to disk (atomically) and restored on re-adoption, letting a reattaching client +//! repaint the screen exactly as it was before the restart. + +use anyhow::{Context, Result, anyhow, bail}; +use std::io::{Read, Write}; +use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +use std::os::unix::net::{UnixListener, UnixStream}; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +/// One-byte commands a server sends to a holder over its control socket after +/// the holder has handed back the master fd. +const HOLDER_CMD_SHUTDOWN: u8 = b'X'; + +/// Magic written into a holder's `meta` file, bumped if the on-disk layout +/// changes so a stale holder from an incompatible build is ignored. +const META_MAGIC: &str = "dosh-holder-1"; + +/// Per-session runtime metadata persisted next to the holder socket. +#[derive(Debug, Clone)] +pub struct HolderMeta { + pub session: String, + pub shell_pid: i32, +} + +/// Root runtime directory holding one subdirectory per persistent session. +/// Lives under `sessions_dir/run` so it shares the session storage location and +/// can be wiped with it. +pub fn runtime_root(sessions_dir: &Path) -> PathBuf { + sessions_dir.join("run") +} + +/// Map a session name to a filesystem-safe directory name. Session names are +/// user-controlled, so hex-encode them rather than trusting them as path +/// components (avoids traversal, slashes, NULs, length issues). +fn session_dir_name(session: &str) -> String { + let mut out = String::with_capacity(session.len() * 2); + for byte in session.as_bytes() { + out.push(char::from_digit((byte >> 4) as u32, 16).unwrap()); + out.push(char::from_digit((byte & 0xf) as u32, 16).unwrap()); + } + out +} + +fn decode_session_dir_name(name: &str) -> Option { + if !name.len().is_multiple_of(2) { + return None; + } + let mut bytes = Vec::with_capacity(name.len() / 2); + let raw = name.as_bytes(); + let mut i = 0; + while i < raw.len() { + let hi = (raw[i] as char).to_digit(16)?; + let lo = (raw[i + 1] as char).to_digit(16)?; + bytes.push(((hi << 4) | lo) as u8); + i += 2; + } + String::from_utf8(bytes).ok() +} + +/// Directory for one session's holder runtime state. +pub fn session_runtime_dir(sessions_dir: &Path, session: &str) -> PathBuf { + runtime_root(sessions_dir).join(session_dir_name(session)) +} + +fn holder_sock_path(dir: &Path) -> PathBuf { + dir.join("holder.sock") +} + +fn meta_path(dir: &Path) -> PathBuf { + dir.join("meta") +} + +fn screen_path(dir: &Path) -> PathBuf { + dir.join("screen") +} + +/// Create the runtime directory tree with private (0700) permissions. +pub fn ensure_runtime_dir(sessions_dir: &Path, session: &str) -> Result { + use std::os::unix::fs::PermissionsExt; + let root = runtime_root(sessions_dir); + std::fs::create_dir_all(&root).with_context(|| format!("create {}", root.display()))?; + let _ = std::fs::set_permissions(&root, std::fs::Permissions::from_mode(0o700)); + let dir = session_runtime_dir(sessions_dir, session); + std::fs::create_dir_all(&dir).with_context(|| format!("create {}", dir.display()))?; + std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700)) + .with_context(|| format!("chmod {}", dir.display()))?; + Ok(dir) +} + +/// Atomically write `data` to `path` (write temp + rename), so a concurrent +/// reader (e.g. a restarting server) never sees a half-written file. +fn atomic_write(path: &Path, data: &[u8]) -> Result<()> { + use std::os::unix::fs::PermissionsExt; + let tmp = path.with_extension("tmp"); + { + let mut file = + std::fs::File::create(&tmp).with_context(|| format!("create {}", tmp.display()))?; + let _ = file.set_permissions(std::fs::Permissions::from_mode(0o600)); + file.write_all(data) + .with_context(|| format!("write {}", tmp.display()))?; + file.sync_all().ok(); + } + std::fs::rename(&tmp, path) + .with_context(|| format!("rename {} -> {}", tmp.display(), path.display()))?; + Ok(()) +} + +/// Persist the vt100 screen snapshot plus recent raw output for a session so a +/// post-restart reattach repaints correctly. `snapshot` is the bytes a fresh +/// attach would receive (alt-screen toggle + `state_formatted`). +pub fn save_screen( + sessions_dir: &Path, + session: &str, + cols: u16, + rows: u16, + output_seq: u64, + snapshot: &[u8], +) -> Result<()> { + let dir = session_runtime_dir(sessions_dir, session); + if !dir.exists() { + // No holder runtime for this session (non-persistent): nothing to do. + return Ok(()); + } + let mut buf = Vec::with_capacity(snapshot.len() + 32); + buf.extend_from_slice(&cols.to_be_bytes()); + buf.extend_from_slice(&rows.to_be_bytes()); + buf.extend_from_slice(&output_seq.to_be_bytes()); + buf.extend_from_slice(&(snapshot.len() as u32).to_be_bytes()); + buf.extend_from_slice(snapshot); + atomic_write(&screen_path(&dir), &buf) +} + +/// A restored screen for a re-adopted session. +pub struct SavedScreen { + pub cols: u16, + pub rows: u16, + pub output_seq: u64, + pub snapshot: Vec, +} + +/// Load a previously persisted screen, if any. +pub fn load_screen(sessions_dir: &Path, session: &str) -> Option { + let dir = session_runtime_dir(sessions_dir, session); + let data = std::fs::read(screen_path(&dir)).ok()?; + if data.len() < 16 { + return None; + } + let cols = u16::from_be_bytes(data[0..2].try_into().ok()?); + let rows = u16::from_be_bytes(data[2..4].try_into().ok()?); + let output_seq = u64::from_be_bytes(data[4..12].try_into().ok()?); + let len = u32::from_be_bytes(data[12..16].try_into().ok()?) as usize; + if data.len() < 16 + len { + return None; + } + Some(SavedScreen { + cols, + rows, + output_seq, + snapshot: data[16..16 + len].to_vec(), + }) +} + +fn write_meta(dir: &Path, meta: &HolderMeta) -> Result<()> { + let body = format!("{META_MAGIC}\n{}\n{}\n", meta.shell_pid, meta.session); + atomic_write(&meta_path(dir), body.as_bytes()) +} + +fn read_meta(dir: &Path) -> Option { + let body = std::fs::read_to_string(meta_path(dir)).ok()?; + let mut lines = body.lines(); + if lines.next()? != META_MAGIC { + return None; + } + let shell_pid: i32 = lines.next()?.parse().ok()?; + let session = lines.next()?.to_string(); + Some(HolderMeta { session, shell_pid }) +} + +/// Spawn a holder process for `session`: re-exec this binary as `dosh-server +/// hold` with the runtime dir, shell, terminal size, and accepted env. The +/// holder daemonizes (setsid, owns the PTY) and listens on its control socket. +/// Returns once the holder has signalled readiness by creating its socket. +#[allow(clippy::too_many_arguments)] +pub fn spawn_holder( + sessions_dir: &Path, + session: &str, + shell: &str, + cols: u16, + rows: u16, + env: &[(String, String)], +) -> Result<()> { + let dir = ensure_runtime_dir(sessions_dir, session)?; + let sock = holder_sock_path(&dir); + // Clear any stale socket left by a crashed holder with the same name. + let _ = std::fs::remove_file(&sock); + + let exe = std::env::current_exe().context("locate current executable")?; + let mut cmd = std::process::Command::new(exe); + cmd.arg("hold") + .arg("--runtime-dir") + .arg(&dir) + .arg("--session") + .arg(session) + .arg("--shell") + .arg(shell) + .arg("--cols") + .arg(cols.to_string()) + .arg("--rows") + .arg(rows.to_string()); + for (name, value) in env { + cmd.arg("--env").arg(format!("{name}={value}")); + } + cmd.stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()); + let mut child = cmd.spawn().context("spawn holder process")?; + + // Wait (briefly) for the holder to come up. The holder forks/setsids before + // creating the socket; once the socket exists we can connect. We also reap + // the short-lived launcher child so it never becomes a zombie. + let deadline = std::time::Instant::now() + Duration::from_secs(5); + loop { + if sock.exists() { + break; + } + if std::time::Instant::now() >= deadline { + let _ = child.kill(); + let _ = child.wait(); + bail!("holder for session {session} did not come up in time"); + } + std::thread::sleep(Duration::from_millis(20)); + } + // The launcher exits immediately after the grandchild setsids; reap it. + let _ = child.wait(); + Ok(()) +} + +/// Connect to a session's holder and receive the PTY master fd over SCM_RIGHTS. +/// Returns the raw fd (caller owns it) and the holder control socket, kept open +/// so the server can later ask the holder to shut down. Returns an error if the +/// holder is gone (caller then degrades to a fresh, non-persistent session). +pub fn adopt_holder(sessions_dir: &Path, session: &str) -> Result<(RawFd, UnixStream)> { + let dir = session_runtime_dir(sessions_dir, session); + let sock = holder_sock_path(&dir); + let mut stream = UnixStream::connect(&sock) + .with_context(|| format!("connect holder socket {}", sock.display()))?; + stream.set_read_timeout(Some(Duration::from_secs(5))).ok(); + let fd = recv_fd(&mut stream).context("receive master fd from holder")?; + Ok((fd, stream)) +} + +/// Ask a holder to terminate its shell and exit, then clean its runtime dir. +/// Used when reaping a truly-abandoned persistent session. +pub fn request_shutdown(sessions_dir: &Path, session: &str, control: Option<&mut UnixStream>) { + if let Some(stream) = control { + let _ = stream.write_all(&[HOLDER_CMD_SHUTDOWN]); + let _ = stream.flush(); + } else { + let dir = session_runtime_dir(sessions_dir, session); + if let Ok(mut stream) = UnixStream::connect(holder_sock_path(&dir)) { + // Drain the fd the holder sends on connect, then send shutdown. + let _ = recv_fd(&mut stream); + let _ = stream.write_all(&[HOLDER_CMD_SHUTDOWN]); + let _ = stream.flush(); + } + } + // Give the holder a moment to tear down, then remove its runtime dir. + std::thread::sleep(Duration::from_millis(50)); + let dir = session_runtime_dir(sessions_dir, session); + let _ = std::fs::remove_dir_all(&dir); +} + +/// Remove a session's runtime directory unconditionally (best effort). +pub fn remove_runtime_dir(sessions_dir: &Path, session: &str) { + let dir = session_runtime_dir(sessions_dir, session); + let _ = std::fs::remove_dir_all(&dir); +} + +/// Scan the runtime root for holders left behind by a previous server. Returns +/// `(session_name, meta)` for each one whose holder process still appears alive. +/// Stale entries (no live process) are cleaned up. +pub fn scan_existing_holders(sessions_dir: &Path) -> Vec { + let root = runtime_root(sessions_dir); + let mut found = Vec::new(); + let Ok(entries) = std::fs::read_dir(&root) else { + return found; + }; + for entry in entries.flatten() { + let dir = entry.path(); + if !dir.is_dir() { + continue; + } + let Some(name) = dir.file_name().and_then(|n| n.to_str()) else { + continue; + }; + // Decode the on-disk name back to a session name; skip junk dirs. + if decode_session_dir_name(name).is_none() { + continue; + } + match read_meta(&dir) { + Some(meta) if process_alive(meta.shell_pid) && holder_sock_path(&dir).exists() => { + found.push(meta); + } + _ => { + // Holder gone or meta unreadable: clean up so we don't try to + // adopt a dead session (degrade to fresh on next attach). + let _ = std::fs::remove_dir_all(&dir); + } + } + } + found +} + +/// Whether a pid refers to a live process (signal 0 probe). +fn process_alive(pid: i32) -> bool { + if pid <= 0 { + return false; + } + unsafe { + libc::kill(pid, 0) == 0 + || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM) + } +} + +// --------------------------------------------------------------------------- +// Holder process entry point +// --------------------------------------------------------------------------- + +/// Run as a holder process. Daemonizes (double-fork + setsid), opens a PTY, +/// spawns the shell as its own child, and serves the control socket: every +/// accepted connection is handed the master fd (SCM_RIGHTS); a subsequent +/// SHUTDOWN byte (or the shell exiting) tears everything down. +/// +/// This function does not return on success — it `exit`s the process. Errors +/// before the daemonization point are returned to the launcher. +#[allow(clippy::too_many_arguments)] +pub fn run_holder( + runtime_dir: &Path, + session: &str, + shell: &str, + cols: u16, + rows: u16, + env: &[(String, String)], +) -> Result<()> { + use portable_pty::{NativePtySystem, PtySize, PtySystem}; + + // Detach from the launching server: own session + process group so a server + // exit (even a process-group kill of the service) does not take us down. + daemonize().context("daemonize holder")?; + + let pty_system = NativePtySystem::default(); + let pair = pty_system + .openpty(PtySize { + rows, + cols, + pixel_width: 0, + pixel_height: 0, + }) + .context("holder open pty")?; + let cmd = crate::pty::build_shell_command(shell, env); + let mut child = pair + .slave + .spawn_command(cmd) + .context("holder spawn shell")?; + drop(pair.slave); + let master_fd = pair + .master + .as_raw_fd() + .ok_or_else(|| anyhow!("holder master has no raw fd"))?; + let shell_pid = child.process_id().map(|p| p as i32).unwrap_or(-1); + + write_meta( + runtime_dir, + &HolderMeta { + session: session.to_string(), + shell_pid, + }, + )?; + + let sock = holder_sock_path(runtime_dir); + let _ = std::fs::remove_file(&sock); + let listener = + UnixListener::bind(&sock).with_context(|| format!("holder bind {}", sock.display()))?; + { + use std::os::unix::fs::PermissionsExt; + let _ = std::fs::set_permissions(&sock, std::fs::Permissions::from_mode(0o600)); + } + listener + .set_nonblocking(false) + .context("holder listener blocking")?; + + // A watcher thread reaps the shell: when it exits, the holder cleans up and + // exits too, so an abandoned shell does not linger forever. + let runtime_owned = runtime_dir.to_path_buf(); + std::thread::spawn(move || { + let _ = child.wait(); + // Shell exited: remove runtime dir and exit the holder process. + let _ = std::fs::remove_dir_all(&runtime_owned); + std::process::exit(0); + }); + + // Accept loop: each connecting server gets the master fd; a SHUTDOWN byte + // from any of them tears the holder down. Multiple servers never run at once + // in practice (one service), but serving repeated connections lets a server + // restart re-adopt cleanly. + for stream in listener.incoming() { + let Ok(mut stream) = stream else { continue }; + if send_fd(&mut stream, master_fd).is_err() { + continue; + } + // Wait for an optional command byte. EOF (server dropped the control + // socket on its own exit) just means "keep running, await re-adoption". + let mut cmd = [0u8; 1]; + match stream.read(&mut cmd) { + Ok(1) if cmd[0] == HOLDER_CMD_SHUTDOWN => { + let _ = std::fs::remove_dir_all(runtime_dir); + std::process::exit(0); + } + _ => { + // Server detached (exit/restart) or sent nothing actionable: + // loop back and wait for the next server to re-adopt us. + } + } + } + Ok(()) +} + +/// Double-fork + `setsid` so the holder runs in its own session, detached from +/// the server's controlling terminal and process group. Without this a +/// `systemctl restart` (which signals the whole service cgroup/process group) +/// could take the holder down with the server. +fn daemonize() -> Result<()> { + // First fork: parent (launcher) returns to reap; child continues. + match unsafe { libc::fork() } { + -1 => bail!("fork: {}", std::io::Error::last_os_error()), + 0 => {} + _ => { + // Parent process: exit so the launcher's wait() returns promptly and + // the grandchild is reparented to init. + std::process::exit(0); + } + } + // New session: detaches from controlling tty and the server's process group. + if unsafe { libc::setsid() } == -1 { + bail!("setsid: {}", std::io::Error::last_os_error()); + } + // Second fork: ensures we are not a session leader, so we can never + // re-acquire a controlling terminal. + match unsafe { libc::fork() } { + -1 => bail!("fork2: {}", std::io::Error::last_os_error()), + 0 => Ok(()), + _ => std::process::exit(0), + } +} + +// --------------------------------------------------------------------------- +// SCM_RIGHTS file-descriptor passing over a Unix domain socket +// --------------------------------------------------------------------------- + +/// Send a single fd over `stream` using SCM_RIGHTS, with one byte of normal data +/// (sendmsg requires at least one iovec byte for the ancillary data to ride on). +fn send_fd(stream: &mut UnixStream, fd: RawFd) -> Result<()> { + let dummy: [u8; 1] = [0]; + let mut iov = libc::iovec { + iov_base: dummy.as_ptr() as *mut libc::c_void, + iov_len: 1, + }; + let mut cmsg_buf = [0u8; cmsg_space_one_fd()]; + let mut msg: libc::msghdr = unsafe { std::mem::zeroed() }; + msg.msg_iov = &mut iov; + msg.msg_iovlen = 1; + msg.msg_control = cmsg_buf.as_mut_ptr() as *mut libc::c_void; + msg.msg_controllen = cmsg_buf.len() as _; + + unsafe { + let cmsg = libc::CMSG_FIRSTHDR(&msg); + if cmsg.is_null() { + bail!("CMSG_FIRSTHDR null"); + } + (*cmsg).cmsg_level = libc::SOL_SOCKET; + (*cmsg).cmsg_type = libc::SCM_RIGHTS; + (*cmsg).cmsg_len = libc::CMSG_LEN(std::mem::size_of::() as u32) as _; + std::ptr::copy_nonoverlapping( + &fd as *const RawFd as *const u8, + libc::CMSG_DATA(cmsg), + std::mem::size_of::(), + ); + let n = libc::sendmsg(stream.as_raw_fd(), &msg, 0); + if n < 0 { + return Err(std::io::Error::last_os_error()).context("sendmsg SCM_RIGHTS"); + } + } + Ok(()) +} + +/// Receive a single fd sent via SCM_RIGHTS. Returns a fresh fd owned by the +/// caller. +fn recv_fd(stream: &mut UnixStream) -> Result { + let mut dummy = [0u8; 1]; + let mut iov = libc::iovec { + iov_base: dummy.as_mut_ptr() as *mut libc::c_void, + iov_len: 1, + }; + let mut cmsg_buf = [0u8; cmsg_space_one_fd()]; + let mut msg: libc::msghdr = unsafe { std::mem::zeroed() }; + msg.msg_iov = &mut iov; + msg.msg_iovlen = 1; + msg.msg_control = cmsg_buf.as_mut_ptr() as *mut libc::c_void; + msg.msg_controllen = cmsg_buf.len() as _; + + unsafe { + let n = libc::recvmsg(stream.as_raw_fd(), &mut msg, 0); + if n < 0 { + return Err(std::io::Error::last_os_error()).context("recvmsg SCM_RIGHTS"); + } + if n == 0 { + bail!("holder closed connection before sending fd"); + } + let cmsg = libc::CMSG_FIRSTHDR(&msg); + if cmsg.is_null() { + bail!("no ancillary data (fd) received from holder"); + } + if (*cmsg).cmsg_level != libc::SOL_SOCKET || (*cmsg).cmsg_type != libc::SCM_RIGHTS { + bail!("unexpected ancillary message from holder"); + } + let mut fd: RawFd = -1; + std::ptr::copy_nonoverlapping( + libc::CMSG_DATA(cmsg), + &mut fd as *mut RawFd as *mut u8, + std::mem::size_of::(), + ); + if fd < 0 { + bail!("invalid fd received from holder"); + } + Ok(fd) + } +} + +/// Space, in bytes, needed for a control-message buffer carrying exactly one fd. +const fn cmsg_space_one_fd() -> usize { + // CMSG_SPACE is not const in libc; this is the equivalent for one RawFd. + // cmsghdr is aligned to size_of::(); add data length rounded up. + let data = std::mem::size_of::(); + let hdr = std::mem::size_of::(); + let align = std::mem::size_of::(); + // round(hdr) + round(data) + ((hdr + align - 1) & !(align - 1)) + ((data + align - 1) & !(align - 1)) +} + +/// Helper used by `adopt_holder` callers to turn the received raw fd into an +/// owned `File`-like object if they need RAII (the server hands it to +/// `pty::adopt_pty_from_fd`, which takes ownership of the raw fd). +pub fn fd_into_file(fd: RawFd) -> std::fs::File { + unsafe { std::fs::File::from_raw_fd(fd) } +} + +/// Convert a `File` back into a raw fd it no longer owns (so it can be handed to +/// `adopt_pty_from_fd`). Currently unused outside tests but kept symmetric. +pub fn file_into_fd(file: std::fs::File) -> RawFd { + file.into_raw_fd() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn session_dir_name_round_trips() { + for name in ["default", "work", "a/b/../c", "weird name", "日本語"] { + let encoded = session_dir_name(name); + assert!(encoded.bytes().all(|b| b.is_ascii_hexdigit())); + assert_eq!(decode_session_dir_name(&encoded).as_deref(), Some(name)); + } + } + + #[test] + fn decode_rejects_non_hex() { + assert!(decode_session_dir_name("zz").is_none()); + assert!(decode_session_dir_name("abc").is_none()); + } + + #[test] + fn save_and_load_screen_round_trips() { + let tmp = tempfile::tempdir().unwrap(); + let sessions_dir = tmp.path(); + ensure_runtime_dir(sessions_dir, "work").unwrap(); + let snap = b"\x1b[?1049lhello world".to_vec(); + save_screen(sessions_dir, "work", 100, 40, 7, &snap).unwrap(); + let loaded = load_screen(sessions_dir, "work").expect("screen restored"); + assert_eq!(loaded.cols, 100); + assert_eq!(loaded.rows, 40); + assert_eq!(loaded.output_seq, 7); + assert_eq!(loaded.snapshot, snap); + } + + #[test] + fn load_screen_absent_is_none() { + let tmp = tempfile::tempdir().unwrap(); + assert!(load_screen(tmp.path(), "missing").is_none()); + } + + #[test] + fn scan_skips_dead_and_junk_entries() { + let tmp = tempfile::tempdir().unwrap(); + let sessions_dir = tmp.path(); + // A meta pointing at a definitely-dead pid is cleaned up, not returned. + let dir = ensure_runtime_dir(sessions_dir, "ghost").unwrap(); + write_meta( + &dir, + &HolderMeta { + session: "ghost".to_string(), + shell_pid: 2_000_000_000, // not a live pid + }, + ) + .unwrap(); + // A junk (non-hex) directory is ignored. + std::fs::create_dir_all(runtime_root(sessions_dir).join("not-hex")).unwrap(); + let found = scan_existing_holders(sessions_dir); + assert!(found.is_empty(), "dead/junk holders must be skipped"); + assert!(!dir.exists(), "dead holder dir should be cleaned up"); + } + + #[test] + fn send_and_recv_fd_round_trips() { + use std::io::Seek; + // Pass a temp file's fd across a socketpair and confirm both ends point + // at the same open file (write via one, read via the other). + let (mut a, mut b) = UnixStream::pair().unwrap(); + let mut file = tempfile::tempfile().unwrap(); + writeln!(file, "shared-fd-marker").unwrap(); + file.flush().unwrap(); + send_fd(&mut a, file.as_raw_fd()).unwrap(); + let received = recv_fd(&mut b).unwrap(); + let mut got = fd_into_file(received); + got.rewind().unwrap(); + let mut contents = String::new(); + got.read_to_string(&mut contents).unwrap(); + assert!(contents.contains("shared-fd-marker")); + } +} diff --git a/src/pty.rs b/src/pty.rs index 2f1361c..5a86237 100644 --- a/src/pty.rs +++ b/src/pty.rs @@ -1,15 +1,36 @@ use anyhow::{Context, Result}; use portable_pty::{Child, CommandBuilder, MasterPty, NativePtySystem, PtySize, PtySystem}; use std::io::{Read, Write}; +use std::os::unix::io::{AsRawFd, FromRawFd, RawFd}; use std::path::Path; use std::sync::{Arc, Mutex}; use std::thread; use tokio::sync::mpsc; +/// Backing for a PTY master held by the server. +/// +/// `Owned` means this process spawned the shell as a child and is responsible +/// for it: dropping the handle kills the shell. This is the original, +/// non-persistent model and stays the default. +/// +/// `Adopted` means the shell lives in a separate holder process and this handle +/// only borrows the master fd (received over a Unix socket via SCM_RIGHTS). +/// Dropping it must NOT kill the shell — it just closes our copy of the fd and +/// stops the reader thread, leaving the holder + shell alive so a server restart +/// can re-adopt them. +enum Backing { + Owned { + child: Mutex>, + _master: Box, + }, + Adopted { + master: Mutex, + }, +} + pub struct PtyHandle { writer: Arc>>, - child: Mutex>, - _master: Box, + backing: Backing, } impl PtyHandle { @@ -21,23 +42,40 @@ impl PtyHandle { } pub fn resize(&self, cols: u16, rows: u16) -> Result<()> { - self._master.resize(PtySize { - rows, - cols, - pixel_width: 0, - pixel_height: 0, - })?; + match &self.backing { + Backing::Owned { _master, .. } => { + _master.resize(PtySize { + rows, + cols, + pixel_width: 0, + pixel_height: 0, + })?; + } + Backing::Adopted { master } => { + let file = master.lock().expect("pty master poisoned"); + resize_fd(file.as_raw_fd(), cols, rows)?; + } + } Ok(()) } + /// True for a handle backed by a separate holder process. Such a handle must + /// be detached (not killed) when the server lets go of a session, so the + /// shell survives a server restart. + pub fn is_persistent(&self) -> bool { + matches!(self.backing, Backing::Adopted { .. }) + } + /// Terminate the shell process backing this PTY and reap it. /// - /// Without this the child shell outlives the session: dropping the master - /// alone is not guaranteed to take the process down, and the `Child` handle - /// used to be discarded at spawn time, which leaked one shell per - /// abandoned session. + /// Only meaningful for an `Owned` backing (the server spawned the shell as a + /// child). For an `Adopted` backing the shell belongs to the holder process, + /// so this is a no-op here; reaping a persistent session is done by asking + /// the holder to shut down (see `persist::request_shutdown`). pub fn kill(&self) { - if let Ok(mut child) = self.child.lock() { + if let Backing::Owned { child, .. } = &self.backing + && let Ok(mut child) = child.lock() + { let _ = child.kill(); let _ = child.wait(); } @@ -46,10 +84,29 @@ impl PtyHandle { impl Drop for PtyHandle { fn drop(&mut self) { - self.kill(); + // Adopted handles must NOT kill the shell: it lives in the holder so it + // can outlive this server. Dropping just closes our fd / stops the + // reader. Owned handles keep the original kill-on-drop behavior. + if let Backing::Owned { .. } = &self.backing { + self.kill(); + } } } +fn resize_fd(fd: RawFd, cols: u16, rows: u16) -> Result<()> { + let winsize = libc::winsize { + ws_row: rows, + ws_col: cols, + ws_xpixel: 0, + ws_ypixel: 0, + }; + let rc = unsafe { libc::ioctl(fd, libc::TIOCSWINSZ, &winsize) }; + if rc != 0 { + return Err(std::io::Error::last_os_error()).context("TIOCSWINSZ"); + } + Ok(()) +} + #[derive(Debug)] pub struct PtyOutput { pub session: String, @@ -74,6 +131,27 @@ pub fn spawn_pty_session( pixel_height: 0, }) .context("open pty")?; + let cmd = build_shell_command(shell, env); + let child = pair.slave.spawn_command(cmd).context("spawn shell")?; + drop(pair.slave); + + let writer = pair.master.take_writer().context("take pty writer")?; + let reader = pair.master.try_clone_reader().context("clone pty reader")?; + spawn_reader_thread(session, reader, tx)?; + + Ok(PtyHandle { + writer: Arc::new(Mutex::new(writer)), + backing: Backing::Owned { + child: Mutex::new(child), + _master: pair.master, + }, + }) +} + +/// Build the [`CommandBuilder`] for a dosh shell, identically for the in-process +/// `spawn_pty_session` and the out-of-process holder, so a persistent session's +/// environment matches a non-persistent one. +pub fn build_shell_command(shell: &str, env: &[(String, String)]) -> CommandBuilder { let mut cmd = CommandBuilder::new(shell); cmd.env("TERM", "xterm-256color"); cmd.env("COLORTERM", "truecolor"); @@ -88,11 +166,39 @@ pub fn spawn_pty_session( } else if let Some(parent) = Path::new(shell).parent() { cmd.env("PWD", parent.as_os_str()); } - let child = pair.slave.spawn_command(cmd).context("spawn shell")?; - drop(pair.slave); + cmd +} - let writer = pair.master.take_writer().context("take pty writer")?; - let mut reader = pair.master.try_clone_reader().context("clone pty reader")?; +/// Build a [`PtyHandle`] from a master fd received from a holder process. +/// +/// `master_fd` is an fd this handle takes ownership of (it is wrapped in a +/// `File` and closed on drop). The shell is NOT a child of this process; it +/// belongs to the holder, so dropping this handle leaves it running. +pub fn adopt_pty_from_fd( + session: String, + master_fd: RawFd, + tx: mpsc::UnboundedSender, +) -> Result { + // Take ownership of the fd. A clone gives us an independent reader so the + // reader thread and the writer/resize side hold separate `File`s and don't + // get closed out from under each other. + let master = unsafe { std::fs::File::from_raw_fd(master_fd) }; + let reader_file = master.try_clone().context("clone master for reader")?; + let writer_file = master.try_clone().context("clone master for writer")?; + spawn_reader_thread(session, Box::new(reader_file), tx)?; + Ok(PtyHandle { + writer: Arc::new(Mutex::new(Box::new(writer_file))), + backing: Backing::Adopted { + master: Mutex::new(master), + }, + }) +} + +fn spawn_reader_thread( + session: String, + mut reader: Box, + tx: mpsc::UnboundedSender, +) -> Result<()> { let reader_session = session.clone(); thread::Builder::new() .name(format!("dosh-pty-{session}")) @@ -127,10 +233,5 @@ pub fn spawn_pty_session( } }) .context("spawn pty reader")?; - - Ok(PtyHandle { - writer: Arc::new(Mutex::new(writer)), - child: Mutex::new(child), - _master: pair.master, - }) + Ok(()) } diff --git a/tests/hostile_network.rs b/tests/hostile_network.rs index 2d8b6ea..206fa68 100644 --- a/tests/hostile_network.rs +++ b/tests/hostile_network.rs @@ -67,6 +67,7 @@ sessions_dir = "{sessions}" secret_path = "{secret}" host_key = "{host_key}" authorized_keys = ["{authorized_keys}"] +persist_sessions = false "#, sessions = dir.path().join("sessions").display(), secret = dir.path().join("secret").display(), diff --git a/tests/integration_smoke.rs b/tests/integration_smoke.rs index f3e51fc..b8709d5 100644 --- a/tests/integration_smoke.rs +++ b/tests/integration_smoke.rs @@ -83,6 +83,7 @@ sessions_dir = "{sessions}" secret_path = "{secret}" host_key = "{host_key}" authorized_keys = ["{authorized_keys}"] +persist_sessions = false "#, sessions = dir.path().join("sessions").display(), secret = dir.path().join("secret").display(), @@ -1800,3 +1801,185 @@ fn native_agent_forwarding_requires_client_opt_in() { "agent proxy dir {agent_dir:?} must not exist without client -A opt-in" ); } + +// --------------------------------------------------------------------------- +// Session persistence across a server restart. +// --------------------------------------------------------------------------- + +/// Write a server config with `persist_sessions = true` so each session's shell +/// runs in a detached holder that survives a server restart. +fn write_persistent_server_config(dir: &tempfile::TempDir, port: u16) -> std::path::PathBuf { + let config = write_server_config(dir, port); + let mut raw = fs::read_to_string(&config).unwrap(); + // The base writer hardcodes `persist_sessions = false`; flip it on. + raw = raw.replace("persist_sessions = false", "persist_sessions = true"); + fs::write(&config, raw).unwrap(); + config +} + +/// Mirror the server's hex encoding of a session name to its runtime dir. +fn session_runtime_dir(sessions_dir: &Path, session: &str) -> PathBuf { + let mut hex = String::new(); + for b in session.as_bytes() { + hex.push_str(&format!("{b:02x}")); + } + sessions_dir.join("run").join(hex) +} + +/// Read a holder's recorded shell pid from its meta file, if present. +fn holder_shell_pid(sessions_dir: &Path, session: &str) -> Option { + let meta = session_runtime_dir(sessions_dir, session).join("meta"); + let body = fs::read_to_string(meta).ok()?; + body.lines().nth(1)?.parse().ok() +} + +/// Tear down a persistent session's holder + shell so a test never leaks them +/// (the TempDir drop would otherwise leave detached processes running). +fn kill_holder(sessions_dir: &Path, session: &str) { + if let Some(pid) = holder_shell_pid(sessions_dir, session) { + let _ = Command::new("kill").arg("-9").arg(pid.to_string()).status(); + } + // Best-effort: also kill any matching holder process for this runtime dir. + let dir = session_runtime_dir(sessions_dir, session); + let _ = Command::new("pkill") + .arg("-9") + .arg("-f") + .arg(format!("dosh-server hold --runtime-dir {}", dir.display())) + .status(); +} + +/// Send one encrypted Input line and give the shell a moment to act. +fn type_line(socket: &UdpSocket, port: u16, ok: &AttachOk, key: &[u8; 32], seq: u64, line: &str) { + let input = Input { + bytes: line.as_bytes().to_vec(), + }; + send_encrypted( + socket, + port, + PacketKind::Input, + ok.client_id, + seq, + 0, + key, + &protocol::to_body(&input).unwrap(), + ); + thread::sleep(Duration::from_millis(150)); +} + +/// Proof of survival: a session's shell + state + screen outlive a full +/// `dosh-server` process restart. +/// +/// We attach, set durable shell state (`cd /tmp`, `export MARK=...`) and paint a +/// recognizable marker on the screen, wait for the screen to be mirrored to disk, +/// then KILL the server process (simulating a crash/upgrade). We restart the +/// server on the same config/sockets, reattach, and assert we land on the SAME +/// shell — the exported variable and working directory are still there, and the +/// restored attach snapshot still shows the pre-restart screen — not a fresh +/// shell. +#[test] +fn session_survives_server_restart_same_shell_and_screen() { + let dir = tempfile::tempdir().unwrap(); + let sessions_dir = dir.path().join("sessions"); + let port = free_udp_port(); + let config = write_persistent_server_config(&dir, port); + + let mut server = start_server(&dir, &config); + + // First attach: establish durable shell state and a screen marker. + let (socket, bootstrap, ok) = direct_attach(&config, port, "read-write"); + let key = bootstrap.session_key; + type_line(&socket, port, &ok, &key, 2, "cd /tmp\n"); + type_line( + &socket, + port, + &ok, + &key, + 3, + "export MARK=zebra_persist_42\n", + ); + // Paint a unique, stable marker into the screen (printf, no trailing newline + // scroll surprises) so the restored snapshot is easy to recognize. + type_line( + &socket, + port, + &ok, + &key, + 4, + "printf 'PRE_RESTART_SCREEN_MARKER\\n'\n", + ); + // Drain output so the parser/screen state and disk mirror catch up. + let pre = collect_frame_text(&socket, &key, 1500); + assert!( + pre.contains("PRE_RESTART_SCREEN_MARKER"), + "expected to see the screen marker before restart, got {pre:?}" + ); + // The periodic flush mirrors the screen every ~2s; wait long enough that the + // pre-restart screen is guaranteed on disk before we crash the server. + thread::sleep(Duration::from_secs(3)); + + let shell_pid_before = holder_shell_pid(&sessions_dir, "default"); + assert!( + shell_pid_before.is_some(), + "a holder shell pid should be recorded for a persistent session" + ); + + // CRASH: kill the server process outright (its in-memory state is lost; only + // the holder + shell + on-disk screen remain). + let _ = server.kill(); + let _ = server.wait(); + thread::sleep(Duration::from_millis(300)); + + // The holder + shell must still be alive after the server died. + let pid = shell_pid_before.unwrap(); + let alive = unsafe { libc::kill(pid, 0) } == 0; + assert!( + alive, + "shell pid {pid} must survive the server crash (holder keeps it alive)" + ); + + // RESTART on the same config/sockets; the server re-adopts the holder. + let mut server = start_server(&dir, &config); + + // Reattach. A fresh bootstrap attach lands on the EXISTING re-adopted + // session (same name), not a new shell. + let (socket2, bootstrap2, ok2) = direct_attach(&config, port, "read-write"); + let key2 = bootstrap2.session_key; + + // 1) Screen survived: the restored attach snapshot still shows the marker. + let snapshot = String::from_utf8_lossy(&ok2.snapshot).to_string(); + + // 2) Shell identity + state survived: ask the SAME shell for its state. + type_line( + &socket2, + port, + &ok2, + &key2, + 2, + "printf 'MARK=%s PWD=%s\\n' \"$MARK\" \"$PWD\"\n", + ); + let post = collect_frame_text(&socket2, &key2, 2000); + + let shell_pid_after = holder_shell_pid(&sessions_dir, "default"); + + // Clean up BEFORE asserting so a failed assert never leaks the holder. + let _ = server.kill(); + let _ = server.wait(); + kill_holder(&sessions_dir, "default"); + + assert_eq!( + shell_pid_after, shell_pid_before, + "re-adopted session must be backed by the same shell pid, not a fresh one" + ); + assert!( + post.contains("MARK=zebra_persist_42"), + "exported variable must survive the restart (same shell), got {post:?}" + ); + assert!( + post.contains("PWD=/tmp"), + "working directory must survive the restart (same shell), got {post:?}" + ); + assert!( + snapshot.contains("PRE_RESTART_SCREEN_MARKER"), + "restored attach snapshot must repaint the pre-restart screen, got {snapshot:?}" + ); +}