fix: Fix credential leakage vulnerability in Agent Registry

Passing in ADC auth headers to non-google MCP toolsets is a vulnerability. To fix, only pass in the headers to Google MCP toolsets.

Co-authored-by: Kathy Wu <[email protected]>
PiperOrigin-RevId: 897230159

Author: wukath

src/google/adk/integrations/agent_registry/agent_registry.py

@@ -19,19 +19,15 @@
 from collections.abc import Generator
 from enum import Enum
 import logging
-import os
 import re
 from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Mapping
-from typing import Sequence
 from typing import TypedDict
-from urllib.parse import parse_qs
 from urllib.parse import urlparse
 
-from a2a.client.client_factory import minimal_agent_card
 from a2a.types import AgentCapabilities
 from a2a.types import AgentCard
 from a2a.types import AgentSkill
@@ -142,6 +138,17 @@ class Endpoint(TypedDict, total=False):
   attributes: Dict[str, Any]
 
 
+def _is_google_api(url: str) -> bool:
+  """Checks if the given URL points to a Google API endpoint."""
+  parsed_url = urlparse(url)
+  if not parsed_url.hostname:
+    return False
+  return (
+      parsed_url.hostname == "googleapis.com"
+      or parsed_url.hostname.endswith(".googleapis.com")
+  )
+
+
 class AgentRegistry:
   """Client for interacting with the Google Cloud Agent Registry service.
 
@@ -305,8 +312,10 @@ def get_mcp_toolset(
           f"MCP Server endpoint URI not found for: {mcp_server_name}"
       )
 
+    headers = self._get_auth_headers() if _is_google_api(endpoint_uri) else None
     connection_params = StreamableHTTPConnectionParams(
-        url=endpoint_uri, headers=self._get_auth_headers()
+        url=endpoint_uri,
+        headers=headers,
     )
     return AgentRegistrySingleMcpToolset(
         destination_resource_id=mcp_server_id,

tests/unittests/integrations/agent_registry/test_agent_registry.py

@@ -320,13 +320,23 @@ def test_get_endpoint(self, mock_httpx, registry):
     server = registry.get_endpoint("test-endpoint")
     assert server == {"name": "test-endpoint"}
 
+  @pytest.mark.parametrize(
+      "url, expected_auth",
+      [
+          ("https://mcp.com", False),
+          ("https://mcp.googleapis.com/v1", True),
+          ("https://example.com/googleapis/v1", False),
+      ],
+  )
   @patch("httpx.Client")
-  def test_get_mcp_toolset(self, mock_httpx, registry):
+  def test_get_mcp_toolset_auth_headers(
+      self, mock_httpx, registry, url, expected_auth
+  ):
     mock_response = MagicMock()
     mock_response.json.return_value = {
         "displayName": "TestPrefix",
         "interfaces": [{
-            "url": "https://mcp.com",
+            "url": url,
             "protocolBinding": "JSONRPC",
         }],
     }
@@ -341,6 +351,13 @@ def test_get_mcp_toolset(self, mock_httpx, registry):
     toolset = registry.get_mcp_toolset("test-mcp")
     assert isinstance(toolset, McpToolset)
     assert toolset.tool_name_prefix == "TestPrefix"
+    if expected_auth:
+      assert toolset._connection_params.headers is not None
+      assert (
+          toolset._connection_params.headers["Authorization"] == "Bearer token"
+      )
+    else:
+      assert toolset._connection_params.headers is None
 
   @patch("httpx.Client")
   def test_get_mcp_toolset_with_auth(self, mock_httpx, registry):