gaze_recognizers/
dictionary.rs

1use std::collections::hash_map::DefaultHasher;
2use std::collections::HashMap;
3use std::hash::{Hash, Hasher};
4use std::sync::{Arc, Mutex};
5
6use aho_corasick::{AhoCorasick, AhoCorasickBuilder};
7use gaze_types::{
8    Candidate, ConflictTier, DetectContext, DictionaryEntry, LocaleTag, PiiClass, Recognizer,
9};
10
11/// Lookup-based [`Recognizer`] for tenant-specific PII.
12///
13/// Matches exact strings from a runtime-supplied dictionary: order IDs, song
14/// titles, customer codes, or any domain-specific identifiers that regex cannot
15/// reliably detect.
16///
17/// Supply dictionaries at runtime via [`DetectContext`] or the CLI's
18/// `--context-json` wiring. Dictionary entries are session-scoped detection
19/// inputs and are not stored in the audit log.
20pub struct DictionaryRecognizer {
21    id: String,
22    class: PiiClass,
23    dictionary_name: String,
24    case_sensitive: bool,
25    token_family: String,
26    locales: Vec<LocaleTag>,
27    score: f32,
28    priority: i32,
29    compiled_dictionaries: Mutex<HashMap<DictionaryCacheKey, Arc<AhoCorasick>>>,
30}
31
32impl DictionaryRecognizer {
33    pub fn new(
34        id: impl Into<String>,
35        class: PiiClass,
36        dictionary_name: impl Into<String>,
37        case_sensitive: bool,
38        token_family: impl Into<String>,
39    ) -> Self {
40        Self::with_rulepack_fields(
41            id,
42            class,
43            dictionary_name,
44            case_sensitive,
45            token_family,
46            vec![LocaleTag::Global],
47            1.0,
48            0,
49        )
50    }
51
52    #[allow(clippy::too_many_arguments)]
53    pub fn with_rulepack_fields(
54        id: impl Into<String>,
55        class: PiiClass,
56        dictionary_name: impl Into<String>,
57        case_sensitive: bool,
58        token_family: impl Into<String>,
59        locales: Vec<LocaleTag>,
60        score: f32,
61        priority: i32,
62    ) -> Self {
63        Self {
64            id: id.into(),
65            class,
66            dictionary_name: dictionary_name.into(),
67            case_sensitive,
68            token_family: token_family.into(),
69            locales,
70            score,
71            priority,
72            compiled_dictionaries: Mutex::new(HashMap::new()),
73        }
74    }
75
76    pub fn dictionary_name(&self) -> &str {
77        &self.dictionary_name
78    }
79
80    pub fn case_sensitive(&self) -> bool {
81        self.case_sensitive
82    }
83
84    fn automaton_for(&self, entry: &DictionaryEntry) -> Arc<AhoCorasick> {
85        let key = DictionaryCacheKey::from_entry(entry);
86        let mut compiled = self
87            .compiled_dictionaries
88            .lock()
89            .expect("dictionary automaton cache poisoned");
90        Arc::clone(compiled.entry(key).or_insert_with(|| {
91            Arc::new(
92                AhoCorasickBuilder::new()
93                    .ascii_case_insensitive(!entry.case_sensitive())
94                    .build(entry.terms())
95                    .expect("DictionaryEntry validates terms before automaton construction"),
96            )
97        }))
98    }
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
102struct DictionaryCacheKey {
103    terms_hash: u64,
104    case_sensitive: bool,
105}
106
107impl DictionaryCacheKey {
108    fn from_entry(entry: &DictionaryEntry) -> Self {
109        let mut hasher = DefaultHasher::new();
110        entry.terms().hash(&mut hasher);
111        Self {
112            terms_hash: hasher.finish(),
113            case_sensitive: entry.case_sensitive(),
114        }
115    }
116}
117
118impl Recognizer for DictionaryRecognizer {
119    fn id(&self) -> &str {
120        &self.id
121    }
122
123    fn supported_class(&self) -> &PiiClass {
124        &self.class
125    }
126
127    fn detect(&self, input: &str, ctx: &DetectContext<'_>) -> Vec<Candidate> {
128        let Some(entry) = ctx.dictionaries.get(&self.dictionary_name) else {
129            return Vec::new();
130        };
131        let automaton = self.automaton_for(entry);
132
133        automaton
134            .find_iter(input)
135            .map(|m| {
136                Candidate::new(
137                    m.start()..m.end(),
138                    self.class.clone(),
139                    self.id.clone(),
140                    self.score,
141                    self.priority,
142                    Some(input[m.start()..m.end()].to_string()),
143                    self.token_family.clone(),
144                    format!(
145                        "dictionary:{}[#{}]",
146                        self.dictionary_name,
147                        m.pattern().as_usize()
148                    ),
149                    ConflictTier::None,
150                    Vec::new(),
151                )
152            })
153            .collect()
154    }
155
156    fn token_family(&self) -> &str {
157        &self.token_family
158    }
159
160    fn locales(&self) -> &[LocaleTag] {
161        &self.locales
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use std::collections::HashMap;
168
169    use gaze::{
170        dictionary_bundle_from_context, ContextDictionary, RecognizerRegistry, TypedContext,
171    };
172    use serde_json::Map;
173
174    use super::*;
175
176    #[test]
177    fn recognizer_detects_dictionary_hits_from_context_bundle() {
178        let ctx = TypedContext {
179            dictionaries: HashMap::from([(
180                "dict_alpha".to_string(),
181                ContextDictionary {
182                    terms: vec!["AAA-12345".to_string()],
183                    case_sensitive: true,
184                },
185            )]),
186            class_map: HashMap::new(),
187            fields: Map::new(),
188        };
189        let bundle = dictionary_bundle_from_context(&ctx);
190        let detect_context = DetectContext::new(&[LocaleTag::Global], &bundle);
191        let recognizer = DictionaryRecognizer::new(
192            "dict/dict_alpha",
193            PiiClass::Custom("class_alpha".to_string()),
194            "dict_alpha",
195            true,
196            "counter",
197        );
198
199        let hits = recognizer.detect("Customer bought AAA-12345", &detect_context);
200        assert_eq!(hits.len(), 1);
201        assert_eq!(hits[0].span, 16..25);
202        assert_eq!(hits[0].class, PiiClass::Custom("class_alpha".to_string()));
203    }
204
205    #[test]
206    fn recognizer_locale_gates_dictionary_hits() {
207        let ctx = TypedContext {
208            dictionaries: HashMap::from([(
209                "songs".to_string(),
210                ContextDictionary {
211                    terms: vec!["Bohemian Rhapsody".to_string()],
212                    case_sensitive: false,
213                },
214            )]),
215            class_map: HashMap::new(),
216            fields: Map::new(),
217        };
218        let bundle = dictionary_bundle_from_context(&ctx);
219        let detect_context = DetectContext::new(&[LocaleTag::EnUs], &bundle);
220        let recognizer = DictionaryRecognizer::with_rulepack_fields(
221            "dict/songs",
222            PiiClass::Custom("song".to_string()),
223            "songs",
224            false,
225            "counter",
226            vec![LocaleTag::DeDe],
227            1.0,
228            0,
229        );
230
231        let registry = RecognizerRegistry::builder().register(recognizer).build();
232        let hits = registry.detect_all("Listening to bohemian rhapsody", &detect_context);
233        assert!(hits.is_empty());
234    }
235
236    #[test]
237    fn dictionary_recognizer_emits_per_term_source() {
238        let ctx = TypedContext {
239            dictionaries: HashMap::from([(
240                "songs".to_string(),
241                ContextDictionary {
242                    terms: vec![
243                        "alpha-one".to_string(),
244                        "bravo-two".to_string(),
245                        "charlie-three".to_string(),
246                    ],
247                    case_sensitive: true,
248                },
249            )]),
250            class_map: HashMap::new(),
251            fields: Map::new(),
252        };
253        let bundle = dictionary_bundle_from_context(&ctx);
254        let detect_context = DetectContext::new(&[LocaleTag::Global], &bundle);
255        let recognizer = DictionaryRecognizer::new(
256            "dict/songs",
257            PiiClass::Custom("song".to_string()),
258            "songs",
259            true,
260            "counter",
261        );
262
263        let hits = recognizer.detect(
264            "first alpha-one, second bravo-two, third charlie-three",
265            &detect_context,
266        );
267
268        assert_eq!(hits.len(), 3);
269        let source_shape = regex::Regex::new(r"^dictionary:[a-z_]+\[#\d+\]$").unwrap();
270        for hit in &hits {
271            assert!(
272                source_shape.is_match(&hit.source),
273                "unexpected source shape: {}",
274                hit.source
275            );
276        }
277        assert_eq!(hits[0].source, "dictionary:songs[#0]");
278        assert_eq!(hits[1].source, "dictionary:songs[#1]");
279        assert_eq!(hits[2].source, "dictionary:songs[#2]");
280    }
281
282    #[test]
283    fn dictionary_recognizer_source_index_matches_automaton_for_duplicate_terms() {
284        let ctx = TypedContext {
285            dictionaries: HashMap::from([(
286                "songs".to_string(),
287                ContextDictionary {
288                    terms: vec![
289                        "same-song".to_string(),
290                        "same-song".to_string(),
291                        "other-song".to_string(),
292                    ],
293                    case_sensitive: true,
294                },
295            )]),
296            class_map: HashMap::new(),
297            fields: Map::new(),
298        };
299        let bundle = dictionary_bundle_from_context(&ctx);
300        let detect_context = DetectContext::new(&[LocaleTag::Global], &bundle);
301        let recognizer = DictionaryRecognizer::new(
302            "dict/songs",
303            PiiClass::Custom("song".to_string()),
304            "songs",
305            true,
306            "counter",
307        );
308
309        let hits = recognizer.detect("same-song then other-song", &detect_context);
310
311        assert_eq!(hits.len(), 2);
312        assert_eq!(hits[0].source, "dictionary:songs[#0]");
313        assert_eq!(hits[1].source, "dictionary:songs[#2]");
314    }
315}
gaze_recognizers/dictionary.rs

gaze_recognizers/
dictionary.rs
