Harden reconnect paths under UDP churn

src/bin/dosh-client.rs

@@ -36,7 +36,10 @@ use dosh::transport::{
     DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig,
     TransportEvent,
 };
-use dosh::udp::{is_transient_udp_error, is_transient_udp_send_error};
+use dosh::udp::{
+    is_transient_udp_error, is_transient_udp_send_error, recv_udp_retrying_transient,
+    send_udp_retrying_transient,
+};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 use std::collections::BTreeMap;
@@ -3789,14 +3792,11 @@ async fn try_native_auth(
             hello: hello.clone(),
         })?,
     )?;
-    socket.send_to(&packet, addr).await?;
+    let auth_timeout = Duration::from_millis(config.native_auth_timeout_ms.max(1));
+    send_udp_retrying_transient(socket, &packet, addr, auth_timeout).await?;
 
     let mut buf = vec![0u8; 65535];
-    let (n, _) = tokio::time::timeout(
-        Duration::from_millis(config.native_auth_timeout_ms.max(1)),
-        socket.recv_from(&mut buf),
-    )
-    .await??;
+    let (n, _) = recv_udp_retrying_transient(socket, &mut buf, auth_timeout).await?;
     let packet = protocol::decode(&buf[..n])?;
     if packet.header.kind != PacketKind::NativeServerHello {
         if packet.header.kind == PacketKind::AttachReject {
@@ -3859,13 +3859,9 @@ async fn try_native_auth(
         CLIENT_TO_SERVER,
         &auth_body,
     )?;
-    socket.send_to(&packet, addr).await?;
+    send_udp_retrying_transient(socket, &packet, addr, auth_timeout).await?;
 
-    let (n, _) = tokio::time::timeout(
-        Duration::from_millis(config.native_auth_timeout_ms.max(1)),
-        socket.recv_from(&mut buf),
-    )
-    .await??;
+    let (n, _) = recv_udp_retrying_transient(socket, &mut buf, auth_timeout).await?;
     let packet = protocol::decode(&buf[..n])?;
     if packet.header.kind != PacketKind::NativeAuthOk {
         if packet.header.kind == PacketKind::AttachReject {
@@ -3947,14 +3943,11 @@ async fn try_native_auth_check(
             hello: hello.clone(),
         })?,
     )?;
-    socket.send_to(&packet, addr).await?;
+    let auth_timeout = Duration::from_millis(config.native_auth_timeout_ms.max(1));
+    send_udp_retrying_transient(socket, &packet, addr, auth_timeout).await?;
 
     let mut buf = vec![0u8; 65535];
-    let (n, _) = tokio::time::timeout(
-        Duration::from_millis(config.native_auth_timeout_ms.max(1)),
-        socket.recv_from(&mut buf),
-    )
-    .await??;
+    let (n, _) = recv_udp_retrying_transient(socket, &mut buf, auth_timeout).await?;
     let packet = protocol::decode(&buf[..n])?;
     if packet.header.kind != PacketKind::NativeServerHello {
         if packet.header.kind == PacketKind::AttachReject {
@@ -4016,13 +4009,9 @@ async fn try_native_auth_check(
         CLIENT_TO_SERVER,
         &protocol::to_body(&NativeUserAuthBody { auth })?,
     )?;
-    socket.send_to(&packet, addr).await?;
+    send_udp_retrying_transient(socket, &packet, addr, auth_timeout).await?;
 
-    let (n, _) = tokio::time::timeout(
-        Duration::from_millis(config.native_auth_timeout_ms.max(1)),
-        socket.recv_from(&mut buf),
-    )
-    .await??;
+    let (n, _) = recv_udp_retrying_transient(socket, &mut buf, auth_timeout).await?;
     let packet = protocol::decode(&buf[..n])?;
     if packet.header.kind != PacketKind::NativeAuthCheckOk {
         if packet.header.kind == PacketKind::AttachReject {
@@ -4182,7 +4171,7 @@ async fn bootstrap_attach(
     let body = protocol::to_body(&req)?;
     let packet =
         protocol::encode_plain(PacketKind::BootstrapAttachRequest, [0u8; 16], 1, 0, &body)?;
-    socket.send_to(&packet, addr).await?;
+    send_udp_retrying_transient(socket, &packet, addr, Duration::from_secs(5)).await?;
     let expected_key_id = protocol::session_key_id(&bootstrap.session_key);
     let ok = recv_response_until(socket, Duration::from_secs(5), |packet| {
         if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] {
@@ -4262,7 +4251,7 @@ async fn try_ticket_attach(
         0,
         &protocol::to_body(&envelope)?,
     )?;
-    socket.send_to(&packet, addr).await?;
+    send_udp_retrying_transient(socket, &packet, addr, Duration::from_millis(700)).await?;
 
     let ok = recv_response_until(socket, Duration::from_millis(700), |packet| {
         if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] {
@@ -4338,7 +4327,7 @@ async fn try_resume_fresh_process(
         CLIENT_TO_SERVER,
         &body,
     )?;
-    socket.send_to(&packet, addr).await?;
+    send_udp_retrying_transient(socket, &packet, addr, Duration::from_millis(700)).await?;
     let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
         if !is_resume_response_for_client(&packet, cached.client_id) {
             return Ok(None);
@@ -4393,7 +4382,7 @@ async fn try_live_resume(
         &body,
     )?;
     *send_seq += 1;
-    socket.send_to(&packet, addr).await?;
+    send_udp_retrying_transient(socket, &packet, addr, Duration::from_millis(700)).await?;
     let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
         if !is_resume_response_for_client(&packet, cached.client_id) {
             return Ok(None);
@@ -4440,14 +4429,9 @@ where
             return Err(anyhow!("timed out waiting for server response"));
         }
         let remaining = wait - elapsed;
-        let (n, _) = match tokio::time::timeout(remaining, socket.recv_from(&mut buf))
+        let (n, _) = recv_udp_retrying_transient(socket, &mut buf, remaining)
             .await
-            .context("timed out waiting for server response")?
-        {
-            Ok(value) => value,
-            Err(err) if is_transient_udp_error(&err) => continue,
-            Err(err) => return Err(err.into()),
-        };
+            .context("timed out waiting for server response")?;
         let Ok(packet) = protocol::decode(&buf[..n]) else {
             continue;
         };

src/bin/dosh-server.rs

@@ -361,6 +361,7 @@ impl NativeAuthRateLimiter {
 struct Session {
     pty: Option<PtyHandle>,
     parser: vt100::Parser,
+    restored_screen: Option<RestoredScreenSnapshot>,
     clients: HashMap<[u8; 16], ClientState>,
     output_seq: u64,
     recent: VecDeque<Vec<u8>>,
@@ -387,6 +388,14 @@ struct Session {
     last_screen_persist_at: Instant,
 }
 
+#[derive(Clone)]
+struct RestoredScreenSnapshot {
+    cols: u16,
+    rows: u16,
+    output_seq: u64,
+    snapshot: Vec<u8>,
+}
+
 #[derive(Clone)]
 struct ClientState {
     endpoint: SocketAddr,
@@ -512,6 +521,7 @@ impl ServerState {
             Session {
                 pty,
                 parser: vt100::Parser::new(rows.max(1), cols.max(1), self.config.scrollback),
+                restored_screen: None,
                 clients: HashMap::new(),
                 output_seq: 0,
                 recent: VecDeque::with_capacity(self.config.scrollback),
@@ -641,15 +651,23 @@ impl ServerState {
                 .unwrap_or((80, 24));
             let mut parser = vt100::Parser::new(rows, cols, self.config.scrollback);
             let mut output_seq = 0;
+            let mut restored_screen = None;
             if let Some(saved) = saved {
                 parser.process(&saved.snapshot);
                 output_seq = saved.output_seq;
+                restored_screen = Some(RestoredScreenSnapshot {
+                    cols,
+                    rows,
+                    output_seq,
+                    snapshot: saved.snapshot,
+                });
             }
             self.sessions.insert(
                 name.clone(),
                 Session {
                     pty: Some(pty),
                     parser,
+                    restored_screen,
                     clients: HashMap::new(),
                     output_seq,
                     recent: VecDeque::with_capacity(self.config.scrollback),
@@ -762,6 +780,19 @@ fn apply_terminal_size(
     Ok(())
 }
 
+fn attach_snapshot(session: &Session, cols: u16, rows: u16) -> Vec<u8> {
+    let cols = cols.max(1);
+    let rows = rows.max(1);
+    if let Some(restored) = &session.restored_screen
+        && restored.output_seq == session.output_seq
+        && restored.cols == cols
+        && restored.rows == rows
+    {
+        return restored.snapshot.clone();
+    }
+    screen_snapshot(session.parser.screen())
+}
+
 fn mode_allows_terminal_updates(mode: &str) -> bool {
     mode != "view-only" && mode != "forward-only"
 }
@@ -1189,7 +1220,7 @@ async fn handle_native_user_auth(
                 apply_terminal_size(session.pty.as_ref(), &mut session.parser, cols, rows)?;
             }
             let client_id = crypto::random_16();
-            let snapshot = screen_snapshot(session.parser.screen());
+            let snapshot = attach_snapshot(session, cols, rows);
             let output_seq = session.output_seq;
             locked.insert_client(
                 &session_name,
@@ -1337,7 +1368,7 @@ async fn handle_bootstrap_attach(
                 )?;
             }
             let client_id = crypto::random_16();
-            let snapshot = screen_snapshot(session.parser.screen());
+            let snapshot = attach_snapshot(session, req.cols, req.rows);
             let output_seq = session.output_seq;
             let bootstrap_session = req.bootstrap.session.clone();
             locked.insert_client(
@@ -1473,7 +1504,7 @@ async fn handle_ticket_attach(
                 )?;
             }
             let client_id = crypto::random_16();
-            let snapshot = screen_snapshot(session.parser.screen());
+            let snapshot = attach_snapshot(session, req.cols, req.rows);
             let output_seq = session.output_seq;
             let ticket_session = req.session.clone();
             locked.insert_client(
@@ -1586,20 +1617,23 @@ async fn handle_resume(
             .sessions
             .get_mut(&req.session)
             .ok_or_else(|| anyhow!("unknown session"))?;
-        let client = session
-            .clients
-            .get_mut(&packet.header.conn_id)
-            .ok_or_else(|| anyhow!("unknown client"))?;
-        if !client.replay.accept(packet.header.seq) {
-            return Ok(());
-        }
-        client.endpoint = peer;
-        client.last_acked = req.last_rendered_seq;
-        client.cols = req.cols;
-        client.rows = req.rows;
-        client.last_seen = Instant::now();
-        client.send_seq += 1;
-        if mode_allows_terminal_updates(&client.mode) {
+        let (send_seq, allow_terminal_updates) = {
+            let client = session
+                .clients
+                .get_mut(&packet.header.conn_id)
+                .ok_or_else(|| anyhow!("unknown client"))?;
+            if !client.replay.accept(packet.header.seq) {
+                return Ok(());
+            }
+            client.endpoint = peer;
+            client.last_acked = req.last_rendered_seq;
+            client.cols = req.cols;
+            client.rows = req.rows;
+            client.last_seen = Instant::now();
+            client.send_seq += 1;
+            (client.send_seq, mode_allows_terminal_updates(&client.mode))
+        };
+        if allow_terminal_updates {
             apply_terminal_size(
                 session.pty.as_ref(),
                 &mut session.parser,
@@ -1607,8 +1641,8 @@ async fn handle_resume(
                 req.rows,
             )?;
         }
-        let snapshot = screen_snapshot(session.parser.screen());
-        (client.send_seq, session.output_seq, snapshot)
+        let snapshot = attach_snapshot(session, req.cols, req.rows);
+        (send_seq, session.output_seq, snapshot)
     };
     let frame = Frame {
         session: session_name,
@@ -3454,6 +3488,7 @@ async fn broadcast_output(
             .get_mut(&output.session)
             .ok_or_else(|| anyhow!("unknown session"))?;
         session.parser.process(&output.bytes);
+        session.restored_screen = None;
         session.output_seq += 1;
         let output_seq = session.output_seq;
         session.recent.push_back(output.bytes.clone());
@@ -4049,6 +4084,34 @@ mod tests {
         assert_eq!(ServerConfig::default().client_timeout_secs, 2_592_000);
     }
 
+    #[test]
+    fn attach_snapshot_prefers_restored_screen_for_same_size_reattach() {
+        let restored = b"\x1b[?1049lRESTORED_SCREEN_MARKER".to_vec();
+        let mut session = Session {
+            pty: None,
+            parser: vt100::Parser::new(24, 80, 100),
+            restored_screen: Some(RestoredScreenSnapshot {
+                cols: 80,
+                rows: 24,
+                output_seq: 7,
+                snapshot: restored.clone(),
+            }),
+            clients: HashMap::new(),
+            output_seq: 7,
+            recent: VecDeque::new(),
+            empty_since: Some(Instant::now()),
+            holder_control: None,
+            persistent: true,
+            bytes_since_persist: 0,
+            last_persisted_seq: 7,
+            last_screen_persist_at: Instant::now(),
+        };
+
+        session.parser.set_size(24, 80);
+        assert_eq!(attach_snapshot(&session, 80, 24), restored);
+        assert_ne!(attach_snapshot(&session, 100, 30), restored);
+    }
+
     #[tokio::test]
     async fn unknown_resume_reject_keeps_client_id() {
         let (pty_tx, _rx) = mpsc::unbounded_channel();
@@ -4399,6 +4462,7 @@ mod tests {
             Session {
                 pty: None,
                 parser: vt100::Parser::new(24, 80, 100),
+                restored_screen: None,
                 clients: HashMap::from([(client_id, client)]),
                 output_seq: 0,
                 recent: VecDeque::new(),
@@ -4463,6 +4527,7 @@ mod tests {
             Session {
                 pty: None,
                 parser: vt100::Parser::new(24, 80, 100),
+                restored_screen: None,
                 clients: HashMap::from([(client_id, client)]),
                 output_seq: 0,
                 recent: VecDeque::new(),
@@ -4539,6 +4604,7 @@ mod tests {
             Session {
                 pty: None,
                 parser: vt100::Parser::new(24, 80, 100),
+                restored_screen: None,
                 clients: HashMap::from([(client_id, client)]),
                 output_seq: 0,
                 recent: VecDeque::new(),
@@ -4679,6 +4745,7 @@ mod tests {
             Session {
                 pty: None,
                 parser: vt100::Parser::new(24, 80, 100),
+                restored_screen: None,
                 clients: HashMap::from([(
                     client_id,
                     ClientState {
@@ -4776,6 +4843,7 @@ mod tests {
             Session {
                 pty: None,
                 parser: vt100::Parser::new(24, 80, 100),
+                restored_screen: None,
                 clients: HashMap::from([(
                     client_id,
                     ClientState {