Skip to main content

contextvm_sdk/core/
validation.rs

1//! Message validation utilities.
2
3use crate::core::constants::MAX_MESSAGE_SIZE;
4use crate::core::types::JsonRpcMessage;
5
6/// Validate that a message's content doesn't exceed the maximum size (1MB).
7pub fn validate_message_size(content: &str) -> bool {
8    content.len() <= MAX_MESSAGE_SIZE
9}
10
11/// Validate size and structure, then parse into a [`JsonRpcMessage`].
12pub fn validate_and_parse(content: &str) -> Option<JsonRpcMessage> {
13    if !validate_message_size(content) {
14        tracing::warn!("Message size validation failed: {} bytes", content.len());
15        return None;
16    }
17
18    let value: serde_json::Value = serde_json::from_str(content).ok()?;
19    validate_message(&value)
20}
21
22/// Validate that a JSON value is a well-formed JSON-RPC 2.0 message.
23///
24/// Checks:
25/// - Has `jsonrpc: "2.0"`
26/// - Is one of: request (id + method), response (id + result), error (id + error), notification (method, no id)
27pub fn validate_message(value: &serde_json::Value) -> Option<JsonRpcMessage> {
28    // Must have jsonrpc field
29    let jsonrpc = value.get("jsonrpc")?.as_str()?;
30    if jsonrpc != "2.0" {
31        return None;
32    }
33
34    serde_json::from_value(value.clone()).ok()
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40    use serde_json::json;
41
42    #[test]
43    fn test_valid_request() {
44        let msg = json!({"jsonrpc": "2.0", "id": 1, "method": "tools/list"});
45        let parsed = validate_message(&msg).unwrap();
46        assert!(parsed.is_request());
47    }
48
49    #[test]
50    fn test_valid_notification() {
51        let msg = json!({"jsonrpc": "2.0", "method": "notifications/initialized"});
52        let parsed = validate_message(&msg).unwrap();
53        assert!(parsed.is_notification());
54    }
55
56    #[test]
57    fn test_valid_response() {
58        let msg = json!({"jsonrpc": "2.0", "id": 1, "result": {"tools": []}});
59        let parsed = validate_message(&msg).unwrap();
60        assert!(parsed.is_response());
61    }
62
63    #[test]
64    fn test_invalid_version() {
65        let msg = json!({"jsonrpc": "1.0", "id": 1, "method": "test"});
66        assert!(validate_message(&msg).is_none());
67    }
68
69    #[test]
70    fn test_size_validation() {
71        assert!(validate_message_size("hello"));
72        let big = "x".repeat(MAX_MESSAGE_SIZE + 1);
73        assert!(!validate_message_size(&big));
74    }
75
76    #[test]
77    fn test_validate_and_parse_valid_request() {
78        let content = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#;
79        let msg = validate_and_parse(content).unwrap();
80        assert!(msg.is_request());
81        assert_eq!(msg.method(), Some("tools/list"));
82    }
83
84    #[test]
85    fn test_validate_and_parse_rejects_oversized() {
86        let padding = "x".repeat(MAX_MESSAGE_SIZE);
87        let content = format!(r#"{{"jsonrpc":"2.0","id":1,"method":"{}"}}"#, padding);
88        assert!(validate_and_parse(&content).is_none());
89    }
90
91    #[test]
92    fn test_validate_and_parse_rejects_invalid_version() {
93        let content = r#"{"jsonrpc":"1.0","id":1,"method":"test"}"#;
94        assert!(validate_and_parse(content).is_none());
95    }
96
97    #[test]
98    fn test_validate_and_parse_rejects_invalid_json() {
99        assert!(validate_and_parse("not json").is_none());
100    }
101}