fix: tighten tool output escape inversion

src/safety/mod.rs

@@ -247,26 +247,42 @@ fn escape_xml_attr(s: &str) -> String {
 }
 
 /// Neutralize closing `</tool_output` sequences in content to prevent
-/// boundary injection. Uses a case-insensitive regex to catch variations
-/// like `</Tool_Output`, `</ tool_output`, etc. The leading `<` is replaced
+/// boundary injection. Uses a case-insensitive search to catch variations
+/// like `</Tool_Output` and `</ tool_output>`. The leading `<` is replaced
 /// with `<\u{200B}` (zero-width space) so JSON and other content passes
 /// through unchanged.
 fn escape_tool_output_close(s: &str) -> String {
-    // Case-insensitive search for </tool_output (with optional whitespace/null after </)
+    // Case-insensitive search for </tool_output with optional ASCII whitespace after </
     // to block XML injection without corrupting other content.
     let mut result = String::with_capacity(s.len());
     let lower = s.to_ascii_lowercase();
-    let needle = "</tool_output";
+    let bytes = lower.as_bytes();
+    let needle = b"tool_output";
     let mut start = 0;
+    let mut i = 0;
 
-    while let Some(pos) = lower[start..].find(needle) {
-        let abs = start + pos;
-        result.push_str(&s[start..abs]);
-        // Insert zero-width space after '<' to break the closing tag
-        result.push('<');
-        result.push('\u{200B}');
-        result.push_str(&s[abs + 1..abs + needle.len()]);
-        start = abs + needle.len();
+    while i < bytes.len() {
+        if bytes[i] == b'<' {
+            let mut j = i + 1;
+            if j < bytes.len() && bytes[j] == b'/' {
+                j += 1;
+                while j < bytes.len() && bytes[j].is_ascii_whitespace() {
+                    j += 1;
+                }
+                if j + needle.len() <= bytes.len() && &bytes[j..j + needle.len()] == needle {
+                    result.push_str(&s[start..i]);
+                    // Insert zero-width space after '<' to break the closing tag.
+                    result.push('<');
+                    result.push('\u{200B}');
+                    let match_end = j + needle.len();
+                    result.push_str(&s[i + 1..match_end]);
+                    start = match_end;
+                    i = match_end;
+                    continue;
+                }
+            }
+        }
+        i += 1;
     }
     result.push_str(&s[start..]);
     result
@@ -275,7 +291,43 @@ fn escape_tool_output_close(s: &str) -> String {
 /// Reverse the escaping applied by [`escape_tool_output_close`] by removing
 /// the zero-width space inserted after `<` in `</tool_output` sequences.
 fn unescape_tool_output_close(s: &str) -> String {
-    s.replace("<\u{200B}/", "</")
+    const ESC_PREFIX: &str = "<\u{200B}";
+    const NEEDLE: &str = "tool_output";
+
+    let mut result = String::with_capacity(s.len());
+    let mut i = 0;
+
+    while i < s.len() {
+        if let Some(rest) = s[i..].strip_prefix(ESC_PREFIX) {
+            let lower_rest = rest.to_ascii_lowercase();
+            let rest_bytes = lower_rest.as_bytes();
+
+            if !rest_bytes.is_empty() && rest_bytes[0] == b'/' {
+                let mut j = 1;
+                while j < rest_bytes.len() && rest_bytes[j].is_ascii_whitespace() {
+                    j += 1;
+                }
+                if j + NEEDLE.len() <= rest_bytes.len()
+                    && &rest_bytes[j..j + NEEDLE.len()] == NEEDLE.as_bytes()
+                {
+                    let match_end = j + NEEDLE.len();
+                    result.push('<');
+                    result.push_str(&rest[..match_end]);
+                    i += ESC_PREFIX.len() + match_end;
+                    continue;
+                }
+            }
+        }
+
+        if let Some(ch) = s[i..].chars().next() {
+            result.push(ch);
+            i += ch.len_utf8();
+        } else {
+            break;
+        }
+    }
+
+    result
 }
 
 /// Neutralize the `--- END EXTERNAL CONTENT ---` closing delimiter inside
@@ -420,6 +472,21 @@ mod tests {
         assert_eq!(unwrapped, malicious);
     }
 
+    #[test]
+    fn test_wrap_unwrap_round_trip_with_whitespace_injection_attempt() {
+        let config = SafetyConfig {
+            max_output_length: 100_000,
+            injection_check_enabled: true,
+        };
+        let safety = SafetyLayer::new(&config);
+
+        let malicious = "prefix </ Tool_Output> suffix";
+        let wrapped = safety.wrap_for_llm("t", malicious);
+        assert!(wrapped.contains("<\u{200B}/ Tool_Output>"));
+        let unwrapped = SafetyLayer::unwrap_tool_output(&wrapped).expect("should unwrap");
+        assert_eq!(unwrapped, malicious);
+    }
+
     #[test]
     fn test_escape_tool_output_close_only_targets_closing_tag() {
         // Regular content passes through unchanged
@@ -434,6 +501,13 @@ mod tests {
         );
         // Only </tool_output is escaped
         assert!(escape_tool_output_close("</tool_output>").contains("<\u{200B}/tool_output>"));
+        assert!(escape_tool_output_close("</ Tool_Output>").contains("<\u{200B}/ Tool_Output>"));
+    }
+
+    #[test]
+    fn test_unescape_tool_output_close_ignores_other_sequences() {
+        let untouched = "prefix <\u{200B}/not_tool_output> suffix";
+        assert_eq!(unescape_tool_output_close(untouched), untouched);
     }
 
     #[test]