Validate RPC method for XSS (#576)

* Validate prompt arguments for XSS

Signed-off-by: Madhav Kandukuri <[email protected]>

* copied from check-prompt-args

Signed-off-by: Madhav Kandukuri <[email protected]>

* Fix tests
Signed-off-by: Madhav Kandukuri <[email protected]>

* flake8 fix

Signed-off-by: Madhav Kandukuri <[email protected]>

* Remove commented code
Signed-off-by: Madhav Kandukuri <[email protected]>

* Minor changes that fix smoketest
Signed-off-by: Madhav Kandukuri <[email protected]>

---------

Signed-off-by: Madhav Kandukuri <[email protected]>

Author: madhav165

mcpgateway/config.py

@@ -506,6 +506,7 @@ def validate_database(self) -> None:
     validation_safe_uri_pattern: str = r"^[a-zA-Z0-9_\-.:/?=&%]+$"
     validation_unsafe_uri_pattern: str = r'[<>"\'\\]'
     validation_tool_name_pattern: str = r"^[a-zA-Z][a-zA-Z0-9._-]*$"  # MCP tool naming
+    validation_tool_method_pattern: str = r"^[a-zA-Z][a-zA-Z0-9_\./-]*$"
 
     # MCP-compliant size limits (configurable via env)
     validation_max_name_length: int = 255
@@ -516,6 +517,8 @@ def validate_database(self) -> None:
     validation_max_url_length: int = 2048
     validation_max_rpc_param_size: int = 262144  # 256KB
 
+    validation_max_method_length: int = 128
+
     # Allowed MIME types
     validation_allowed_mime_types: List[str] = [
         "text/plain",

mcpgateway/main.py

@@ -85,6 +85,7 @@
     ResourceCreate,
     ResourceRead,
     ResourceUpdate,
+    RPCRequest,
     ServerCreate,
     ServerRead,
     ServerUpdate,
@@ -131,7 +132,6 @@
 from mcpgateway.utils.verify_credentials import require_auth, require_auth_override
 from mcpgateway.validation.jsonrpc import (
     JSONRPCError,
-    validate_request,
 )
 
 # Import the admin routes from the new module
@@ -1970,12 +1970,13 @@ async def handle_rpc(request: Request, db: Session = Depends(get_db), user: str
     try:
         logger.debug(f"User {user} made an RPC request")
         body = await request.json()
-        validate_request(body)
         method = body["method"]
         # rpc_id = body.get("id")
         params = body.get("params", {})
         cursor = params.get("cursor")  # Extract cursor parameter
 
+        RPCRequest(jsonrpc="2.0", method=method, params=params)  # Validate the request body against the RPCRequest model
+
         if method == "tools/list":
             tools = await tool_service.list_tools(db, cursor=cursor)
             result = [t.model_dump(by_alias=True, exclude_none=True) for t in tools]
@@ -2030,6 +2031,8 @@ async def handle_rpc(request: Request, db: Session = Depends(get_db), user: str
     except JSONRPCError as e:
         return e.to_dict()
     except Exception as e:
+        if isinstance(e, ValueError):
+            return JSONResponse(content={"message": "Method invalid"}, status_code=422)
         logger.error(f"RPC error: {str(e)}")
         return {
             "jsonrpc": "2.0",

mcpgateway/schemas.py

@@ -2122,9 +2122,10 @@ def validate_method(cls, v: str) -> str:
         Raises:
             ValueError: When value is not safe
         """
-        if not re.match(r"^[a-zA-Z][a-zA-Z0-9_\.]*$", v):
+        SecurityValidator.validate_no_xss(v, "RPC method name")
+        if not re.match(settings.validation_tool_method_pattern, v):
             raise ValueError("Invalid method name format")
-        if len(v) > 128:  # MCP method name limit
+        if len(v) > settings.validation_max_method_length:
             raise ValueError("Method name too long")
         return v
 

tests/security/test_input_validation.py

@@ -802,7 +802,6 @@ def test_rpc_request_validation(self):
         # Invalid method names
         invalid_methods = [
             "method with spaces",
-            "method-with-dash",
             "method@special",
             "<script>alert('XSS')</script>",
             "9method",  # Starts with number

tests/unit/mcpgateway/test_main.py

@@ -239,9 +239,9 @@ def test_static_files(self, test_client):
 class TestProtocolEndpoints:
     """Tests for MCP protocol operations: initialize, ping, notifications, etc."""
 
-    @patch("mcpgateway.main.validate_request")
+    # @patch("mcpgateway.main.validate_request")
     @patch("mcpgateway.main.session_registry.handle_initialize_logic")
-    def test_initialize_endpoint(self, mock_handle_initialize, _mock_validate, test_client, auth_headers):
+    def test_initialize_endpoint(self, mock_handle_initialize, test_client, auth_headers):
         """Test MCP protocol initialization."""
         mock_capabilities = ServerCapabilities(
             prompts={"listChanged": True},
@@ -271,8 +271,8 @@ def test_initialize_endpoint(self, mock_handle_initialize, _mock_validate, test_
         assert body["protocolVersion"] == PROTOCOL_VERSION
         mock_handle_initialize.assert_called_once()
 
-    @patch("mcpgateway.main.validate_request")
-    def test_ping_endpoint(self, _mock_validate, test_client, auth_headers):
+    # @patch("mcpgateway.main.validate_request")
+    def test_ping_endpoint(self, test_client, auth_headers):
         """Test MCP ping endpoint."""
         req = {"jsonrpc": "2.0", "method": "ping", "id": "test-id"}
         response = test_client.post("/protocol/ping", json=req, headers=auth_headers)
@@ -807,8 +807,8 @@ def test_rpc_tool_invocation(self, mock_invoke_tool, test_client, auth_headers):
         mock_invoke_tool.assert_called_once_with(db=ANY, name="test_tool", arguments={"param": "value"})
 
     @patch("mcpgateway.main.prompt_service.get_prompt")
-    @patch("mcpgateway.main.validate_request")
-    def test_rpc_prompt_get(self, _mock_validate, mock_get_prompt, test_client, auth_headers):
+    # @patch("mcpgateway.main.validate_request")
+    def test_rpc_prompt_get(self, mock_get_prompt, test_client, auth_headers):
         """Test prompt retrieval via JSON-RPC."""
         mock_get_prompt.return_value = {
             "messages": [{"role": "user", "content": {"type": "text", "text": "Rendered prompt"}}],
@@ -829,8 +829,8 @@ def test_rpc_prompt_get(self, _mock_validate, mock_get_prompt, test_client, auth
         mock_get_prompt.assert_called_once_with(ANY, "test_prompt", {"param": "value"})
 
     @patch("mcpgateway.main.tool_service.list_tools")
-    @patch("mcpgateway.main.validate_request")
-    def test_rpc_list_tools(self, _mock_validate, mock_list_tools, test_client, auth_headers):
+    # @patch("mcpgateway.main.validate_request")
+    def test_rpc_list_tools(self, mock_list_tools, test_client, auth_headers):
         """Test listing tools via JSON-RPC."""
         mock_tool = MagicMock()
         mock_tool.model_dump.return_value = MOCK_TOOL_READ
@@ -849,26 +849,26 @@ def test_rpc_list_tools(self, _mock_validate, mock_list_tools, test_client, auth
         assert isinstance(body, list)
         mock_list_tools.assert_called_once()
 
-    @patch("mcpgateway.main.validate_request")
-    def test_rpc_invalid_request(self, mock_validate, test_client, auth_headers):
+    @patch("mcpgateway.main.RPCRequest")
+    def test_rpc_invalid_request(self, mock_rpc_request, test_client, auth_headers):
         """Test RPC error handling for invalid requests."""
-        mock_validate.side_effect = Exception("Invalid request")
+        mock_rpc_request.side_effect = ValueError("Invalid method")
 
         req = {"jsonrpc": "1.0", "id": "test-id", "method": "invalid_method"}
         response = test_client.post("/rpc/", json=req, headers=auth_headers)
 
-        assert response.status_code == 200
+        assert response.status_code == 422
         body = response.json()
-        assert "error" in body and "Invalid request" in body["error"]["data"]
+        assert "Method invalid" in body.get("message")
 
     def test_rpc_invalid_json(self, test_client, auth_headers):
         """Test RPC error handling for malformed JSON."""
         headers = auth_headers
         headers["content-type"] = "application/json"
         response = test_client.post("/rpc/", content="invalid json", headers=headers)
-        assert response.status_code == 200  # Returns error response, not HTTP error
+        assert response.status_code == 422  # Returns error response, not HTTP error
         body = response.json()
-        assert "error" in body
+        assert "Method invalid" in body.get("message")
 
     @patch("mcpgateway.main.logging_service.set_level")
     def test_set_log_level_endpoint(self, mock_set_level, test_client, auth_headers):