diff --git a/application/agents/tools/mcp_tool.py b/application/agents/tools/mcp_tool.py index 59299def..d0fffcdf 100644 --- a/application/agents/tools/mcp_tool.py +++ b/application/agents/tools/mcp_tool.py @@ -24,6 +24,7 @@ from application.api.user.tasks import mcp_oauth_status_task, mcp_oauth_task from application.cache import get_redis_instance from application.core.mongo_db import MongoDB from application.core.settings import settings +from application.core.url_validation import SSRFError, validate_url from application.security.encryption import decrypt_credentials logger = logging.getLogger(__name__) @@ -61,7 +62,8 @@ class MCPTool(Tool): """ self.config = config self.user_id = user_id - self.server_url = config.get("server_url", "") + raw_url = config.get("server_url", "") + self.server_url = self._validate_server_url(raw_url) if raw_url else "" self.transport_type = config.get("transport_type", "auto") self.auth_type = config.get("auth_type", "none") self.timeout = config.get("timeout", 30) @@ -87,6 +89,18 @@ class MCPTool(Tool): if self.server_url and self.auth_type != "oauth": self._setup_client() + @staticmethod + def _validate_server_url(server_url: str) -> str: + """Validate server_url to prevent SSRF to internal networks. + + Raises: + ValueError: If the URL points to a private/internal address. + """ + try: + return validate_url(server_url) + except SSRFError as exc: + raise ValueError(f"Invalid MCP server URL: {exc}") from exc + def _resolve_redirect_uri(self, configured_redirect_uri: Optional[str]) -> str: if configured_redirect_uri: return configured_redirect_uri.rstrip("/") diff --git a/application/api/user/tools/mcp.py b/application/api/user/tools/mcp.py index 3f9b2873..c38031ce 100644 --- a/application/api/user/tools/mcp.py +++ b/application/api/user/tools/mcp.py @@ -14,6 +14,7 @@ from application.api.user.tools.routes import transform_actions from application.cache import get_redis_instance from application.core.mongo_db import MongoDB from application.core.settings import settings +from application.core.url_validation import SSRFError, validate_url from application.security.encryption import decrypt_credentials, encrypt_credentials from application.utils import check_required_fields @@ -63,6 +64,21 @@ def _extract_auth_credentials(config): return auth_credentials +def _validate_mcp_server_url(config: dict) -> None: + """Validate the server_url in an MCP config to prevent SSRF. + + Raises: + ValueError: If the URL is missing or points to a blocked address. + """ + server_url = (config.get("server_url") or "").strip() + if not server_url: + raise ValueError("server_url is required") + try: + validate_url(server_url) + except SSRFError as exc: + raise ValueError(f"Invalid server URL: {exc}") from exc + + @tools_mcp_ns.route("/mcp_server/test") class TestMCPServerConfig(Resource): @api.expect( @@ -97,6 +113,8 @@ class TestMCPServerConfig(Resource): 400, ) + _validate_mcp_server_url(config) + auth_credentials = _extract_auth_credentials(config) test_config = config.copy() test_config["auth_credentials"] = auth_credentials @@ -114,6 +132,11 @@ class TestMCPServerConfig(Resource): result["message"] = "Connection test failed" return make_response(jsonify(result), 200) + except ValueError as e: + return make_response( + jsonify({"success": False, "error": str(e)}), + 400, + ) except Exception as e: current_app.logger.error(f"Error testing MCP server: {e}", exc_info=True) return make_response( @@ -165,6 +188,8 @@ class MCPServerSave(Resource): 400, ) + _validate_mcp_server_url(config) + auth_credentials = _extract_auth_credentials(config) auth_type = config.get("auth_type", "none") mcp_config = config.copy() @@ -279,6 +304,11 @@ class MCPServerSave(Resource): "tools_count": len(transformed_actions), } return make_response(jsonify(response_data), 200) + except ValueError as e: + return make_response( + jsonify({"success": False, "error": str(e)}), + 400, + ) except Exception as e: current_app.logger.error(f"Error saving MCP server: {e}", exc_info=True) return make_response( diff --git a/application/api/user/tools/routes.py b/application/api/user/tools/routes.py index 316bea36..70b76c1e 100644 --- a/application/api/user/tools/routes.py +++ b/application/api/user/tools/routes.py @@ -8,6 +8,7 @@ from application.agents.tools.spec_parser import parse_spec from application.agents.tools.tool_manager import ToolManager from application.api import api from application.api.user.base import user_tools_collection +from application.core.url_validation import SSRFError, validate_url from application.security.encryption import decrypt_credentials, encrypt_credentials from application.utils import check_required_fields, validate_function_name @@ -238,6 +239,16 @@ class CreateTool(Resource): if missing_fields: return missing_fields try: + if data["name"] == "mcp_tool": + server_url = (data.get("config", {}).get("server_url") or "").strip() + if server_url: + try: + validate_url(server_url) + except SSRFError as exc: + return make_response( + jsonify({"success": False, "message": f"Invalid server URL: {exc}"}), + 400, + ) tool_instance = tool_manager.tools.get(data["name"]) if not tool_instance: return make_response( @@ -423,6 +434,16 @@ class UpdateToolConfig(Resource): return make_response(jsonify({"success": False}), 404) tool_name = tool_doc.get("name") + if tool_name == "mcp_tool": + server_url = (data["config"].get("server_url") or "").strip() + if server_url: + try: + validate_url(server_url) + except SSRFError as exc: + return make_response( + jsonify({"success": False, "message": f"Invalid server URL: {exc}"}), + 400, + ) tool_instance = tool_manager.tools.get(tool_name) config_requirements = ( tool_instance.get_config_requirements() if tool_instance else {} diff --git a/tests/agents/tools/test_mcp_tool.py b/tests/agents/tools/test_mcp_tool.py index 1ca3546f..4e1beb82 100644 --- a/tests/agents/tools/test_mcp_tool.py +++ b/tests/agents/tools/test_mcp_tool.py @@ -33,6 +33,8 @@ def _patch_mcp_globals(monkeypatch): monkeypatch.setattr(mcp_mod, "mongo", mock_mongo) monkeypatch.setattr(mcp_mod, "db", mock_db) monkeypatch.setattr(mcp_mod, "_mcp_clients_cache", {}) + # Bypass DNS-resolving URL validation for tests using fake hostnames. + monkeypatch.setattr(mcp_mod, "validate_url", lambda u, **kw: u) @pytest.fixture @@ -136,6 +138,47 @@ class TestMCPToolInit: }) assert tool.custom_headers == {"X-Custom": "val"} + def test_rejects_metadata_ip(self, monkeypatch): + from application.agents.tools.mcp_tool import MCPTool + from application.core.url_validation import validate_url as real_validate_url + import application.agents.tools.mcp_tool as mcp_mod + + monkeypatch.setattr(mcp_mod, "validate_url", real_validate_url) + with pytest.raises(ValueError, match="Invalid MCP server URL"): + MCPTool(config={"server_url": "http://169.254.169.254/latest/meta-data", "auth_type": "none"}) + + def test_rejects_localhost(self, monkeypatch): + from application.agents.tools.mcp_tool import MCPTool + from application.core.url_validation import validate_url as real_validate_url + import application.agents.tools.mcp_tool as mcp_mod + + monkeypatch.setattr(mcp_mod, "validate_url", real_validate_url) + with pytest.raises(ValueError, match="Invalid MCP server URL"): + MCPTool(config={"server_url": "http://localhost:8080/mcp", "auth_type": "none"}) + + def test_rejects_private_ip(self, monkeypatch): + from application.agents.tools.mcp_tool import MCPTool + from application.core.url_validation import validate_url as real_validate_url + import application.agents.tools.mcp_tool as mcp_mod + + monkeypatch.setattr(mcp_mod, "validate_url", real_validate_url) + with pytest.raises(ValueError, match="Invalid MCP server URL"): + MCPTool(config={"server_url": "http://10.0.0.1/mcp", "auth_type": "none"}) + + def test_accepts_public_url(self): + tool = _make_tool({ + "server_url": "https://mcp.example.com/api", + "auth_type": "none", + }) + assert tool.server_url == "https://mcp.example.com/api" + + def test_empty_server_url_allowed(self): + from application.agents.tools.mcp_tool import MCPTool + + with patch.object(MCPTool, "_setup_client"): + tool = MCPTool(config={"server_url": "", "auth_type": "none"}) + assert tool.server_url == "" + # ===================================================================== # Redirect URI Resolution