Enhanced Validation Missing in GatewayCreate closes #694 (#695)

Signed-off-by: Mihai Criveti <[email protected]>

Author: crivetimihai

mcpgateway/schemas.py

@@ -1818,20 +1818,7 @@ def create_auth_value(cls, v, info):
         if (auth_type is None) or (auth_type == ""):
             return v  # If no auth_type is provided, no need to create auth_value
 
-        # If custom headers, use all headers
-        if auth_type == "authheaders":
-            auth_headers = data.get("auth_headers")
-            if auth_headers and isinstance(auth_headers, list):
-                # Convert list of {key, value} to dict
-                header_dict = {h["key"]: h["value"] for h in auth_headers if h.get("key")}
-                return encode_auth(header_dict)
-            # Fallback to old single key/value
-            header_key = data.get("auth_header_key")
-            header_value = data.get("auth_header_value")
-            if header_key and header_value:
-                return encode_auth({header_key: header_value})
-
-        # Otherwise, use the default logic
+        # Process the auth fields and generate auth_value based on auth_type
         auth_value = cls._process_auth_fields(info)
         return auth_value
 
@@ -1902,10 +1889,44 @@ def _process_auth_fields(info: ValidationInfo) -> Optional[Dict[str, Any]]:
             # Support both new multi-headers format and legacy single header format
             auth_headers = data.get("auth_headers")
             if auth_headers and isinstance(auth_headers, list):
-                # New multi-headers format
-                header_dict = {h["key"]: h["value"] for h in auth_headers if h.get("key")}
+                # New multi-headers format with enhanced validation
+                header_dict = {}
+                duplicate_keys = set()
+
+                for header in auth_headers:
+                    if not isinstance(header, dict):
+                        continue
+
+                    key = header.get("key")
+                    value = header.get("value", "")
+
+                    # Skip headers without keys
+                    if not key:
+                        continue
+
+                    # Track duplicate keys (last value wins)
+                    if key in header_dict:
+                        duplicate_keys.add(key)
+
+                    # Validate header key format (basic HTTP header validation)
+                    if not all(c.isalnum() or c in "-_" for c in key.replace(" ", "")):
+                        raise ValueError(f"Invalid header key format: '{key}'. Header keys should contain only alphanumeric characters, hyphens, and underscores.")
+
+                    # Store header (empty values are allowed)
+                    header_dict[key] = value
+
+                # Ensure at least one valid header
                 if not header_dict:
-                    raise ValueError("For 'headers' auth, at least one header must be provided.")
+                    raise ValueError("For 'headers' auth, at least one valid header with a key must be provided.")
+
+                # Warn about duplicate keys (optional - could log this instead)
+                if duplicate_keys:
+                    logging.warning(f"Duplicate header keys detected (last value used): {', '.join(duplicate_keys)}")
+
+                # Check for excessive headers (prevent abuse)
+                if len(header_dict) > 100:
+                    raise ValueError("Maximum of 100 headers allowed per gateway.")
+
                 return encode_auth(header_dict)
 
             # Legacy single header format (backward compatibility)
@@ -2027,17 +2048,7 @@ def create_auth_value(cls, v, info):
         if (auth_type is None) or (auth_type == ""):
             return v  # If no auth_type is provided, no need to create auth_value
 
-        # If custom headers, use all headers
-        if auth_type == "authheaders":
-            auth_headers = data.get("auth_headers")
-            if auth_headers and isinstance(auth_headers, list):
-                header_dict = {h["key"]: h["value"] for h in auth_headers if h.get("key")}
-                return encode_auth(header_dict)
-            header_key = data.get("auth_header_key")
-            header_value = data.get("auth_header_value")
-            if header_key and header_value:
-                return encode_auth({header_key: header_value})
-
+        # Process the auth fields and generate auth_value based on auth_type
         auth_value = cls._process_auth_fields(info)
         return auth_value
 

tests/unit/mcpgateway/test_multi_auth_headers.py

@@ -111,14 +111,25 @@ async def test_gateway_update_add_multi_headers(self):
         assert decoded["X-New-Header"] == "new_value"
 
     @pytest.mark.asyncio
-    async def test_special_characters_in_headers(self):
-        """Test headers with special characters."""
+    async def test_special_characters_in_headers_rejected(self):
+        """Test headers with invalid special characters are rejected."""
         auth_headers = [{"key": "X-Special-!@#", "value": "value-with-特殊字符"}, {"key": "Content-Type", "value": "application/json; charset=utf-8"}]
 
