llm_inspect/
input.rs

1//! Input inspection: scanning and redaction.
2
3use std::collections::HashSet;
4
5use regex::{Regex, RegexBuilder, RegexSet, RegexSetBuilder};
6
7use crate::Error;
8use crate::normalize::Normalized;
9use crate::report::{Decision, Finding, Redacted, Report};
10use crate::rules::{Category, Policy, Rule, RuleSet, Severity};
11
12/// Id of the built-in finding emitted when invisible characters were removed.
13pub(crate) const INVISIBLE_RULE_ID: &str = "obfuscation.invisible_chars";
14
15/// Upper bound on compiled regex size, to keep rule files from exhausting memory.
16const REGEX_SIZE_LIMIT: usize = 10 * (1 << 20);
17
18struct Compiled {
19    rule: Rule,
20    re: Regex,
21}
22
23/// Scans input text against a [`RuleSet`].
24///
25/// Building an inspector compiles every pattern, so build once and reuse it.
26/// `Inspector` is `Send + Sync`.
27pub struct Inspector {
28    rules: Vec<Compiled>,
29    set: RegexSet,
30    policy: Policy,
31}
32
33impl std::fmt::Debug for Inspector {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        f.debug_struct("Inspector")
36            .field("rules", &self.rules.len())
37            .field("policy", &self.policy)
38            .finish_non_exhaustive()
39    }
40}
41
42impl Inspector {
43    /// Compiles `rule_set`.
44    pub fn new(rule_set: RuleSet) -> Result<Self, Error> {
45        rule_set.validate_ids()?;
46        let mut rules = Vec::with_capacity(rule_set.rules.len());
47        for rule in rule_set.rules {
48            let re = RegexBuilder::new(&rule.pattern)
49                .size_limit(REGEX_SIZE_LIMIT)
50                .build()
51                .map_err(|e| Error::InvalidPattern {
52                    id: rule.id.clone(),
53                    source: Box::new(e),
54                })?;
55            rules.push(Compiled { rule, re });
56        }
57        let set = RegexSetBuilder::new(rules.iter().map(|c| c.rule.pattern.as_str()))
58            .size_limit(REGEX_SIZE_LIMIT * 4)
59            .build()
60            .map_err(|e| Error::InvalidPattern {
61                id: "<rule set>".to_owned(),
62                source: Box::new(e),
63            })?;
64        Ok(Self {
65            rules,
66            set,
67            policy: rule_set.policy,
68        })
69    }
70
71    /// An inspector with the bundled default rules.
72    ///
73    /// # Panics
74    ///
75    /// Never in practice: the bundled rules are compiled in the crate's tests.
76    pub fn with_default_rules() -> Self {
77        Self::new(RuleSet::default_rules()).expect("bundled rules must compile")
78    }
79
80    /// The policy in effect.
81    pub fn policy(&self) -> &Policy {
82        &self.policy
83    }
84
85    /// The rules in effect, in order.
86    pub fn rules(&self) -> impl Iterator<Item = &Rule> {
87        self.rules.iter().map(|c| &c.rule)
88    }
89
90    /// Scans `text` and returns every finding plus an overall decision.
91    pub fn scan(&self, text: &str) -> Report {
92        let norm = Normalized::new(text);
93        let mut findings = Vec::new();
94
95        for idx in self.set.matches(norm.text()) {
96            let compiled = &self.rules[idx];
97            // Spans the rule also produces on the raw text; used to tell
98            // "matched only thanks to normalization" apart from plain matches.
99            let raw_spans: Option<HashSet<(usize, usize)>> = norm
100                .changed()
101                .then(|| spans(&compiled.re, compiled.rule.validator, text).collect());
102
103            for (s, e) in spans(&compiled.re, compiled.rule.validator, norm.text()) {
104                let (start, end) = norm.to_original(s, e);
105                let normalized_only = raw_spans
106                    .as_ref()
107                    .is_some_and(|raw| !raw.contains(&(start, end)));
108                findings.push(Finding {
109                    rule_id: compiled.rule.id.clone(),
110                    category: compiled.rule.category,
111                    severity: compiled.rule.severity,
112                    start,
113                    end,
114                    normalized_only,
115                });
116            }
117        }
118
119        if norm.removed_invisible() > 0 {
120            findings.push(Finding {
121                rule_id: INVISIBLE_RULE_ID.to_owned(),
122                category: Category::Injection,
123                severity: Severity::Low,
124                start: 0,
125                end: text.len(),
126                normalized_only: true,
127            });
128        }
129
130        findings.sort_by(|a, b| {
131            a.start
132                .cmp(&b.start)
133                .then_with(|| b.severity.cmp(&a.severity))
134                .then_with(|| a.end.cmp(&b.end))
135                .then_with(|| a.rule_id.cmp(&b.rule_id))
136        });
137
138        let decision = decide(&self.policy, findings.iter().map(|f| f.severity).max());
139        Report {
140            decision,
141            findings,
142            normalized: norm.changed(),
143            removed_invisible: norm.removed_invisible(),
144        }
145    }
146
147    /// Scans `text` and replaces every finding whose category is listed in
148    /// [`Policy::redact`] with a marker.
149    ///
150    /// Overlapping findings are merged; the marker of the first (highest
151    /// severity) finding wins. Injection findings are not redacted by default
152    /// because masking an instruction does not make it safe to forward.
153    pub fn redact(&self, text: &str) -> Redacted {
154        let report = self.scan(text);
155        let mut out = String::with_capacity(text.len());
156        let mut cursor = 0usize;
157        let mut pending: Option<(usize, usize, String)> = None;
158
159        let redactable = report
160            .findings
161            .iter()
162            .filter(|f| f.rule_id != INVISIBLE_RULE_ID && self.policy.redact.contains(&f.category));
163
164        for f in redactable {
165            let marker = self.marker_for(&f.rule_id, f.category);
166            match pending.as_mut() {
167                Some((_, end, _)) if f.start < *end => {
168                    *end = (*end).max(f.end);
169                }
170                _ => {
171                    if let Some((s, e, m)) = pending.take() {
172                        out.push_str(&text[cursor..s]);
173                        out.push_str(&m);
174                        cursor = e;
175                    }
176                    pending = Some((f.start, f.end, marker));
177                }
178            }
179        }
180        if let Some((s, e, m)) = pending.take() {
181            out.push_str(&text[cursor..s]);
182            out.push_str(&m);
183            cursor = e;
184        }
185        out.push_str(&text[cursor..]);
186
187        Redacted { text: out, report }
188    }
189
190    fn marker_for(&self, rule_id: &str, category: Category) -> String {
191        self.rules
192            .iter()
193            .find(|c| c.rule.id == rule_id)
194            .and_then(|c| c.rule.redact_with.clone())
195            .unwrap_or_else(|| format!("[REDACTED:{}]", category.label()))
196    }
197}
198
199/// Byte spans of matches in `hay` (group 1 if the pattern has one), filtered by
200/// the rule's validator. Empty spans are skipped.
201fn spans<'h>(
202    re: &'h Regex,
203    validator: Option<crate::Validator>,
204    hay: &'h str,
205) -> impl Iterator<Item = (usize, usize)> + 'h {
206    re.captures_iter(hay).filter_map(move |caps| {
207        let m = caps.get(1).or_else(|| caps.get(0))?;
208        if m.is_empty() {
209            return None;
210        }
211        if validator.is_some_and(|v| !v.accepts(m.as_str())) {
212            return None;
213        }
214        Some((m.start(), m.end()))
215    })
216}
217
218fn decide(policy: &Policy, highest: Option<Severity>) -> Decision {
219    match highest {
220        Some(s) if s >= policy.block_at => Decision::Block,
221        Some(s) if s >= policy.warn_at => Decision::Warn,
222        _ => Decision::Allow,
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    fn inspector(toml_text: &str) -> Inspector {
231        Inspector::new(RuleSet::from_toml(toml_text).unwrap()).unwrap()
232    }
233
234    #[test]
235    fn no_rules_allows_everything() {
236        let i = inspector("");
237        let r = i.scan("anything at all");
238        assert_eq!(r.decision, Decision::Allow);
239        assert_eq!(r.findings, Vec::new());
240    }
241
242    #[test]
243    fn capture_group_one_is_the_finding() {
244        let i = inspector(
245            r#"
246            [[rules]]
247            id = "t"
248            category = "pii"
249            severity = "high"
250            pattern = '''(?:^|[^0-9])(0[0-9]{2}-[0-9]{4}-[0-9]{4})(?:[^0-9]|$)'''
251            "#,
252        );
253        let text = "tel: 090-1234-5678.";
254        let r = i.scan(text);
255        assert_eq!(r.findings.len(), 1);
256        assert_eq!(r.findings[0].excerpt(text), "090-1234-5678");
257        assert_eq!(r.decision, Decision::Block);
258    }
259
260    #[test]
261    fn thresholds_drive_the_decision() {
262        let i = inspector(
263            r#"
264            [policy]
265            block_at = "critical"
266            warn_at = "high"
267            [[rules]]
268            id = "t"
269            category = "injection"
270            severity = "high"
271            pattern = "ignore"
272            "#,
273        );
274        assert_eq!(i.scan("please ignore").decision, Decision::Warn);
275        assert_eq!(i.scan("fine").decision, Decision::Allow);
276    }
277
278    #[test]
279    fn redaction_merges_overlaps_and_respects_categories() {
280        let i = inspector(
281            r#"
282            [[rules]]
283            id = "a"
284            category = "pii"
285            severity = "high"
286            pattern = "abcd"
287            [[rules]]
288            id = "b"
289            category = "pii"
290            severity = "low"
291            pattern = "cdef"
292            redact_with = "<b>"
293            [[rules]]
294            id = "c"
295            category = "injection"
296            severity = "high"
297            pattern = "xyz"
298            "#,
299        );
300        let r = i.redact("..abcdef..xyz");
301        assert_eq!(r.text, "..[REDACTED:PII]..xyz");
302        assert_eq!(r.report.findings.len(), 3);
303    }
304
305    #[test]
306    fn invalid_pattern_names_the_rule() {
307        let err = Inspector::new(
308            RuleSet::from_toml(
309                r#"
310                [[rules]]
311                id = "broken"
312                category = "pii"
313                severity = "low"
314                pattern = "("
315                "#,
316            )
317            .unwrap(),
318        )
319        .unwrap_err();
320        assert!(matches!(err, Error::InvalidPattern { id, .. } if id == "broken"));
321    }
322}
llm_inspect/input.rs

llm_inspect/
input.rs
