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

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")
        );
    }
}