From 9d0396dbf918062bb802b195b2657d66aef7ae0f Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Sat, 27 Jun 2026 11:46:26 -0400 Subject: [PATCH] Fix reconnect stale keys and TUI status rendering --- src/bin/dosh-client.rs | 52 +++++++++++++++++++++++++++++++++----- src/bin/dosh-server.rs | 27 ++++++++------------ tests/integration_smoke.rs | 46 +++++++++++++++++++++++++++++++-- 3 files changed, 100 insertions(+), 25 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index d921b68..e13b27f 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -3480,7 +3480,11 @@ async fn run_terminal( // on how long the link has been silent (recomputed after any // reconnect attempt above may have reset `last_packet_at`). if !forward_only { - let action = disconnect_status.on_tick(last_packet_at.elapsed()); + let action = if predictor.alternate_screen { + disconnect_status.on_suppressed() + } else { + disconnect_status.on_tick(last_packet_at.elapsed()) + }; disconnect_status.apply(action)?; } } @@ -4854,8 +4858,10 @@ const DISCONNECT_STATUS_THRESHOLD_SECS: u64 = 2; /// When server packets stop arriving we paint one blue line on the terminal's /// top row — e.g. `[dosh] reconnecting — last contact 3s ago` — using /// save/restore cursor (ESC 7 / ESC 8) so the real application cursor never -/// moves and full-screen TUIs are not corrupted. The bar is refreshed on the -/// status timer tick and cleared the instant a packet resumes. +/// moves. The run loop suppresses the bar while the server is in the alternate +/// screen because a saved cursor does not make writing row 1 non-destructive for +/// full-screen TUIs. The bar is refreshed on the status timer tick and cleared +/// the instant a packet resumes. /// /// The decision logic and the rendered text are pure and unit-tested; only /// [`DisconnectStatus::emit`] / [`DisconnectStatus::emit_clear`] touch stdout. @@ -4878,6 +4884,8 @@ enum StatusAction { Show(String), /// Erase the status line (link recovered or feature disabled). Clear, + /// Drop tracked status state without touching the terminal. + Forget, } impl DisconnectStatus { @@ -4936,6 +4944,18 @@ impl DisconnectStatus { } } + /// Pure decision for periods where drawing the top-row status is unsafe + /// (alternate-screen TUIs). Forget any tracked bar without writing a clear + /// sequence, because clearing row 1 would corrupt the TUI just like showing + /// the bar would. + fn on_suppressed(&self) -> StatusAction { + if self.shown { + StatusAction::Forget + } else { + StatusAction::None + } + } + /// Record that an action was applied to the screen, updating internal state. fn applied(&mut self, action: &StatusAction) { match action { @@ -4947,6 +4967,10 @@ impl DisconnectStatus { self.shown = false; self.drawn_text.clear(); } + StatusAction::Forget => { + self.shown = false; + self.drawn_text.clear(); + } StatusAction::None => {} } } @@ -4956,6 +4980,7 @@ impl DisconnectStatus { match &action { StatusAction::Show(text) => self.emit(text)?, StatusAction::Clear => self.emit_clear()?, + StatusAction::Forget => {} StatusAction::None => {} } self.applied(&action); @@ -6143,9 +6168,24 @@ mod tests { assert_eq!(status.on_tick(Duration::from_secs(10)), StatusAction::Clear); } - /// The rendered overlay must be non-destructive: save cursor (ESC 7), park on - /// the top row, paint the blue reconnect bar, then restore cursor (ESC 8) so - /// the application's cursor and full-screen TUIs are never disturbed. + #[test] + fn disconnect_status_suppresses_without_touching_alternate_screen() { + let mut status = DisconnectStatus::new(true); + let show = status.on_tick(Duration::from_secs(3)); + status.applied(&show); + assert!(status.shown); + + let suppressed = status.on_suppressed(); + assert_eq!(suppressed, StatusAction::Forget); + status.applied(&suppressed); + assert!(!status.shown); + assert_eq!(status.drawn_text, ""); + } + + /// The rendered overlay must preserve cursor position: save cursor (ESC 7), + /// park on the top row, paint the blue reconnect bar, then restore cursor + /// (ESC 8). The run loop suppresses this renderer entirely in full-screen + /// alternate-screen TUIs. #[test] fn disconnect_status_overlay_is_save_restore_top_blue_bar() { let bytes = render_status_overlay("[dosh] reconnecting — last contact 3s ago", 24); diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 10f19cd..0f74072 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -674,8 +674,9 @@ impl ServerState { /// previous epoch's key must still decrypt instead of being dropped as fatal /// (spec: "stale encrypted packets ... ignored, not fatal"). We pick whichever /// of the current or retained-previous key matches the packet's - /// `session_key_id`, falling back to the current key when neither matches (so - /// a genuinely stale/forged packet still fails the AEAD check, not silently). + /// `session_key_id`. If neither key id matches, the client is on a stale + /// epoch and should be forced through the normal reconnect path instead of + /// waiting for a silent timeout. fn lookup_client_key_for( &self, client_id: &[u8; 16], @@ -692,7 +693,7 @@ impl ServerState { { return Some((previous, session_name.clone())); } - Some((client.session_key, session_name.clone())) + None } /// O(1) mutable access to a live client via the conn_id index. @@ -836,7 +837,7 @@ async fn handle_packet( | PacketKind::StreamEof | PacketKind::StreamWindowAdjust | PacketKind::RekeyAck - if find_client_key(state, &packet.header.conn_id).is_err() => + if find_client_decrypt_key(state, &packet.header).is_err() => { send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client").await } @@ -2822,6 +2823,7 @@ fn screen_snapshot(screen: &vt100::Screen) -> Vec { let mut bytes = Vec::new(); if screen.alternate_screen() { bytes.extend_from_slice(b"\x1b[?1049h"); + bytes.extend_from_slice(b"\x1b[H\x1b[2J"); } else { bytes.extend_from_slice(b"\x1b[?1049l"); } @@ -3172,19 +3174,10 @@ fn cleanup_disconnected_clients(state: &Arc>) { } } -fn find_client_key( - state: &Arc>, - client_id: &[u8; 16], -) -> Result<([u8; 32], String)> { - let locked = state.lock().expect("server state poisoned"); - locked - .lookup_client(client_id) - .ok_or_else(|| anyhow!("unknown client")) -} - -/// Like [`find_client_key`] but selects the key (current or retained previous -/// epoch) matching the packet header's `session_key_id`, so a packet sealed -/// under the pre-rekey key during the grace window still decrypts. +/// 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 +/// treated as an unknown client and triggers a client reconnect. fn find_client_decrypt_key( state: &Arc>, header: &protocol::Header, diff --git a/tests/integration_smoke.rs b/tests/integration_smoke.rs index 27cdf7d..bd19f0e 100644 --- a/tests/integration_smoke.rs +++ b/tests/integration_smoke.rs @@ -620,6 +620,48 @@ fn disconnected_client_times_out_and_gets_reject() { assert_eq!(reject.reason, "unknown client"); } +#[test] +fn known_client_with_stale_key_id_gets_reject_not_timeout() { + let dir = tempfile::tempdir().unwrap(); + let port = free_udp_port(); + let config = write_server_config(&dir, port); + let mut server = start_server(&dir, &config); + let (_bootstrap, ok) = { + let (_socket, bootstrap, ok) = direct_attach(&config, port, "read-write"); + (bootstrap, ok) + }; + let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + socket + .set_read_timeout(Some(Duration::from_secs(2))) + .unwrap(); + let stale_key = crypto::random_32(); + let packet = protocol::encode_encrypted( + PacketKind::Ping, + ok.client_id, + 99, + 0, + &stale_key, + CLIENT_TO_SERVER, + b"", + ) + .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(); + let reject: AttachReject = protocol::from_body(&packet.body).unwrap(); + + let _ = server.kill(); + let _ = server.wait(); + + assert_eq!(packet.header.kind, PacketKind::AttachReject); + assert_eq!(packet.header.conn_id, ok.client_id); + assert_eq!(reject.reason, "unknown client"); +} + #[test] fn local_attach_only_smoke() { let dir = tempfile::tempdir().unwrap(); @@ -1525,8 +1567,8 @@ fn resume_snapshot_preserves_alternate_screen_mode() { assert!(resume_frame.snapshot); assert!( - resume_frame.bytes.starts_with(b"\x1b[?1049h"), - "snapshot did not enter alternate screen first: {:?}", + resume_frame.bytes.starts_with(b"\x1b[?1049h\x1b[H\x1b[2J"), + "snapshot did not enter and clear alternate screen first: {:?}", String::from_utf8_lossy(&resume_frame.bytes) ); }