Big Brain v0.4.0

privacy.rs

Rust · 173 lines · 7.8 KB
//! Masking before text leaves the machine. Entries are already refused if they look like they hold
//! a secret; this goes further for remote models (embeddings, the librarian, `--deep`): emails, IP
//! addresses and the user name in home-folder paths are replaced with placeholders, plus any
//! patterns the user adds (account-number shapes, internal host names). Models on this machine
//! (LM Studio on localhost) get the text unchanged.
//!
//! All patterns are matched against the original text and the matches merged, so one rule's
//! placeholder can't hide text from another rule. A pattern that doesn't compile makes masking
//! fail closed: nothing is sent until the setting is fixed.

use anyhow::{bail, Result};
use regex::{Regex, RegexBuilder};
use std::sync::LazyLock;

/// Longest user pattern accepted, and the compiled-size cap.
const MAX_PATTERN_LEN: usize = 300;
const MAX_REGEX_SIZE: usize = 1 << 20;
/// Work bounds: the most text masked in one call, and the most custom patterns. Past either,
/// masking fails closed (nothing is sent). Patterns come from the user's own config, so these
/// bound accidents (a slow pattern on a huge text), not attacks.
const MAX_TEXT: usize = 64 * 1024;
const MAX_PATTERNS: usize = 50;

/// Candidate IPv6 tokens: hex digits, colons and dots with at least two colons; each is checked by
/// actually parsing it, so every valid form (`::dead:beef`, `2001:db8::1`, `::ffff:192.0.2.1`) is caught.
static V6_CANDIDATE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"[0-9A-Fa-f:.]*:[0-9A-Fa-f:.]*:[0-9A-Fa-f:.]*").unwrap());

static BUILTIN: LazyLock<Vec<(Regex, &'static str)>> = LazyLock::new(|| {
    [
        (r"(?i)\b[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}\b", "<email>"),
        (r"\b(?:25[0-5]|2[0-4]\d|1?\d?\d)(?:\.(?:25[0-5]|2[0-4]\d|1?\d?\d)){3}\b", "<ip>"),
        // Home folders: keep the shape, drop the user-name component. A name may contain spaces, so
        // everything up to the next slash or line end goes (over-masking is the safe side).
        (r"/home/[^/\r\n]+", "/home/<user>"),
        (r"(?i)/Users/[^/\r\n]+", "/Users/<user>"),
        // Windows, with either slash and any case: `C:\Users\<name>`, `c:/users/<name>`.
        // Also escaped, as in JSON or source code: `C:\\Users\\<name>`.
        (r"(?i)\b[A-Z]:(\\+|/)Users(\\+|/)[^\\/\r\n]+", r"C:\Users\<user>"),
    ]
    .into_iter()
    .map(|(r, rep)| (Regex::new(r).unwrap(), rep))
    .collect()
});

#[derive(Clone, Default)]
pub struct Masker {
    on: bool,
    extra: Vec<Regex>,
    /// A configuration problem: masking refuses to run (fail closed) until it is fixed.
    broken: Option<String>,
}

impl Masker {
    /// From the `[privacy]` settings.
    pub fn new(cfg: &crate::store::Privacy) -> Masker {
        let mut m = Masker { on: cfg.mask, extra: vec![], broken: None };
        if cfg.patterns.len() > MAX_PATTERNS {
            m.broken = Some(format!("more than {MAX_PATTERNS} privacy patterns"));
            return m;
        }
        // Errors name a pattern by its position, never its text: patterns are often internal host
        // names or account shapes, and these messages reach agents.
        for (i, p) in cfg.patterns.iter().enumerate() {
            let n = i + 1;
            if p.len() > MAX_PATTERN_LEN {
                m.broken = Some(format!("privacy pattern #{n} is longer than {MAX_PATTERN_LEN} bytes"));
                break;
            }
            match RegexBuilder::new(p).size_limit(MAX_REGEX_SIZE).build() {
                Ok(r) => m.extra.push(r),
                Err(_) => {
                    m.broken = Some(format!("privacy pattern #{n} isn't a valid regular expression (`brain doctor` tests it)"));
                    break;
                }
            }
        }
        m
    }

