Render Dosh frames in stream order

src/bin/dosh-client.rs

@@ -14,6 +14,7 @@ use dosh::protocol::{
     TicketAttachOkEnvelope,
 };
 use serde::{Deserialize, Serialize};
+use std::collections::BTreeMap;
 use std::fs;
 use std::io::{Read, Write};
 use std::net::{SocketAddr, ToSocketAddrs};
@@ -604,6 +605,7 @@ async fn run_terminal(
     let mut status_tick = tokio::time::interval(Duration::from_secs(1));
     let mut resize_tick = tokio::time::interval(Duration::from_millis(250));
     let mut last_size = size().unwrap_or((80, 24));
+    let mut frame_buffer = FrameBuffer::default();
     if let Some(frame) = first_frame {
         render_frame(&frame)?;
         if frame.closed {
@@ -675,13 +677,16 @@ async fn run_terminal(
                     PacketKind::Frame | PacketKind::ResumeOk => {
                         let plain = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT)?;
                         let frame: Frame = protocol::from_body(&plain)?;
+                        let frames = frame_buffer.accept(frame, &mut cred.last_rendered_seq);
+                        for frame in frames {
                             render_frame(&frame)?;
-                        cred.last_rendered_seq = frame.output_seq;
-                        send_ack(&socket, addr, &cred, &mut send_seq).await?;
                             if frame.closed {
-                            break;
+                                send_ack(&socket, addr, &cred, &mut send_seq).await?;
+                                return Ok(());
                             }
                         }
+                        send_ack(&socket, addr, &cred, &mut send_seq).await?;
+                    }
                     PacketKind::Pong => {}
                     PacketKind::AttachReject => {
                         let reject: AttachReject = protocol::from_body(&packet.body)?;
@@ -712,6 +717,32 @@ async fn run_terminal(
     Ok(())
 }
 
+#[derive(Default)]
+struct FrameBuffer {
+    pending: BTreeMap<u64, Frame>,
+}
+
+impl FrameBuffer {
+    fn accept(&mut self, frame: Frame, last_rendered_seq: &mut u64) -> Vec<Frame> {
+        if frame.snapshot {
+            self.pending.clear();
+            *last_rendered_seq = frame.output_seq;
+            return vec![frame];
+        }
+        if frame.output_seq <= *last_rendered_seq {
+            return Vec::new();
+        }
+        self.pending.entry(frame.output_seq).or_insert(frame);
+
+        let mut ready = Vec::new();
+        while let Some(frame) = self.pending.remove(&(*last_rendered_seq + 1)) {
+            *last_rendered_seq = frame.output_seq;
+            ready.push(frame);
+        }
+        ready
+    }
+}
+
 async fn send_resize(
     socket: &UdpSocket,
     addr: SocketAddr,
@@ -803,9 +834,6 @@ struct RawMode;
 impl RawMode {
     fn enter() -> Result<Self> {
         enable_raw_mode()?;
-        let mut stdout = std::io::stdout();
-        stdout.write_all(TERMINAL_CLEANUP)?;
-        stdout.flush()?;
         Ok(Self)
     }
 }
@@ -837,7 +865,8 @@ const TERMINAL_CLEANUP: &[u8] = concat!(
 
 #[cfg(test)]
 mod tests {
-    use super::{parse_ssh_config_hostname, ssh_destination_host};
+    use super::{FrameBuffer, parse_ssh_config_hostname, ssh_destination_host};
+    use dosh::protocol::Frame;
 
     #[test]
     fn parses_ssh_config_hostname() {
@@ -861,4 +890,53 @@ mod tests {
     fn strips_user_from_fallback_udp_host() {
         assert_eq!(ssh_destination_host("palav@example.com"), "example.com");
     }
+
+    fn test_frame(seq: u64, snapshot: bool) -> Frame {
+        Frame {
+            session: "test".to_string(),
+            output_seq: seq,
+            bytes: format!("frame-{seq}").into_bytes(),
+            snapshot,
+            closed: false,
+        }
+    }
+
+    #[test]
+    fn frame_buffer_renders_only_contiguous_frames() {
+        let mut buffer = FrameBuffer::default();
+        let mut last = 10;
+
+        assert!(buffer.accept(test_frame(12, false), &mut last).is_empty());
+        assert_eq!(last, 10);
+
+        let ready = buffer.accept(test_frame(11, false), &mut last);
+        assert_eq!(ready.len(), 2);
+        assert_eq!(ready[0].output_seq, 11);
+        assert_eq!(ready[1].output_seq, 12);
+        assert_eq!(last, 12);
+    }
+
+    #[test]
+    fn frame_buffer_ignores_duplicate_frames() {
+        let mut buffer = FrameBuffer::default();
+        let mut last = 4;
+
+        assert_eq!(buffer.accept(test_frame(5, false), &mut last).len(), 1);
+        assert!(buffer.accept(test_frame(5, false), &mut last).is_empty());
+        assert_eq!(last, 5);
+    }
+
+    #[test]
+    fn frame_buffer_accepts_snapshot_as_complete_state() {
+        let mut buffer = FrameBuffer::default();
+        let mut last = 10;
+
+        assert!(buffer.accept(test_frame(12, false), &mut last).is_empty());
+        let ready = buffer.accept(test_frame(20, true), &mut last);
+
+        assert_eq!(ready.len(), 1);
+        assert_eq!(ready[0].output_seq, 20);
+        assert_eq!(last, 20);
+        assert!(buffer.accept(test_frame(12, false), &mut last).is_empty());
+    }
 }