Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 46 Packages Projects Releases 51 Wiki Activity

Merge persistence: server-restart-survivable sessions (holder + SCM_RIGHTS fd passing)

Browse Source
This commit is contained in:
DuProcess
2026-06-14 14:21:17 -04:00
parent a88d912d2b 8b1af51bc6
commit 14b7e75025
11 changed files with 1347 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dosh-server.rs
+358 -30
View File
@@ -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<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>,
},
}
#[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<std::path::PathBuf>) -> 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<std::path::PathBuf>) -> 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<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)]
@@ -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<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) {
@@ -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<u8>)> = 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<UdpSocket>,
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<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.
let reaped: Vec<Session> = {
// 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<Mutex<ServerState>>) {
.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<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::*;
@@ -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`.