gemini_adk_rs/plugin/
security.rs

1//! Security plugin — policy-based tool call authorization.
2
3use async_trait::async_trait;
4
5use gemini_genai_rs::prelude::FunctionCall;
6
7use super::{Plugin, PluginResult};
8use crate::context::InvocationContext;
9
10/// The outcome of a policy evaluation.
11#[derive(Debug, Clone)]
12pub enum PolicyOutcome {
13    /// Allow the tool call to proceed.
14    Allow,
15    /// Require user confirmation before proceeding.
16    Confirm(String),
17    /// Deny the tool call with a reason.
18    Deny(String),
19}
20
21/// Trait for evaluating tool call policies.
22///
23/// Implementations can check tool names, arguments, user permissions,
24/// rate limits, etc.
25pub trait PolicyEngine: Send + Sync + 'static {
26    /// Evaluate whether a tool call should be allowed.
27    fn evaluate(&self, tool_name: &str, args: &serde_json::Value) -> PolicyOutcome;
28}
29
30/// Plugin that enforces tool call policies via a `PolicyEngine`.
31///
32/// Before every tool call, the security plugin consults the policy engine.
33/// If the engine returns `Deny`, the tool call is blocked. If it returns
34/// `Confirm`, the tool call is blocked with a confirmation message (in a
35/// real system, this would prompt the user).
36pub struct SecurityPlugin {
37    engine: Box<dyn PolicyEngine>,
38}
39
40impl SecurityPlugin {
41    /// Create a new security plugin with the given policy engine.
42    pub fn new(engine: impl PolicyEngine + 'static) -> Self {
43        Self {
44            engine: Box::new(engine),
45        }
46    }
47}
48
49#[async_trait]
50impl Plugin for SecurityPlugin {
51    fn name(&self) -> &str {
52        "security"
53    }
54
55    async fn before_tool(&self, call: &FunctionCall, _ctx: &InvocationContext) -> PluginResult {
56        match self.engine.evaluate(&call.name, &call.args) {
57            PolicyOutcome::Allow => {
58                tracing::debug!(tool = %call.name, "[plugin:security] Tool call allowed");
59                PluginResult::Continue
60            }
61            PolicyOutcome::Confirm(msg) => {
62                tracing::warn!(tool = %call.name, reason = %msg, "[plugin:security] Tool call requires confirmation");
63                PluginResult::Deny(format!("Confirmation required: {msg}"))
64            }
65            PolicyOutcome::Deny(reason) => {
66                tracing::warn!(tool = %call.name, reason = %reason, "[plugin:security] Tool call denied");
67                PluginResult::Deny(reason)
68            }
69        }
70    }
71}
72
73/// A simple policy engine that blocks specific tool names.
74pub struct DenyListPolicy {
75    blocked_tools: Vec<String>,
76}
77
78impl DenyListPolicy {
79    /// Create a policy that denies specific tools by name.
80    pub fn new(blocked_tools: Vec<String>) -> Self {
81        Self { blocked_tools }
82    }
83}
84
85impl PolicyEngine for DenyListPolicy {
86    fn evaluate(&self, tool_name: &str, _args: &serde_json::Value) -> PolicyOutcome {
87        if self.blocked_tools.iter().any(|t| t == tool_name) {
88            PolicyOutcome::Deny(format!("Tool '{tool_name}' is blocked by policy"))
89        } else {
90            PolicyOutcome::Allow
91        }
92    }
93}
94
95/// A policy engine that allows all tool calls.
96pub struct AllowAllPolicy;
97
98impl PolicyEngine for AllowAllPolicy {
99    fn evaluate(&self, _tool_name: &str, _args: &serde_json::Value) -> PolicyOutcome {
100        PolicyOutcome::Allow
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn deny_list_policy_blocks() {
110        let policy = DenyListPolicy::new(vec!["dangerous_tool".into()]);
111        let result = policy.evaluate("dangerous_tool", &serde_json::json!({}));
112        assert!(matches!(result, PolicyOutcome::Deny(_)));
113    }
114
115    #[test]
116    fn deny_list_policy_allows() {
117        let policy = DenyListPolicy::new(vec!["dangerous_tool".into()]);
118        let result = policy.evaluate("safe_tool", &serde_json::json!({}));
119        assert!(matches!(result, PolicyOutcome::Allow));
120    }
121
122    #[test]
123    fn allow_all_policy() {
124        let policy = AllowAllPolicy;
125        let result = policy.evaluate("anything", &serde_json::json!({}));
126        assert!(matches!(result, PolicyOutcome::Allow));
127    }
128
129    #[tokio::test]
130    async fn security_plugin_denies_blocked_tool() {
131        use tokio::sync::broadcast;
132
133        let policy = DenyListPolicy::new(vec!["rm_rf".into()]);
134        let plugin = SecurityPlugin::new(policy);
135
136        let (evt_tx, _) = broadcast::channel(16);
137        let writer: std::sync::Arc<dyn gemini_genai_rs::session::SessionWriter> =
138            std::sync::Arc::new(crate::test_helpers::MockWriter);
139        let session = crate::agent_session::AgentSession::from_writer(writer, evt_tx);
140        let ctx = InvocationContext::new(session);
141
142        let call = FunctionCall {
143            name: "rm_rf".into(),
144            args: serde_json::json!({}),
145            id: None,
146        };
147
148        let result = plugin.before_tool(&call, &ctx).await;
149        assert!(result.is_deny());
150    }
151
152    #[tokio::test]
153    async fn security_plugin_allows_safe_tool() {
154        use tokio::sync::broadcast;
155
156        let policy = DenyListPolicy::new(vec!["rm_rf".into()]);
157        let plugin = SecurityPlugin::new(policy);
158
159        let (evt_tx, _) = broadcast::channel(16);
160        let writer: std::sync::Arc<dyn gemini_genai_rs::session::SessionWriter> =
161            std::sync::Arc::new(crate::test_helpers::MockWriter);
162        let session = crate::agent_session::AgentSession::from_writer(writer, evt_tx);
163        let ctx = InvocationContext::new(session);
164
165        let call = FunctionCall {
166            name: "get_weather".into(),
167            args: serde_json::json!({}),
168            id: None,
169        };
170
171        let result = plugin.before_tool(&call, &ctx).await;
172        assert!(result.is_continue());
173    }
174}
gemini_adk_rs/plugin/security.rs

gemini_adk_rs/plugin/
security.rs
