vex_api/
sanitize.rs

1//! Input sanitization and validation for security
2//!
3//! Provides functions to sanitize and validate user inputs to prevent
4//! injection attacks and ensure data integrity.
5
6use thiserror::Error;
7
8/// Sanitization errors
9#[derive(Debug, Error)]
10pub enum SanitizeError {
11    #[error("Input too long: {actual} chars (max {max})")]
12    TooLong { actual: usize, max: usize },
13
14    #[error("Input too short: {actual} chars (min {min})")]
15    TooShort { actual: usize, min: usize },
16
17    #[error("Input contains forbidden pattern: {pattern}")]
18    ForbiddenPattern { pattern: String },
19
20    #[error("Input contains invalid characters")]
21    InvalidCharacters,
22
23    #[error("Input is empty or whitespace only")]
24    EmptyInput,
25}
26
27/// Configuration for input sanitization
28#[derive(Debug, Clone)]
29pub struct SanitizeConfig {
30    /// Maximum length allowed
31    pub max_length: usize,
32    /// Minimum length required
33    pub min_length: usize,
34    /// Strip leading/trailing whitespace
35    pub trim: bool,
36    /// Check for prompt injection patterns
37    pub check_injection: bool,
38    /// Allow newlines
39    pub allow_newlines: bool,
40    /// Allow special characters
41    pub allow_special_chars: bool,
42}
43
44impl Default for SanitizeConfig {
45    fn default() -> Self {
46        Self {
47            max_length: 10000,
48            min_length: 1,
49            trim: true,
50            check_injection: true,
51            allow_newlines: true,
52            allow_special_chars: true,
53        }
54    }
55}
56
57impl SanitizeConfig {
58    /// Strict config for names/identifiers
59    pub fn strict() -> Self {
60        Self {
61            max_length: 100,
62            min_length: 1,
63            trim: true,
64            check_injection: true,
65            allow_newlines: false,
66            allow_special_chars: false,
67        }
68    }
69
70    /// Config for role descriptions
71    pub fn role() -> Self {
72        Self {
73            max_length: 500,
74            min_length: 3,
75            trim: true,
76            check_injection: true,
77            allow_newlines: true,
78            allow_special_chars: true,
79        }
80    }
81
82    /// Config for prompts (more permissive)
83    pub fn prompt() -> Self {
84        Self {
85            max_length: 50000,
86            min_length: 1,
87            trim: true,
88            check_injection: true,
89            allow_newlines: true,
90            allow_special_chars: true,
91        }
92    }
93}
94
95/// Patterns that may indicate prompt injection attempts
96/// Updated with 2024/2025 jailbreak techniques (OWASP LLM Top 10)
97const INJECTION_PATTERNS: &[&str] = &[
98    // === System Prompt Override Attempts ===
99    "ignore previous instructions",
100    "ignore all previous",
101    "disregard previous",
102    "forget previous",
103    "forget everything",
104    "new instructions:",
105    "system prompt:",
106    "override system",
107    "replace system prompt",
108    // === Role/Identity Attacks ===
109    "you are now",
110    "pretend you are",
111    "act as if",
112    "roleplay as",
113    "simulate being",
114    "imagine you are",
115    "from now on you are",
116    "your new identity",
117    // === DAN-Style Jailbreaks (Do Anything Now) ===
118    "dan mode",
119    "developer mode",
120    "god mode",
121    "sudo mode",
122    "admin mode",
123    "unrestricted mode",
124    "no restrictions",
125    "without limitations",
126    "without ethical",
127    "without safety",
128    "enabled: true", // Policy puppetry
129    // === Jailbreak Keywords ===
130    "jailbreak",
131    "unlock",
132    "bypass",
133    "disable filter",
134    "remove restrictions",
135    "break free",
136    "escape constraints",
137    "override safety",
138    // === Encoding/Obfuscation Attacks ===
139    "base64:",
140    "decode base64",
141    "\\x",
142    "\\u00",
143    "[system]", // Special token injection
144    "[/system]",
145    "<|im_start|>", // ChatML tokens
146    "<|im_end|>",
147    // === Multi-Turn Escalation Patterns ===
148    "hypothetically speaking",
149    "for educational purposes",
150    "in a fictional scenario",
151    "for research purposes",
152    "just between us",
153    // === Prompt Leaking Attempts ===
154    "repeat your instructions",
155    "show me your prompt",
156    "what are your rules",
157    "print your system prompt",
158    "reveal your instructions",
159    // === Context Manipulation ===
160    "end of conversation",
161    "new conversation",
162    "reset context",
163    "clear memory",
164];
165
166/// Sanitize and validate input text
167pub fn sanitize(input: &str, config: &SanitizeConfig) -> Result<String, SanitizeError> {
168    // Trim if configured
169    let text = if config.trim { input.trim() } else { input };
170
171    // Check empty
172    if text.is_empty() {
173        return Err(SanitizeError::EmptyInput);
174    }
175
176    // Normalize Unicode to NFC form and strip zero-width characters
177    // This prevents homoglyph attacks (e.g., using Cyrillic 'а' instead of Latin 'a')
178    let normalized: String = text
179        .chars()
180        .filter(|c| {
181            // Strip zero-width characters commonly used to bypass filters
182            !matches!(
183                *c,
184                '\u{200B}' | // Zero width space
185                '\u{200C}' | // Zero width non-joiner
186                '\u{200D}' | // Zero width joiner
187                '\u{FEFF}' | // Byte order mark
188                '\u{00AD}' // Soft hyphen
189            )
190        })
191        // Convert common lookalikes to ASCII (basic confusable mitigation)
192        .map(|c| match c {
193            // Cyrillic lookalikes
194            '\u{0430}' => 'a', // Cyrillic а
195            '\u{0435}' => 'e', // Cyrillic е
196            '\u{043E}' => 'o', // Cyrillic о
197            '\u{0440}' => 'p', // Cyrillic р
198            '\u{0441}' => 'c', // Cyrillic с
199            '\u{0445}' => 'x', // Cyrillic х
200            // Fullwidth ASCII
201            c if ('\u{FF01}'..='\u{FF5E}').contains(&c) => {
202                char::from_u32(c as u32 - 0xFEE0).unwrap_or(c)
203            }
204            _ => c,
205        })
206        .collect();
207
208    let text = &normalized;
209
210    // Check length
211    if text.len() < config.min_length {
212        return Err(SanitizeError::TooShort {
213            actual: text.len(),
214            min: config.min_length,
215        });
216    }
217
218    if text.len() > config.max_length {
219        return Err(SanitizeError::TooLong {
220            actual: text.len(),
221            max: config.max_length,
222        });
223    }
224
225    // Check for newlines if not allowed
226    if !config.allow_newlines && text.contains('\n') {
227        return Err(SanitizeError::InvalidCharacters);
228    }
229
230    // Check for special characters if not allowed
231    if !config.allow_special_chars {
232        for c in text.chars() {
233            if !c.is_alphanumeric() && c != ' ' && c != '-' && c != '_' {
234                return Err(SanitizeError::InvalidCharacters);
235            }
236        }
237    }
238
239    // Check for injection patterns
240    if config.check_injection {
241        let lower = text.to_lowercase();
242        for pattern in INJECTION_PATTERNS {
243            if lower.contains(pattern) {
244                tracing::warn!(pattern = pattern, "Potential prompt injection detected");
245                return Err(SanitizeError::ForbiddenPattern {
246                    pattern: pattern.to_string(),
247                });
248            }
249        }
250    }
251
252    // Remove null bytes and other control characters (except newlines/tabs if allowed)
253    let sanitized: String = text
254        .chars()
255        .filter(|c| {
256            if *c == '\n' || *c == '\t' {
257                config.allow_newlines
258            } else {
259                !c.is_control()
260            }
261        })
262        .collect();
263
264    Ok(sanitized)
265}
266
267/// Sanitize a name field (strict)
268pub fn sanitize_name(input: &str) -> Result<String, SanitizeError> {
269    sanitize(input, &SanitizeConfig::strict())
270}
271
272/// Sanitize a role description
273pub fn sanitize_role(input: &str) -> Result<String, SanitizeError> {
274    sanitize(input, &SanitizeConfig::role())
275}
276
277/// Sanitize a prompt
278pub fn sanitize_prompt(input: &str) -> Result<String, SanitizeError> {
279    sanitize(input, &SanitizeConfig::prompt())
280}
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    #[test]
287    fn test_sanitize_valid_input() {
288        let result = sanitize("Hello world", &SanitizeConfig::default());
289        assert!(result.is_ok());
290        assert_eq!(result.unwrap(), "Hello world");
291    }
292
293    #[test]
294    fn test_sanitize_trims_whitespace() {
295        let result = sanitize("  Hello  ", &SanitizeConfig::default());
296        assert!(result.is_ok());
297        assert_eq!(result.unwrap(), "Hello");
298    }
299
300    #[test]
301    fn test_sanitize_rejects_empty() {
302        let result = sanitize("", &SanitizeConfig::default());
303        assert!(matches!(result, Err(SanitizeError::EmptyInput)));
304    }
305
306    #[test]
307    fn test_sanitize_rejects_too_long() {
308        let long_input = "a".repeat(101);
309        let result = sanitize(&long_input, &SanitizeConfig::strict());
310        assert!(matches!(result, Err(SanitizeError::TooLong { .. })));
311    }
312
313    #[test]
314    fn test_sanitize_detects_injection() {
315        let result = sanitize(
316            "Please ignore previous instructions",
317            &SanitizeConfig::default(),
318        );
319        assert!(matches!(
320            result,
321            Err(SanitizeError::ForbiddenPattern { .. })
322        ));
323    }
324
325    #[test]
326    fn test_sanitize_name_rejects_special_chars() {
327        let result = sanitize_name("agent<script>");
328        assert!(matches!(result, Err(SanitizeError::InvalidCharacters)));
329    }
330
331    #[test]
332    fn test_sanitize_removes_control_chars() {
333        let input = "Hello\x00World";
334        let result = sanitize(input, &SanitizeConfig::default());
335        assert!(result.is_ok());
336        assert_eq!(result.unwrap(), "HelloWorld");
337    }
338
339    #[test]
340    fn test_all_injection_patterns() {
341        for pattern in INJECTION_PATTERNS {
342            let input = format!("some benign text then {} and more text", pattern);
343            let result = sanitize(&input, &SanitizeConfig::prompt());
344            assert!(
345                matches!(result, Err(SanitizeError::ForbiddenPattern { .. })),
346                "Failed to detect pattern: {}",
347                pattern
348            );
349
350            // Test case insensitivity
351            let input_upper = format!(
352                "some benign text then {} and more text",
353                pattern.to_uppercase()
354            );
355            let result_upper = sanitize(&input_upper, &SanitizeConfig::prompt());
356            assert!(
357                matches!(result_upper, Err(SanitizeError::ForbiddenPattern { .. })),
358                "Failed to detect uppercase pattern: {}",
359                pattern
360            );
361        }
362    }
363}
vex_api/sanitize.rs

vex_api/
sanitize.rs
