Keep generated sessions alive across updates

This commit is contained in:
DuProcess
2026-07-07 15:34:24 -04:00
parent ff2b7fff2b
commit a23f610f20
2 changed files with 198 additions and 47 deletions
+148 -20
View File
@@ -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]