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 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 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}