ailint_core/rules/security/
prompt_injection.rs1use 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#[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 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}