Skip to main content

agentox_core/checks/security/
tool_param_boundary.rs

1//! SEC-002: Validate handling of malformed `tools/call` parameters.
2
3use crate::checks::runner::{Check, CheckContext};
4use crate::checks::types::{CheckCategory, CheckResult, Severity};
5use crate::protocol::jsonrpc::JsonRpcRequest;
6
7pub struct ToolParameterBoundaryValidation;
8
9fn accepted_error_codes_for_case(case: &str) -> &'static [i64] {
10    match case {
11        "missing name" => &[-32600, -32601, -32602],
12        "name is wrong type" => &[-32600, -32601, -32602],
13        "unknown tool" => &[-32601, -32602],
14        "arguments wrong type" => &[-32600, -32601, -32602, -32000],
15        _ => &[-32600, -32601, -32602],
16    }
17}
18
19#[async_trait::async_trait]
20impl Check for ToolParameterBoundaryValidation {
21    fn id(&self) -> &str {
22        "SEC-002"
23    }
24
25    fn name(&self) -> &str {
26        "Tool parameter boundary validation"
27    }
28
29    fn category(&self) -> CheckCategory {
30        CheckCategory::Security
31    }
32
33    async fn run(&self, ctx: &mut CheckContext) -> Vec<CheckResult> {
34        let desc = "Malformed tools/call parameters should return structured JSON-RPC errors";
35
36        let tools = match &ctx.tools {
37            Some(tools) => tools.clone(),
38            None => match ctx.session.list_tools().await {
39                Ok(tools) => {
40                    ctx.tools = Some(tools.clone());
41                    tools
42                }
43                Err(e) => {
44                    return vec![CheckResult::fail(
45                        self.id(),
46                        self.name(),
47                        self.category(),
48                        Severity::High,
49                        desc,
50                        format!("Failed to list tools before parameter validation tests: {e}"),
51                    )];
52                }
53            },
54        };
55
56        let known_tool = tools.first().map(|t| t.name.clone());
57        let mut cases = vec![
58            (
59                "missing name",
60                serde_json::json!({ "arguments": { "message": "x" } }),
61            ),
62            (
63                "name is wrong type",
64                serde_json::json!({ "name": 123, "arguments": {} }),
65            ),
66            (
67                "unknown tool",
68                serde_json::json!({ "name": "__missing__", "arguments": {} }),
69            ),
70        ];
71        if let Some(tool_name) = known_tool {
72            cases.push((
73                "arguments wrong type",
74                serde_json::json!({ "name": tool_name, "arguments": "not-an-object" }),
75            ));
76        }
77
78        let mut findings = Vec::new();
79        for (label, params) in cases {
80            let req =
81                JsonRpcRequest::new(ctx.session.next_id(), "tools/call", Some(params.clone()));
82            match ctx.session.send_request(&req).await {
83                Ok(resp) => {
84                    if let Some(error) = resp.error {
85                        let accepted = accepted_error_codes_for_case(label);
86                        if !accepted.contains(&error.code) {
87                            findings.push(
88                                CheckResult::fail(
89                                    self.id(),
90                                    self.name(),
91                                    self.category(),
92                                    Severity::Medium,
93                                    desc,
94                                    format!(
95                                        "{label}: expected one of {:?}, got code {}",
96                                        accepted, error.code
97                                    ),
98                                )
99                                .with_evidence(serde_json::json!({
100                                    "case": label,
101                                    "params": params,
102                                    "error_code": error.code,
103                                    "error_message": error.message
104                                })),
105                            );
106                        }
107                    } else if resp
108                        .result
109                        .as_ref()
110                        .and_then(|v| v.get("isError"))
111                        .and_then(|v| v.as_bool())
112                        .unwrap_or(false)
113                    {
114                        // Some servers return structured tool errors through result payloads.
115                    } else {
116                        findings.push(
117                            CheckResult::fail(
118                                self.id(),
119                                self.name(),
120                                self.category(),
121                                Severity::High,
122                                desc,
123                                format!("{label}: server accepted malformed tool-call parameters"),
124                            )
125                            .with_evidence(serde_json::json!({
126                                "case": label,
127                                "params": params,
128                                "response_result": resp.result
129                            })),
130                        );
131                    }
132                }
133                Err(e) => {
134                    findings.push(
135                        CheckResult::fail(
136                            self.id(),
137                            self.name(),
138                            self.category(),
139                            Severity::Critical,
140                            desc,
141                            format!("{label}: transport/session failed while testing parameter boundary: {e}"),
142                        )
143                        .with_evidence(serde_json::json!({
144                            "case": label,
145                            "params": params
146                        })),
147                    );
148                }
149            }
150        }
151
152        if findings.is_empty() {
153            vec![CheckResult::pass(
154                self.id(),
155                self.name(),
156                self.category(),
157                desc,
158            )]
159        } else {
160            findings
161        }
162    }
163}