From c56e8f9fe1a26df0fc3f5f6a82e173fc89f4b802 Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Fri, 12 Jun 2026 13:39:27 -0400 Subject: [PATCH] Add guarded local typing prediction --- src/bin/dosh-client.rs | 118 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 114 insertions(+), 4 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index c8de101..4de8ea0 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -41,6 +41,8 @@ struct Args { #[arg(long)] attach_only: bool, #[arg(long)] + predict: bool, + #[arg(long)] ssh_port: Option, #[arg(long)] ssh_auth_command: Option, @@ -138,7 +140,7 @@ async fn main() -> Result<()> { detach_once(&socket, &cred, 2).await?; return Ok(()); } - return run_terminal(socket, cred, Some(frame)).await; + return run_terminal(socket, cred, Some(frame), args.predict).await; } Err(err) => { log_debug( @@ -156,7 +158,7 @@ async fn main() -> Result<()> { detach_once(&socket, &cred, 2).await?; return Ok(()); } - return run_terminal(socket, cred, Some(frame)).await; + return run_terminal(socket, cred, Some(frame), args.predict).await; } Err(err) => { log_debug( @@ -213,7 +215,7 @@ async fn main() -> Result<()> { detach_once(&socket, &cred, 2).await?; return Ok(()); } - run_terminal(socket, cred, Some(first)).await + run_terminal(socket, cred, Some(first), args.predict).await } fn run_update(config: &dosh::config::ClientConfig) -> Result<()> { @@ -597,6 +599,7 @@ async fn run_terminal( socket: UdpSocket, mut cred: CachedCredential, first_frame: Option, + predict: bool, ) -> Result<()> { let _raw = RawMode::enter()?; let addr = resolve_addr(&cred.udp_host, cred.udp_port)?; @@ -606,8 +609,10 @@ async fn run_terminal( 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(); + let mut predictor = Predictor::new(predict && cred.mode != "view-only"); if let Some(frame) = first_frame { render_frame(&frame)?; + predictor.observe_output(&frame.bytes); if frame.closed { return Ok(()); } @@ -642,6 +647,7 @@ async fn run_terminal( break; } if cred.mode != "view-only" { + predictor.observe_input(&bytes)?; let body = protocol::to_body(&Input { bytes })?; let packet = protocol::encode_encrypted( PacketKind::Input, @@ -679,7 +685,9 @@ async fn run_terminal( let frame: Frame = protocol::from_body(&plain)?; let frames = frame_buffer.accept(frame, &mut cred.last_rendered_seq); for frame in frames { + predictor.clear_pending()?; render_frame(&frame)?; + predictor.observe_output(&frame.bytes); if frame.closed { send_ack(&socket, addr, &cred, &mut send_seq).await?; return Ok(()); @@ -717,6 +725,88 @@ async fn run_terminal( Ok(()) } +struct Predictor { + enabled: bool, + alternate_screen: bool, + pending: Vec, +} + +impl Predictor { + fn new(enabled: bool) -> Self { + Self { + enabled, + alternate_screen: false, + pending: Vec::new(), + } + } + + fn observe_input(&mut self, bytes: &[u8]) -> Result<()> { + if !self.enabled || self.alternate_screen { + return Ok(()); + } + if !Self::predictable(bytes) { + self.clear_pending()?; + return Ok(()); + } + self.pending.extend_from_slice(bytes); + self.draw_pending() + } + + fn observe_output(&mut self, bytes: &[u8]) { + if contains_bytes(bytes, b"\x1b[?1049h") + || contains_bytes(bytes, b"\x1b[?1047h") + || contains_bytes(bytes, b"\x1b[?47h") + { + self.alternate_screen = true; + self.pending.clear(); + } + if contains_bytes(bytes, b"\x1b[?1049l") + || contains_bytes(bytes, b"\x1b[?1047l") + || contains_bytes(bytes, b"\x1b[?47l") + { + self.alternate_screen = false; + self.pending.clear(); + } + } + + fn predictable(bytes: &[u8]) -> bool { + !bytes.is_empty() && bytes.iter().all(|byte| matches!(byte, 0x20..=0x7e)) + } + + fn draw_pending(&self) -> Result<()> { + if self.pending.is_empty() { + return Ok(()); + } + let mut stdout = std::io::stdout(); + stdout.write_all(b"\x1b7\x1b[4m")?; + stdout.write_all(&self.pending)?; + stdout.write_all(b"\x1b[24m\x1b8")?; + stdout.flush()?; + Ok(()) + } + + fn clear_pending(&mut self) -> Result<()> { + if self.pending.is_empty() { + return Ok(()); + } + let mut stdout = std::io::stdout(); + stdout.write_all(b"\x1b7")?; + for _ in 0..self.pending.len() { + stdout.write_all(b" ")?; + } + stdout.write_all(b"\x1b8")?; + stdout.flush()?; + self.pending.clear(); + Ok(()) + } +} + +fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool { + haystack + .windows(needle.len()) + .any(|window| window == needle) +} + #[derive(Default)] struct FrameBuffer { pending: BTreeMap, @@ -868,7 +958,7 @@ const TERMINAL_CLEANUP: &[u8] = concat!( #[cfg(test)] mod tests { - use super::{FrameBuffer, parse_ssh_config_hostname, ssh_destination_host}; + use super::{FrameBuffer, Predictor, parse_ssh_config_hostname, ssh_destination_host}; use dosh::protocol::Frame; #[test] @@ -951,4 +1041,24 @@ mod tests { assert!(buffer.accept(test_frame(19, true), &mut last).is_empty()); assert_eq!(last, 20); } + + #[test] + fn predictor_only_accepts_plain_printable_input() { + assert!(Predictor::predictable(b"hello world")); + assert!(!Predictor::predictable(b"")); + assert!(!Predictor::predictable(b"\n")); + assert!(!Predictor::predictable(b"\x1b[A")); + assert!(!Predictor::predictable("snowman: \u{2603}".as_bytes())); + } + + #[test] + fn predictor_tracks_alternate_screen_state() { + let mut predictor = Predictor::new(true); + + predictor.observe_output(b"\x1b[?1049h"); + assert!(predictor.alternate_screen); + + predictor.observe_output(b"\x1b[?1049l"); + assert!(!predictor.alternate_screen); + } }