llm_kernel/safety/
sanitize.rs

1//! Output sanitization — secret masking and control character removal.
2
3use regex::Regex;
4use std::sync::LazyLock;
5
6// Single compiled pattern covering all secret types in one pass:
7//   group 1 — Bearer/Basic token value
8//   group 2 — key=value prefix (password=, token=, ...)
9//   group 3 — key=value value
10//   groups 4–6 — standalone sk-*, AKIA*, gh[posu]_* tokens
11static SECRET_RE: LazyLock<Regex> = LazyLock::new(|| {
12    Regex::new(
13        r"(?i:bearer |basic )(\S+)|((?:password|token|key|secret|api_key|apikey|access_token|private_key)=)(\S+)|(sk-\S*)|(AKIA[A-Za-z0-9]+)|(gh[posu]_[A-Za-z0-9]+)"
14    ).expect("SECRET_RE is valid")
15});
16
17/// Mask known secret patterns in a string in a single pass.
18///
19/// Handles `Bearer`/`Basic` auth headers (case-insensitive), `sk-*` API keys,
20/// AWS keys (`AKIA...`), GitHub tokens (`ghp_`/`gho_`/`ghs_`/`ghu_`), and
21/// `password=`, `token=`, `key=`, `secret=`, `api_key=`, `apikey=`,
22/// `access_token=`, `private_key=` values. All occurrences are masked.
23pub fn mask_secrets(input: &str) -> String {
24    SECRET_RE
25        .replace_all(input, |caps: &regex::Captures| -> String {
26            if caps.get(1).is_some() {
27                // Bearer/Basic: full match = "<prefix> <value>"; group 1 = value only
28                let full = caps.get(0).unwrap().as_str();
29                let val = caps.get(1).unwrap().as_str();
30                let prefix = &full[..full.len() - val.len()];
31                format!("{prefix}****")
32            } else if let Some(key_prefix) = caps.get(2) {
33                // key=value: group 2 = "key=", group 3 = value
34                format!("{}****", key_prefix.as_str())
35            } else {
36                // standalone: sk-*, AKIA*, gh[posu]_*
37                "****".to_string()
38            }
39        })
40        .into_owned()
41}
42
43/// Remove ANSI escape sequences from text.
44pub fn strip_ansi(input: &str) -> String {
45    let mut result = String::with_capacity(input.len());
46    let mut chars = input.chars().peekable();
47    while let Some(ch) = chars.next() {
48        if ch == '\x1B' {
49            // Skip ESC and the following sequence
50            if chars.peek() == Some(&'[') {
51                chars.next(); // consume '['
52                // Skip parameter bytes (0x30-0x3F), intermediate bytes (0x20-0x2F), final byte (0x40-0x7E)
53                while let Some(&next) = chars.peek() {
54                    let cp = next as u32;
55                    if (0x30..=0x3F).contains(&cp) || (0x20..=0x2F).contains(&cp) {
56                        chars.next();
57                    } else if (0x40..=0x7E).contains(&cp) {
58                        chars.next(); // consume final byte
59                        break;
60                    } else {
61                        break;
62                    }
63                }
64            }
65            continue;
66        }
67        result.push(ch);
68    }
69    result
70}
71
72/// Sanitize output for safe display by removing:
73///
74/// - Bidi override characters (U+202A–U+202E) — invisible text direction attacks
75/// - Plane-14 tag characters (U+E0000–U+E007F) — LLM injection vectors
76/// - Line/paragraph separators (U+2028, U+2029)
77/// - Null bytes (U+0000)
78/// - C1 control characters (U+0080–U+009F) except common whitespace
79pub fn sanitize_output(input: &str) -> String {
80    input
81        .chars()
82        .filter(|c| {
83            let cp = *c as u32;
84            // Null byte
85            if cp == 0 {
86                return false;
87            }
88            // C1 controls (except \u{0085} NEL which is sometimes valid)
89            if (0x80..=0x9F).contains(&cp) {
90                return false;
91            }
92            // Bidi overrides
93            if (0x202A..=0x202E).contains(&cp) {
94                return false;
95            }
96            // Line/paragraph separators
97            if cp == 0x2028 || cp == 0x2029 {
98                return false;
99            }
100            // Plane-14 tags
101            if (0xE0000..=0xE007F).contains(&cp) {
102                return false;
103            }
104            true
105        })
106        .collect()
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn mask_bearer_token() {
115        let input = "Authorization: Bearer sk-longtoken123456";
116        let masked = mask_secrets(input);
117        assert!(masked.contains("****"));
118    }
119
120    #[test]
121    fn mask_password_value() {
122        let input = "password=supersecret other=value";
123        let masked = mask_secrets(input);
124        assert!(masked.contains("password=****"));
125        assert!(masked.contains("other=value"));
126    }
127
128    #[test]
129    fn mask_token_value() {
130        let masked = mask_secrets("token=abc123");
131        assert!(masked.contains("token=****"));
132    }
133
134    #[test]
135    fn mask_key_value() {
136        let masked = mask_secrets("key=my-api-key-here");
137        assert!(masked.contains("key=****"));
138    }
139
140    #[test]
141    fn mask_secret_value() {
142        let masked = mask_secrets("secret=data rest=ok");
143        assert!(masked.contains("secret=****"));
144        assert!(masked.contains("rest=ok"));
145    }
146
147    #[test]
148    fn no_secrets_unchanged() {
149        let input = "hello world 123";
150        assert_eq!(mask_secrets(input), input);
151    }
152
153    #[test]
154    fn mask_multiple_passwords() {
155        let masked = mask_secrets("password=a password=b");
156        assert!(masked.contains("password=****"), "got: {masked}");
157    }
158
159    #[test]
160    fn mask_multiple_bearer() {
161        let masked = mask_secrets("Bearer token1 and Bearer token2");
162        assert_eq!(masked, "Bearer **** and Bearer ****");
163    }
164
165    #[test]
166    fn mask_standalone_sk_key() {
167        let masked = mask_secrets("key is sk-proj-abc123 here");
168        assert!(masked.contains("****"), "got: {masked}");
169        assert!(!masked.contains("sk-proj"), "got: {masked}");
170    }
171
172    #[test]
173    fn mask_api_key_value() {
174        let masked = mask_secrets("api_key=my-secret-key other=ok");
175        assert!(masked.contains("api_key=****"));
176        assert!(masked.contains("other=ok"));
177    }
178
179    #[test]
180    fn mask_apikey_nounderscore() {
181        let masked = mask_secrets("apikey=abc123");
182        assert!(masked.contains("apikey=****"));
183    }
184
185    #[test]
186    fn mask_access_token_value() {
187        let masked = mask_secrets("access_token=eyJhbGciOiJI");
188        assert!(masked.contains("access_token=****"));
189    }
190
191    #[test]
192    fn mask_private_key_value() {
193        let masked = mask_secrets("private_key=-----BEGIN RSA");
194        assert!(masked.contains("private_key=****"));
195    }
196
197    #[test]
198    fn mask_basic_auth() {
199        let masked = mask_secrets("Authorization: Basic dXNlcjpwYXNz");
200        assert!(masked.contains("Basic ****"));
201    }
202
203    #[test]
204    fn mask_bearer_case_insensitive() {
205        let masked = mask_secrets("auth: bearer token123 here");
206        assert!(masked.contains("bearer ****"));
207    }
208
209    #[test]
210    fn mask_aws_access_key() {
211        let masked = mask_secrets("key=AKIAIOSFODNN7EXAMPLE end");
212        assert!(masked.contains("key=****"));
213    }
214
215    #[test]
216    fn mask_standalone_aws_key() {
217        let masked = mask_secrets("found AKIAIOSFODNN7EXAMPLE in config");
218        assert!(!masked.contains("AKIAIOSFODNN7EXAMPLE"), "got: {masked}");
219        assert!(masked.contains("****"));
220    }
221
222    #[test]
223    fn mask_github_pat() {
224        let masked = mask_secrets("token ghp_ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefgh end");
225        assert!(!masked.contains("ghp_"), "got: {masked}");
226        assert!(masked.contains("****"));
227    }
228
229    #[test]
230    fn mask_github_oauth() {
231        let masked = mask_secrets("using gho_ABCDEF1234567890 here");
232        assert!(!masked.contains("gho_"), "got: {masked}");
233    }
234
235    #[test]
236    fn sanitize_removes_bidi() {
237        let input = "hello\u{202E}world";
238        let clean = sanitize_output(input);
239        assert_eq!(clean, "helloworld");
240    }
241
242    #[test]
243    fn sanitize_removes_plane14() {
244        let input = "text\u{E0001}hidden";
245        let clean = sanitize_output(input);
246        assert_eq!(clean, "texthidden");
247    }
248
249    #[test]
250    fn sanitize_removes_null() {
251        let clean = sanitize_output("a\u{0000}b");
252        assert_eq!(clean, "ab");
253    }
254
255    #[test]
256    fn sanitize_removes_line_sep() {
257        let clean = sanitize_output("a\u{2028}b");
258        assert_eq!(clean, "ab");
259    }
260
261    #[test]
262    fn sanitize_preserves_normal() {
263        let input = "Hello, 世界! 🎉";
264        assert_eq!(sanitize_output(input), input);
265    }
266
267    #[test]
268    fn sanitize_removes_c1_controls() {
269        let clean = sanitize_output("a\u{0080}b\u{009F}c");
270        assert_eq!(clean, "abc");
271    }
272
273    #[test]
274    fn strip_ansi_removes_color_codes() {
275        let input = "\x1B[31mHello\x1B[0m \x1B[1;32mWorld\x1B[0m";
276        let clean = strip_ansi(input);
277        assert_eq!(clean, "Hello World");
278    }
279
280    #[test]
281    fn strip_ansi_preserves_plain_text() {
282        let input = "Hello, 世界! 🎉";
283        assert_eq!(strip_ansi(input), input);
284    }
285
286    #[test]
287    fn strip_ansi_handles_complex_sequence() {
288        // 256-color and RGB sequences
289        let input = "\x1B[38;5;196mRed\x1B[0m \x1B[38;2;0;255;0mGreen\x1B[0m";
290        let clean = strip_ansi(input);
291        assert_eq!(clean, "Red Green");
292    }
293}
llm_kernel/safety/sanitize.rs

llm_kernel/safety/
sanitize.rs
