Add Mosh parity hardening

src/bin/dosh-server.rs

@@ -203,12 +203,15 @@ async fn serve(config_path: Option<std::path::PathBuf>) -> Result<()> {
     let retransmit_state = Arc::clone(&state);
     let retransmit_socket = Arc::clone(&socket);
     tokio::spawn(async move {
-        let mut interval = tokio::time::interval(Duration::from_millis(100));
+        let mut interval = tokio::time::interval(Duration::from_millis(20));
         loop {
             interval.tick().await;
             if let Err(err) = retransmit_pending(&retransmit_state, &retransmit_socket).await {
                 eprintln!("retransmit error: {err:#}");
             }
+            if let Err(err) = flush_paced_snapshots(&retransmit_state, &retransmit_socket).await {
+                eprintln!("paced snapshot error: {err:#}");
+            }
             if let Err(err) = maybe_rekey_clients(&retransmit_state, &retransmit_socket).await {
                 eprintln!("rekey error: {err:#}");
             }
@@ -370,6 +373,8 @@ struct ClientState {
     rows: u16,
     last_seen: Instant,
     pending: VecDeque<PendingFrame>,
+    last_frame_sent: Instant,
+    paced_snapshot_due: bool,
     allowed_forwardings: Vec<ForwardingRequest>,
     stream_writers: HashMap<u64, mpsc::Sender<Vec<u8>>>,
     opened_streams: HashSet<u64>,
@@ -1122,6 +1127,8 @@ async fn handle_native_user_auth(
                     rows,
                     last_seen: Instant::now(),
                     pending: VecDeque::new(),
+                    last_frame_sent: Instant::now() - Duration::from_secs(3600),
+                    paced_snapshot_due: false,
                     allowed_forwardings: req.auth.requested_forwardings.clone(),
                     stream_writers: HashMap::new(),
                     opened_streams: HashSet::new(),
@@ -1271,6 +1278,8 @@ async fn handle_bootstrap_attach(
                 rows: req.rows,
                 last_seen: Instant::now(),
                 pending: VecDeque::new(),
+                last_frame_sent: Instant::now() - Duration::from_secs(3600),
+                paced_snapshot_due: false,
                 allowed_forwardings: Vec::new(),
                 stream_writers: HashMap::new(),
                 opened_streams: HashSet::new(),
@@ -1394,6 +1403,8 @@ async fn handle_ticket_attach(
                 rows: req.rows,
                 last_seen: Instant::now(),
                 pending: VecDeque::new(),
+                last_frame_sent: Instant::now() - Duration::from_secs(3600),
+                paced_snapshot_due: false,
                 allowed_forwardings: Vec::new(),
                 stream_writers: HashMap::new(),
                 opened_streams: HashSet::new(),
@@ -2630,6 +2641,8 @@ async fn broadcast_output(
         let mut locked = state.lock().expect("server state poisoned");
         let scrollback = locked.config.scrollback;
         let retransmit_window = locked.config.retransmit_window;
+        let frame_interval = Duration::from_millis(locked.config.output_frame_interval_ms);
+        let now = Instant::now();
         let session = locked
             .sessions
             .get_mut(&output.session)
@@ -2664,6 +2677,15 @@ async fn broadcast_output(
         }
         let mut sends = Vec::new();
         for (client_id, client) in session.clients.iter_mut() {
+            let terminal_control =
+                output.bytes.contains(&0x1b) || session.parser.screen().alternate_screen();
+            if !terminal_control
+                && !frame_interval.is_zero()
+                && now.duration_since(client.last_frame_sent) < frame_interval
+            {
+                client.paced_snapshot_due = true;
+                continue;
+            }
             client.send_seq += 1;
             // Count traffic toward the packet-count rekey trigger (spec §11).
             client.epoch_packets = client.epoch_packets.saturating_add(1);
@@ -2702,6 +2724,8 @@ async fn broadcast_output(
                 last_sent: Instant::now(),
                 attempts: 0,
             });
+            client.last_frame_sent = now;
+            client.paced_snapshot_due = false;
             sends.push((client.endpoint, packet));
         }
         sends
@@ -2852,6 +2876,76 @@ async fn retransmit_pending(
     Ok(())
 }
 
+async fn flush_paced_snapshots(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+) -> Result<()> {
+    let sends = {
+        let mut locked = state.lock().expect("server state poisoned");
+        let frame_interval = Duration::from_millis(locked.config.output_frame_interval_ms);
+        if frame_interval.is_zero() {
+            return Ok(());
+        }
+        let retransmit_window = locked.config.retransmit_window;
+        let now = Instant::now();
+        let mut sends = Vec::new();
+        for (session_name, session) in locked.sessions.iter_mut() {
+            let due = session.clients.values().any(|client| {
+                client.paced_snapshot_due
+                    && now.duration_since(client.last_frame_sent) >= frame_interval
+            });
+            if !due {
+                continue;
+            }
+            let snapshot = screen_snapshot(session.parser.screen());
+            let output_seq = session.output_seq;
+            for (client_id, client) in session.clients.iter_mut() {
+                if !client.paced_snapshot_due
+                    || now.duration_since(client.last_frame_sent) < frame_interval
+                {
+                    continue;
+                }
+                client.send_seq += 1;
+                client.epoch_packets = client.epoch_packets.saturating_add(1);
+                let frame = Frame {
+                    session: session_name.clone(),
+                    output_seq,
+                    bytes: snapshot.clone(),
+                    snapshot: true,
+                    closed: false,
+                };
+                let body = protocol::to_body(&frame)?;
+                let packet = protocol::encode_encrypted(
+                    PacketKind::Frame,
+                    *client_id,
+                    client.send_seq,
+                    client.last_acked,
+                    &client.session_key,
+                    SERVER_TO_CLIENT,
+                    &body,
+                )?;
+                while client.pending.len() >= retransmit_window {
+                    client.pending.pop_front();
+                }
+                client.pending.push_back(PendingFrame {
+                    output_seq,
+                    packet: packet.clone(),
+                    last_sent: now,
+                    attempts: 0,
+                });
+                client.last_frame_sent = now;
+                client.paced_snapshot_due = false;
+                sends.push((client.endpoint, packet));
+            }
+        }
+        sends
+    };
+    for (endpoint, packet) in sends {
+        socket.send_to(&packet, endpoint).await?;
+    }
+    Ok(())
+}
+
 /// How long the previous epoch's key is retained after a rekey so in-flight
 /// pre-rekey packets still decrypt instead of being dropped as fatal.
 const REKEY_PREVIOUS_KEY_GRACE_SECS: u64 = 5;
@@ -3185,6 +3279,11 @@ mod tests {
         assert!(limiter.buckets.is_empty());
     }
 
+    #[test]
+    fn default_client_timeout_supports_long_sleep() {
+        assert_eq!(ServerConfig::default().client_timeout_secs, 2_592_000);
+    }
+
     fn test_client_state(session_key: [u8; 32]) -> ClientState {
         ClientState {
             endpoint: "127.0.0.1:9".parse().unwrap(),
@@ -3197,6 +3296,8 @@ mod tests {
             rows: 24,
             last_seen: Instant::now(),
             pending: VecDeque::new(),
+            last_frame_sent: Instant::now() - Duration::from_secs(3600),
+            paced_snapshot_due: false,
             allowed_forwardings: Vec::new(),
             stream_writers: HashMap::new(),
             opened_streams: HashSet::new(),
@@ -3397,6 +3498,8 @@ mod tests {
                         rows: 24,
                         last_seen: Instant::now(),
                         pending: VecDeque::new(),
+                        last_frame_sent: Instant::now() - Duration::from_secs(3600),
+                        paced_snapshot_due: false,
                         allowed_forwardings: Vec::new(),
                         stream_writers: HashMap::new(),
                         opened_streams: HashSet::new(),
@@ -3490,6 +3593,8 @@ mod tests {
                         rows: 24,
                         last_seen: Instant::now(),
                         pending: VecDeque::new(),
+                        last_frame_sent: Instant::now() - Duration::from_secs(3600),
+                        paced_snapshot_due: false,
                         allowed_forwardings: Vec::new(),
                         stream_writers: HashMap::new(),
                         opened_streams: HashSet::from([42]),
@@ -3557,6 +3662,89 @@ mod tests {
         assert_eq!(frame.bytes, b"terminal-priority");
     }
 
+    #[tokio::test]
+    async fn output_pacing_coalesces_fast_frames_into_snapshot() {
+        let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
+        let config = ServerConfig {
+            output_frame_interval_ms: 1000,
+            ..ServerConfig::default()
+        };
+        let mut state = ServerState::new(config, [0u8; 32], pty_tx);
+        state
+            .ensure_session("test", 80, 24, "forward-only", &[])
+            .unwrap();
+        let client_id = [42u8; 16];
+        let session_key = [7u8; 32];
+        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+        let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
+        let mut client = test_client_state(session_key);
+        client.endpoint = receiver.local_addr().unwrap();
+        state.insert_client("test", client_id, client);
+        let state = Arc::new(Mutex::new(state));
+
+        broadcast_output(
+            &state,
+            &sender,
+            PtyOutput {
+                session: "test".to_string(),
+                bytes: b"first".to_vec(),
+                exited: false,
+            },
+        )
+        .await
+        .unwrap();
+        broadcast_output(
+            &state,
+            &sender,
+            PtyOutput {
+                session: "test".to_string(),
+                bytes: b"second".to_vec(),
+                exited: false,
+            },
+        )
+        .await
+        .unwrap();
+
+        let mut buf = [0u8; 4096];
+        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
+            .await
+            .unwrap()
+            .unwrap();
+        let packet = protocol::decode(&buf[..n]).unwrap();
+        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
+        let frame: Frame = protocol::from_body(&plain).unwrap();
+        assert!(!frame.snapshot);
+        assert_eq!(frame.bytes, b"first");
+        assert!(
+            tokio::time::timeout(Duration::from_millis(30), receiver.recv_from(&mut buf))
+                .await
+                .is_err(),
+            "second frame should be coalesced"
+        );
+
+        {
+            let mut locked = state.lock().unwrap();
+            let client = locked
+                .sessions
+                .get_mut("test")
+                .unwrap()
+                .clients
+                .get_mut(&client_id)
+                .unwrap();
+            client.last_frame_sent = Instant::now() - Duration::from_secs(2);
+        }
+        flush_paced_snapshots(&state, &sender).await.unwrap();
+        let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
+            .await
+            .unwrap()
+            .unwrap();
+        let packet = protocol::decode(&buf[..n]).unwrap();
+        let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
+        let frame: Frame = protocol::from_body(&plain).unwrap();
+        assert!(frame.snapshot);
+        assert_eq!(frame.output_seq, 2);
+    }
+
     #[test]
     fn remote_non_loopback_bind_requires_explicit_config() {
         let config = ServerConfig::default();