veil_engine/
lib.rs

1#![doc = include_str!("../README.md")]
2#![warn(missing_docs)]
3
4mod allowlist;
5mod builder;
6mod detect;
7mod error;
8mod kind;
9mod ner;
10mod placeholder;
11mod rules;
12mod scan_cache;
13mod stream;
14mod vault;
15mod walk;
16
17pub use builder::EngineBuilder;
18pub use detect::Source;
19pub use error::{Error, RecognizerError, VaultError};
20pub use kind::{kinds, EntityKind, InvalidKind};
21pub use ner::{Entity, EntityRecognizer};
22pub use placeholder::{PlaceholderFormat, Token, DEFAULT_FORMAT};
23pub use rules::{
24    CompiledRules, Detectors, DictionaryEntry, IpOptions, NatIdOptions, NerOptions, NoOptions, Options, PhoneOptions,
25    RegexRule, RulesError, RulesFile, Setting,
26};
27pub use scan_cache::{ScanCacheStats, DEFAULT_BUDGET as DEFAULT_SCAN_CACHE_BUDGET};
28pub use stream::StreamRestorer;
29pub use vault::{MemoryVault, Vault};
30pub use walk::{KeyAction, KeyPolicy, WalkAll};
31
32/// This crate's version as compiled in, such as `"0.2.0"`, so an application
33/// can show which release of it is running.
34pub const VERSION: &str = env!("CARGO_PKG_VERSION");
35
36use std::borrow::Cow;
37use std::collections::{BTreeMap, HashMap, HashSet};
38use std::fmt;
39use std::sync::Arc;
40
41use arc_swap::ArcSwap;
42use scan_cache::{ScanCache, Scanned};
43use serde_json::Value;
44use walk::StrKind;
45
46/// Identifies which placeholder mapping a request uses.
47#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
48pub struct ScopeId(String);
49
50impl ScopeId {
51    /// A scope with exactly this ID.
52    pub fn new(id: impl Into<String>) -> Self {
53        ScopeId(id.into())
54    }
55
56    /// The scope for a named profile, `profile:{name}`. Every conversation using
57    /// the profile shares its placeholders.
58    pub fn profile(name: &str) -> Self {
59        ScopeId(format!("profile:{name}"))
60    }
61
62    /// A scope derived from `hint` (for example a conversation's ID or first
63    /// message) with the secret `key`: HMAC-SHA256, so the scope reveals
64    /// nothing about the hint. The same key and hint give the same scope.
65    pub fn keyed(key: &[u8], hint: &str) -> Self {
66        use hmac::{Hmac, Mac};
67        let mut mac = Hmac::<sha2::Sha256>::new_from_slice(key).expect("HMAC accepts any key length");
68        mac.update(hint.as_bytes());
69        let digest = mac.finalize().into_bytes();
70        let hex: String = digest[..16].iter().map(|b| format!("{b:02x}")).collect();
71        ScopeId(format!("conv:{hex}"))
72    }
73
74    /// The scope's ID.
75    pub fn as_str(&self) -> &str {
76        &self.0
77    }
78}
79
80impl fmt::Display for ScopeId {
81    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82        f.write_str(&self.0)
83    }
84}
85
86/// A detected value. Offsets are bytes into the scanned text.
87#[derive(Clone, Debug, PartialEq, Eq)]
88#[non_exhaustive]
89pub struct Match {
90    /// Byte offset where the value starts.
91    pub start: usize,
92    /// Byte offset just past the value.
93    pub end: usize,
94    /// What kind of value it is.
95    pub kind: EntityKind,
96    /// Which detector found it.
97    pub source: Source,
98}
99
100/// Counts from one pass. Never holds values.
101#[derive(Clone, Debug, Default, PartialEq, Eq)]
102#[non_exhaustive]
103pub struct Report {
104    /// Values replaced (redact mode) or found (audit mode).
105    pub redacted: usize,
106    /// [`redacted`](Self::redacted), by kind name.
107    pub by_kind: BTreeMap<String, usize>,
108    /// Placeholders put back to their real values.
109    pub restored: usize,
110    /// Placeholder-shaped text whose number this scope doesn't know, left as it
111    /// is. Often a placeholder the model made up.
112    pub unknown_placeholders: usize,
113}
114
115impl Report {
116    /// Adds `other`'s counts to these.
117    pub fn merge(&mut self, other: &Report) {
118        self.redacted += other.redacted;
119        for (k, v) in &other.by_kind {
120            *self.by_kind.entry(k.clone()).or_default() += v;
121        }
122        self.restored += other.restored;
123        self.unknown_placeholders += other.unknown_placeholders;
124    }
125
126    fn count(&mut self, kind: &EntityKind) {
127        self.redacted += 1;
128        *self.by_kind.entry(kind.as_str().to_owned()).or_default() += 1;
129    }
130}
131
132/// The result of redacting one string.
133#[derive(Clone, Debug, PartialEq, Eq)]
134#[non_exhaustive]
135pub struct Redacted {
136    /// The text with every match replaced by its placeholder.
137    pub text: String,
138    /// Offsets refer to the original text.
139    pub matches: Vec<Match>,
140    /// Counts for this text.
141    pub report: Report,
142}
143
144/// The result of restoring one string.
145#[derive(Clone, Debug, PartialEq, Eq)]
146#[non_exhaustive]
147pub struct Restored {
148    /// The text with every known placeholder replaced by its real value.
149    pub text: String,
150    /// Counts for this text.
151    pub report: Report,
152}
153
154/// Told about the values live traffic matched, e.g. to show recent matches.
155/// Not called for restores, or for [`Engine::scan`].
156pub trait Recorder: Send + Sync {
157    /// `value` matched `times` times in one request, in order of first
158    /// appearance. `placeholder` is what replaced it, or `None` when it was
159    /// only found and sent unchanged ([`Engine::scan_json`], audit mode).
160    fn matched(&self, kind: &EntityKind, value: &str, placeholder: Option<&str>, times: u64);
161}
162
163/// Finds sensitive values, replaces them with placeholders and puts the real
164/// values back. Build one with [`Engine::builder`].
165pub struct Engine {
166    rules: ArcSwap<CompiledRules>,
167    vault: Arc<dyn Vault>,
168    recorder: Option<Arc<dyn Recorder>>,
169    scans: ScanCache,
170}
171
172/// The values one call has seen. Each distinct value is looked up in the
173/// vault once, however often it appears, and reported to the recorder once.
174#[derive(Default)]
175struct Pass {
176    index: HashMap<EntityKind, HashMap<String, usize>>,
177    seen: Vec<Seen>,
178    reserved: HashSet<(EntityKind, u64)>,
179}
180
181struct Seen {
182    kind: EntityKind,
183    value: String,
184    placeholder: Option<String>,
185    times: u64,
186}
187
188impl Pass {
189    fn find(&self, kind: &EntityKind, value: &str) -> Option<usize> {
190        self.index.get(kind).and_then(|m| m.get(value)).copied()
191    }
192
193    fn add(&mut self, kind: &EntityKind, value: &str, placeholder: Option<String>) -> usize {
194        let i = self.seen.len();
195        self.seen.push(Seen { kind: kind.clone(), value: value.to_owned(), placeholder, times: 0 });
196        self.index.entry(kind.clone()).or_default().insert(value.to_owned(), i);
197        i
198    }
199}
200
201impl Default for Engine {
202    /// The default rules and placeholder format, with an in-memory vault.
203    fn default() -> Self {
204        Engine::builder().build().expect("default rules compile")
205    }
206}
207
208impl Engine {
209    /// Starts building an engine. With nothing set, it uses the default
210    /// rules, the `[{TYPE}_{N}]` placeholder format and an in-memory vault.
211    pub fn builder() -> EngineBuilder {
212        EngineBuilder::default()
213    }
214
215    /// How the cache of recent scans is doing.
216    pub fn cache_stats(&self) -> ScanCacheStats {
217        self.scans.stats()
218    }
219
220    /// Swaps in new rules. Calls already running finish with the old ones.
221    /// Scans made with the old rules are dropped: they can't be reused.
222    pub fn set_rules(&self, rules: CompiledRules) {
223        self.rules.store(Arc::new(rules));
224        self.scans.clear();
225    }
226
227    /// The rules in use now.
228    pub fn rules(&self) -> Arc<CompiledRules> {
229        self.rules.load_full()
230    }
231
232    /// The vault placeholders are kept in.
233    pub fn vault(&self) -> &Arc<dyn Vault> {
234        &self.vault
235    }
236
237    /// Finds sensitive values without changing anything or touching the vault.
238    pub fn scan(&self, text: &str) -> Result<Vec<Match>, Error> {
239        scan_with(&self.rules(), text)
240    }
241
242    /// Replaces every detected value in `text` with its placeholder.
243    pub fn redact_text(&self, scope: &ScopeId, text: &str) -> Result<Redacted, Error> {
244        let rules = self.rules();
245        let mut pass = Pass::default();
246        let (out, report, scanned) = self.redact_with(&rules, scope, text, &mut pass)?;
247        self.report(pass);
248        Ok(Redacted { text: out.into_owned(), matches: scanned.matches.clone(), report })
249    }
250
251    /// Redacts every string in `body` that `policy` does not skip.
252    pub fn redact_json(&self, scope: &ScopeId, body: &mut Value, policy: &dyn KeyPolicy) -> Result<Report, Error> {
253        let rules = self.rules();
254        let mut report = Report::default();
255        let mut pass = Pass::default();
256        walk::walk(body, policy, &mut |s, kind| self.redact_string(&rules, scope, s, kind, &mut report, &mut pass))?;
257        self.report(pass);
258        Ok(report)
259    }
260
261    /// Counts what would be redacted in `body` without changing it (audit mode).
262    /// Each match is reported to the recorder as sent unchanged.
263    pub fn scan_json(&self, body: &Value, policy: &dyn KeyPolicy) -> Result<Report, Error> {
264        let rules = self.rules();
265        let mut report = Report::default();
266        let mut pass = Pass::default();
267        let mut copy = body.clone();
268        walk::walk(&mut copy, policy, &mut |s, kind| {
269            let text: Cow<'_, str> = match kind {
270                StrKind::Json => {
271                    serde_json::from_str::<Value>(s).map_or(Cow::Borrowed(s.as_str()), |v| json_strings(&v).into())
272                }
273                StrKind::Plain => Cow::Borrowed(s.as_str()),
274            };
275            let scanned = self.scanned(&rules, &text)?;
276            for m in &scanned.matches {
277                report.count(&m.kind);
278                if self.recorder.is_some() {
279                    let value = &text[m.start..m.end];
280                    let i = pass.find(&m.kind, value).unwrap_or_else(|| pass.add(&m.kind, value, None));
281                    pass.seen[i].times += 1;
282                }
283            }
284            Ok(())
285        })?;
286        self.report(pass);
287        Ok(report)
288    }
289
290    /// Replaces every known placeholder in `text` with its real value.
291    /// Unknown placeholders are left as they are and counted.
292    pub fn restore_text(&self, scope: &ScopeId, text: &str) -> Restored {
293        self.restore_with(scope, text, |v| v.to_owned())
294    }
295
296    /// Like [`Engine::restore_text`], but passes each real value through
297    /// `escape` first, e.g. to restore inside raw JSON text.
298    pub fn restore_with(&self, scope: &ScopeId, text: &str, escape: impl Fn(&str) -> String) -> Restored {
299        self.restore_after(scope, None, text, escape)
300    }
301
302    /// Restores `text` that follows the character `before`, which decides
303    /// whether a bracketless placeholder at the very start stands alone.
304    pub(crate) fn restore_after(
305        &self,
306        scope: &ScopeId,
307        before: Option<char>,
308        text: &str,
309        escape: impl Fn(&str) -> String,
310    ) -> Restored {
311        let rules = self.rules();
312        let tokens: Vec<(usize, usize, String, u64, bool)> = if rules.fuzzy_restore {
313            rules
314                .placeholder
315                .fuzzy_tokens_after(before, text)
316                .into_iter()
317                .map(|t| (t.start, t.end, t.kind, t.n, t.exact))
318                .collect()
319        } else {
320            rules.placeholder.tokens(text).map(|t| (t.start, t.end, t.kind.to_owned(), t.n, true)).collect()
321        };
322        let mut report = Report::default();
323        let mut out = String::with_capacity(text.len());
324        let mut last = 0;
325        for (start, end, kind, n, exact) in tokens {
326            let value = EntityKind::new(&kind).ok().and_then(|kind| self.vault.lookup(scope, &kind, n).ok().flatten());
327            out.push_str(&text[last..start]);
328            match value {
329                Some(v) => {
330                    out.push_str(&escape(&v));
331                    report.restored += 1;
332                }
333                None => {
334                    out.push_str(&text[start..end]);
335                    // Near misses that aren't in the vault are just text.
336                    if exact {
337                        report.unknown_placeholders += 1;
338                    }
339                }
340            }
341            last = end;
342        }
343        out.push_str(&text[last..]);
344        Restored { text: out, report }
345    }
346
347    /// Restores every string in `value` that `policy` does not skip.
348    pub fn restore_json(&self, scope: &ScopeId, value: &mut Value, policy: &dyn KeyPolicy) -> Report {
349        let mut report = Report::default();
350        let result = walk::walk(value, policy, &mut |s, kind| {
351            let r = match kind {
352                StrKind::Plain => self.restore_text(scope, s),
353                StrKind::Json => self.restore_with(scope, s, json_escape),
354            };
355            *s = r.text;
356            report.merge(&r.report);
357            Ok(())
358        });
359        debug_assert!(result.is_ok(), "restoration never fails");
360        report
361    }
362
363    fn redact_string(
364        &self,
365        rules: &CompiledRules,
366        scope: &ScopeId,
367        s: &mut String,
368        kind: StrKind,
369        report: &mut Report,
370        pass: &mut Pass,
371    ) -> Result<(), Error> {
372        if kind == StrKind::Json {
373            if let Ok(mut inner) = serde_json::from_str::<Value>(s) {
374                walk::walk(&mut inner, &WalkAll, &mut |s, k| self.redact_string(rules, scope, s, k, report, pass))?;
375                *s = inner.to_string();
376                return Ok(());
377            }
378        }
379        let (out, r, _) = self.redact_with(rules, scope, s, pass)?;
380        report.merge(&r);
381        if let Cow::Owned(out) = out {
382            *s = out;
383        }
384        Ok(())
385    }
386
387    /// Redacts one string. Returns it unchanged (borrowed) when nothing matched.
388    fn redact_with<'t>(
389        &self,
390        rules: &CompiledRules,
391        scope: &ScopeId,
392        text: &'t str,
393        pass: &mut Pass,
394    ) -> Result<(Cow<'t, str>, Report, Arc<Scanned>), Error> {
395        let scanned = self.scanned(rules, text)?;
396        // Placeholder-shaped text the vault doesn't know was written by the
397        // user. Reserve its number so a real value never gets it later.
398        for (kind, n) in &scanned.tokens {
399            if pass.reserved.insert((kind.clone(), *n)) && self.vault.lookup(scope, kind, *n)?.is_none() {
400                self.vault.reserve(scope, kind, *n)?;
401            }
402        }
403        let mut report = Report::default();
404        if scanned.matches.is_empty() {
405            return Ok((Cow::Borrowed(text), report, scanned));
406        }
407        let mut out = String::with_capacity(text.len());
408        let mut last = 0;
409        for m in &scanned.matches {
410            out.push_str(&text[last..m.start]);
411            let value = &text[m.start..m.end];
412            let i = match pass.find(&m.kind, value) {
413                Some(i) => i,
414                None => {
415                    let n = self.vault.get_or_create(scope, &m.kind, value)?;
416                    pass.add(&m.kind, value, Some(rules.placeholder.format(&m.kind, n)))
417                }
418            };
419            let seen = &mut pass.seen[i];
420            seen.times += 1;
421            out.push_str(seen.placeholder.as_deref().expect("redacted values have a placeholder"));
422            report.count(&m.kind);
423            last = m.end;
424        }
425        out.push_str(&text[last..]);
426        Ok((Cow::Owned(out), report, scanned))
427    }
428
429    /// Scans `text`, or reuses the scan of the same text with the same rules.
430    fn scanned(&self, rules: &CompiledRules, text: &str) -> Result<Arc<Scanned>, Error> {
431        if !self.scans.enabled() || text.len() < scan_cache::MIN_LEN {
432            return Ok(Arc::new(scan_full(rules, text)?));
433        }
434        let key = ScanCache::key(rules.id, text);
435        if let Some(hit) = self.scans.get(&key) {
436            return Ok(hit);
437        }
438        let scanned = Arc::new(scan_full(rules, text)?);
439        self.scans.put(key, scanned.clone());
440        Ok(scanned)
441    }
442
443    /// Tells the recorder what this call matched.
444    fn report(&self, pass: Pass) {
445        if let Some(r) = &self.recorder {
446            for s in pass.seen {
447                r.matched(&s.kind, &s.value, s.placeholder.as_deref(), s.times);
448            }
449        }
450    }
451}
452
453fn scan_full(rules: &CompiledRules, text: &str) -> Result<Scanned, Error> {
454    let tokens: Vec<Token<'_>> = rules.placeholder.tokens(text).collect();
455    let matches = scan_tokens(rules, text, &tokens)?;
456    let tokens = tokens.iter().filter_map(|t| EntityKind::new(t.kind).ok().map(|k| (k, t.n))).collect();
457    Ok(Scanned { matches, tokens })
458}
459
460fn scan_with(rules: &CompiledRules, text: &str) -> Result<Vec<Match>, Error> {
461    let tokens: Vec<Token<'_>> = rules.placeholder.tokens(text).collect();
462    scan_tokens(rules, text, &tokens)
463}
464
465fn scan_tokens(rules: &CompiledRules, text: &str, tokens: &[Token<'_>]) -> Result<Vec<Match>, Error> {
466    let protected: Vec<(usize, usize)> = tokens.iter().map(|t| (t.start, t.end)).collect();
467    let mut candidates = Vec::new();
468    for d in &rules.detectors {
469        d.find(text, &mut candidates);
470    }
471    if let Some(ner) = &rules.ner {
472        if text.chars().any(char::is_alphabetic) {
473            let recognizer = ner.recognizer.as_ref().ok_or(Error::RecognizerMissing)?;
474            for e in recognizer.recognize(text)? {
475                let valid = e.start < e.end
476                    && e.end <= text.len()
477                    && text.is_char_boundary(e.start)
478                    && text.is_char_boundary(e.end);
479                if valid && e.score >= ner.threshold && ner.kinds.contains(&e.kind) {
480                    candidates.push(detect::Candidate {
481                        start: e.start,
482                        end: e.end,
483                        kind: e.kind,
484                        source: Source::Ner,
485                        priority: detect::P_NER,
486                    });
487                }
488            }
489        }
490    }
491    candidates.retain(|c| protected.iter().all(|&(s, e)| c.end <= s || c.start >= e));
492    Ok(detect::resolve(candidates)
493        .into_iter()
494        .filter(|c| !rules.allowlist.allows(&text[c.start..c.end]))
495        .map(|c| Match { start: c.start, end: c.end, kind: c.kind, source: c.source })
496        .collect())
497}
498
499/// All string values in a JSON value joined by newlines, for scanning.
500fn json_strings(v: &Value) -> String {
501    let mut out = String::new();
502    let mut copy = v.clone();
503    let _ = walk::walk(&mut copy, &WalkAll, &mut |s, _| {
504        out.push_str(s);
505        out.push('\n');
506        Ok(())
507    });
508    out
509}
510
511fn json_escape(v: &str) -> String {
512    let quoted = Value::String(v.to_owned()).to_string();
513    quoted[1..quoted.len() - 1].to_owned()
514}
515
516#[cfg(test)]
517mod tests {
518    use super::*;
519    use serde_json::json;
520
521    fn engine_with(rules: RulesFile) -> Engine {
522        Engine::builder().rules(rules).build().unwrap()
523    }
524
525    fn engine() -> Engine {
526        engine_with(RulesFile::default())
527    }
528
529    fn scope() -> ScopeId {
530        ScopeId::profile("default")
531    }
532
533    #[test]
534    fn redacts_and_restores() {
535        let e = engine();
536        let r = e.redact_text(&scope(), "Mail ana@example.com or bob@test.org, ana@example.com.").unwrap();
537        assert_eq!(r.text, "Mail [EMAIL_1] or [EMAIL_2], [EMAIL_1].");
538        assert_eq!(r.report.redacted, 3);
539        assert_eq!(r.report.by_kind["EMAIL"], 3);
540        let Restored { text: restored, report } = e.restore_text(&scope(), &r.text);
541        assert_eq!(restored, "Mail ana@example.com or bob@test.org, ana@example.com.");
542        assert_eq!(report.restored, 3);
543    }
544
545    #[test]
546    fn several_kinds_in_one_text() {
547        let e = engine();
548        let text = "Card 4111 1111 1111 1111, IBAN GB82 WEST 1234 5698 7654 32, call +44 20 7946 0958, key sk-ant-api03-abcdefghijklmnopqrstuvwxyz";
549        let r = e.redact_text(&scope(), text).unwrap();
550        assert_eq!(r.text, "Card [CARD_1], IBAN [IBAN_1], call [PHONE_1], key [SECRET_1]");
551        assert_eq!(e.restore_text(&scope(), &r.text).text, text);
552    }
553
554    #[test]
555    fn scopes_do_not_share_placeholders() {
556        let e = engine();
557        let redacted = e.redact_text(&ScopeId::profile("a"), "ana@example.com").unwrap().text;
558        let Restored { text: restored, report } = e.restore_text(&ScopeId::profile("b"), &redacted);
559        assert_eq!(restored, "[EMAIL_1]");
560        assert_eq!(report.unknown_placeholders, 1);
561    }
562
563    #[test]
564    fn keyed_scopes_never_change() {
565        // Stored mappings are filed under these IDs: a change would orphan
566        // every existing user's placeholders.
567        assert_eq!(ScopeId::keyed(b"fixed key", "conversation 1").as_str(), "conv:994b52d98b13f5df97e24363ad999bbc");
568        assert_ne!(ScopeId::keyed(b"other key", "conversation 1"), ScopeId::keyed(b"fixed key", "conversation 1"));
569    }
570
571    #[test]
572    fn user_written_placeholders_are_never_reassigned() {
573        let e = engine();
574        let text = "Template uses [EMAIL_1] and [EMAIL_2]; real one is ana@example.com";
575        let r = e.redact_text(&scope(), text).unwrap();
576        assert_eq!(r.text, "Template uses [EMAIL_1] and [EMAIL_2]; real one is [EMAIL_3]");
577        // The literal tokens stay literal on the way back.
578        assert_eq!(e.restore_text(&scope(), &r.text).text, text);
579    }
580
581    #[test]
582    fn redaction_is_idempotent() {
583        let e = engine();
584        let once = e.redact_text(&scope(), "ana@example.com and 4111111111111111").unwrap().text;
585        let twice = e.redact_text(&scope(), &once).unwrap().text;
586        assert_eq!(once, twice);
587    }
588
589    #[test]
590    fn allowlist_applies_after_overlap_resolution() {
591        let mut rules = RulesFile::default();
592        rules.allowlist.push("support@acme.com".into());
593        rules.dictionary.push(DictionaryEntry { term: "acme".into(), kind: "CLIENT".into(), match_case: false });
594        let e = engine_with(rules);
595        assert_eq!(
596            e.redact_text(&scope(), "write support@acme.com about ACME").unwrap().text,
597            "write support@acme.com about [CLIENT_1]"
598        );
599        assert_eq!(e.redact_text(&scope(), "sales@acme.com").unwrap().text, "[EMAIL_1]");
600    }
601
602    #[test]
603    fn the_default_rules_allow_nothing() {
604        assert!(RulesFile::default().allowlist.is_empty());
605        assert_eq!(engine().redact_text(&scope(), "noreply@anthropic.com").unwrap().text, "[EMAIL_1]");
606    }
607
608    #[test]
609    fn dictionary_and_custom_regex() {
610        let mut rules = RulesFile::default();
611        rules.dictionary.push(DictionaryEntry {
612            term: "Project Halcyon".into(),
613            kind: "CODENAME".into(),
614            match_case: true,
615        });
616        rules.dictionary.push(DictionaryEntry {
617            term: "Acme Holdings".into(),
618            kind: "CLIENT".into(),
619            match_case: false,
620        });
621        rules.regex.push(RegexRule { name: "tickets".into(), kind: "TICKET".into(), pattern: r"\bTCK-\d{6}\b".into() });
622        let e = engine_with(rules);
623        let r = e
624            .redact_text(&scope(), "Project Halcyon for ACME HOLDINGS, see TCK-004211 not TCK-12; project halcyon")
625            .unwrap();
626        assert_eq!(r.text, "[CODENAME_1] for [CLIENT_1], see [TICKET_1] not TCK-12; project halcyon");
627        assert_eq!(r.matches[0].source, Source::Dictionary);
628        assert_eq!(r.matches[2].source, Source::Regex("tickets".into()));
629    }
630
631    #[test]
632    fn redact_json_respects_policy_and_json_strings() {
633        struct Policy;
634        impl KeyPolicy for Policy {
635            fn action(&self, _: &serde_json::Map<String, Value>, key: &str) -> KeyAction {
636                match key {
637                    "model" => KeyAction::Skip,
638                    "arguments" => KeyAction::JsonString,
639                    _ => KeyAction::Walk,
640                }
641            }
642        }
643        let e = engine();
644        let mut body = json!({
645            "model": "a@b.co",
646            "messages": [{"content": "hi a@b.co"}],
647            "tool": {"arguments": "{\"to\":\"a@b.co\",\"n\":1}"},
648            "broken": {"arguments": "{\"to\": \"a@b.co\""}
649        });
650        let report = e.redact_json(&scope(), &mut body, &Policy).unwrap();
651        assert_eq!(body["model"], "a@b.co");
652        assert_eq!(body["messages"][0]["content"], "hi [EMAIL_1]");
653        assert_eq!(body["tool"]["arguments"], "{\"to\":\"[EMAIL_1]\",\"n\":1}");
654        assert_eq!(body["broken"]["arguments"], "{\"to\": \"[EMAIL_1]\"");
655        assert_eq!(report.redacted, 3);
656
657        let audit = e
658            .scan_json(
659                &json!({"model": "a@b.co", "x": "c@d.co", "tool": {"arguments": "{\"to\":\"e@f.co\"}"}}),
660                &Policy,
661            )
662            .unwrap();
663        assert_eq!(audit.redacted, 2);
664    }
665
666    #[test]
667    fn restore_json_escapes_inside_json_strings() {
668        let e = engine();
669        // A value with a quote: only valid inside JSON text if escaped.
670        e.redact_text(&scope(), "\"ana\"@example.com").unwrap();
671        let mut rules = RulesFile::default();
672        rules.dictionary.push(DictionaryEntry { term: "say \"hi\"".into(), kind: "PHRASE".into(), match_case: false });
673        let e = engine_with(rules);
674        let r = e.redact_text(&scope(), "say \"hi\" now").unwrap();
675        assert_eq!(r.text, "[PHRASE_1] now");
676        struct Args;
677        impl KeyPolicy for Args {
678            fn action(&self, _: &serde_json::Map<String, Value>, _: &str) -> KeyAction {
679                KeyAction::JsonString
680            }
681        }
682        let mut v = json!({"arguments": "{\"text\":\"[PHRASE_1]\"}"});
683        e.restore_json(&scope(), &mut v, &Args);
684        let inner: Value = serde_json::from_str(v["arguments"].as_str().unwrap()).unwrap();
685        assert_eq!(inner["text"], "say \"hi\"");
686    }
687
688    #[test]
689    fn rules_can_be_swapped_live() {
690        let e = engine();
691        assert_eq!(e.redact_text(&scope(), "see https://x.example/a").unwrap().text, "see https://x.example/a");
692        let mut rules = RulesFile::default();
693        rules.detectors.url = Setting::on();
694        e.set_rules(CompiledRules::compile(&rules, &Options::default(), None).unwrap());
695        assert_eq!(e.redact_text(&scope(), "see https://x.example/a").unwrap().text, "see [URL_1]");
696    }
697
698    /// Finds fixed words, with fixed scores.
699    struct FakeNer(Vec<(&'static str, &'static str, f32)>);
700
701    impl EntityRecognizer for FakeNer {
702        fn recognize(&self, text: &str) -> Result<Vec<Entity>, RecognizerError> {
703            if text.contains("FAIL") {
704                return Err(RecognizerError::new("model crashed"));
705            }
706            Ok(self
707                .0
708                .iter()
709                .flat_map(|&(word, kind, score)| {
710                    text.match_indices(word).map(move |(i, _)| Entity {
711                        start: i,
712                        end: i + word.len(),
713                        kind: EntityKind::new(kind).unwrap(),
714                        score,
715                    })
716                })
717                .collect())
718        }
719    }
720
721    fn ner_engine(rules: &RulesFile) -> Engine {
722        let ner = FakeNer(vec![
723            ("Ana Pop", "PERSON", 0.97),
724            ("Acme", "ORG", 0.9),
725            ("Jordan", "PERSON", 0.4),
726            ("Paris", "LOCATION", 0.95),
727        ]);
728        Engine::builder().rules(rules.clone()).recognizer(Arc::new(ner)).build().unwrap()
729    }
730
731    #[test]
732    fn ner_redacts_names_above_the_threshold() {
733        let mut rules = RulesFile::default();
734        rules.detectors.ner = Setting::on();
735        let e = ner_engine(&rules);
736        let text = "Ana Pop (ana@example.com) of Acme met Jordan in Paris";
737        let r = e.redact_text(&scope(), text).unwrap();
738        assert_eq!(r.text, "[PERSON_1] ([EMAIL_1]) of [ORG_1] met Jordan in [LOCATION_1]");
739        assert!(r.matches.iter().any(|m| m.source == Source::Ner));
740        assert_eq!(e.restore_text(&scope(), &r.text).text, text);
741
742        rules.detectors.ner.options.kinds = vec!["PERSON".into()];
743        rules.detectors.ner.options.threshold = 0.3;
744        rules.allowlist.push("Ana Pop".into());
745        let e = ner_engine(&rules);
746        assert_eq!(e.redact_text(&scope(), text).unwrap().text, "Ana Pop ([EMAIL_1]) of Acme met [PERSON_1] in Paris");
747    }
748
749    #[test]
750    fn dictionary_beats_ner_on_the_same_span() {
751        let mut rules = RulesFile::default();
752        rules.detectors.ner = Setting::on();
753        rules.dictionary.push(DictionaryEntry { term: "Acme".into(), kind: "CLIENT".into(), match_case: true });
754        assert_eq!(ner_engine(&rules).redact_text(&scope(), "Acme").unwrap().text, "[CLIENT_1]");
755    }
756
757    #[test]
758    fn ner_fails_closed() {
759        let mut rules = RulesFile::default();
760        rules.detectors.ner = Setting::on();
761        // Turned on without a model: refuse rather than send names.
762        let compiled = CompiledRules::compile(&rules, &Options::default(), None).unwrap();
763        assert!(compiled.ner_missing());
764        let e = Engine::builder().compiled_rules(compiled).build().unwrap();
765        assert!(matches!(e.redact_text(&scope(), "Ana Pop"), Err(Error::RecognizerMissing)));
766        // Text with no letters can't hold a name, so it doesn't need the model.
767        assert_eq!(e.redact_text(&scope(), "12 + 3").unwrap().text, "12 + 3");
768        assert!(matches!(ner_engine(&rules).redact_text(&scope(), "FAIL"), Err(Error::Recognizer(_))));
769    }
770
771    #[test]
772    fn fuzzy_restore_is_opt_in() {
773        let compile = |fuzzy| {
774            let options = Options { fuzzy_restore: fuzzy, ..Options::default() };
775            CompiledRules::compile(&RulesFile::default(), &options, None).unwrap()
776        };
777        let e = Engine::builder().compiled_rules(compile(false)).build().unwrap();
778        e.redact_text(&scope(), "ana@example.com and +44 20 7946 0958").unwrap();
779        let model = r"Wrote to [email_1], [ PHONE 1 ], EMAIL_1 and [EMAIL\_1]; [email_9] is unknown.";
780        let Restored { text: exact, report } = e.restore_text(&scope(), model);
781        assert_eq!(exact, model);
782        assert_eq!(report.restored, 0);
783        e.set_rules(compile(true));
784        let Restored { text: fuzzy, report } = e.restore_text(&scope(), model);
785        assert_eq!(
786            fuzzy,
787            "Wrote to ana@example.com, +44 20 7946 0958, ana@example.com and ana@example.com; [email_9] is unknown."
788        );
789        assert_eq!((report.restored, report.unknown_placeholders), (4, 0));
790    }
791
792    #[test]
793    fn ner_options_are_checked() {
794        let mut rules = RulesFile::default();
795        rules.detectors.ner = Setting::on();
796        rules.detectors.ner.options.kinds = vec!["DATE".into()];
797        assert!(matches!(CompiledRules::compile(&rules, &Options::default(), None), Err(RulesError::Ner(_))));
798        rules.detectors.ner.options = NerOptions { threshold: 1.5, ..NerOptions::default() };
799        assert!(matches!(CompiledRules::compile(&rules, &Options::default(), None), Err(RulesError::Ner(_))));
800    }
801
802    mod props {
803        use super::*;
804        use proptest::prelude::*;
805
806        proptest! {
807            #[test]
808            fn redact_then_restore_is_identity(text in "[a-z0-9 .,@\\[\\]_A-Z\\n+-]{0,80}") {
809                let e = engine();
810                let r = e.redact_text(&scope(), &text).unwrap();
811                prop_assert_eq!(e.restore_text(&scope(), &r.text).text, text);
812            }
813
814            #[test]
815            fn redaction_is_deterministic(text in "[a-z0-9 .,@\\[\\]_A-Z\\n+-]{0,80}") {
816                let e = engine();
817                let a = e.redact_text(&scope(), &text).unwrap().text;
818                let b = e.redact_text(&scope(), &text).unwrap().text;
819                prop_assert_eq!(a, b);
820            }
821        }
822    }
823}
veil_engine/lib.rs

veil_engine/
lib.rs
