Skip to main content

llm_inspect/
rules.rs

1//! Rule definitions and the TOML format they are loaded from.
2//!
3//! A rule file looks like this:
4//!
5//! ```toml
6//! [policy]
7//! block_at = "high"      # findings at or above this severity block the input
8//! warn_at = "medium"     # findings at or above this severity (below block_at) warn
9//! redact = ["pii", "secret"]
10//!
11//! [[rules]]
12//! id = "pii.email"
13//! category = "pii"
14//! severity = "medium"
15//! description = "Email address"
16//! pattern = '''[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}'''
17//! ```
18//!
19//! If a pattern contains a capture group, group 1 is reported as the finding;
20//! otherwise the whole match is. This lets a rule anchor on context (for example
21//! "not preceded by a digit") without the regex crate's lack of look-around.
22
23use serde::{Deserialize, Serialize};
24
25use crate::Error;
26
27/// The bundled default rule set, as TOML.
28pub const DEFAULT_RULES_TOML: &str = include_str!("../rules/default.toml");
29
30/// What kind of problem a rule detects.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
32#[serde(rename_all = "snake_case")]
33#[non_exhaustive]
34pub enum Category {
35    /// Personally identifiable information (email, phone, card numbers...).
36    Pii,
37    /// Credentials and keys.
38    Secret,
39    /// Prompt-injection or jailbreak attempts, including obfuscation signals.
40    Injection,
41}
42
43impl Category {
44    /// Upper-case label used in default redaction markers.
45    pub fn label(self) -> &'static str {
46        match self {
47            Category::Pii => "PII",
48            Category::Secret => "SECRET",
49            Category::Injection => "INJECTION",
50        }
51    }
52}
53
54/// How serious a finding is. Ordered from least to most severe.
55#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
56#[serde(rename_all = "snake_case")]
57pub enum Severity {
58    /// Informational; worth logging.
59    Low,
60    /// Likely sensitive; usually warn.
61    Medium,
62    /// Sensitive or dangerous; usually block.
63    High,
64    /// Credentials or clear attacks; block.
65    Critical,
66}
67
68/// Extra validation applied to a regex match before it becomes a finding.
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
70#[serde(rename_all = "snake_case")]
71#[non_exhaustive]
72pub enum Validator {
73    /// The digits of the match must pass the Luhn checksum (payment cards).
74    Luhn,
75}
76
77impl Validator {
78    /// Returns `true` if `matched` passes this validator.
79    pub fn accepts(self, matched: &str) -> bool {
80        match self {
81            Validator::Luhn => luhn(matched),
82        }
83    }
84}
85
86/// Luhn checksum over the ASCII digits of `s`; non-digits are ignored.
87/// Requires 13 to 19 digits, the range used by payment card numbers.
88fn luhn(s: &str) -> bool {
89    let digits: Vec<u32> = s.chars().filter_map(|c| c.to_digit(10)).collect();
90    if !(13..=19).contains(&digits.len()) {
91        return false;
92    }
93    let sum: u32 = digits
94        .iter()
95        .rev()
96        .enumerate()
97        .map(|(i, &d)| {
98            if i % 2 == 1 {
99                let dd = d * 2;
100                if dd > 9 { dd - 9 } else { dd }
101            } else {
102                d
103            }
104        })
105        .sum();
106    sum % 10 == 0
107}
108
109/// A single detection rule.
110#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
111pub struct Rule {
112    /// Stable identifier, e.g. `pii.email`. Must be unique within a rule set.
113    pub id: String,
114    /// Category of the finding.
115    pub category: Category,
116    /// Severity of the finding.
117    pub severity: Severity,
118    /// Human-readable description.
119    #[serde(default)]
120    pub description: String,
121    /// Regular expression (regex crate syntax). Group 1, if present, is the finding.
122    pub pattern: String,
123    /// Optional extra validation of the matched text.
124    #[serde(default, skip_serializing_if = "Option::is_none")]
125    pub validator: Option<Validator>,
126    /// Replacement used by [`Inspector::redact`](crate::Inspector::redact).
127    /// Defaults to `[REDACTED:<CATEGORY>]`.
128    #[serde(default, skip_serializing_if = "Option::is_none")]
129    pub redact_with: Option<String>,
130}
131
132/// Decision thresholds and redaction settings.
133#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
134#[serde(default)]
135pub struct Policy {
136    /// Findings at or above this severity make the decision [`Block`](crate::Decision::Block).
137    pub block_at: Severity,
138    /// Findings at or above this severity (and below `block_at`) make the decision
139    /// [`Warn`](crate::Decision::Warn).
140    pub warn_at: Severity,
141    /// Categories that [`Inspector::redact`](crate::Inspector::redact) masks.
142    pub redact: Vec<Category>,
143}
144
145impl Default for Policy {
146    fn default() -> Self {
147        Self {
148            block_at: Severity::High,
149            warn_at: Severity::Medium,
150            redact: vec![Category::Pii, Category::Secret],
151        }
152    }
153}
154
155/// A policy plus a list of rules; the unit that is loaded from TOML.
156#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
157#[serde(default)]
158pub struct RuleSet {
159    /// Decision thresholds.
160    pub policy: Policy,
161    /// Detection rules.
162    pub rules: Vec<Rule>,
163}
164
165impl RuleSet {
166    /// The bundled default rules.
167    ///
168    /// # Panics
169    ///
170    /// Never in practice: the bundled TOML is validated by the crate's tests.
171    pub fn default_rules() -> Self {
172        Self::from_toml(DEFAULT_RULES_TOML).expect("bundled rules must parse")
173    }
174
175    /// Parses a rule set from TOML. Patterns are compiled later by
176    /// [`Inspector::new`](crate::Inspector::new).
177    pub fn from_toml(toml_text: &str) -> Result<Self, Error> {
178        Ok(toml::from_str(toml_text)?)
179    }
180
181    /// Appends the rules of `other` and takes its policy.
182    pub fn merge(&mut self, other: RuleSet) {
183        self.policy = other.policy;
184        self.rules.extend(other.rules);
185    }
186
187    /// Checks that ids are non-empty and unique.
188    pub(crate) fn validate_ids(&self) -> Result<(), Error> {
189        let mut seen = std::collections::HashSet::new();
190        for (i, rule) in self.rules.iter().enumerate() {
191            if rule.id.is_empty() {
192                return Err(Error::EmptyRuleId(i));
193            }
194            if !seen.insert(rule.id.as_str()) {
195                return Err(Error::DuplicateRuleId(rule.id.clone()));
196            }
197        }
198        Ok(())
199    }
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn luhn_accepts_known_test_numbers() {
208        assert!(luhn("4111 1111 1111 1111"));
209        assert!(luhn("5555-5555-5555-4444"));
210        assert!(luhn("378282246310005"));
211    }
212
213    #[test]
214    fn luhn_rejects_bad_checksum_and_wrong_length() {
215        assert!(!luhn("4111 1111 1111 1112"));
216        assert!(!luhn("1234"));
217        assert!(!luhn("12345678901234567890"));
218    }
219
220    #[test]
221    fn default_rules_parse_and_have_unique_ids() {
222        let rs = RuleSet::default_rules();
223        assert!(rs.rules.len() >= 10);
224        rs.validate_ids().unwrap();
225    }
226
227    #[test]
228    fn severity_is_ordered() {
229        assert!(Severity::Low < Severity::Medium);
230        assert!(Severity::Medium < Severity::High);
231        assert!(Severity::High < Severity::Critical);
232    }
233
234    #[test]
235    fn duplicate_ids_are_rejected() {
236        let rs = RuleSet::from_toml(
237            r#"
238            [[rules]]
239            id = "a"
240            category = "pii"
241            severity = "low"
242            pattern = "x"
243            [[rules]]
244            id = "a"
245            category = "pii"
246            severity = "low"
247            pattern = "y"
248            "#,
249        )
250        .unwrap();
251        assert!(matches!(rs.validate_ids(), Err(Error::DuplicateRuleId(id)) if id == "a"));
252    }
253}