Skip to main content

cloakrs_core/
scanner.rs

1//! Scanner orchestration.
2
3use crate::masker::deduplicate;
4use crate::{
5    apply_mask, CloakError, Confidence, EntityType, Locale, MaskStrategy, PiiEntity, Recognizer,
6    RecognizerRegistry, Result,
7};
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::time::Instant;
11
12/// Builder for configuring a [`Scanner`].
13///
14/// # Examples
15///
16/// ```
17/// use cloakrs_core::{Locale, Scanner};
18///
19/// let scanner = Scanner::builder().locale(Locale::US).build();
20/// assert!(scanner.is_err());
21/// ```
22pub struct ScannerBuilder {
23    registry: RecognizerRegistry,
24    locale: Locale,
25    strategy: Option<MaskStrategy>,
26    min_confidence: Confidence,
27    allow_list: Vec<String>,
28    deny_list: Vec<String>,
29}
30
31impl Default for ScannerBuilder {
32    fn default() -> Self {
33        Self {
34            registry: RecognizerRegistry::new(),
35            locale: Locale::Universal,
36            strategy: Some(MaskStrategy::default()),
37            min_confidence: Confidence::ZERO,
38            allow_list: Vec::new(),
39            deny_list: Vec::new(),
40        }
41    }
42}
43
44impl ScannerBuilder {
45    /// Creates a scanner builder.
46    #[must_use]
47    pub fn new() -> Self {
48        Self::default()
49    }
50
51    /// Creates a scanner builder from an existing recognizer registry.
52    #[must_use]
53    pub fn from_registry(registry: RecognizerRegistry) -> Self {
54        Self {
55            registry,
56            ..Self::default()
57        }
58    }
59
60    /// Sets the scanner locale.
61    #[must_use]
62    pub fn locale(mut self, locale: Locale) -> Self {
63        self.locale = locale;
64        self
65    }
66
67    /// Sets the masking strategy.
68    #[must_use]
69    pub fn strategy(mut self, strategy: MaskStrategy) -> Self {
70        self.strategy = Some(strategy);
71        self
72    }
73
74    /// Disables masked text generation.
75    #[must_use]
76    pub fn without_masking(mut self) -> Self {
77        self.strategy = None;
78        self
79    }
80
81    /// Adds a recognizer.
82    #[must_use]
83    pub fn recognizer<R>(mut self, recognizer: R) -> Self
84    where
85        R: Recognizer + 'static,
86    {
87        self.registry.register(recognizer);
88        self
89    }
90
91    /// Adds a boxed recognizer.
92    #[must_use]
93    pub fn boxed_recognizer(mut self, recognizer: Box<dyn Recognizer>) -> Self {
94        self.registry.register_boxed(recognizer);
95        self
96    }
97
98    /// Sets the minimum confidence threshold.
99    pub fn min_confidence(mut self, confidence: f64) -> Result<Self> {
100        self.min_confidence = Confidence::new(confidence)?;
101        Ok(self)
102    }
103
104    /// Adds literal values that should never be reported or masked.
105    #[must_use]
106    pub fn allow_list<I, S>(mut self, values: I) -> Self
107    where
108        I: IntoIterator<Item = S>,
109        S: Into<String>,
110    {
111        self.allow_list.extend(
112            values
113                .into_iter()
114                .map(Into::into)
115                .filter(|value| !value.is_empty()),
116        );
117        self
118    }
119
120    /// Adds literal values that should always be reported and masked.
121    #[must_use]
122    pub fn deny_list<I, S>(mut self, values: I) -> Self
123    where
124        I: IntoIterator<Item = S>,
125        S: Into<String>,
126    {
127        self.deny_list.extend(
128            values
129                .into_iter()
130                .map(Into::into)
131                .filter(|value| !value.is_empty()),
132        );
133        self
134    }
135
136    /// Builds a scanner.
137    pub fn build(self) -> Result<Scanner> {
138        if self.registry.is_empty() {
139            return Err(CloakError::NoRecognizers);
140        }
141
142        Ok(Scanner {
143            registry: self.registry,
144            locale: self.locale,
145            strategy: self.strategy,
146            min_confidence: self.min_confidence,
147            allow_list: self.allow_list,
148            deny_list: self.deny_list,
149        })
150    }
151}
152
153/// The main scanner that ties detection and masking together.
154pub struct Scanner {
155    registry: RecognizerRegistry,
156    locale: Locale,
157    strategy: Option<MaskStrategy>,
158    min_confidence: Confidence,
159    allow_list: Vec<String>,
160    deny_list: Vec<String>,
161}
162
163impl Scanner {
164    /// Creates a scanner builder.
165    #[must_use]
166    pub fn builder() -> ScannerBuilder {
167        ScannerBuilder::new()
168    }
169
170    /// Scans text and returns findings, optional masked text, and stats.
171    pub fn scan(&self, text: &str) -> Result<ScanResult> {
172        let started = Instant::now();
173        let mut findings = self.registry.scan_locale(text, &self.locale);
174        findings.extend(deny_list_findings(text, &self.deny_list));
175        let allow_spans = literal_spans(text, &self.allow_list);
176        findings.retain(|finding| !allow_spans.iter().any(|span| span.overlaps(finding.span)));
177        findings.retain(|finding| finding.confidence >= self.min_confidence);
178        findings = deduplicate_for_reporting(&findings);
179        findings.sort_by_key(|finding| finding.span.start);
180
181        let masked_text = self
182            .strategy
183            .as_ref()
184            .map(|strategy| apply_mask(text, &findings, strategy))
185            .transpose()?;
186
187        let stats = ScanStats::from_findings(&findings, started.elapsed().as_millis(), text.len());
188
189        Ok(ScanResult {
190            findings,
191            masked_text,
192            stats,
193        })
194    }
195}
196
197fn deny_list_findings(text: &str, deny_list: &[String]) -> Vec<PiiEntity> {
198    literal_spans(text, deny_list)
199        .into_iter()
200        .map(|span| PiiEntity {
201            entity_type: EntityType::Custom("DenyList".to_string()),
202            span,
203            text: text[span.start..span.end].to_string(),
204            confidence: Confidence::ONE,
205            recognizer_id: "deny_list_v1".to_string(),
206        })
207        .collect()
208}
209
210fn literal_spans(text: &str, values: &[String]) -> Vec<crate::Span> {
211    let mut spans = Vec::new();
212    for value in values {
213        if value.is_empty() {
214            continue;
215        }
216        let mut search_start = 0;
217        while let Some(offset) = text[search_start..].find(value) {
218            let start = search_start + offset;
219            let end = start + value.len();
220            spans.push(crate::Span::new(start, end));
221            search_start = end;
222        }
223    }
224    spans
225}
226
227fn deduplicate_for_reporting(findings: &[PiiEntity]) -> Vec<PiiEntity> {
228    let mut sorted = findings.to_vec();
229    sorted.sort_by_key(|finding| (finding.span.start, std::cmp::Reverse(finding.span.end)));
230
231    let mut keep: Vec<PiiEntity> = Vec::with_capacity(sorted.len());
232    for finding in sorted {
233        if keep
234            .iter()
235            .any(|kept| should_preserve_nested_url_query(kept, &finding))
236        {
237            keep.push(finding);
238            continue;
239        }
240
241        if let Some(overlap_index) = keep
242            .iter()
243            .rposition(|kept| finding.span.overlaps(kept.span))
244        {
245            if should_keep_existing_url_query(&keep[overlap_index], &finding) {
246                continue;
247            }
248            if should_replace_with_url_query(&keep[overlap_index], &finding) {
249                keep[overlap_index] = finding;
250                continue;
251            }
252
253            let merged = deduplicate(&[keep[overlap_index].clone(), finding])
254                .into_iter()
255                .next()
256                .unwrap_or_else(|| keep[overlap_index].clone());
257            keep[overlap_index] = merged;
258            continue;
259        }
260        keep.push(finding);
261    }
262
263    keep
264}
265
266fn should_preserve_nested_url_query(outer: &PiiEntity, inner: &PiiEntity) -> bool {
267    outer.entity_type == EntityType::Url
268        && inner.recognizer_id.starts_with("url_query_")
269        && inner.span.start >= outer.span.start
270        && inner.span.end <= outer.span.end
271}
272
273fn should_keep_existing_url_query(existing: &PiiEntity, incoming: &PiiEntity) -> bool {
274    existing.span == incoming.span
275        && existing.entity_type == incoming.entity_type
276        && existing.recognizer_id.starts_with("url_query_")
277        && !incoming.recognizer_id.starts_with("url_query_")
278}
279
280fn should_replace_with_url_query(existing: &PiiEntity, incoming: &PiiEntity) -> bool {
281    existing.span == incoming.span
282        && existing.entity_type == incoming.entity_type
283        && !existing.recognizer_id.starts_with("url_query_")
284        && incoming.recognizer_id.starts_with("url_query_")
285}
286
287/// Result of scanning text.
288#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
289pub struct ScanResult {
290    /// All PII entities found.
291    pub findings: Vec<PiiEntity>,
292    /// Text with PII masked, when masking is enabled.
293    pub masked_text: Option<String>,
294    /// Statistics about the scan.
295    pub stats: ScanStats,
296}
297
298/// Statistics about a scan.
299#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
300pub struct ScanStats {
301    /// Total number of findings.
302    pub total_findings: usize,
303    /// Count of findings per entity type.
304    pub findings_by_type: HashMap<EntityType, usize>,
305    /// Wall-clock scan duration in milliseconds.
306    pub scan_duration_ms: u64,
307    /// Number of input bytes scanned.
308    pub bytes_scanned: usize,
309}
310
311impl ScanStats {
312    fn from_findings(findings: &[PiiEntity], duration_ms: u128, bytes_scanned: usize) -> Self {
313        let mut findings_by_type = HashMap::new();
314        for finding in findings {
315            *findings_by_type
316                .entry(finding.entity_type.clone())
317                .or_insert(0) += 1;
318        }
319
320        Self {
321            total_findings: findings.len(),
322            findings_by_type,
323            scan_duration_ms: duration_ms.try_into().unwrap_or(u64::MAX),
324            bytes_scanned,
325        }
326    }
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332    use crate::Span;
333
334    struct EmailRecognizer;
335
336    impl Recognizer for EmailRecognizer {
337        fn id(&self) -> &str {
338            "email_test_v1"
339        }
340
341        fn entity_type(&self) -> EntityType {
342            EntityType::Email
343        }
344
345        fn supported_locales(&self) -> &[Locale] {
346            &[]
347        }
348
349        fn scan(&self, text: &str) -> Vec<PiiEntity> {
350            let Some(start) = text
351                .find('@')
352                .and_then(|at| text[..at].rfind(' ').map(|space| space + 1).or(Some(0)))
353            else {
354                return Vec::new();
355            };
356            let end = text[start..]
357                .find(' ')
358                .map_or(text.len(), |offset| start + offset);
359            vec![PiiEntity {
360                entity_type: EntityType::Email,
361                span: Span::new(start, end),
362                text: text[start..end].to_string(),
363                confidence: Confidence::new(0.95).unwrap(),
364                recognizer_id: self.id().to_string(),
365            }]
366        }
367    }
368
369    #[test]
370    fn test_scanner_builder_without_recognizers_errors() {
371        assert!(Scanner::builder().build().is_err());
372    }
373
374    #[test]
375    fn test_scanner_scan_returns_findings_and_masked_text() {
376        let scanner = Scanner::builder()
377            .recognizer(EmailRecognizer)
378            .build()
379            .unwrap();
380        let result = scanner.scan("Contact user@example.com").unwrap();
381        assert_eq!(result.findings.len(), 1);
382        assert_eq!(result.masked_text.as_deref(), Some("Contact [EMAIL]"));
383    }
384
385    #[test]
386    fn test_scanner_without_masking_returns_no_masked_text() {
387        let scanner = Scanner::builder()
388            .recognizer(EmailRecognizer)
389            .without_masking()
390            .build()
391            .unwrap();
392        let result = scanner.scan("Contact user@example.com").unwrap();
393        assert!(result.masked_text.is_none());
394    }
395
396    #[test]
397    fn test_scanner_min_confidence_filters_low_confidence_findings() {
398        let scanner = Scanner::builder()
399            .recognizer(EmailRecognizer)
400            .min_confidence(1.0)
401            .unwrap()
402            .build()
403            .unwrap();
404        let result = scanner.scan("Contact user@example.com").unwrap();
405        assert!(result.findings.is_empty());
406    }
407
408    #[test]
409    fn test_scanner_allow_list_suppresses_overlapping_findings() {
410        let scanner = Scanner::builder()
411            .recognizer(EmailRecognizer)
412            .allow_list(["user@example.com"])
413            .build()
414            .unwrap();
415        let result = scanner.scan("Contact user@example.com").unwrap();
416        assert!(result.findings.is_empty());
417        assert_eq!(
418            result.masked_text.as_deref(),
419            Some("Contact user@example.com")
420        );
421    }
422
423    #[test]
424    fn test_scanner_deny_list_adds_custom_finding() {
425        let scanner = Scanner::builder()
426            .recognizer(EmailRecognizer)
427            .deny_list(["PRJ-12345"])
428            .build()
429            .unwrap();
430        let result = scanner.scan("ticket PRJ-12345").unwrap();
431        assert_eq!(result.findings.len(), 1);
432        assert_eq!(
433            result.findings[0].entity_type,
434            EntityType::Custom("DenyList".to_string())
435        );
436        assert_eq!(result.masked_text.as_deref(), Some("ticket [DENYLIST]"));
437    }
438}