use atd_protocol::ToolDefinition;
pub trait Middleware: Send + Sync {
fn name(&self) -> &'static str;
fn on_result(&self, tool_id: &str, tool_def: &ToolDefinition, result: &mut serde_json::Value);
}
fn walk_strings(value: &mut serde_json::Value, f: &mut impl FnMut(&mut String)) {
match value {
serde_json::Value::String(s) => f(s),
serde_json::Value::Array(arr) => {
for v in arr.iter_mut() {
walk_strings(v, f);
}
}
serde_json::Value::Object(obj) => {
for (_k, v) in obj.iter_mut() {
walk_strings(v, f);
}
}
_ => {}
}
}
pub struct RedactPathsMiddleware {
patterns: Vec<(regex::Regex, String)>,
}
impl RedactPathsMiddleware {
pub fn new(patterns: Vec<(regex::Regex, String)>) -> Self {
Self { patterns }
}
pub fn with_home_default() -> Self {
let patterns = match std::env::var("HOME") {
Ok(home) if !home.is_empty() => {
let escaped = regex::escape(&home);
match regex::Regex::new(&escaped) {
Ok(re) => vec![(re, "<redacted:home>".to_string())],
Err(_) => vec![],
}
}
_ => vec![],
};
Self { patterns }
}
}
impl Middleware for RedactPathsMiddleware {
fn name(&self) -> &'static str {
"redact_paths"
}
fn on_result(
&self,
_tool_id: &str,
_tool_def: &ToolDefinition,
result: &mut serde_json::Value,
) {
if self.patterns.is_empty() {
return;
}
let patterns = &self.patterns;
walk_strings(result, &mut |s| {
for (re, rep) in patterns {
*s = re.replace_all(s, rep.as_str()).into_owned();
}
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use atd_protocol::{
BindingProtocol, SafetyLevel, ToolBinding, ToolCapability, ToolResources, ToolSafety,
ToolTrust, ToolVisibility, TrustLevel,
};
fn tool_def() -> ToolDefinition {
ToolDefinition {
id: "test:mw".into(),
name: "mw".into(),
description: "middleware test fixture".into(),
version: "0.0.0".into(),
capability: ToolCapability {
domain: "test".into(),
actions: vec![],
tags: vec![],
intent_examples: vec![],
},
input_schema: serde_json::json!({}),
output_schema: serde_json::json!({}),
bindings: vec![ToolBinding {
protocol: BindingProtocol::Cli,
config: serde_json::json!({}),
}],
safety: ToolSafety {
level: SafetyLevel::Read,
dry_run: false,
side_effects: vec![],
data_sensitivity: None,
},
resources: ToolResources {
timeout_ms: 1000,
max_concurrent: 1,
rate_limit_per_min: None,
estimated_tokens: None,
},
trust: ToolTrust {
publisher: "test".into(),
trust_level: TrustLevel::L0Unverified,
signature: None,
},
visibility: ToolVisibility::Read,
required_capabilities: vec![],
tier: None,
errors: vec![],
}
}
fn mw_with(pattern: &str, rep: &str) -> RedactPathsMiddleware {
let re = regex::Regex::new(pattern).unwrap();
RedactPathsMiddleware::new(vec![(re, rep.to_string())])
}
#[test]
fn redacts_pattern_in_top_level_string() {
let mw = mw_with(r"/home/[^/]+", "<redacted>");
let def = tool_def();
let mut v = serde_json::json!({"path": "/home/alice/x.txt"});
mw.on_result("test:mw", &def, &mut v);
assert_eq!(v["path"], "<redacted>/x.txt");
}
#[test]
fn redacts_in_nested_object() {
let mw = mw_with(r"secret", "***");
let def = tool_def();
let mut v = serde_json::json!({
"outer": {"inner": "this is a secret value"}
});
mw.on_result("t", &def, &mut v);
assert_eq!(v["outer"]["inner"], "this is a *** value");
}
#[test]
fn redacts_in_array_elements() {
let mw = mw_with(r"password=\w+", "password=<redacted>");
let def = tool_def();
let mut v = serde_json::json!({
"entries": ["password=hunter2", "normal line", "password=correct horse"]
});
mw.on_result("t", &def, &mut v);
let arr = v["entries"].as_array().unwrap();
assert_eq!(arr[0], "password=<redacted>");
assert_eq!(arr[1], "normal line");
assert_eq!(arr[2], "password=<redacted> horse");
}
#[test]
fn leaves_non_string_leaves_untouched() {
let mw = mw_with(r"\d+", "N");
let def = tool_def();
let mut v = serde_json::json!({
"num": 42,
"bool": true,
"null": null,
"str_with_num": "port 42"
});
mw.on_result("t", &def, &mut v);
assert_eq!(v["num"], 42);
assert_eq!(v["bool"], true);
assert_eq!(v["null"], serde_json::Value::Null);
assert_eq!(v["str_with_num"], "port N");
}
#[test]
fn applies_multiple_patterns_in_order() {
let p1 = (regex::Regex::new(r"aaa").unwrap(), "bbb".to_string());
let p2 = (regex::Regex::new(r"bbb").unwrap(), "ccc".to_string());
let mw = RedactPathsMiddleware::new(vec![p1, p2]);
let def = tool_def();
let mut v = serde_json::json!({"x": "aaa"});
mw.on_result("t", &def, &mut v);
assert_eq!(v["x"], "ccc");
}
#[test]
fn name_is_stable() {
let mw = RedactPathsMiddleware::new(vec![]);
assert_eq!(mw.name(), "redact_paths");
}
#[test]
fn empty_middleware_is_a_noop() {
let mw = RedactPathsMiddleware::new(vec![]);
let def = tool_def();
let mut v = serde_json::json!({"x": "unchanged"});
mw.on_result("t", &def, &mut v);
assert_eq!(v["x"], "unchanged");
}
#[test]
fn with_home_default_handles_home_path_or_is_noop_when_unset() {
let prev = std::env::var_os("HOME");
unsafe {
std::env::set_var("HOME", "/tmp/fakehome-sp12");
}
let mw = RedactPathsMiddleware::with_home_default();
let def = tool_def();
let mut v = serde_json::json!({"p": "/tmp/fakehome-sp12/secret"});
mw.on_result("t", &def, &mut v);
assert_eq!(v["p"], "<redacted:home>/secret");
unsafe {
std::env::remove_var("HOME");
}
let mw2 = RedactPathsMiddleware::with_home_default();
let mut v2 = serde_json::json!({"p": "/tmp/anything"});
mw2.on_result("t", &def, &mut v2);
assert_eq!(v2["p"], "/tmp/anything");
if let Some(h) = prev {
unsafe {
std::env::set_var("HOME", h);
}
}
}
}