Merge pull request #600 from TS0713/handle-invalid-gateway-url

handle-invalid-gateway-url

Author: madhav165

.env.example

@@ -240,6 +240,13 @@ MAX_PROMPT_SIZE=102400
 # Timeout for rendering prompt templates (in seconds)
 PROMPT_RENDER_TIMEOUT=10
 
+#####################################
+# Gateway Validation
+#####################################
+
+# Timeout for gateway validation (in seconds)
+GATEWAY_VALIDATION_TIMEOUT=5
+
 #####################################
 # Health Checks
 #####################################

mcpgateway/config.py

@@ -280,6 +280,9 @@ def _parse_federation_peers(cls, v):
     health_check_timeout: int = 10  # seconds
     unhealthy_threshold: int = 5  # after this many failures, mark as Offline
 
+    # Validation Gateway URL
+    gateway_validation_timeout: int = 5  # seconds
+
     filelock_name: str = "gateway_service_leader.lock"
 
     # Default Roots

mcpgateway/services/gateway_service.py

@@ -244,6 +244,57 @@ def __init__(self) -> None:
         else:
             self._redis_client = None
 
+    async def _validate_gateway_url(self, url: str, headers: dict, transport_type: str, timeout: Optional[int] = None):
+        """
+        Validate if the given URL is a live Server-Sent Events (SSE) endpoint.
+
+        Args:
+            url (str): The full URL of the endpoint to validate.
+            headers (dict): Headers to be included in the requests (e.g., Authorization).
+            transport_type (str): SSE or STREAMABLEHTTP
+            timeout (int, optional): Timeout in seconds. Defaults to settings.gateway_validation_timeout.
+
+        Returns:
+            bool: True if the endpoint is reachable and supports SSE/StreamableHTTP, otherwise False.
+        """
+        if timeout is None:
+            timeout = settings.gateway_validation_timeout
+        validation_client = ResilientHttpClient(client_args={"timeout": settings.gateway_validation_timeout, "verify": not settings.skip_ssl_verify})
+        try:
+            async with validation_client.client.stream("GET", url, headers=headers, timeout=timeout) as response:
+                response_headers = dict(response.headers)
+                location = response_headers.get("location")
+                content_type = response_headers.get("content-type")
+                if response.status_code in (401, 403):
+                    logger.debug(f"Authentication failed for {url} with status {response.status_code}")
+                    return False
+
+                if transport_type == "STREAMABLEHTTP":
+                    if location:
+                        async with validation_client.client.stream("GET", location, headers=headers, timeout=timeout) as response_redirect:
+                            response_headers = dict(response_redirect.headers)
+                            mcp_session_id = response_headers.get("mcp-session-id")
+                            content_type = response_headers.get("content-type")
+                            if response_redirect.status_code in (401, 403):
+                                logger.debug(f"Authentication failed at redirect location {location}")
+                                return False
+                            if mcp_session_id is not None and mcp_session_id != "":
+                                if content_type is not None and content_type != "" and "application/json" in content_type:
+                                    return True
+
+                elif transport_type == "SSE":
+                    if content_type is not None and content_type != "" and "text/event-stream" in content_type:
+                        return True
+                return False
+        except httpx.UnsupportedProtocol as e:
+            logger.debug(f"Gateway URL Unsupported Protocol for {url}: {str(e)}", exc_info=True)
+            return False
+        except Exception as e:
+            logger.debug(f"Gateway validation failed for {url}: {str(e)}", exc_info=True)
+            return False
+        finally:
+            await validation_client.aclose()
+
     async def initialize(self) -> None:
         """Initialize the service and start health check if this instance is the leader.
 
@@ -844,13 +895,11 @@ async def forward_request(self, gateway: DbGateway, method: str, params: Optiona
 
             # Update last seen timestamp
             gateway.last_seen = datetime.now(timezone.utc)
-
-            if "error" in result:
-                raise GatewayError(f"Gateway error: {result['error'].get('message')}")
-            return result.get("result")
-
-        except Exception as e:
-            raise GatewayConnectionError(f"Failed to forward request to {gateway.name}: {str(e)}")
+        except Exception:
+            raise GatewayConnectionError(f"Failed to forward request to {gateway.name}")
+        if "error" in result:
+            raise GatewayError(f"Gateway error: {result['error'].get('message')}")
+        return result.get("result")
 
     async def _handle_gateway_failure(self, gateway: str) -> None:
         """Tracks and handles gateway failures during health checks.