    pub fn off() -> Masker {
        Masker::default()
    }

    /// The masked text, or an error if the masking settings are broken (nothing should be sent).
    pub fn mask(&self, text: &str) -> Result<String> {
        if let Some(why) = &self.broken {
            bail!("{why}; fix [privacy] patterns in config.toml (nothing is sent until then)");
        }
        if !self.on && self.extra.is_empty() {
            return Ok(text.to_string());
        }
        if text.len() > MAX_TEXT {
            bail!("text too long to mask safely ({} KB); not sent", text.len() / 1024);
        }
        // Collect every match on the ORIGINAL text, then merge overlaps and replace once.
        let mut spans: Vec<(usize, usize, &str)> = vec![];
        if self.on {
            for (re, rep) in BUILTIN.iter() {
                for m in re.find_iter(text) {
                    spans.push((m.start(), m.end(), rep));
                }
            }
            let word = |c: Option<char>| c.is_some_and(|c| c.is_alphanumeric() || c == '_');
            for m in V6_CANDIDATE.find_iter(text) {
                // Punctuation can cling to either end (`…2001:db8::1: down`, `(2001:db8::1.)`):
                // try the token with trailing dots/colons trimmed, then check what really surrounds
                // the address, not the raw match. `d::f` inside `std::fs` is not standalone.
                let raw = m.as_str();
                let first = m.start() + (raw.len() - raw.trim_start_matches('.').len());
                // The address may also start after a label's single colon (`ip:2001:db8::1`,
                // `id:2001:db8::1`): try each such start, left to right. Never inside `::`.
                let mut starts = vec![first];
                let bytes = text.as_bytes();
                for i in first..m.end() {
                    if bytes[i] == b':' && bytes.get(i + 1) != Some(&b':') && (i == 0 || bytes[i - 1] != b':') {
                        starts.push(i + 1);
                    }
                }
                'starts: for start in starts {
                    if start >= m.end() || word(text[..start].chars().next_back()) {
                        continue;
                    }
                    // Cut back at `.`/`:` until a standalone address remains
                    // (`2001:db8::1.example`, `2001:db8::1:failed` → `2001:db8::1`).
                    let mut t = &text[start..m.end()];
                    loop {
                        let end = start + t.len();
                        if t.parse::<std::net::Ipv6Addr>().is_ok() && !word(text[end..].chars().next()) {
                            spans.push((start, end, "<ip>"));
                            break 'starts;
                        }
                        match t.rfind(['.', ':']) {
                            Some(i) if i >= 2 => t = &t[..i],
                            _ => break,
                        }
                    }
                }
            }
        }
        // User patterns run on the whole text, so they can match across lines.
        for re in &self.extra {
            for m in re.find_iter(text) {
                if m.end() > m.start() {
                    spans.push((m.start(), m.end(), "<masked>"));
                }
            }
        }
        spans.sort_by_key(|s| (s.0, std::cmp::Reverse(s.1)));
        let mut out = String::with_capacity(text.len());
        let mut pos = 0;
        let mut i = 0;
        while i < spans.len() {
            let (start, mut end, rep) = spans[i];
            if start < pos {
                i += 1;
                continue;
            }
            // Swallow every span that overlaps this one.
            let mut j = i + 1;
            while j < spans.len() && spans[j].0 < end {
                end = end.max(spans[j].1);
                j += 1;
            }
            out.push_str(&text[pos..start]);
            // A lone match gets its own placeholder; overlapping matches from different rules become one.
            out.push_str(if j == i + 1 { rep } else { "<masked>" });
            pos = end;
            i = j;
        }
        out.push_str(&text[pos..]);
        Ok(out)
    }
}