From a23f610f20508c8805872d307723225b78361420 Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:34:24 -0400 Subject: [PATCH] Keep generated sessions alive across updates --- src/bin/dosh-server.rs | 77 +++++++++++------ tests/integration_smoke.rs | 168 ++++++++++++++++++++++++++++++++----- 2 files changed, 198 insertions(+), 47 deletions(-) diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 86873ac..b65248b 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -52,6 +52,7 @@ const STREAM_INITIAL_WINDOW: usize = 1024 * 1024; /// have accumulated since the last mirror. Keeps the atomic write off the /// per-packet hot path while bounding how much fresh output a crash can lose. const SCREEN_PERSIST_BYTES: usize = 4096; +const IMPLICIT_EMPTY_SESSION_GRACE_SECS: u64 = 300; /// Sentinel `target_host` sent to the client on a server-initiated `StreamOpen` /// that carries an SSH-agent connection (the client splices it into its local @@ -512,15 +513,16 @@ impl ServerState { Ok(()) } - /// Whether a session named `name` should be backed by a persistent holder. + /// Whether a session should be backed by a persistent holder. /// - /// With persistence enabled, explicitly named sessions get a detached - /// holder. Implicit `term--` sessions are one-off terminals - /// created by `dosh host`; users cannot intentionally reattach to them by - /// name, so persisting them only leaves stale holders that can be - /// accidentally re-adopted after restarts. - fn should_persist_session(&self, name: &str) -> bool { - self.config.persist_sessions && !protocol::is_implicit_session_name(name) + /// With persistence enabled, every interactive session gets a detached holder + /// so an update/restart of `dosh-server` behaves like a transient network + /// break. Named sessions can sit empty for the normal client timeout. + /// Generated one-off sessions are also persisted, but cleanup gives them a + /// shorter empty grace period so invisible abandoned holders do not linger + /// indefinitely. + fn should_persist_session(&self, _name: &str) -> bool { + self.config.persist_sessions } /// Spawn (or, in the persistent case, spawn-and-adopt) a PTY for a session. @@ -608,20 +610,6 @@ impl ServerState { if self.sessions.contains_key(&name) { continue; } - if !self.should_persist_session(&name) { - eprintln!( - "discarding stale implicit persistent session {name} (shell pid {})", - meta.shell_pid - ); - persist::request_shutdown(&sessions_dir, &name, None); - if meta.shell_pid > 0 { - let _ = unsafe { libc::kill(meta.shell_pid, libc::SIGHUP) }; - let _ = unsafe { libc::kill(meta.shell_pid, libc::SIGTERM) }; - std::thread::sleep(Duration::from_millis(50)); - let _ = unsafe { libc::kill(meta.shell_pid, libc::SIGKILL) }; - } - continue; - } let (pty, control) = match self.adopt_persistent_pty(&name) { Ok(pair) => pair, Err(err) => { @@ -671,6 +659,7 @@ impl ServerState { fn insert_client(&mut self, session_name: &str, client_id: [u8; 16], client: ClientState) { if let Some(session) = self.sessions.get_mut(session_name) { session.clients.insert(client_id, client); + session.empty_since = None; self.client_index .insert(client_id, session_name.to_string()); } @@ -3842,9 +3831,9 @@ fn cleanup_disconnected_clients(state: &Arc>) { .iter() .filter(|(name, session)| { !prewarm.contains(name.as_str()) - && session - .empty_since - .is_some_and(|since| now.duration_since(since) >= timeout) + && session.empty_since.is_some_and(|since| { + now.duration_since(since) >= empty_session_timeout(name, timeout) + }) }) .map(|(name, _)| name.clone()) .collect(); @@ -3863,6 +3852,14 @@ fn cleanup_disconnected_clients(state: &Arc>) { } } +fn empty_session_timeout(name: &str, configured_timeout: Duration) -> Duration { + if protocol::is_implicit_session_name(name) { + configured_timeout.min(Duration::from_secs(IMPLICIT_EMPTY_SESSION_GRACE_SECS)) + } else { + configured_timeout + } +} + /// Select the client key (current or retained previous epoch) matching the /// packet header's `session_key_id`, so a packet sealed under the pre-rekey key /// during the grace window still decrypts while an expired/stale epoch is @@ -3909,7 +3906,7 @@ mod tests { use super::*; #[test] - fn persists_named_sessions_when_enabled() { + fn persists_terminal_sessions_when_enabled() { let (pty_tx, _rx) = mpsc::unbounded_channel(); let config = ServerConfig { persist_sessions: true, @@ -3919,7 +3916,7 @@ mod tests { let state = ServerState::new(config, [0u8; 32], pty_tx); assert!(state.should_persist_session("default")); // prewarmed assert!(state.should_persist_session("work")); // explicitly named - assert!(!state.should_persist_session("term-1781470634216-76685")); // one-off terminal + assert!(state.should_persist_session("term-1781470634216-76685")); // one-off update-survivable terminal } #[test] @@ -4074,6 +4071,12 @@ mod tests { Some((key, "work".to_string())) ); assert!(state.client_mut(&client_id).is_some()); + state.sessions.get_mut("work").unwrap().empty_since = Some(Instant::now()); + state.insert_client("work", [7u8; 16], test_client_state([9u8; 32])); + assert!( + state.sessions["work"].empty_since.is_none(), + "reattach must clear empty-since so cleanup cannot reap a live session" + ); // Removing purges both the session map and the index. assert!(state.remove_client_everywhere(&client_id)); @@ -4113,6 +4116,26 @@ mod tests { assert!(locked.lookup_client(&client_id).is_none()); } + #[test] + fn implicit_empty_session_timeout_is_bounded_for_update_reconnect() { + let configured = Duration::from_secs(2_592_000); + let implicit = protocol::generate_implicit_session_name(); + assert_eq!( + empty_session_timeout(&implicit, configured), + Duration::from_secs(IMPLICIT_EMPTY_SESSION_GRACE_SECS) + ); + assert_eq!( + empty_session_timeout("work", configured), + configured, + "named sessions keep the normal long timeout" + ); + assert_eq!( + empty_session_timeout(&implicit, Duration::from_secs(1)), + Duration::from_secs(1), + "tests/admins can still configure a shorter timeout" + ); + } + #[tokio::test] async fn forged_plaintext_detach_does_not_remove_client() { let (pty_tx, _pty_rx) = mpsc::unbounded_channel(); diff --git a/tests/integration_smoke.rs b/tests/integration_smoke.rs index f17a640..be40288 100644 --- a/tests/integration_smoke.rs +++ b/tests/integration_smoke.rs @@ -18,7 +18,8 @@ use dosh::crypto; use dosh::native::{host_public_key, load_or_create_host_key, ssh_ed25519_public_blob, trust_host}; use dosh::protocol::{ self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, - PacketKind, Resize, ResumeRequest, SERVER_TO_CLIENT, + PacketKind, Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody, TicketAttachEnvelope, + TicketAttachOkEnvelope, }; use ed25519_dalek::{Signer, SigningKey, VerifyingKey}; @@ -530,6 +531,75 @@ fn direct_attach_session( (socket, bootstrap, ok) } +fn ticket_attach_session( + socket: &UdpSocket, + port: u16, + bootstrap: &dosh::auth::BootstrapResponse, + session: &str, + mode: &str, +) -> AttachOk { + let client_nonce = crypto::random_12(); + let request_key = crypto::hkdf32( + &bootstrap.attach_ticket_psk, + &client_nonce, + b"dosh/ticket-attach-request/v1", + ) + .unwrap(); + let body = protocol::to_body(&TicketAttachBody { + session: session.to_string(), + mode: mode.to_string(), + cols: 80, + rows: 24, + requested_env: Vec::new(), + }) + .unwrap(); + let ciphertext = crypto::seal( + &request_key, + &client_nonce, + b"dosh-ticket-attach-request-v1", + &body, + ) + .unwrap(); + let envelope = TicketAttachEnvelope { + ticket: bootstrap.attach_ticket.clone(), + client_nonce, + ciphertext, + }; + let packet = protocol::encode_plain( + PacketKind::TicketAttachRequest, + [0u8; 16], + 1, + 0, + &protocol::to_body(&envelope).unwrap(), + ) + .unwrap(); + socket + .send_to(&packet, format!("127.0.0.1:{port}")) + .unwrap(); + let mut buf = [0u8; 65535]; + let (n, _) = socket.recv_from(&mut buf).unwrap(); + let packet = protocol::decode(&buf[..n]).unwrap(); + assert_eq!(packet.header.kind, PacketKind::AttachOk); + let envelope: TicketAttachOkEnvelope = protocol::from_body(&packet.body).unwrap(); + let mut salt = Vec::with_capacity(24); + salt.extend_from_slice(&client_nonce); + salt.extend_from_slice(&envelope.server_nonce); + let response_key = crypto::hkdf32( + &bootstrap.attach_ticket_psk, + &salt, + b"dosh/ticket-attach-ok/v1", + ) + .unwrap(); + let plain = crypto::open( + &response_key, + &envelope.server_nonce, + b"dosh-ticket-attach-ok-v1", + &envelope.ciphertext, + ) + .unwrap(); + protocol::from_body(&plain).unwrap() +} + #[allow(clippy::too_many_arguments)] fn send_encrypted( socket: &UdpSocket, @@ -2733,7 +2803,7 @@ fn session_survives_server_restart_same_shell_and_screen() { } #[test] -fn implicit_session_does_not_create_restart_holder() { +fn implicit_session_survives_update_restart_via_ticket_reattach() { let dir = tempfile::tempdir().unwrap(); let sessions_dir = dir.path().join("sessions"); let port = free_udp_port(); @@ -2743,21 +2813,34 @@ fn implicit_session_does_not_create_restart_holder() { let mut server = start_server(&dir, &config); let (socket, bootstrap, ok) = direct_attach_session(&config, port, "read-write", &session); let key = bootstrap.session_key; + type_line(&socket, port, &ok, &key, 2, "cd /tmp\n"); type_line( &socket, port, &ok, &key, - 2, + 3, "export IMPLICIT_MARK=implicit_restart_42\n", ); - let _ = collect_frame_text(&socket, &key, 1000); + type_line( + &socket, + port, + &ok, + &key, + 4, + "printf 'IMPLICIT_SCREEN_MARKER\\n'\n", + ); + let pre = collect_frame_text(&socket, &key, 1500); + assert!( + pre.contains("IMPLICIT_SCREEN_MARKER"), + "implicit marker missing before restart: {pre:?}" + ); thread::sleep(Duration::from_secs(3)); let shell_pid_before = holder_shell_pid(&sessions_dir, &session); assert!( - shell_pid_before.is_none(), - "implicit sessions must not get persistent holders when persistence is enabled" + shell_pid_before.is_some(), + "implicit sessions must get temporary holders so active clients survive updates" ); let _ = server.kill(); @@ -2765,34 +2848,78 @@ fn implicit_session_does_not_create_restart_holder() { thread::sleep(Duration::from_millis(300)); let mut server = start_server(&dir, &config); - let (socket2, bootstrap2, ok2) = direct_attach_session(&config, port, "read-write", &session); - let key2 = bootstrap2.session_key; + + let resume = ResumeRequest { + session: session.clone(), + last_rendered_seq: ok.initial_seq, + cols: 80, + rows: 24, + }; + send_encrypted( + &socket, + port, + PacketKind::ResumeRequest, + ok.client_id, + 5, + 0, + &key, + &protocol::to_body(&resume).unwrap(), + ); + let mut buf = [0u8; 65535]; + let deadline = std::time::Instant::now() + Duration::from_secs(2); + loop { + assert!( + std::time::Instant::now() < deadline, + "old live resume did not receive unknown-client reject" + ); + let (n, _) = socket.recv_from(&mut buf).unwrap(); + let packet = protocol::decode(&buf[..n]).unwrap(); + if packet.header.kind == PacketKind::Frame { + continue; + } + assert_eq!(packet.header.kind, PacketKind::AttachReject); + assert_eq!(packet.header.conn_id, ok.client_id); + break; + } + + let ok2 = ticket_attach_session(&socket, port, &bootstrap, &session, "read-write"); + let key2 = ok2.session_key; + let snapshot = String::from_utf8_lossy(&ok2.snapshot).to_string(); type_line( - &socket2, + &socket, port, &ok2, &key2, 2, - "printf 'IMPLICIT_MARK=%s\\n' \"$IMPLICIT_MARK\"\n", + "printf 'IMPLICIT_MARK=%s PWD=%s\\n' \"$IMPLICIT_MARK\" \"$PWD\"\n", ); - let post = collect_frame_text(&socket2, &key2, 2000); + let post = collect_frame_text(&socket, &key2, 2000); let shell_pid_after = holder_shell_pid(&sessions_dir, &session); let _ = server.kill(); let _ = server.wait(); + kill_holder(&sessions_dir, &session); assert!( - shell_pid_after.is_none(), - "implicit restart attach must stay ephemeral" + snapshot.contains("IMPLICIT_SCREEN_MARKER"), + "ticket reattach snapshot must repaint the exact pre-update screen, got {snapshot:?}" + ); + assert_eq!( + shell_pid_after, shell_pid_before, + "implicit update reconnect must use the same holder shell" ); assert!( - post.contains("IMPLICIT_MARK=") && !post.contains("implicit_restart_42"), - "implicit session must start fresh after server restart, got {post:?}" + post.contains("IMPLICIT_MARK=implicit_restart_42"), + "implicit session env must survive update reconnect, got {post:?}" + ); + assert!( + post.contains("PWD=/tmp"), + "implicit session cwd must survive update reconnect, got {post:?}" ); } #[test] -fn startup_discards_legacy_implicit_holders() { +fn startup_readopts_implicit_holders_for_update_reconnect() { let dir = tempfile::tempdir().unwrap(); let sessions_dir = dir.path().join("sessions"); let port = free_udp_port(); @@ -2812,13 +2939,14 @@ fn startup_discards_legacy_implicit_holders() { let _ = server.kill(); let _ = server.wait(); assert!( - holder_shell_pid(&sessions_dir, &session).is_none(), - "startup must remove old implicit holder metadata" + holder_shell_pid(&sessions_dir, &session).is_some(), + "startup must keep live implicit holders so active update reconnects can reattach" ); assert!( - unsafe { libc::kill(shell_pid, 0) } != 0, - "startup must shut down old implicit holder shell" + unsafe { libc::kill(shell_pid, 0) } == 0, + "startup must not shut down a live implicit holder before reconnect grace" ); + kill_holder(&sessions_dir, &session); } #[test]