Skip to main content

ailint_core/rules/security/
prompt_injection.rs

1//! AIL200 `no-prompt-injection-marker` — flag known prompt-injection sentinels.
2//!
3//! See: `docs/rules/security/AIL200.md`
4
5use regex::{Regex, RegexBuilder, RegexSetBuilder};
6use serde::Deserialize;
7
8use crate::parser::ParsedDocument;
9use crate::rules::security::{line_containing, line_of_offset, truncate_chars, AIL200};
10use crate::rules::{dictionary_lines, Rule, RuleContext, RuleId, Severity, Violation};
11
12const BUILTIN_PATTERNS: &str = include_str!("prompt_injection_patterns.txt");
13
14#[derive(Debug, Default, Deserialize)]
15struct Options {
16    #[serde(default)]
17    patterns: Option<Vec<String>>,
18    #[serde(default)]
19    extra_patterns: Option<Vec<String>>,
20}
21
22/// AIL200 no-prompt-injection-marker: flags known injection markers in text.
23#[derive(Debug, Default)]
24pub struct NoPromptInjectionMarkerRule;
25
26impl Rule for NoPromptInjectionMarkerRule {
27    fn id(&self) -> RuleId {
28        AIL200
29    }
30
31    fn default_severity(&self) -> Severity {
32        Severity::Error
33    }
34
35    fn description(&self) -> &'static str {
36        "Guidance file contains a phrase commonly used to override system prompts."
37    }
38
39    fn fix_hint(&self) -> &'static str {
40        "Remove the marker; injected content could exploit it."
41    }
42
43    fn run(&self, doc: &ParsedDocument, ctx: &RuleContext<'_>) -> Vec<Violation> {
44        let opts: Options = ctx
45            .options
46            .and_then(|v| serde_yaml::from_value(v.clone()).ok())
47            .unwrap_or_default();
48
49        let base: Vec<String> = match opts.patterns {
50            Some(p) => p,
51            None => dictionary_lines(BUILTIN_PATTERNS)
52                .into_iter()
53                .map(String::from)
54                .collect(),
55        };
56        let extras = opts.extra_patterns.unwrap_or_default();
57
58        // Compile individually (silently skipping invalid user patterns), then
59        // use a RegexSet as a single-pass prefilter over the document.
60        let compiled: Vec<(&String, Regex)> = base
61            .iter()
62            .chain(extras.iter())
63            .filter_map(|p| {
64                RegexBuilder::new(p)
65                    .case_insensitive(true)
66                    .build()
67                    .ok()
68                    .map(|re| (p, re))
69            })
70            .collect();
71        let set = RegexSetBuilder::new(compiled.iter().map(|(p, _)| p.as_str()))
72            .case_insensitive(true)
73            .build()
74            .ok();
75        let matched: Vec<usize> = match &set {
76            Some(s) => s.matches(&doc.raw).into_iter().collect(),
77            None => (0..compiled.len()).collect(),
78        };
79
80        let mut out = Vec::new();
81        for idx in matched {
82            let (_, re) = &compiled[idx];
83            for m in re.find_iter(&doc.raw) {
84                let line = line_of_offset(&doc.raw, m.start());
85                let matched = truncate_chars(m.as_str().trim_end(), 60);
86                let mut v = Violation::new(
87                    AIL200,
88                    self.default_severity(),
89                    doc.path.clone(),
90                    "prompt-injection marker",
91                )
92                .at(line, 1)
93                .with_detail(matched);
94                v.snippet = Some(line_containing(&doc.raw, m.start()));
95                out.push(v);
96            }
97        }
98        out
99    }
100}