1use once_cell::sync::Lazy;
21use regex::Regex;
22use serde::{Deserialize, Serialize};
23
24struct RedactRule {
25 pattern: Regex,
26 replacement: &'static str,
27 category: &'static str,
28}
29
30fn compile(rules: &[(&str, &'static str, &'static str)]) -> Vec<RedactRule> {
31 rules
32 .iter()
33 .filter_map(|(p, r, c)| {
34 Regex::new(p).ok().map(|pattern| RedactRule {
35 pattern,
36 replacement: r,
37 category: c,
38 })
39 })
40 .collect()
41}
42
43fn secret_rules() -> Vec<RedactRule> {
45 compile(&[
46 (
48 r"sk-(?:proj-|live-)?[A-Za-z0-9_\-]{20,}",
49 "sk-<redacted>",
50 "openai_key",
51 ),
52 (
54 r"sk-ant-[A-Za-z0-9_\-]{20,}",
55 "sk-ant-<redacted>",
56 "anthropic_key",
57 ),
58 (
60 r"(?i)Bearer\s+[A-Za-z0-9._\-]{16,}",
61 "Bearer <redacted>",
62 "bearer_token",
63 ),
64 (r"gh[pousr]_[A-Za-z0-9]{20,}", "gh<redacted>", "github_pat"),
66 (r"AKIA[0-9A-Z]{16}", "AKIA<redacted>", "aws_key"),
68 (r"ASIA[0-9A-Z]{16}", "ASIA<redacted>", "aws_key"),
69 (
71 r"xox[baprs]-[A-Za-z0-9\-]{10,}",
72 "xox<redacted>",
73 "slack_token",
74 ),
75 (r"AIza[0-9A-Za-z_\-]{35}", "AIza<redacted>", "google_key"),
77 (
79 r"(?s)-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----",
80 "<redacted-private-key>",
81 "private_key",
82 ),
83 (
85 r#"(?i)api[_-]?key["']?\s*[:=]\s*["']?[A-Za-z0-9_\-]{24,}"#,
86 "api_key=<redacted>",
87 "api_key_assignment",
88 ),
89 ])
90}
91
92fn pii_rules() -> Vec<RedactRule> {
94 compile(&[
95 (
97 r"([A-Za-z0-9._%+\-]+)@([A-Za-z0-9.\-]+\.[A-Za-z]{2,})",
98 "<redacted>@$2",
99 "email",
100 ),
101 (
103 r"\+?\d{1,3}[\s\-]?\(?\d{2,4}\)?[\s\-]?\d{3,4}[\s\-]?\d{3,4}",
104 "<redacted-phone>",
105 "phone",
106 ),
107 ])
108}
109
110static SECRET_RULES: Lazy<Vec<RedactRule>> = Lazy::new(secret_rules);
111static PII_RULES: Lazy<Vec<RedactRule>> = Lazy::new(pii_rules);
112
113fn apply_rules(input: &str, rules: &[RedactRule]) -> String {
114 let mut out = input.to_string();
115 for rule in rules {
116 if rule.pattern.is_match(&out) {
117 out = rule
118 .pattern
119 .replace_all(&out, rule.replacement)
120 .into_owned();
121 }
122 }
123 out
124}
125
126pub fn redact(input: &str) -> String {
129 let secrets = apply_rules(input, &SECRET_RULES);
130 apply_rules(&secrets, &PII_RULES)
131}
132
133pub fn redact_in_place(buf: &mut String) {
135 let new = redact(buf);
136 if new != *buf {
137 *buf = new;
138 }
139}
140
141pub fn redact_secrets(input: &str) -> String {
145 apply_rules(input, &SECRET_RULES)
146}
147
148pub fn scan_secrets(input: &str) -> Vec<&'static str> {
151 let mut found: Vec<&'static str> = Vec::new();
152 for rule in SECRET_RULES.iter() {
153 if rule.pattern.is_match(input) && !found.contains(&rule.category) {
154 found.push(rule.category);
155 }
156 }
157 found
158}
159
160#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
162#[serde(rename_all = "lowercase")]
163pub enum OutputGuardrailMode {
164 #[default]
166 Off,
167 Redact,
169 Block,
171}
172
173#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
175pub struct OutputGuardrailConfig {
176 #[serde(default)]
177 pub mode: OutputGuardrailMode,
178}
179
180pub struct OutputGuardrailOutcome {
182 pub text: String,
184 pub categories: Vec<&'static str>,
186 pub blocked: bool,
188}
189
190pub const OUTPUT_BLOCKED_NOTICE: &str =
192 "[message withheld by output guardrail: a credential-like string was detected]";
193
194pub fn apply_output_guardrail(
198 text: &str,
199 mode: OutputGuardrailMode,
200) -> Option<OutputGuardrailOutcome> {
201 if mode == OutputGuardrailMode::Off {
202 return None;
203 }
204 let categories = scan_secrets(text);
205 if categories.is_empty() {
206 return None;
207 }
208 match mode {
209 OutputGuardrailMode::Off => None,
210 OutputGuardrailMode::Redact => Some(OutputGuardrailOutcome {
211 text: redact_secrets(text),
212 categories,
213 blocked: false,
214 }),
215 OutputGuardrailMode::Block => Some(OutputGuardrailOutcome {
216 text: OUTPUT_BLOCKED_NOTICE.to_string(),
217 categories,
218 blocked: true,
219 }),
220 }
221}
222
223#[cfg(test)]
224mod tests {
225 use super::*;
226
227 #[test]
228 fn openai_keys_redacted() {
229 let out = redact("key is sk-proj-abcdef1234567890ABCDEF here");
230 assert!(!out.contains("abcdef"));
231 assert!(out.contains("sk-<redacted>"));
232 }
233
234 #[test]
235 fn bearer_header_redacted() {
236 let out = redact("Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.payload.sig");
237 assert!(!out.contains("eyJhbG"));
238 assert!(out.contains("Bearer <redacted>"));
239 }
240
241 #[test]
242 fn github_pat_redacted() {
243 let out = redact("token=ghp_abcdef1234567890ABCDEFghij");
244 assert!(!out.contains("abcdef"));
245 }
246
247 #[test]
248 fn aws_key_redacted() {
249 let out = redact("AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE");
250 assert!(out.contains("AKIA<redacted>"));
251 }
252
253 #[test]
254 fn email_partially_masked() {
255 let out = redact("send to alice@example.com please");
256 assert!(out.contains("<redacted>@example.com"));
257 }
258
259 #[test]
260 fn passthrough_for_plain_text() {
261 let input = "hello world, nothing sensitive here";
262 assert_eq!(redact(input), input);
263 }
264
265 #[test]
266 fn multiple_secrets_in_one_string() {
267 let input = format!(
268 "sk-live-1234567890abcdefghij and Bearer {}",
269 "xyzabcdefghijk1234567890",
270 );
271 let out = redact(&input);
272 assert!(!out.contains("1234567890abcdefghij"));
273 assert!(out.contains("sk-<redacted>"));
274 }
275
276 #[test]
277 fn redact_secrets_keeps_pii() {
278 let out = redact_secrets("email alice@example.com, key sk-live-1234567890abcdefghij");
280 assert!(out.contains("alice@example.com"));
281 assert!(out.contains("sk-<redacted>"));
282 }
283
284 #[test]
285 fn scan_secrets_reports_categories() {
286 let cats = scan_secrets("ghp_abcdef1234567890ABCDEFghij and AKIAIOSFODNN7EXAMPLE");
287 assert!(cats.contains(&"github_pat"));
288 assert!(cats.contains(&"aws_key"));
289 }
290
291 #[test]
292 fn private_key_block_redacted() {
293 let pem = "-----BEGIN RSA PRIVATE KEY-----\nMIIEabc\n-----END RSA PRIVATE KEY-----";
294 assert!(scan_secrets(pem).contains(&"private_key"));
295 assert!(redact_secrets(pem).contains("<redacted-private-key>"));
296 }
297
298 #[test]
299 fn guardrail_off_is_noop() {
300 assert!(
301 apply_output_guardrail("sk-live-1234567890abcdefghij", OutputGuardrailMode::Off)
302 .is_none()
303 );
304 }
305
306 #[test]
307 fn guardrail_passes_clean_text() {
308 assert!(
309 apply_output_guardrail("just a normal reply", OutputGuardrailMode::Redact).is_none()
310 );
311 }
312
313 #[test]
314 fn guardrail_redacts_secret() {
315 let out = apply_output_guardrail(
316 "here is ghp_abcdef1234567890ABCDEFghij",
317 OutputGuardrailMode::Redact,
318 )
319 .expect("should fire");
320 assert!(!out.blocked);
321 assert!(out.text.contains("gh<redacted>"));
322 assert!(out.categories.contains(&"github_pat"));
323 }
324
325 #[test]
326 fn guardrail_blocks_secret() {
327 let out = apply_output_guardrail(
328 "here is ghp_abcdef1234567890ABCDEFghij",
329 OutputGuardrailMode::Block,
330 )
331 .expect("should fire");
332 assert!(out.blocked);
333 assert_eq!(out.text, OUTPUT_BLOCKED_NOTICE);
334 }
335}