Improve terminal input and client cleanup

src/bin/dosh-client.rs

@@ -555,7 +555,10 @@ 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 remote_terminal = vt100::Parser::new(last_size.1.max(1), last_size.0.max(1), 0);
+    let mut input_translator = UserInputTranslator::default();
     if let Some(frame) = first_frame {
+        remote_terminal.process(&frame.bytes);
         render_frame(&frame)?;
         if frame.closed {
             return Ok(());
@@ -584,23 +587,32 @@ async fn run_terminal(
     let mut recv_buf = vec![0u8; 65535];
     loop {
         tokio::select! {
-            Some(bytes) = stdin_rx.recv() => {
-                if bytes == [0x1d] {
-                    break;
-                }
-                if cred.mode != "view-only" {
-                    let body = protocol::to_body(&Input { bytes })?;
-                    let packet = protocol::encode_encrypted(
-                        PacketKind::Input,
-                        cred.client_id,
-                        send_seq,
-                        cred.last_rendered_seq,
-                        &cred.session_key,
-                        CLIENT_TO_SERVER,
-                        &body,
-                    )?;
-                    send_seq += 1;
-                    socket.send_to(&packet, addr).await?;
+            stdin_msg = stdin_rx.recv() => {
+                match stdin_msg {
+                    Some(bytes) => {
+                        if bytes == [0x1d] {
+                            break;
+                        }
+                        if cred.mode != "view-only" {
+                            let bytes = input_translator.translate(
+                                &bytes,
+                                remote_terminal.screen().application_cursor(),
+                            );
+                            let body = protocol::to_body(&Input { bytes })?;
+                            let packet = protocol::encode_encrypted(
+                                PacketKind::Input,
+                                cred.client_id,
+                                send_seq,
+                                cred.last_rendered_seq,
+                                &cred.session_key,
+                                CLIENT_TO_SERVER,
+                                &body,
+                            )?;
+                            send_seq += 1;
+                            socket.send_to(&packet, addr).await?;
+                        }
+                    }
+                    None => break,
                 }
             }
             _ = resize_tick.tick() => {
@@ -608,6 +620,7 @@ async fn run_terminal(
                     && (cols, rows) != last_size
                 {
                     last_size = (cols, rows);
+                    remote_terminal.set_size(rows.max(1), cols.max(1));
                     if cred.mode != "view-only" {
                         send_resize(&socket, addr, &cred, &mut send_seq, cols, rows).await?;
                     }
@@ -625,6 +638,7 @@ 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)?;
+                        remote_terminal.process(&frame.bytes);
                         render_frame(&frame)?;
                         cred.last_rendered_seq = frame.output_seq;
                         send_ack(&socket, addr, &cred, &mut send_seq).await?;
@@ -664,9 +678,58 @@ async fn run_terminal(
             }
         }
     }
+    let _ = detach_once(&socket, &cred, send_seq).await;
     Ok(())
 }
 
+#[derive(Default)]
+struct UserInputTranslator {
+    state: UserInputState,
+}
+
+#[derive(Default)]
+enum UserInputState {
+    #[default]
+    Ground,
+    Esc,
+    Ss3,
+}
+
+impl UserInputTranslator {
+    fn translate(&mut self, bytes: &[u8], application_cursor: bool) -> Vec<u8> {
+        let mut out = Vec::with_capacity(bytes.len());
+        for &byte in bytes {
+            match self.state {
+                UserInputState::Ground => {
+                    if byte == 0x1b {
+                        self.state = UserInputState::Esc;
+                    }
+                    out.push(byte);
+                }
+                UserInputState::Esc => {
+                    if byte == b'O' {
+                        self.state = UserInputState::Ss3;
+                    } else {
+                        self.state = UserInputState::Ground;
+                        out.push(byte);
+                    }
+                }
+                UserInputState::Ss3 => {
+                    self.state = UserInputState::Ground;
+                    if !application_cursor && matches!(byte, b'A'..=b'D') {
+                        out.push(b'[');
+                        out.push(byte);
+                    } else {
+                        out.push(b'O');
+                        out.push(byte);
+                    }
+                }
+            }
+        }
+        out
+    }
+}
+
 async fn send_resize(
     socket: &UdpSocket,
     addr: SocketAddr,
@@ -824,3 +887,31 @@ const TERMINAL_CLEANUP: &[u8] = concat!(
     "\x1b[?1049l"
 )
 .as_bytes();
+
+#[cfg(test)]
+mod tests {
+    use super::UserInputTranslator;
+
+    #[test]
+    fn translates_application_cursor_arrows_for_normal_remote_mode() {
+        let mut translator = UserInputTranslator::default();
+        assert_eq!(translator.translate(b"\x1bOA", false), b"\x1b[A");
+        assert_eq!(translator.translate(b"\x1bOB", false), b"\x1b[B");
+        assert_eq!(translator.translate(b"\x1bOC", false), b"\x1b[C");
+        assert_eq!(translator.translate(b"\x1bOD", false), b"\x1b[D");
+    }
+
+    #[test]
+    fn preserves_application_cursor_arrows_for_remote_application_mode() {
+        let mut translator = UserInputTranslator::default();
+        assert_eq!(translator.translate(b"\x1bOA", true), b"\x1bOA");
+    }
+
+    #[test]
+    fn preserves_split_escape_state_between_reads() {
+        let mut translator = UserInputTranslator::default();
+        assert_eq!(translator.translate(b"\x1b", false), b"\x1b");
+        assert_eq!(translator.translate(b"O", false), b"");
+        assert_eq!(translator.translate(b"A", false), b"[A");
+    }
+}