Fix reconnect stale keys and TUI status rendering

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);

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<u8> {
     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<Mutex<ServerState>>) {
     }
 }
 
-fn find_client_key(
-    state: &Arc<Mutex<ServerState>>,
-    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<Mutex<ServerState>>,
     header: &protocol::Header,

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)
     );
 }