@@ -1115,9 +1164,10 @@ async def _initialize_gateway(self, url: str, authentication: Optional[Dict[str,
             >>> import asyncio
             >>> async def test_params():
             ...     try:
-            ...         await service._initialize_gateway("invalid://url")
+            ...         await service._initialize_gateway("hello//")
             ...     except Exception as e:
-            ...         return "Failed" in str(e) or "GatewayConnectionError" in str(type(e).__name__)
+            ...         return isinstance(e, GatewayConnectionError) or "Failed" in str(e)
+
             >>> asyncio.run(test_params())
             True
 
@@ -1172,21 +1222,23 @@ async def connect_to_sse_server(server_url: str, authentication: Optional[Dict[s
                 # Store the context managers so they stay alive
                 decoded_auth = decode_auth(authentication)
 
-                # Use async with for both sse_client and ClientSession
-                async with sse_client(url=server_url, headers=decoded_auth) as streams:
-                    async with ClientSession(*streams) as session:
-                        # Initialize the session
-                        response = await session.initialize()
-                        capabilities = response.capabilities.model_dump(by_alias=True, exclude_none=True)
+                if await self._validate_gateway_url(url=server_url, headers=decoded_auth, transport_type="SSE"):
+                    # Use async with for both sse_client and ClientSession
+                    async with sse_client(url=server_url, headers=decoded_auth) as streams:
+                        async with ClientSession(*streams) as session:
+                            # Initialize the session
+                            response = await session.initialize()
+                            capabilities = response.capabilities.model_dump(by_alias=True, exclude_none=True)
 
-                        response = await session.list_tools()
-                        tools = response.tools
-                        tools = [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools]
+                            response = await session.list_tools()
+                            tools = response.tools
+                            tools = [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools]
 
-                        tools = [ToolCreate.model_validate(tool) for tool in tools]
-                        logger.info(f"{tools[0]=}")
+                            tools = [ToolCreate.model_validate(tool) for tool in tools]
+                            logger.info(f"{tools[0]=}")
 
-                return capabilities, tools
+                    return capabilities, tools
+                raise GatewayConnectionError(f"Failed to initialize gateway at {url}")
 
             async def connect_to_streamablehttp_server(server_url: str, authentication: Optional[Dict[str, str]] = None):
                 """
@@ -1203,25 +1255,26 @@ async def connect_to_streamablehttp_server(server_url: str, authentication: Opti
                     authentication = {}
                 # Store the context managers so they stay alive
                 decoded_auth = decode_auth(authentication)
-
-                # Use async with for both streamablehttp_client and ClientSession
-                async with streamablehttp_client(url=server_url, headers=decoded_auth) as (read_stream, write_stream, _get_session_id):
-                    async with ClientSession(read_stream, write_stream) as session:
-                        # Initialize the session
-                        response = await session.initialize()
-                        # if get_session_id:
-                        #     session_id = get_session_id()
-                        #     if session_id:
-                        #         print(f"Session ID: {session_id}")
-                        capabilities = response.capabilities.model_dump(by_alias=True, exclude_none=True)
-                        response = await session.list_tools()
-                        tools = response.tools
-                        tools = [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools]
-                        tools = [ToolCreate.model_validate(tool) for tool in tools]
-                        for tool in tools:
-                            tool.request_type = "STREAMABLEHTTP"
-
-                return capabilities, tools
+                if await self._validate_gateway_url(url=server_url, headers=decoded_auth, transport_type="STREAMABLEHTTP"):
+                    # Use async with for both streamablehttp_client and ClientSession
+                    async with streamablehttp_client(url=server_url, headers=decoded_auth) as (read_stream, write_stream, _get_session_id):
+                        async with ClientSession(read_stream, write_stream) as session:
+                            # Initialize the session
+                            response = await session.initialize()
+                            # if get_session_id:
+                            #     session_id = get_session_id()
+                            #     if session_id:
+                            #         print(f"Session ID: {session_id}")
+                            capabilities = response.capabilities.model_dump(by_alias=True, exclude_none=True)
+                            response = await session.list_tools()
+                            tools = response.tools
+                            tools = [tool.model_dump(by_alias=True, exclude_none=True) for tool in tools]
+                            tools = [ToolCreate.model_validate(tool) for tool in tools]
+                            for tool in tools:
+                                tool.request_type = "STREAMABLEHTTP"
+
+                    return capabilities, tools
+                raise GatewayConnectionError(f"Failed to initialize gateway at {url}")
 
             capabilities = {}
             tools = []
@@ -1232,7 +1285,8 @@ async def connect_to_streamablehttp_server(server_url: str, authentication: Opti
 
             return capabilities, tools
         except Exception as e:
-            raise GatewayConnectionError(f"Failed to initialize gateway at {url}: {str(e)}")
+            logger.debug(f"Gateway initialization failed for {url}: {str(e)}", exc_info=True)
+            raise GatewayConnectionError(f"Failed to initialize gateway at {url}")
 
     def _get_gateways(self, include_inactive: bool = True) -> list[DbGateway]:
         """Sync function for database operations (runs in thread).

mcpgateway/utils/verify_credentials.py

@@ -145,7 +145,7 @@ async def verify_jwt_token(token: str) -> dict:
 
         # Log warning for non-expiring tokens
         if "exp" not in unverified:
-            logger.warning("JWT token without expiration accepted. " "Consider enabling REQUIRE_TOKEN_EXPIRATION for better security. " f"Token sub: {unverified.get('sub', 'unknown')}")
+            logger.warning(f"JWT token without expiration accepted. Consider enabling REQUIRE_TOKEN_EXPIRATION for better security. Token sub: {unverified.get('sub', 'unknown')}")
 
         # Full validation
         options = {}