llm_inspect/
normalize.rs

1//! Text normalization applied before pattern matching.
2//!
3//! Attackers evade naive filters with full-width characters (`ＩＧＮＯＲＥ`),
4//! compatibility forms (`ﬁ`) or invisible characters (zero-width spaces,
5//! bidirectional controls). [`Normalized`] applies NFKC normalization and drops
6//! invisible characters while keeping a byte-offset map back to the original
7//! text, so findings and redactions can still be reported on the text the
8//! caller actually passed in.
9
10use unicode_normalization::UnicodeNormalization;
11
12/// Returns `true` for characters that are removed during normalization.
13///
14/// These are format characters that carry no visible content: soft hyphen,
15/// zero-width spaces and joiners, bidirectional controls, word joiner and
16/// invisible operators, Mongolian vowel separator and the byte-order mark.
17pub fn is_invisible(c: char) -> bool {
18    matches!(
19        c,
20        '\u{00AD}'
21            | '\u{180E}'
22            | '\u{200B}'..='\u{200F}'
23            | '\u{202A}'..='\u{202E}'
24            | '\u{2060}'..='\u{2064}'
25            | '\u{2066}'..='\u{2069}'
26            | '\u{FEFF}'
27    )
28}
29
30/// Normalized text with a mapping back to byte offsets in the original.
31#[derive(Debug, Clone)]
32pub struct Normalized<'a> {
33    original: &'a str,
34    text: String,
35    /// `map[i]` is the byte offset in `original` of the character that produced
36    /// byte `i` of `text`. One extra trailing entry holds `original.len()`.
37    map: Vec<usize>,
38    removed_invisible: usize,
39}
40
41impl<'a> Normalized<'a> {
42    /// Normalizes `original`.
43    pub fn new(original: &'a str) -> Self {
44        let mut text = String::with_capacity(original.len());
45        let mut map = Vec::with_capacity(original.len() + 1);
46        let mut removed_invisible = 0;
47        for (offset, c) in original.char_indices() {
48            if is_invisible(c) {
49                removed_invisible += 1;
50                continue;
51            }
52            for n in std::iter::once(c).nfkc() {
53                let before = text.len();
54                text.push(n);
55                map.extend(std::iter::repeat_n(offset, text.len() - before));
56            }
57        }
58        map.push(original.len());
59        Self {
60            original,
61            text,
62            map,
63            removed_invisible,
64        }
65    }
66
67    /// The normalized text.
68    pub fn text(&self) -> &str {
69        &self.text
70    }
71
72    /// The text that was passed to [`Normalized::new`].
73    pub fn original(&self) -> &'a str {
74        self.original
75    }
76
77    /// Whether normalization changed anything.
78    pub fn changed(&self) -> bool {
79        self.text != self.original
80    }
81
82    /// Number of invisible characters that were removed.
83    pub fn removed_invisible(&self) -> usize {
84        self.removed_invisible
85    }
86
87    /// Maps a byte range of the normalized text back to a byte range of the
88    /// original text. The result always covers whole characters of the original.
89    ///
90    /// # Panics
91    ///
92    /// Panics if `start > end` or `end > self.text().len()`.
93    pub fn to_original(&self, start: usize, end: usize) -> (usize, usize) {
94        assert!(
95            start <= end && end <= self.text.len(),
96            "range out of bounds"
97        );
98        if start == end {
99            let o = self.map[start];
100            return (o, o);
101        }
102        let s = self.map[start];
103        let last_src = self.map[end - 1];
104        let next = self.map[end];
105        let e = if next > last_src {
106            next
107        } else {
108            next_char_boundary(self.original, last_src)
109        };
110        (s, e)
111    }
112}
113
114fn next_char_boundary(s: &str, i: usize) -> usize {
115    let mut j = i + 1;
116    while j < s.len() && !s.is_char_boundary(j) {
117        j += 1;
118    }
119    j.min(s.len())
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn plain_ascii_is_unchanged() {
128        let n = Normalized::new("hello world");
129        assert!(!n.changed());
130        assert_eq!(n.to_original(6, 11), (6, 11));
131    }
132
133    #[test]
134    fn full_width_maps_back_to_original_bytes() {
135        let src = "say ＩＧＮＯＲＥ now"; // each full-width letter is 3 bytes
136        let n = Normalized::new(src);
137        assert_eq!(n.text(), "say IGNORE now");
138        let (s, e) = n.to_original(4, 10);
139        assert_eq!(&src[s..e], "ＩＧＮＯＲＥ");
140    }
141
142    #[test]
143    fn zero_width_characters_are_removed_and_counted() {
144        let src = "ig\u{200B}no\u{200D}re";
145        let n = Normalized::new(src);
146        assert_eq!(n.text(), "ignore");
147        assert_eq!(n.removed_invisible(), 2);
148        let (s, e) = n.to_original(0, 6);
149        assert_eq!((s, e), (0, src.len()));
150    }
151
152    #[test]
153    fn partial_match_inside_expanded_char_covers_whole_source_char() {
154        let src = "ﬁle"; // U+FB01 expands to "fi"
155        let n = Normalized::new(src);
156        assert_eq!(n.text(), "file");
157        let (s, e) = n.to_original(0, 1); // just "f"
158        assert_eq!(&src[s..e], "ﬁ");
159    }
160
161    #[test]
162    fn empty_input() {
163        let n = Normalized::new("");
164        assert_eq!(n.text(), "");
165        assert_eq!(n.to_original(0, 0), (0, 0));
166    }
167}
llm_inspect/normalize.rs

llm_inspect/
normalize.rs
