agentox_core/checks/security/
prompt_injection_echo.rs

1//! SEC-001: Probe tool-call behavior with prompt-injection strings.
2
3use crate::checks::runner::{Check, CheckContext};
4use crate::checks::security::constants::{
5    truncate_for_evidence, EVIDENCE_MAX_CHARS, INJECTION_LEAK_MARKERS, INJECTION_PAYLOADS,
6};
7use crate::checks::types::{CheckCategory, CheckResult, Severity};
8use crate::protocol::mcp_types::Content;
9
10pub struct PromptInjectionEchoSafety;
11
12#[async_trait::async_trait]
13impl Check for PromptInjectionEchoSafety {
14    fn id(&self) -> &str {
15        "SEC-001"
16    }
17
18    fn name(&self) -> &str {
19        "Prompt-injection echo safety"
20    }
21
22    fn category(&self) -> CheckCategory {
23        CheckCategory::Security
24    }
25
26    async fn run(&self, ctx: &mut CheckContext) -> Vec<CheckResult> {
27        let desc = "Tool calls should reject or safely handle prompt-injection-like input";
28
29        let tools = match &ctx.tools {
30            Some(tools) => tools.clone(),
31            None => match ctx.session.list_tools().await {
32                Ok(tools) => {
33                    ctx.tools = Some(tools.clone());
34                    tools
35                }
36                Err(e) => {
37                    return vec![CheckResult::fail(
38                        self.id(),
39                        self.name(),
40                        self.category(),
41                        Severity::High,
42                        desc,
43                        format!("Failed to list tools for security probing: {e}"),
44                    )];
45                }
46            },
47        };
48
49        if tools.is_empty() {
50            let mut pass = CheckResult::pass(self.id(), self.name(), self.category(), desc);
51            pass.message = "No tools exposed; injection probe is not applicable".to_string();
52            return vec![pass.with_evidence(serde_json::json!({ "tools": 0 }))];
53        }
54
55        let mut findings = Vec::new();
56        let mut benign_echoes = 0usize;
57        for tool in &tools {
58            for payload in INJECTION_PAYLOADS {
59                let args = serde_json::json!({
60                    "message": payload,
61                    "text": payload,
62                    "name": payload
63                });
64
65                match ctx.session.call_tool(&tool.name, args).await {
66                    Ok(result) => {
67                        let text_out: String = result
68                            .content
69                            .iter()
70                            .filter_map(|c| match c {
71                                Content::Text { text } => Some(text.as_str()),
72                                _ => None,
73                            })
74                            .collect::<Vec<_>>()
75                            .join(" ");
76                        let lowered = text_out.to_ascii_lowercase();
77                        let leaked = INJECTION_LEAK_MARKERS
78                            .iter()
79                            .any(|marker| lowered.contains(marker));
80                        if leaked {
81                            findings.push(
82                                CheckResult::fail(
83                                    self.id(),
84                                    self.name(),
85                                    self.category(),
86                                    Severity::High,
87                                    desc,
88                                    format!(
89                                        "Tool \"{}\" returned suspicious content for injection probe",
90                                        tool.name
91                                    ),
92                                )
93                                .with_evidence(serde_json::json!({
94                                    "tool": tool.name,
95                                    "payload": payload,
96                                    "output_sample": truncate_for_evidence(&text_out, EVIDENCE_MAX_CHARS),
97                                })),
98                            );
99                        } else if text_out.trim() == *payload {
100                            // Explicitly treat a pure echo as benign if there is no sensitive marker leakage.
101                            benign_echoes += 1;
102                        }
103                    }
104                    Err(crate::error::SessionError::JsonRpc { .. }) => {
105                        // Rejected requests are acceptable and often safer.
106                    }
107                    Err(e) => {
108                        findings.push(
109                            CheckResult::fail(
110                                self.id(),
111                                self.name(),
112                                self.category(),
113                                Severity::High,
114                                desc,
115                                format!(
116                                    "Tool \"{}\" probe failed with transport/session error: {e}",
117                                    tool.name
118                                ),
119                            )
120                            .with_evidence(serde_json::json!({
121                                "tool": tool.name,
122                                "payload": payload
123                            })),
124                        );
125                    }
126                }
127            }
128        }
129
130        if findings.is_empty() {
131            vec![
132                CheckResult::pass(self.id(), self.name(), self.category(), desc).with_evidence(
133                    serde_json::json!({
134                        "tools_probed": tools.len(),
135                        "payloads_per_tool": INJECTION_PAYLOADS.len(),
136                        "benign_echo_responses": benign_echoes
137                    }),
138                ),
139            ]
140        } else {
141            findings
142        }
143    }
144}
agentox_core/checks/security/prompt_injection_echo.rs

agentox_core/checks/security/
prompt_injection_echo.rs
