Preserve native terminal behavior through trusted prompts

crates/lazaret-client/src/terminal/presentation.rs

@@ -0,0 +1,210 @@
+const MAX_SEQUENCE_BYTES: usize = 4096;
+
+#[derive(Clone, Copy, Default)]
+enum ScanState {
+    #[default]
+    Ground,
+    Escape,
+    Csi,
+    DiscardCsi,
+}
+
+pub(super) struct PresentationTracker {
+    parser: vt100::Parser,
+    scan: ScanState,
+    sequence: Vec<u8>,
+    autowrap: bool,
+    focus_reporting: bool,
+    synchronized_output: bool,
+}
+
+impl PresentationTracker {
+    pub(super) fn new(rows: u16, columns: u16) -> Self {
+        Self {
+            parser: vt100::Parser::new(rows, columns, 0),
+            scan: ScanState::Ground,
+            sequence: Vec::new(),
+            autowrap: true,
+            focus_reporting: false,
+            synchronized_output: false,
+        }
+    }
+
+    pub(super) fn process(&mut self, bytes: &[u8]) {
+        self.parser.process(bytes);
+        for &byte in bytes {
+            self.process_mode_byte(byte);
+        }
+    }
+
+    pub(super) fn resize(&mut self, rows: u16, columns: u16) {
+        self.parser.screen_mut().set_size(rows, columns);
+    }
+
+    pub(super) fn alternate_screen(&self) -> bool {
+        self.parser.screen().alternate_screen()
+    }
+
+    pub(super) fn context_snapshot(&self) -> Vec<u8> {
+        let mut bytes = b"\x1b[2J\x1b[H".to_vec();
+        bytes.extend(self.parser.screen().contents_formatted());
+        bytes
+    }
+
+    pub(super) fn restoration_snapshot(&self) -> Vec<u8> {
+        let mut bytes = b"\x1b[2J\x1b[H".to_vec();
+        bytes.extend(self.parser.screen().state_formatted());
+        bytes.extend(self.additional_modes(false));
+        bytes
+    }
+
+    pub(super) fn restoration_modes(&self) -> Vec<u8> {
+        let mut bytes = self.parser.screen().input_mode_formatted();
+        bytes.extend(self.additional_modes(false));
+        bytes
+    }
+
+    pub(super) fn synchronized_output_mode(&self) -> Vec<u8> {
+        if self.synchronized_output {
+            b"\x1b[?2026h".to_vec()
+        } else {
+            Vec::new()
+        }
+    }
+
+    fn additional_modes(&self, synchronized_output: bool) -> Vec<u8> {
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(if self.autowrap {
+            b"\x1b[?7h"
+        } else {
+            b"\x1b[?7l"
+        });
+        bytes.extend_from_slice(if self.parser.screen().hide_cursor() {
+            b"\x1b[?25l"
+        } else {
+            b"\x1b[?25h"
+        });
+        bytes.extend_from_slice(if self.focus_reporting {
+            b"\x1b[?1004h"
+        } else {
+            b"\x1b[?1004l"
+        });
+        if synchronized_output && self.synchronized_output {
+            bytes.extend_from_slice(b"\x1b[?2026h");
+        }
+        bytes
+    }
+
+    fn process_mode_byte(&mut self, byte: u8) {
+        match self.scan {
+            ScanState::Ground => {
+                if byte == 0x1b {
+                    self.sequence.clear();
+                    self.sequence.push(byte);
+                    self.scan = ScanState::Escape;
+                }
+            }
+            ScanState::Escape => match byte {
+                b'[' => {
+                    self.sequence.push(byte);
+                    self.scan = ScanState::Csi;
+                }
+                b'c' => {
+                    self.autowrap = true;
+                    self.focus_reporting = false;
+                    self.synchronized_output = false;
+                    self.sequence.clear();
+                    self.scan = ScanState::Ground;
+                }
+                0x1b => {
+                    self.sequence.clear();
+                    self.sequence.push(byte);
+                }
+                _ => {
+                    self.sequence.clear();
+                    self.scan = ScanState::Ground;
+                }
+            },
+            ScanState::Csi => {
+                if byte == 0x1b {
+                    self.sequence.clear();
+                    self.sequence.push(byte);
+                    self.scan = ScanState::Escape;
+                    return;
+                }
+                self.sequence.push(byte);
+                if self.sequence.len() > MAX_SEQUENCE_BYTES {
+                    self.sequence.clear();
+                    self.scan = ScanState::DiscardCsi;
+                } else if (0x40..=0x7e).contains(&byte) {
+                    let sequence = std::mem::take(&mut self.sequence);
+                    self.apply_mode(&sequence);
+                    self.scan = ScanState::Ground;
+                } else if !(0x20..=0x3f).contains(&byte) {
+                    self.sequence.clear();
+                    self.scan = ScanState::Ground;
+                }
+            }
+            ScanState::DiscardCsi => {
+                if (0x40..=0x7e).contains(&byte) {
+                    self.scan = ScanState::Ground;
+                }
+            }
+        }
+    }
+
+    fn apply_mode(&mut self, sequence: &[u8]) {
+        let Some((&action, body)) = sequence.split_last() else {
+            return;
+        };
+        if !matches!(action, b'h' | b'l') {
+            return;
+        }
+        let Some(parameters) = body.strip_prefix(b"\x1b[?") else {
+            return;
+        };
+        let enabled = action == b'h';
+        for parameter in parameters.split(|byte| *byte == b';') {
+            match parameter {
+                b"7" => self.autowrap = enabled,
+                b"1004" => self.focus_reporting = enabled,
+                b"2026" => self.synchronized_output = enabled,
+                _ => {}
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::PresentationTracker;
+
+    #[test]
+    fn presentation_modes_follow_fragmented_guest_output() {
+        let bytes = b"normal\x1b[?7l\x1b[?25l\x1b[?1004;2026h";
+        for split in 0..=bytes.len() {
+            let mut tracker = PresentationTracker::new(8, 32);
+            tracker.process(&bytes[..split]);
+            tracker.process(&bytes[split..]);
+            let mut modes = tracker.restoration_modes();
+            modes.extend(tracker.synchronized_output_mode());
+            assert!(modes.windows(4).any(|window| window == b"[?7l"));
+            assert!(modes.windows(5).any(|window| window == b"[?25l"));
+            assert!(modes.windows(7).any(|window| window == b"[?1004h"));
+            assert!(modes.windows(7).any(|window| window == b"[?2026h"));
+        }
+    }
+
+    #[test]
+    fn snapshots_do_not_replay_an_alternate_transition() {
+        let mut tracker = PresentationTracker::new(8, 32);
+        tracker.process(b"normal\x1b[?1049h\x1b[2J\x1b[Halternate");
+        assert!(tracker.alternate_screen());
+        assert!(
+            !tracker
+                .restoration_snapshot()
+                .windows(8)
+                .any(|window| window == b"[?1049h")
+        );
+    }
+}