Add guarded local typing prediction

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<u16>,
     #[arg(long)]
     ssh_auth_command: Option<String>,
@@ -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<Frame>,
+    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<u8>,
+}
+
+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<u64, Frame>,
@@ -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);
+    }
 }