返回 CodeWhale
regex_cache.rs
根目录 / crates / runtime / src / regex_cache.rs
1 //! Compiled-regex cache for patterns a user or model supplies (tool search,
2 //! purge). Those patterns are untrusted input, so compilation is bounded:
3 //! the pattern length is capped, the compiled program and lazy DFA have
4 //! explicit size limits, compilation runs outside the cache lock, and a
5 //! rejected pattern is remembered so repeating it costs a lookup.
6
7 use std::num::NonZeroUsize;
8 use std::sync::{Mutex, OnceLock};
9
10 use lru::LruCache;
11 use regex::{Regex, RegexBuilder};
12
13 const DEFAULT_USER_REGEX_CACHE_CAPACITY: usize = 64;
14 /// Longest pattern accepted, in bytes. Search and purge patterns are short.
15 pub const MAX_USER_REGEX_PATTERN_BYTES: usize = 4 * 1024;
16 /// Compiled program limit (the `regex` default is 10 MiB).
17 const USER_REGEX_SIZE_LIMIT: usize = 1024 * 1024;
18 /// Lazy DFA cache limit per matcher.
19 const USER_REGEX_DFA_SIZE_LIMIT: usize = 1024 * 1024;
20
21 static USER_REGEX_CACHE: OnceLock<UserRegexCache> = OnceLock::new();
22
23 pub fn compile_user_regex(pattern: &str) -> Result<Regex, regex::Error> {
24 user_regex_cache().compile(pattern)
25 }
26
27 fn user_regex_cache() -> &'static UserRegexCache {
28 USER_REGEX_CACHE.get_or_init(UserRegexCache::new)
29 }
30
31 fn compile_bounded(pattern: &str) -> Result<Regex, regex::Error> {
32 if pattern.len() > MAX_USER_REGEX_PATTERN_BYTES {
33 return Err(regex::Error::Syntax(format!(
34 "pattern is {} bytes; the limit is {MAX_USER_REGEX_PATTERN_BYTES}",
35 pattern.len()
36 )));
37 }
38 RegexBuilder::new(pattern)
39 .size_limit(USER_REGEX_SIZE_LIMIT)
40 .dfa_size_limit(USER_REGEX_DFA_SIZE_LIMIT)
41 .build()
42 }
43
44 struct UserRegexCache {
45 /// Compiled patterns and rejected ones alike, so a repeated bad pattern
46 /// is not recompiled.
47 inner: Mutex<LruCache<String, Result<Regex, regex::Error>>>,
48 }
49
50 impl UserRegexCache {
51 fn new() -> Self {
52 Self::with_capacity(
53 NonZeroUsize::new(DEFAULT_USER_REGEX_CACHE_CAPACITY).expect("non-zero capacity"),
54 )
55 }
56
57 fn with_capacity(capacity: NonZeroUsize) -> Self {
58 Self {
59 inner: Mutex::new(LruCache::new(capacity)),
60 }
61 }
62
63 fn compile(&self, pattern: &str) -> Result<Regex, regex::Error> {
64 // An over-long pattern is refused before it is hashed or stored.
65 if pattern.len() > MAX_USER_REGEX_PATTERN_BYTES {
66 return compile_bounded(pattern);
67 }
68 if let Ok(mut cache) = self.inner.lock()
69 && let Some(cached) = cache.get(pattern)
70 {
71 return cached.clone();
72 }
73 // Compile without the lock: one slow pattern must not stall every
74 // other caller. Two racing misses both compile; the result is equal.
75 let compiled = compile_bounded(pattern);
76 if let Ok(mut cache) = self.inner.lock() {
77 cache.put(pattern.to_string(), compiled.clone());
78 }
79 compiled
80 }
81
82 #[cfg(test)]
83 fn len(&self) -> usize {
84 self.inner.lock().expect("cache lock").len()
85 }
86
87 #[cfg(test)]
88 fn contains(&self, pattern: &str) -> bool {
89 self.inner.lock().expect("cache lock").contains(pattern)
90 }
91 }
92
93 #[cfg(test)]
94 mod tests {
95 use super::*;
96
97 #[test]
98 fn repeated_pattern_uses_one_cache_entry() {
99 let cache = UserRegexCache::with_capacity(NonZeroUsize::new(2).unwrap());
100
101 let first = cache.compile("alpha|beta").expect("regex compiles");
102 let second = cache.compile("alpha|beta").expect("regex cache hit");
103
104 assert!(first.is_match("alpha"));
105 assert!(second.is_match("beta"));
106 assert_eq!(cache.len(), 1);
107 }
108
109 #[test]
110 fn capacity_evicts_least_recently_used_pattern() {
111 let cache = UserRegexCache::with_capacity(NonZeroUsize::new(2).unwrap());
112
113 cache.compile("one").expect("one compiles");
114 cache.compile("two").expect("two compiles");
115 cache.compile("one").expect("one is refreshed");
116 cache.compile("three").expect("three compiles");
117
118 assert!(cache.contains("one"));
119 assert!(!cache.contains("two"));
120 assert!(cache.contains("three"));
121 }
122
123 #[test]
124 fn invalid_pattern_is_remembered_as_invalid() {
125 let cache = UserRegexCache::with_capacity(NonZeroUsize::new(2).unwrap());
126
127 assert!(cache.compile("[").is_err());
128 assert!(cache.contains("["));
129 assert!(cache.compile("[").is_err());
130
131 assert_eq!(cache.len(), 1);
132 }
133
134 #[test]
135 fn oversized_patterns_are_refused_without_being_cached() {
136 let cache = UserRegexCache::with_capacity(NonZeroUsize::new(2).unwrap());
137
138 let long = "a".repeat(MAX_USER_REGEX_PATTERN_BYTES + 1);
139 assert!(cache.compile(&long).is_err());
140 assert_eq!(cache.len(), 0);
141
142 // Within the length cap, a pattern whose compiled form exceeds the
143 // program limit is refused rather than built.
144 let explosive = r"(?:\w{100}){100}";
145 assert!(explosive.len() <= MAX_USER_REGEX_PATTERN_BYTES);
146 assert!(matches!(
147 cache.compile(explosive),
148 Err(regex::Error::CompiledTooBig(_))
149 ));
150 }
151 }
152
152 lines RUST