claude-utils 0.2.0

use serde_json::Value;
use crate::mcp::protocol::*;

pub struct RequestValidator;

impl RequestValidator {
    /// Validate JSON-RPC request structure
    pub fn validate_jsonrpc_request(request: &JsonRpcRequest) -> Result<(), String> {
        // Validate JSON-RPC version
        if request.jsonrpc != "2.0" {
            return Err("Invalid JSON-RPC version".to_string());
        }

        // Validate method name
        if request.method.is_empty() {
            return Err("Method name cannot be empty".to_string());
        }

        // Validate method format (should not contain spaces or special chars)
        if !request.method.chars().all(|c| c.is_alphanumeric() || c == '.' || c == '/' || c == '_') {
            return Err("Invalid method name format".to_string());
        }

        Ok(())
    }

    /// Validate tool call parameters
    pub fn validate_tool_call(tool_request: &ToolCallRequest) -> Result<(), String> {
        // Validate tool name
        if tool_request.name.is_empty() {
            return Err("Tool name cannot be empty".to_string());
        }

        // Validate known tools
        match tool_request.name.as_str() {
            "clipboard.get" => Self::validate_clipboard_get_args(tool_request.arguments.as_ref()),
            "clipboard.set" => Self::validate_clipboard_set_args(tool_request.arguments.as_ref()),
            _ => Err(format!("Unknown tool: {}", tool_request.name)),
        }
    }

    fn validate_clipboard_get_args(args: Option<&Value>) -> Result<(), String> {
        if let Some(args) = args {
            if let Some(format) = args.get("format").and_then(|f| f.as_str()) {
                match format {
                    "auto" | "text" | "image" => Ok(()),
                    _ => Err(format!("Invalid format: {format}. Must be 'auto', 'text', or 'image'")),
                }
            } else if args.get("format").is_some() {
                Err("Format must be a string".to_string())
            } else {
                Ok(())
            }
        } else {
            Ok(())
        }
    }

    fn validate_clipboard_set_args(args: Option<&Value>) -> Result<(), String> {
        match args {
            None => Err("Missing required arguments for clipboard.set".to_string()),
            Some(args) => {
                // Validate required fields
                let content_type = args.get("type")
                    .and_then(|t| t.as_str())
                    .ok_or("Missing required field 'type'".to_string())?;

                let data = args.get("data")
                    .and_then(|d| d.as_str())
                    .ok_or("Missing required field 'data'".to_string())?;

                // Validate content type
                match content_type {
                    "text/plain" => {
                        // Validate text data isn't too large (10MB limit)
                        if data.len() > 10 * 1024 * 1024 {
                            return Err("Text data too large (max 10MB)".to_string());
                        }
                        Ok(())
                    }
                    "image/png" => {
                        // Validate base64 format
                        use base64::{Engine as _, engine::general_purpose::STANDARD};
                        if STANDARD.decode(data).is_err() {
                            return Err("Invalid base64 image data".to_string());
                        }
                        // Check size (50MB limit for images)
                        let decoded_size = (data.len() * 3) / 4;
                        if decoded_size > 50 * 1024 * 1024 {
                            return Err("Image data too large (max 50MB)".to_string());
                        }
                        Ok(())
                    }
                    _ => Err(format!("Unsupported content type: {content_type}")),
                }
            }
        }
    }

    /// Sanitize request data to prevent injection attacks
    pub fn sanitize_request(request: &mut JsonRpcRequest) {
        // Remove any potentially dangerous characters from method name
        request.method = request.method
            .chars()
            .filter(|c| c.is_alphanumeric() || *c == '.' || *c == '/' || *c == '_')
            .collect();

        // Limit method name length
        if request.method.len() > 100 {
            request.method.truncate(100);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn test_validate_jsonrpc_request() {
        let valid_request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            method: "tools.list".to_string(),
            params: None,
            id: Some(json!(1)),
        };
        assert!(RequestValidator::validate_jsonrpc_request(&valid_request).is_ok());

        let invalid_version = JsonRpcRequest {
            jsonrpc: "1.0".to_string(),
            method: "tools.list".to_string(),
            params: None,
            id: Some(json!(1)),
        };
        assert!(RequestValidator::validate_jsonrpc_request(&invalid_version).is_err());
    }

    #[test]
    fn test_validate_clipboard_set() {
        let valid_text = ToolCallRequest {
            name: "clipboard.set".to_string(),
            arguments: Some(json!({
                "type": "text/plain",
                "data": "Hello, world!"
            })),
        };
        assert!(RequestValidator::validate_tool_call(&valid_text).is_ok());

        let missing_data = ToolCallRequest {
            name: "clipboard.set".to_string(),
            arguments: Some(json!({
                "type": "text/plain"
            })),
        };
        assert!(RequestValidator::validate_tool_call(&missing_data).is_err());
    }
}