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 run(&self, doc: &ParsedDocument, ctx: &RuleContext<'_>) -> Vec<Violation> {
36        let opts: Options = ctx
37            .options
38            .and_then(|v| serde_yaml::from_value(v.clone()).ok())
39            .unwrap_or_default();
40
41        let base: Vec<String> = match opts.patterns {
42            Some(p) => p,
43            None => dictionary_lines(BUILTIN_PATTERNS)
44                .into_iter()
45                .map(String::from)
46                .collect(),
47        };
48        let extras = opts.extra_patterns.unwrap_or_default();
49
50        // Compile individually (silently skipping invalid user patterns), then
51        // use a RegexSet as a single-pass prefilter over the document.
52        let compiled: Vec<(&String, Regex)> = base
53            .iter()
54            .chain(extras.iter())
55            .filter_map(|p| {
56                RegexBuilder::new(p)
57                    .case_insensitive(true)
58                    .build()
59                    .ok()
60                    .map(|re| (p, re))
61            })
62            .collect();
63        let set = RegexSetBuilder::new(compiled.iter().map(|(p, _)| p.as_str()))
64            .case_insensitive(true)
65            .build()
66            .ok();
67        let matched: Vec<usize> = match &set {
68            Some(s) => s.matches(&doc.raw).into_iter().collect(),
69            None => (0..compiled.len()).collect(),
70        };
71
72        let mut out = Vec::new();
73        for idx in matched {
74            let (_, re) = &compiled[idx];
75            for m in re.find_iter(&doc.raw) {
76                let line = line_of_offset(&doc.raw, m.start());
77                let matched = truncate_chars(m.as_str().trim_end(), 60);
78                let mut v = Violation::new(
79                    AIL200,
80                    self.default_severity(),
81                    doc.path.clone(),
82                    format!("possible prompt-injection marker: '{}'", matched),
83                )
84                .at(line, 1);
85                v.snippet = Some(line_containing(&doc.raw, m.start()));
86                out.push(v);
87            }
88        }
89        out
90    }
91}