shell_tunnel/output/
sanitizer.rs

1//! Output sanitization for stripping ANSI escape codes.
2
3use vte::{Params, Parser, Perform};
4
5/// Output sanitizer using VTE parser.
6pub struct OutputSanitizer;
7
8impl OutputSanitizer {
9    /// Strip ANSI escape codes from raw bytes.
10    ///
11    /// Returns clean UTF-8 text with all control sequences removed.
12    pub fn strip_ansi(input: &[u8]) -> String {
13        let mut extractor = PlainTextExtractor::new();
14        let mut parser = Parser::new();
15
16        parser.advance(&mut extractor, input);
17
18        extractor.into_string()
19    }
20
21    /// Strip ANSI codes from a string.
22    pub fn strip_ansi_str(input: &str) -> String {
23        Self::strip_ansi(input.as_bytes())
24    }
25}
26
27/// VTE performer that extracts plain text.
28struct PlainTextExtractor {
29    output: Vec<u8>,
30}
31
32impl PlainTextExtractor {
33    fn new() -> Self {
34        Self { output: Vec::new() }
35    }
36
37    fn into_string(self) -> String {
38        String::from_utf8_lossy(&self.output).into_owned()
39    }
40}
41
42impl Perform for PlainTextExtractor {
43    fn print(&mut self, c: char) {
44        let mut buf = [0u8; 4];
45        let encoded = c.encode_utf8(&mut buf);
46        self.output.extend_from_slice(encoded.as_bytes());
47    }
48
49    fn execute(&mut self, byte: u8) {
50        // Handle control characters
51        match byte {
52            // Newline, carriage return, tab
53            0x0A | 0x0D | 0x09 => self.output.push(byte),
54            // Ignore other control characters
55            _ => {}
56        }
57    }
58
59    fn hook(&mut self, _params: &Params, _intermediates: &[u8], _ignore: bool, _action: char) {
60        // Ignore DCS sequences
61    }
62
63    fn put(&mut self, _byte: u8) {
64        // Ignore DCS data
65    }
66
67    fn unhook(&mut self) {
68        // Ignore DCS end
69    }
70
71    fn osc_dispatch(&mut self, _params: &[&[u8]], _bell_terminated: bool) {
72        // Ignore OSC sequences
73    }
74
75    fn csi_dispatch(
76        &mut self,
77        _params: &Params,
78        _intermediates: &[u8],
79        _ignore: bool,
80        _action: char,
81    ) {
82        // Ignore CSI sequences (cursor movement, colors, etc.)
83    }
84
85    fn esc_dispatch(&mut self, _intermediates: &[u8], _ignore: bool, _byte: u8) {
86        // Ignore escape sequences
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn test_plain_text() {
96        let input = b"hello world";
97        let output = OutputSanitizer::strip_ansi(input);
98        assert_eq!(output, "hello world");
99    }
100
101    #[test]
102    fn test_strip_color_codes() {
103        // "\x1b[31mred\x1b[0m" - red text with reset
104        let input = b"\x1b[31mred\x1b[0m";
105        let output = OutputSanitizer::strip_ansi(input);
106        assert_eq!(output, "red");
107    }
108
109    #[test]
110    fn test_strip_bold() {
111        // "\x1b[1mbold\x1b[0m"
112        let input = b"\x1b[1mbold\x1b[0m";
113        let output = OutputSanitizer::strip_ansi(input);
114        assert_eq!(output, "bold");
115    }
116
117    #[test]
118    fn test_preserve_newlines() {
119        let input = b"line1\nline2\nline3";
120        let output = OutputSanitizer::strip_ansi(input);
121        assert_eq!(output, "line1\nline2\nline3");
122    }
123
124    #[test]
125    fn test_strip_cursor_movement() {
126        // "\x1b[2J\x1b[H" - clear screen and home cursor
127        let input = b"\x1b[2J\x1b[Hcontent";
128        let output = OutputSanitizer::strip_ansi(input);
129        assert_eq!(output, "content");
130    }
131
132    #[test]
133    fn test_complex_sequence() {
134        // Mix of colors, bold, cursor movement
135        let input = b"\x1b[32m\x1b[1mGreen Bold\x1b[0m Normal \x1b[34mBlue\x1b[0m";
136        let output = OutputSanitizer::strip_ansi(input);
137        assert_eq!(output, "Green Bold Normal Blue");
138    }
139
140    #[test]
141    fn test_osc_title() {
142        // "\x1b]0;Window Title\x07" - set window title
143        let input = b"\x1b]0;Window Title\x07actual content";
144        let output = OutputSanitizer::strip_ansi(input);
145        assert_eq!(output, "actual content");
146    }
147
148    #[test]
149    fn test_strip_ansi_str() {
150        let input = "\x1b[31mcolored\x1b[0m";
151        let output = OutputSanitizer::strip_ansi_str(input);
152        assert_eq!(output, "colored");
153    }
154
155    #[test]
156    fn test_preserve_tabs() {
157        let input = b"col1\tcol2\tcol3";
158        let output = OutputSanitizer::strip_ansi(input);
159        assert_eq!(output, "col1\tcol2\tcol3");
160    }
161
162    #[test]
163    fn test_empty_input() {
164        let output = OutputSanitizer::strip_ansi(b"");
165        assert_eq!(output, "");
166    }
167
168    #[test]
169    fn test_only_escape_codes() {
170        let input = b"\x1b[31m\x1b[0m\x1b[2J";
171        let output = OutputSanitizer::strip_ansi(input);
172        assert_eq!(output, "");
173    }
174}
shell_tunnel/output/sanitizer.rs

shell_tunnel/output/
sanitizer.rs
