Skip to main content

claude_utils/mcp/
validator.rs

1use serde_json::Value;
2use crate::mcp::protocol::*;
3
4pub struct RequestValidator;
5
6impl RequestValidator {
7    /// Validate JSON-RPC request structure
8    pub fn validate_jsonrpc_request(request: &JsonRpcRequest) -> Result<(), String> {
9        // Validate JSON-RPC version
10        if request.jsonrpc != "2.0" {
11            return Err("Invalid JSON-RPC version".to_string());
12        }
13
14        // Validate method name
15        if request.method.is_empty() {
16            return Err("Method name cannot be empty".to_string());
17        }
18
19        // Validate method format (should not contain spaces or special chars)
20        if !request.method.chars().all(|c| c.is_alphanumeric() || c == '.' || c == '/' || c == '_') {
21            return Err("Invalid method name format".to_string());
22        }
23
24        Ok(())
25    }
26
27    /// Validate tool call parameters
28    pub fn validate_tool_call(tool_request: &ToolCallRequest) -> Result<(), String> {
29        // Validate tool name
30        if tool_request.name.is_empty() {
31            return Err("Tool name cannot be empty".to_string());
32        }
33
34        // Validate known tools
35        match tool_request.name.as_str() {
36            "clipboard.get" => Self::validate_clipboard_get_args(tool_request.arguments.as_ref()),
37            "clipboard.set" => Self::validate_clipboard_set_args(tool_request.arguments.as_ref()),
38            _ => Err(format!("Unknown tool: {}", tool_request.name)),
39        }
40    }
41
42    fn validate_clipboard_get_args(args: Option<&Value>) -> Result<(), String> {
43        if let Some(args) = args {
44            if let Some(format) = args.get("format").and_then(|f| f.as_str()) {
45                match format {
46                    "auto" | "text" | "image" => Ok(()),
47                    _ => Err(format!("Invalid format: {format}. Must be 'auto', 'text', or 'image'")),
48                }
49            } else if args.get("format").is_some() {
50                Err("Format must be a string".to_string())
51            } else {
52                Ok(())
53            }
54        } else {
55            Ok(())
56        }
57    }
58
59    fn validate_clipboard_set_args(args: Option<&Value>) -> Result<(), String> {
60        match args {
61            None => Err("Missing required arguments for clipboard.set".to_string()),
62            Some(args) => {
63                // Validate required fields
64                let content_type = args.get("type")
65                    .and_then(|t| t.as_str())
66                    .ok_or("Missing required field 'type'".to_string())?;
67
68                let data = args.get("data")
69                    .and_then(|d| d.as_str())
70                    .ok_or("Missing required field 'data'".to_string())?;
71
72                // Validate content type
73                match content_type {
74                    "text/plain" => {
75                        // Validate text data isn't too large (10MB limit)
76                        if data.len() > 10 * 1024 * 1024 {
77                            return Err("Text data too large (max 10MB)".to_string());
78                        }
79                        Ok(())
80                    }
81                    "image/png" => {
82                        // Validate base64 format
83                        use base64::{Engine as _, engine::general_purpose::STANDARD};
84                        if STANDARD.decode(data).is_err() {
85                            return Err("Invalid base64 image data".to_string());
86                        }
87                        // Check size (50MB limit for images)
88                        let decoded_size = (data.len() * 3) / 4;
89                        if decoded_size > 50 * 1024 * 1024 {
90                            return Err("Image data too large (max 50MB)".to_string());
91                        }
92                        Ok(())
93                    }
94                    _ => Err(format!("Unsupported content type: {content_type}")),
95                }
96            }
97        }
98    }
99
100    /// Sanitize request data to prevent injection attacks
101    pub fn sanitize_request(request: &mut JsonRpcRequest) {
102        // Remove any potentially dangerous characters from method name
103        request.method = request.method
104            .chars()
105            .filter(|c| c.is_alphanumeric() || *c == '.' || *c == '/' || *c == '_')
106            .collect();
107
108        // Limit method name length
109        if request.method.len() > 100 {
110            request.method.truncate(100);
111        }
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use serde_json::json;
119
120    #[test]
121    fn test_validate_jsonrpc_request() {
122        let valid_request = JsonRpcRequest {
123            jsonrpc: "2.0".to_string(),
124            method: "tools.list".to_string(),
125            params: None,
126            id: Some(json!(1)),
127        };
128        assert!(RequestValidator::validate_jsonrpc_request(&valid_request).is_ok());
129
130        let invalid_version = JsonRpcRequest {
131            jsonrpc: "1.0".to_string(),
132            method: "tools.list".to_string(),
133            params: None,
134            id: Some(json!(1)),
135        };
136        assert!(RequestValidator::validate_jsonrpc_request(&invalid_version).is_err());
137    }
138
139    #[test]
140    fn test_validate_clipboard_set() {
141        let valid_text = ToolCallRequest {
142            name: "clipboard.set".to_string(),
143            arguments: Some(json!({
144                "type": "text/plain",
145                "data": "Hello, world!"
146            })),
147        };
148        assert!(RequestValidator::validate_tool_call(&valid_text).is_ok());
149
150        let missing_data = ToolCallRequest {
151            name: "clipboard.set".to_string(),
152            arguments: Some(json!({
153                "type": "text/plain"
154            })),
155        };
156        assert!(RequestValidator::validate_tool_call(&missing_data).is_err());
157    }
158}