diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index e430ee9..98a8068 100644 --- a/src/bin/dosh-client.rs +++ b/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,12 +677,15 @@ 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)?; - render_frame(&frame)?; - cred.last_rendered_seq = frame.output_seq; - send_ack(&socket, addr, &cred, &mut send_seq).await?; - if frame.closed { - break; + let frames = frame_buffer.accept(frame, &mut cred.last_rendered_seq); + for frame in frames { + render_frame(&frame)?; + if frame.closed { + send_ack(&socket, addr, &cred, &mut send_seq).await?; + return Ok(()); + } } + send_ack(&socket, addr, &cred, &mut send_seq).await?; } PacketKind::Pong => {} PacketKind::AttachReject => { @@ -712,6 +717,32 @@ async fn run_terminal( Ok(()) } +#[derive(Default)] +struct FrameBuffer { + pending: BTreeMap, +} + +impl FrameBuffer { + fn accept(&mut self, frame: Frame, last_rendered_seq: &mut u64) -> Vec { + 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 { 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()); + } }