use std::collections::{HashMap, HashSet};
use tokio::sync::RwLock;
pub struct TaintLedger {
untrusted_tools: HashSet<String>,
tainted: RwLock<HashMap<Option<String>, HashSet<String>>>,
}
impl TaintLedger {
pub fn new(untrusted_tools: HashSet<String>) -> Self {
Self {
untrusted_tools,
tainted: RwLock::new(HashMap::new()),
}
}
pub fn untrusted_tools(&self) -> &HashSet<String> {
&self.untrusted_tools
}
pub async fn record_result(
&self,
tenant: Option<&str>,
action: &car_ir::Action,
written_keys: impl IntoIterator<Item = String>,
) {
let written: Vec<String> = written_keys.into_iter().collect();
if written.is_empty() {
return;
}
let tool_is_untrusted = action
.tool
.as_ref()
.map(|t| self.untrusted_tools.contains(t))
.unwrap_or(false);
let reads = action.effective_read_set();
let mut guard = self.tainted.write().await;
let key = tenant.map(str::to_string);
let entry = guard.entry(key).or_default();
let tainted = tool_is_untrusted || reads.iter().any(|k| entry.contains(k));
if tainted {
for k in written {
entry.insert(k);
}
} else {
for k in &written {
entry.remove(k);
}
}
}
pub async fn untrusted_action_ids(
&self,
tenant: Option<&str>,
actions: &[car_ir::Action],
) -> HashSet<String> {
let guard = self.tainted.read().await;
let Some(keys) = guard.get(&tenant.map(str::to_string)) else {
return HashSet::new();
};
if keys.is_empty() {
return HashSet::new();
}
actions
.iter()
.filter(|a| a.effective_read_set().iter().any(|k| keys.contains(k)))
.map(|a| a.id.clone())
.collect()
}
pub async fn tainted_keys(&self, tenant: Option<&str>) -> HashSet<String> {
self.tainted
.read()
.await
.get(&tenant.map(str::to_string))
.cloned()
.unwrap_or_default()
}
}
#[cfg(test)]
mod tests {
use super::*;
use car_ir::{Action, ActionType};
fn tool_action(id: &str, tool: &str) -> Action {
let mut a = Action::new(ActionType::ToolCall);
a.id = id.to_string();
a.tool = Some(tool.to_string());
a.max_retries = 0;
a
}
fn ledger() -> TaintLedger {
TaintLedger::new(["fetch_web".to_string()].into_iter().collect())
}
#[tokio::test]
async fn a_trusted_tool_reading_a_tainted_key_is_untrusted() {
let led = ledger();
led.record_result(None, &tool_action("a1", "fetch_web"), ["page".to_string()])
.await;
assert!(led.tainted_keys(None).await.contains("page"));
let mut summarize = tool_action("a2", "summarize");
summarize.state_dependencies.push("page".to_string());
let ids = led.untrusted_action_ids(None, &[summarize]).await;
assert!(
ids.contains("a2"),
"an action reading a tainted key must be reported untrusted: {ids:?}"
);
}
#[tokio::test]
async fn taint_flows_through_a_trusted_intermediate() {
let led = ledger();
led.record_result(None, &tool_action("a1", "fetch_web"), ["page".to_string()])
.await;
let mut summarize = tool_action("a2", "summarize");
summarize.state_dependencies.push("page".to_string());
led.record_result(None, &summarize, ["digest".to_string()])
.await;
assert!(led.tainted_keys(None).await.contains("digest"));
let mut pay = tool_action("a3", "send_payment");
pay.state_dependencies.push("digest".to_string());
assert!(led.untrusted_action_ids(None, &[pay]).await.contains("a3"));
}
#[tokio::test]
async fn a_trusted_overwrite_clears_the_taint() {
let led = ledger();
led.record_result(None, &tool_action("a1", "fetch_web"), ["page".to_string()])
.await;
assert!(led.tainted_keys(None).await.contains("page"));
led.record_result(None, &tool_action("a2", "summarize"), ["page".to_string()])
.await;
assert!(
!led.tainted_keys(None).await.contains("page"),
"a trusted overwrite must clear the key"
);
let mut reader = tool_action("a3", "send_payment");
reader.state_dependencies.push("page".to_string());
assert!(led.untrusted_action_ids(None, &[reader]).await.is_empty());
}
#[tokio::test]
async fn taint_does_not_cross_tenants() {
let led = ledger();
led.record_result(
Some("tenant-a"),
&tool_action("a1", "fetch_web"),
["page".to_string()],
)
.await;
let mut reader = tool_action("b1", "summarize");
reader.state_dependencies.push("page".to_string());
assert!(led
.untrusted_action_ids(Some("tenant-a"), std::slice::from_ref(&reader))
.await
.contains("b1"));
assert!(led
.untrusted_action_ids(Some("tenant-b"), std::slice::from_ref(&reader))
.await
.is_empty());
assert!(led
.untrusted_action_ids(None, std::slice::from_ref(&reader))
.await
.is_empty());
}
#[tokio::test]
async fn an_action_writing_nothing_changes_nothing() {
let led = ledger();
led.record_result(None, &tool_action("a1", "fetch_web"), Vec::new())
.await;
assert!(led.tainted_keys(None).await.is_empty());
}
#[tokio::test]
async fn an_untainted_ledger_reports_no_untrusted_ids() {
let led = ledger();
let mut reader = tool_action("a1", "summarize");
reader.state_dependencies.push("page".to_string());
assert!(led.untrusted_action_ids(None, &[reader]).await.is_empty());
}
}