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

llm_kernel/safety/
sanitize.rs