+        with pytest.raises(ValidationError) as exc_info:
+            GatewayCreate(name="Test Gateway", url="http://example.com", auth_type="authheaders", auth_headers=auth_headers)
+
+        assert "Invalid header key format" in str(exc_info.value)
+        assert "X-Special-!@#" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_valid_special_characters_in_values(self):
+        """Test headers with special characters in values (allowed) but valid keys."""
+        auth_headers = [{"key": "X-Special-Header", "value": "value-with-特殊字符"}, {"key": "Content-Type", "value": "application/json; charset=utf-8"}]
+
         gateway = GatewayCreate(name="Test Gateway", url="http://example.com", auth_type="authheaders", auth_headers=auth_headers)
 
         decoded = decode_auth(gateway.auth_value)
-        assert decoded["X-Special-!@#"] == "value-with-特殊字符"
+        assert decoded["X-Special-Header"] == "value-with-特殊字符"
         assert decoded["Content-Type"] == "application/json; charset=utf-8"
 
     @pytest.mark.asyncio
@@ -169,3 +180,77 @@ async def test_authorization_header_in_multi_headers(self):
         decoded = decode_auth(gateway.auth_value)
         assert decoded["Authorization"] == "Bearer token123"
         assert decoded["X-API-Key"] == "secret"
+
+    @pytest.mark.asyncio
+    async def test_gateway_create_invalid_header_key_format(self):
+        """Test creating gateway with invalid header key format."""
+        auth_headers = [{"key": "Invalid@Key!", "value": "secret123"}]
+
+        with pytest.raises(ValidationError) as exc_info:
+            GatewayCreate(name="Test Gateway", url="http://example.com", auth_type="authheaders", auth_headers=auth_headers)
+
+        assert "Invalid header key format" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_gateway_create_excessive_headers(self):
+        """Test creating gateway with more than 100 headers."""
+        auth_headers = [{"key": f"X-Header-{i}", "value": f"value-{i}"} for i in range(101)]
+
+        with pytest.raises(ValidationError) as exc_info:
+            GatewayCreate(name="Test Gateway", url="http://example.com", auth_type="authheaders", auth_headers=auth_headers)
+
+        assert "Maximum of 100 headers allowed" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_gateway_create_duplicate_keys_with_warning(self, caplog):
+        """Test creating gateway with duplicate header keys logs warning."""
+        auth_headers = [
+            {"key": "X-API-Key", "value": "first_value"},
+            {"key": "X-API-Key", "value": "second_value"},  # Duplicate
+            {"key": "X-Client-ID", "value": "client123"}
+        ]
+
+        gateway = GatewayCreate(name="Test Gateway", url="http://example.com", auth_type="authheaders", auth_headers=auth_headers)
+
+        # Check that duplicate warning was logged
+        assert "Duplicate header keys detected" in caplog.text
+        assert "X-API-Key" in caplog.text
+
+        # Check that last value wins
+        decoded = decode_auth(gateway.auth_value)
+        assert decoded["X-API-Key"] == "second_value"
+        assert decoded["X-Client-ID"] == "client123"
+
+    @pytest.mark.asyncio
+    async def test_gateway_create_mixed_valid_invalid_keys(self):
+        """Test creating gateway with mixed valid and invalid header keys."""
+        auth_headers = [
+            {"key": "Valid-Header", "value": "test123"},
+            {"key": "Invalid@Key!", "value": "should_fail"}  # This should fail validation
+        ]
+
+        with pytest.raises(ValidationError) as exc_info:
+            GatewayCreate(name="Test Gateway", url="http://example.com", auth_type="authheaders", auth_headers=auth_headers)
+
+        assert "Invalid header key format" in str(exc_info.value)
+        assert "Invalid@Key!" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_gateway_create_edge_case_header_keys(self):
+        """Test creating gateway with edge case header keys."""
+        # Test valid edge cases
+        auth_headers = [
+            {"key": "X-API-Key", "value": "test1"},  # Standard format
+            {"key": "X_API_KEY", "value": "test2"},  # Underscores allowed
+            {"key": "API-Key-123", "value": "test3"},  # Numbers and hyphens
+            {"key": "UPPERCASE", "value": "test4"},  # Uppercase
+            {"key": "lowercase", "value": "test5"}   # Lowercase
+        ]
+
+        gateway = GatewayCreate(name="Test Gateway", url="http://example.com", auth_type="authheaders", auth_headers=auth_headers)
+
+        decoded = decode_auth(gateway.auth_value)
+        assert len(decoded) == 5
+        assert decoded["X-API-Key"] == "test1"
+        assert decoded["X_API_KEY"] == "test2"
+        assert decoded["API-Key-123"] == "test3"