use leviath_core::taint::{
GateDecision, GateDecisionSource, GateEvent, SecurityConfig, TaintLevel, ToolClassification,
builtin_tool_classification,
};
use std::collections::HashMap;
use crate::components::ContextWindow;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GateResolution {
AllowOnce,
AlwaysAllow,
Deny,
}
#[async_trait::async_trait]
pub trait GatePrompt: Send + Sync {
async fn resolve(&self, decision: &GateDecision) -> GateResolution;
}
pub type ScriptRuleChecker = dyn Fn(&str, Option<&str>, TaintLevel) -> Option<String> + Send + Sync;
#[derive(Debug, Clone, bevy_ecs::component::Component)]
pub struct TaintGate {
config: SecurityConfig,
tool_overrides: HashMap<String, ToolClassification>,
audit_log: Vec<GateEvent>,
}
impl TaintGate {
pub fn new(config: SecurityConfig) -> Self {
Self {
config,
tool_overrides: HashMap::new(),
audit_log: Vec::new(),
}
}
pub fn disabled() -> Self {
Self {
config: SecurityConfig {
taint_tracking: false,
},
tool_overrides: HashMap::new(),
audit_log: Vec::new(),
}
}
pub fn config(&self) -> &SecurityConfig {
&self.config
}
pub fn set_tool_classification(
&mut self,
tool_name: String,
classification: ToolClassification,
) {
self.tool_overrides.insert(tool_name, classification);
}
pub fn tool_classification(&self, tool_name: &str) -> ToolClassification {
self.tool_overrides
.get(tool_name)
.cloned()
.unwrap_or_else(|| builtin_tool_classification(tool_name))
}
pub fn apply_mcp_overrides(
&mut self,
overrides: &HashMap<String, leviath_core::policy::McpToolOverride>,
) {
for (tool_name, over) in overrides {
let mut classification = self.tool_classification(tool_name);
if let Some(sensitivity) = over.sensitivity {
classification.sensitivity = sensitivity;
}
if let Some(clearance) = over.clearance {
classification.clearance = clearance;
}
if let Some(direction) = over.direction.as_deref() {
match leviath_core::taint::ToolDirection::from_str_loose(direction) {
Some(parsed) => classification.direction = parsed,
None => tracing::warn!(
tool = %tool_name,
direction = %direction,
"ignoring unrecognized direction in [mcp_overrides]"
),
}
}
self.set_tool_classification(tool_name.clone(), classification);
}
}
pub fn check_traditional(
&mut self,
agent_id: &str,
tool_name: &str,
window: &ContextWindow,
) -> GateDecision {
if !self.config.taint_tracking {
self.log_event(
agent_id,
tool_name,
TaintLevel::Public,
TaintLevel::Public,
true,
GateDecisionSource::TaintDisabled,
);
return GateDecision::Allowed;
}
let classification = self.tool_classification(tool_name);
if !classification.is_outbound() {
self.log_event(
agent_id,
tool_name,
TaintLevel::Public,
classification.clearance,
true,
GateDecisionSource::AutoAllow,
);
return GateDecision::Allowed;
}
let taint = window.overall_taint().unwrap_or(TaintLevel::Public);
if classification.check_clearance(taint) {
self.log_event(
agent_id,
tool_name,
taint,
classification.clearance,
true,
GateDecisionSource::AutoAllow,
);
GateDecision::Allowed
} else {
let source_regions: Vec<String> = window
.taint_summary()
.into_iter()
.filter(|(_, level)| *level > classification.clearance)
.map(|(name, _)| name)
.collect();
self.log_event(
agent_id,
tool_name,
taint,
classification.clearance,
false,
GateDecisionSource::AutoBlock,
);
GateDecision::Blocked {
taint_level: taint,
clearance: classification.clearance,
source_regions,
tool_name: tool_name.to_string(),
}
}
}
pub fn check_with_policy(
&mut self,
agent_id: &str,
tool_name: &str,
window: &ContextWindow,
target: Option<&str>,
policy: &leviath_core::PolicyConfig,
script_checker: Option<&ScriptRuleChecker>,
) -> GateDecision {
let decision = self.check_traditional(agent_id, tool_name, window);
if decision.is_allowed() {
return decision;
}
let (taint, clearance) = decision
.blocked_levels()
.expect("infallible: a non-Allowed GateDecision is always Blocked");
if let Some(rule_idx) = policy.check_allowlist(tool_name, target, taint) {
self.log_event(
agent_id,
tool_name,
taint,
clearance,
true,
GateDecisionSource::AllowlistRule {
rule_index: rule_idx,
},
);
return GateDecision::Allowed;
}
if let Some(checker) = script_checker
&& let Some(script_name) = checker(tool_name, target, taint)
{
self.log_event(
agent_id,
tool_name,
taint,
clearance,
true,
GateDecisionSource::ScriptedRule { script_name },
);
return GateDecision::Allowed;
}
decision
}
pub fn record_allow(
&mut self,
agent_id: &str,
tool_name: &str,
taint: TaintLevel,
clearance: TaintLevel,
source: GateDecisionSource,
) {
self.log_event(agent_id, tool_name, taint, clearance, true, source);
}
pub fn record_deny(
&mut self,
agent_id: &str,
tool_name: &str,
taint: TaintLevel,
clearance: TaintLevel,
source: GateDecisionSource,
) {
self.log_event(agent_id, tool_name, taint, clearance, false, source);
}
pub fn apply_resolution(
&mut self,
agent_id: &str,
tool_name: &str,
tool_id: &str,
taint: TaintLevel,
clearance: TaintLevel,
resolution: GateResolution,
) -> Option<(String, String)> {
match resolution {
GateResolution::AllowOnce => {
self.record_allow(
agent_id,
tool_name,
taint,
clearance,
GateDecisionSource::UserAllowOnce,
);
None
}
GateResolution::AlwaysAllow => {
self.record_allow(
agent_id,
tool_name,
taint,
clearance,
GateDecisionSource::UserAlwaysAllow,
);
let mut cls = self.tool_classification(tool_name);
cls.clearance = TaintLevel::Private;
self.set_tool_classification(tool_name.to_string(), cls);
None
}
GateResolution::Deny => {
self.record_deny(
agent_id,
tool_name,
taint,
clearance,
GateDecisionSource::UserDenied,
);
Some((
tool_id.to_string(),
format!(
"[blocked] Tool '{}' would send data at {} sensitivity, above its {} \
clearance. Denied by user.",
tool_name, taint, clearance
),
))
}
}
}
pub fn audit_log(&self) -> &[GateEvent] {
&self.audit_log
}
#[allow(clippy::too_many_arguments)]
fn log_event(
&mut self,
agent_id: &str,
tool_name: &str,
taint_level: TaintLevel,
clearance: TaintLevel,
allowed: bool,
decision_source: GateDecisionSource,
) {
self.audit_log.push(GateEvent {
timestamp: chrono::Utc::now().timestamp(),
agent_id: agent_id.to_string(),
tool_name: tool_name.to_string(),
taint_level,
clearance,
allowed,
decision_source,
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use leviath_core::taint::ToolDirection;
use leviath_core::{Region, RegionKind};
fn make_window_with_taint(taint: TaintLevel) -> ContextWindow {
let mut window = ContextWindow::new(10000);
let region =
Region::new("conv".to_string(), RegionKind::Temporary, 5000).with_taint_tracking();
window.add_region(region);
if taint != TaintLevel::Public {
window
.add_tainted_to_region("conv", "data".to_string(), 10, taint)
.unwrap();
}
window
}
#[test]
fn gate_disabled_always_allows() {
let mut gate = TaintGate::disabled();
assert!(!gate.config().taint_tracking);
let window = make_window_with_taint(TaintLevel::Private);
let decision = gate.check_traditional("agent-1", "shell", &window);
assert!(decision.is_allowed());
assert_eq!(gate.audit_log().len(), 1);
assert_eq!(
gate.audit_log()[0].decision_source,
GateDecisionSource::TaintDisabled
);
}
#[test]
fn gate_allows_non_outbound_tool() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Private);
let decision = gate.check_traditional("agent-1", "read_file", &window);
assert!(decision.is_allowed());
}
#[test]
fn gate_allows_outbound_when_taint_within_clearance() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Public);
let decision = gate.check_traditional("agent-1", "shell", &window);
assert!(decision.is_allowed());
}
#[test]
fn gate_blocks_outbound_when_taint_exceeds_clearance() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Private);
let decision = gate.check_traditional("agent-1", "shell", &window);
assert!(!decision.is_allowed());
assert_eq!(
decision,
GateDecision::Blocked {
taint_level: TaintLevel::Private,
clearance: TaintLevel::Public,
source_regions: vec!["conv".to_string()],
tool_name: "shell".to_string(),
}
);
}
#[test]
fn gate_uses_tool_override() {
let mut gate = TaintGate::new(SecurityConfig::default());
gate.set_tool_classification(
"shell".to_string(),
ToolClassification::new(
TaintLevel::Public,
ToolDirection::Outbound,
TaintLevel::Private, ),
);
let window = make_window_with_taint(TaintLevel::Private);
let decision = gate.check_traditional("agent-1", "shell", &window);
assert!(decision.is_allowed());
}
#[test]
fn apply_mcp_overrides_replaces_only_the_set_fields() {
let mut gate = TaintGate::new(SecurityConfig::default());
let before = gate.tool_classification("srv.notify");
let overrides = std::collections::HashMap::from([(
"srv.notify".to_string(),
leviath_core::policy::McpToolOverride {
sensitivity: Some(TaintLevel::Private),
direction: None,
clearance: None,
},
)]);
gate.apply_mcp_overrides(&overrides);
let after = gate.tool_classification("srv.notify");
assert_eq!(after.sensitivity, TaintLevel::Private);
assert_eq!(after.direction, before.direction);
assert_eq!(after.clearance, before.clearance);
}
#[test]
fn apply_mcp_overrides_parses_direction_and_clearance() {
let mut gate = TaintGate::new(SecurityConfig::default());
let overrides = std::collections::HashMap::from([(
"srv.post".to_string(),
leviath_core::policy::McpToolOverride {
sensitivity: None,
direction: Some("outbound".to_string()),
clearance: Some(TaintLevel::Internal),
},
)]);
gate.apply_mcp_overrides(&overrides);
let after = gate.tool_classification("srv.post");
assert_eq!(after.direction, ToolDirection::Outbound);
assert_eq!(after.clearance, TaintLevel::Internal);
}
#[test]
fn apply_mcp_overrides_keeps_direction_on_unrecognized_string() {
let mut gate = TaintGate::new(SecurityConfig::default());
let before = gate.tool_classification("srv.odd");
let overrides = std::collections::HashMap::from([(
"srv.odd".to_string(),
leviath_core::policy::McpToolOverride {
sensitivity: None,
direction: Some("sideways".to_string()),
clearance: None,
},
)]);
gate.apply_mcp_overrides(&overrides);
assert_eq!(
gate.tool_classification("srv.odd").direction,
before.direction
);
}
#[test]
fn session_approval_still_wins_over_an_mcp_override() {
let mut gate = TaintGate::new(SecurityConfig::default());
let overrides = std::collections::HashMap::from([(
"srv.send".to_string(),
leviath_core::policy::McpToolOverride {
sensitivity: None,
direction: Some("outbound".to_string()),
clearance: Some(TaintLevel::Public),
},
)]);
gate.apply_mcp_overrides(&overrides);
gate.set_tool_classification(
"srv.send".to_string(),
ToolClassification::new(
TaintLevel::Public,
ToolDirection::Outbound,
TaintLevel::Private,
),
);
assert_eq!(
gate.tool_classification("srv.send").clearance,
TaintLevel::Private
);
}
#[test]
fn gate_blocked_identifies_source_regions() {
let mut gate = TaintGate::new(SecurityConfig::default());
let mut window = ContextWindow::new(10000);
let r1 =
Region::new("clean".to_string(), RegionKind::Temporary, 5000).with_taint_tracking();
let r2 =
Region::new("dirty".to_string(), RegionKind::Temporary, 5000).with_taint_tracking();
window.add_region(r1);
window.add_region(r2);
window
.add_tainted_to_region("clean", "ok".to_string(), 5, TaintLevel::Public)
.unwrap();
window
.add_tainted_to_region("dirty", "secret".to_string(), 5, TaintLevel::Private)
.unwrap();
let decision = gate.check_traditional("agent-1", "shell", &window);
assert_eq!(
decision,
GateDecision::Blocked {
taint_level: TaintLevel::Private,
clearance: TaintLevel::Public,
source_regions: vec!["dirty".to_string()],
tool_name: "shell".to_string(),
}
);
}
#[test]
fn gate_audit_log_records_events() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Public);
gate.check_traditional("agent-1", "shell", &window);
gate.check_traditional("agent-1", "read_file", &window);
assert_eq!(gate.audit_log().len(), 2);
assert!(gate.audit_log()[0].allowed);
assert!(gate.audit_log()[1].allowed);
}
#[test]
fn gate_record_allow() {
let mut gate = TaintGate::new(SecurityConfig::default());
gate.record_allow(
"agent-1",
"shell",
TaintLevel::Private,
TaintLevel::Public,
GateDecisionSource::UserAllowOnce,
);
assert_eq!(gate.audit_log().len(), 1);
assert!(gate.audit_log()[0].allowed);
assert_eq!(
gate.audit_log()[0].decision_source,
GateDecisionSource::UserAllowOnce
);
}
#[test]
fn gate_tool_classification_returns_override() {
let mut gate = TaintGate::new(SecurityConfig::default());
let custom = ToolClassification::new(
TaintLevel::Private,
ToolDirection::Outbound,
TaintLevel::Private,
);
gate.set_tool_classification("my_tool".to_string(), custom.clone());
assert_eq!(gate.tool_classification("my_tool"), custom);
}
#[test]
fn gate_tool_classification_falls_back_to_builtin() {
let gate = TaintGate::new(SecurityConfig::default());
let tc = gate.tool_classification("read_file");
assert_eq!(tc.direction, ToolDirection::Inbound);
}
#[test]
fn gate_new_and_config() {
let config = SecurityConfig {
taint_tracking: true,
};
let gate = TaintGate::new(config.clone());
assert!(gate.config().taint_tracking);
assert!(gate.config().taint_tracking);
}
#[test]
fn gate_with_policy_allows_via_allowlist() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Private);
let policy = leviath_core::PolicyConfig {
allowlist: vec![leviath_core::AllowlistRule {
tool: "shell".into(),
to: vec![],
channel: vec![],
max_sensitivity: TaintLevel::Private,
}],
mcp_overrides: Default::default(),
};
let decision = gate.check_with_policy("agent-1", "shell", &window, None, &policy, None);
assert!(decision.is_allowed());
let last = gate.audit_log().last().unwrap();
assert!(last.allowed);
assert_eq!(
last.decision_source,
GateDecisionSource::AllowlistRule { rule_index: 0 }
);
}
#[test]
fn gate_with_policy_allows_via_scripted_rule() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Private);
let policy = leviath_core::PolicyConfig::default();
let checker = |tool: &str, _target: Option<&str>, _taint: TaintLevel| -> Option<String> {
(tool == "shell").then(|| "company_rule.rhai".to_string())
};
let decision =
gate.check_with_policy("agent-1", "shell", &window, None, &policy, Some(&checker));
assert!(decision.is_allowed());
let last = gate.audit_log().last().unwrap();
assert_eq!(
last.decision_source,
GateDecisionSource::ScriptedRule {
script_name: "company_rule.rhai".to_string()
}
);
}
#[test]
fn gate_with_policy_blocks_when_no_rule_matches() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Private);
let policy = leviath_core::PolicyConfig::default();
let decision = gate.check_with_policy("agent-1", "shell", &window, None, &policy, None);
assert!(!decision.is_allowed());
}
#[test]
fn gate_with_policy_passes_through_when_already_allowed() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Public);
let policy = leviath_core::PolicyConfig::default();
let decision = gate.check_with_policy("agent-1", "shell", &window, None, &policy, None);
assert!(decision.is_allowed());
}
#[test]
fn gate_with_policy_target_pattern_matching() {
let mut gate = TaintGate::new(SecurityConfig::default());
gate.set_tool_classification(
"send_email".to_string(),
ToolClassification::new(
TaintLevel::Public,
ToolDirection::Outbound,
TaintLevel::Public,
),
);
let window = make_window_with_taint(TaintLevel::Private);
let policy = leviath_core::PolicyConfig {
allowlist: vec![leviath_core::AllowlistRule {
tool: "send_email".into(),
to: vec!["megan@*".into()],
channel: vec![],
max_sensitivity: TaintLevel::Private,
}],
mcp_overrides: Default::default(),
};
let decision = gate.check_with_policy(
"agent-1",
"send_email",
&window,
Some("megan@work.com"),
&policy,
None,
);
assert!(decision.is_allowed());
let decision2 = gate.check_with_policy(
"agent-1",
"send_email",
&window,
Some("bob@work.com"),
&policy,
None,
);
assert!(!decision2.is_allowed());
}
#[test]
fn gate_check_traditional_with_internal_taint() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Internal);
let decision = gate.check_traditional("agent-1", "shell", &window);
assert!(!decision.is_allowed());
assert_eq!(
decision,
GateDecision::Blocked {
taint_level: TaintLevel::Internal,
clearance: TaintLevel::Public,
source_regions: vec!["conv".to_string()],
tool_name: "shell".to_string(),
}
);
}
#[test]
fn gate_audit_log_records_blocked_events() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Private);
gate.check_traditional("agent-1", "shell", &window);
assert_eq!(gate.audit_log().len(), 1);
assert!(!gate.audit_log()[0].allowed);
assert_eq!(gate.audit_log()[0].tool_name, "shell");
assert_eq!(gate.audit_log()[0].taint_level, TaintLevel::Private);
assert_eq!(
gate.audit_log()[0].decision_source,
GateDecisionSource::AutoBlock,
);
}
#[test]
fn gate_with_policy_scripted_rule_non_matching_tool() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Private);
let policy = leviath_core::PolicyConfig::default();
let checker = |_tool: &str, _target: Option<&str>, _taint: TaintLevel| -> Option<String> {
None };
let decision =
gate.check_with_policy("agent-1", "shell", &window, None, &policy, Some(&checker));
assert!(!decision.is_allowed());
}
#[test]
fn gate_with_policy_non_outbound_skips_policy_check() {
let mut gate = TaintGate::new(SecurityConfig::default());
let window = make_window_with_taint(TaintLevel::Private);
let policy = leviath_core::PolicyConfig::default();
let decision = gate.check_with_policy("agent-1", "read_file", &window, None, &policy, None);
assert!(decision.is_allowed());
}
#[test]
fn gate_multiple_tool_overrides() {
let mut gate = TaintGate::new(SecurityConfig::default());
gate.set_tool_classification(
"tool_a".to_string(),
ToolClassification::new(
TaintLevel::Public,
ToolDirection::Outbound,
TaintLevel::Internal,
),
);
gate.set_tool_classification(
"tool_b".to_string(),
ToolClassification::new(
TaintLevel::Public,
ToolDirection::Outbound,
TaintLevel::Private,
),
);
let window = make_window_with_taint(TaintLevel::Private);
let decision_a = gate.check_traditional("agent-1", "tool_a", &window);
assert!(!decision_a.is_allowed());
let decision_b = gate.check_traditional("agent-1", "tool_b", &window);
assert!(decision_b.is_allowed());
}
#[test]
fn apply_resolution_allow_once_records_and_executes() {
let mut gate = TaintGate::new(SecurityConfig::default());
let out = gate.apply_resolution(
"a",
"shell",
"call1",
TaintLevel::Private,
TaintLevel::Public,
GateResolution::AllowOnce,
);
assert!(out.is_none()); let allow = gate
.audit_log()
.iter()
.find(|e| e.allowed)
.expect("an allowed event should be logged");
assert_eq!(allow.decision_source, GateDecisionSource::UserAllowOnce);
}
#[test]
fn apply_resolution_always_allow_raises_clearance() {
let mut gate = TaintGate::new(SecurityConfig::default());
let out = gate.apply_resolution(
"a",
"shell",
"call1",
TaintLevel::Private,
TaintLevel::Public,
GateResolution::AlwaysAllow,
);
assert!(out.is_none());
assert_eq!(
gate.tool_classification("shell").clearance,
TaintLevel::Private
);
let allow = gate
.audit_log()
.iter()
.find(|e| e.allowed)
.expect("an allowed event should be logged");
assert_eq!(allow.decision_source, GateDecisionSource::UserAlwaysAllow);
}
#[test]
fn apply_resolution_deny_returns_blocked_result() {
let mut gate = TaintGate::new(SecurityConfig::default());
let out = gate.apply_resolution(
"a",
"shell",
"call1",
TaintLevel::Private,
TaintLevel::Public,
GateResolution::Deny,
);
let (id, msg) = out.expect("deny yields a blocked result");
assert_eq!(id, "call1");
assert!(msg.contains("[blocked]") && msg.contains("shell"));
let deny = gate
.audit_log()
.iter()
.find(|e| !e.allowed)
.expect("a denied event should be logged");
assert_eq!(deny.decision_source, GateDecisionSource::UserDenied);
}
}