pulseengine_mcp_auth/security/
mod.rs

1//! Security features for MCP request/response processing
2//!
3//! This module provides comprehensive security validation, sanitization,
4//! and protection features for MCP protocol messages.
5
6pub mod request_security;
7
8pub use request_security::{
9    InputSanitizer, RequestLimitsConfig, RequestSecurityConfig, RequestSecurityValidator,
10    SecuritySeverity, SecurityValidationError, SecurityViolation, SecurityViolationType,
11};
12
13#[cfg(test)]
14mod tests {
15    use super::*;
16    use pulseengine_mcp_protocol::Request;
17    use serde_json::json;
18
19    #[test]
20    fn test_security_module_exports() {
21        // Test that all security types are accessible
22
23        let config = RequestSecurityConfig::default();
24        assert!(config.limits.max_request_size > 0);
25        assert!(config.limits.max_parameters > 0);
26
27        let _sanitizer = InputSanitizer::new();
28        // InputSanitizer should be creatable
29
30        let violation = SecurityViolation {
31            violation_type: SecurityViolationType::SizeLimit,
32            severity: SecuritySeverity::High,
33            description: "Test violation".to_string(),
34            field: None,
35            value: None,
36            timestamp: chrono::Utc::now(),
37        };
38
39        assert_eq!(violation.violation_type, SecurityViolationType::SizeLimit);
40        assert_eq!(violation.severity, SecuritySeverity::High);
41    }
42
43    #[test]
44    fn test_security_severity_ordering() {
45        // Test that severity levels are properly ordered
46        assert!(SecuritySeverity::Critical > SecuritySeverity::High);
47        assert!(SecuritySeverity::High > SecuritySeverity::Medium);
48        assert!(SecuritySeverity::Medium > SecuritySeverity::Low);
49        assert!(SecuritySeverity::Medium > SecuritySeverity::Low);
50    }
51
52    #[test]
53    fn test_security_violation_types() {
54        let violation_types = vec![
55            SecurityViolationType::SizeLimit,
56            SecurityViolationType::ParameterLimit,
57            SecurityViolationType::InjectionAttempt,
58            SecurityViolationType::MaliciousContent,
59            SecurityViolationType::InvalidFormat,
60            SecurityViolationType::RateLimit,
61            SecurityViolationType::UnauthorizedMethod,
62        ];
63
64        for violation_type in violation_types {
65            let violation = SecurityViolation {
66                violation_type: violation_type.clone(),
67                severity: SecuritySeverity::Medium,
68                description: format!("Test {:?}", violation_type),
69                field: None,
70                value: None,
71                timestamp: chrono::Utc::now(),
72            };
73
74            assert_eq!(violation.violation_type, violation_type);
75            assert!(!violation.description.is_empty());
76        }
77    }
78
79    #[tokio::test]
80    async fn test_request_security_validator() {
81        let config = RequestSecurityConfig::default();
82        let validator = RequestSecurityValidator::new(config);
83
84        // Test valid request
85        let valid_request = Request {
86            jsonrpc: "2.0".to_string(),
87            method: "tools/list".to_string(),
88            id: json!(1),
89            params: json!({}),
90        };
91
92        let result = validator.validate_request(&valid_request, None).await;
93        assert!(result.is_ok());
94
95        // Test request with too many parameters
96        let large_params = (0..1000)
97            .map(|i| (format!("param_{}", i), json!(i)))
98            .collect::<serde_json::Map<_, _>>();
99        let large_request = Request {
100            jsonrpc: "2.0".to_string(),
101            method: "tools/call".to_string(),
102            id: json!(2),
103            params: json!(large_params),
104        };
105
106        let result = validator.validate_request(&large_request, None).await;
107        // Should detect too many parameters (depending on limits)
108        if result.is_err() {
109            match result.unwrap_err() {
110                SecurityValidationError::TooManyParameters { current, limit } => {
111                    assert!(current > limit);
112                }
113                _ => panic!("Expected TooManyParameters error"),
114            }
115        }
116    }
117
118    #[test]
119    fn test_input_sanitizer() {
120        let sanitizer = InputSanitizer::new();
121
122        // Test normal input
123        let normal_input = "hello world";
124        let sanitized = sanitizer.sanitize_string(normal_input);
125        assert_eq!(sanitized, normal_input);
126
127        // Test input with potential issues
128        let suspicious_input = "<script>alert('xss')</script>";
129        let sanitized = sanitizer.sanitize_string(suspicious_input);
130        // Should be sanitized (exact behavior depends on implementation)
131        assert!(sanitized != suspicious_input || sanitized.is_empty());
132
133        // Test very long input
134        let long_input = "a".repeat(10000);
135        let sanitized = sanitizer.sanitize_string(&long_input);
136        // Should be truncated or rejected
137        assert!(sanitized.len() <= long_input.len());
138    }
139
140    #[test]
141    fn test_request_limits_config() {
142        let config = RequestLimitsConfig {
143            max_request_size: 1024,
144            max_parameters: 10,
145            max_parameter_size: 512,
146            max_string_length: 100,
147            max_array_length: 50,
148            max_object_depth: 5,
149            max_object_keys: 20,
150        };
151
152        assert_eq!(config.max_request_size, 1024);
153        assert_eq!(config.max_parameters, 10);
154        assert_eq!(config.max_string_length, 100);
155        assert_eq!(config.max_array_length, 50);
156        assert_eq!(config.max_object_depth, 5);
157    }
158
159    #[test]
160    fn test_security_config_presets() {
161        let permissive = RequestSecurityConfig::permissive();
162        let default = RequestSecurityConfig::default();
163        let strict = RequestSecurityConfig::strict();
164
165        // Strict should have lower limits than default
166        assert!(strict.limits.max_request_size <= default.limits.max_request_size);
167        assert!(strict.limits.max_parameters <= default.limits.max_parameters);
168
169        // Permissive should have higher limits than default
170        assert!(permissive.limits.max_request_size >= default.limits.max_request_size);
171        assert!(permissive.limits.max_parameters >= default.limits.max_parameters);
172    }
173
174    #[test]
175    fn test_security_validation_error_types() {
176        let errors = vec![
177            SecurityValidationError::RequestTooLarge {
178                current: 1000,
179                limit: 500,
180            },
181            SecurityValidationError::TooManyParameters {
182                current: 50,
183                limit: 20,
184            },
185            SecurityValidationError::InjectionDetected {
186                param: "test_param".to_string(),
187            },
188            SecurityValidationError::MaliciousContent {
189                reason: "test malicious content".to_string(),
190            },
191        ];
192
193        for error in errors {
194            let error_string = error.to_string();
195            assert!(!error_string.is_empty());
196            assert!(error_string.len() > 5);
197        }
198    }
199
200    #[tokio::test]
201    async fn test_security_integration() {
202        // Test that security components work together
203
204        let config = RequestSecurityConfig::strict();
205        let validator = RequestSecurityValidator::new(config);
206        let sanitizer = InputSanitizer::new();
207
208        // Create a potentially problematic request
209        let suspicious_request = Request {
210            jsonrpc: "2.0".to_string(),
211            method: "tools/call".to_string(),
212            id: json!(1),
213            params: json!({
214                "name": "test_tool",
215                "arguments": {
216                    "input": "<script>alert('xss')</script>",
217                    "data": "x".repeat(10000), // Very long string
218                }
219            }),
220        };
221
222        // Validate the request
223        let validation_result = validator.validate_request(&suspicious_request, None).await;
224
225        // If validation passes, sanitize the input
226        if let Ok(_) = validation_result {
227            if let Some(args) = suspicious_request.params.get("arguments") {
228                if let Some(input) = args.get("input").and_then(|v| v.as_str()) {
229                    let sanitized = sanitizer.sanitize_string(input);
230                    assert!(sanitized != input || sanitized.is_empty());
231                }
232            }
233        }
234        // If validation fails, that's also acceptable for strict config
235    }
236}
pulseengine_mcp_auth/security/mod.rs

pulseengine_mcp_auth/security/
mod.rs
