fix: address review feedback and add tests

- Move context ownership warning to __init__ (was unreachable)
- Remove unreachable owner-is-None guard in _check_context_ownership
- Change _BLOCKED_NETWORKS to tuple for immutability
- Restore dropped docstrings in card_resolver and
default_request_handler
- Fix non-ASCII chars in comments (use -> and --)
- Add tests/utils/test_url_validation.py with 26 SSRF validation tests

Author: amit-raut

src/a2a/client/card_resolver.py

@@ -46,6 +46,13 @@ def __init__(
         base_url: str,
         agent_card_path: str = AGENT_CARD_WELL_KNOWN_PATH,
     ) -> None:
+        """Initializes the A2ACardResolver.
+
+        Args:
+            httpx_client: An async HTTP client instance (e.g., httpx.AsyncClient).
+            base_url: The base URL of the agent's host.
+            agent_card_path: The path to the agent card endpoint, relative to the base URL.
+        """
         self.base_url = base_url.rstrip('/')
         self.agent_card_path = agent_card_path.lstrip('/')
         self.httpx_client = httpx_client
@@ -56,6 +63,27 @@ async def get_agent_card(
         http_kwargs: dict[str, Any] | None = None,
         signature_verifier: Callable[[AgentCard], None] | None = None,
     ) -> AgentCard:
+        """Fetches an agent card from a specified path relative to the base_url.
+
+        If relative_card_path is None, it defaults to the resolver's configured
+        agent_card_path (for the public agent card).
+
+        Args:
+            relative_card_path: Optional path to the agent card endpoint,
+                relative to the base URL. If None, uses the default public
+                agent card path. Use `'/'` for an empty path.
+            http_kwargs: Optional dictionary of keyword arguments to pass to the
+                underlying httpx.get request.
+            signature_verifier: A callable used to verify the agent card's signatures.
+
+        Returns:
+            An `AgentCard` object representing the agent's capabilities.
+
+        Raises:
+            A2AClientHTTPError: If an HTTP error occurs during the request.
+            A2AClientJSONError: If the response body cannot be decoded as JSON,
+                validated against the AgentCard schema, or fails SSRF URL validation.
+        """
         if not relative_card_path:
             path_segment = self.agent_card_path
         else:
@@ -77,7 +105,7 @@ async def get_agent_card(
             )
             agent_card = AgentCard.model_validate(agent_card_data)
 
-            # ---- FIX: A2A-SSRF-01 — validate card.url before returning ----
+            # ---- FIX: A2A-SSRF-01 -- validate card.url before returning ----
             # Without this check, any caller who controls the card endpoint
             # can redirect all subsequent RPC calls to an internal address.
             try:

src/a2a/server/request_handlers/default_request_handler.py

@@ -6,15 +6,15 @@
 Root cause of vulnerability:
   _setup_message_execution() uses params.message.context_id directly without
   any ownership check. An attacker who knows a victim's contextId can send a
-  new task under that context — task_manager.get_task() returns None for the
+  new task under that context -- task_manager.get_task() returns None for the
   new task_id, so the original task-level check is never reached.
 
 Fix design:
-  DefaultRequestHandler maintains a _context_owners dict (context_id → owner)
+  DefaultRequestHandler maintains a _context_owners dict (context_id -> owner)
   in memory. When a get_caller_id extractor is configured:
     1. On first message for a context_id: record caller as owner.
     2. On subsequent messages for same context_id: verify caller matches owner.
-  If get_caller_id is None (default): no ownership tracking — backward compatible.
+  If get_caller_id is None (default): no ownership tracking -- backward compatible.
 
 Target file: src/a2a/server/request_handlers/default_request_handler.py
 """
@@ -118,7 +118,7 @@ def __init__(  # noqa: PLR0913
                 fingerprint). When provided, the handler tracks which caller
                 created each contextId and rejects messages from different
                 callers attempting to join that context (A2A-INJ-01 fix).
-                If None (default), no ownership tracking is performed —
+                If None (default), no ownership tracking is performed --
                 backward compatible with existing deployments.
 
                 Example::
@@ -147,8 +147,15 @@ def get_caller_id(ctx: ServerCallContext | None) -> str | None:
         )
         # ---- NEW (fix for A2A-INJ-01) ----
         self._get_caller_id: CallerIdExtractor | None = get_caller_id
-        # Maps context_id → owner identity; populated on first message per context.
+        # Maps context_id -> owner identity; populated on first message per context.
         self._context_owners: dict[str, str] = {}
+        if get_caller_id is None:
+            logger.warning(
+                'DefaultRequestHandler initialized without get_caller_id: '
+                'context ownership is not enforced. Cross-user context injection '
+                '(A2A-INJ-01 / CWE-639) is possible. Provide a get_caller_id '
+                'extractor to enable ownership checks.'
+            )
         # ----------------------------------
         self._running_agents = {}
         self._running_agents_lock = asyncio.Lock()
@@ -168,7 +175,10 @@ async def on_get_task(
     async def on_cancel_task(
         self, params: TaskIdParams, context: ServerCallContext | None = None
     ) -> Task | None:
-        """Default handler for 'tasks/cancel'."""
+        """Default handler for 'tasks/cancel'.
+
+        Attempts to cancel the task managed by the `AgentExecutor`.
+        """
         task: Task | None = await self.task_store.get(params.id, context)
         if not task:
             raise ServerError(error=TaskNotFoundError())
@@ -225,6 +235,12 @@ async def on_cancel_task(
     async def _run_event_stream(
         self, request: RequestContext, queue: EventQueue
     ) -> None:
+        """Runs the agent's `execute` method and closes the queue afterwards.
+
+        Args:
+            request: The request context for the agent.
+            queue: The event queue for the agent to publish to.
+        """
         await self.agent_executor.execute(request, queue)
         await queue.close()
 
@@ -236,32 +252,13 @@ def _check_context_ownership(
         """Enforce context ownership when get_caller_id is configured.
 
         Called before any message is processed for an existing context_id.
+        Only invoked when context_id is already present in _context_owners,
+        which guarantees _get_caller_id is not None and owner is not None.
         Raises ServerError(InvalidParamsError) if the caller does not own
         the context.
         """
-        if self._get_caller_id is None:
-            # Ownership tracking not configured — log warning and allow.
-            # Operators should configure get_caller_id in production.
-            logger.warning(
-                'Context ownership not enforced for context_id=%s: '
-                'no get_caller_id configured on DefaultRequestHandler. '
-                'This allows cross-user context injection (A2A-INJ-01 / CWE-639). '
-                'Provide a get_caller_id extractor to enable ownership checks.',
-                context_id,
-            )
-            return
-
-        caller = self._get_caller_id(context)
-        owner = self._context_owners.get(context_id)
-
-        if owner is None:
-            # Context exists in the store but ownership was not recorded
-            # (e.g. created before this patch was deployed). Skip check.
-            logger.debug(
-                'context_id=%s has no recorded owner; skipping ownership check.',
-                context_id,
-            )
-            return
+        caller = self._get_caller_id(context)  # type: ignore[misc]
+        owner = self._context_owners[context_id]
 
         if caller is None:
             raise ServerError(
@@ -308,10 +305,10 @@ async def _setup_message_execution(
     ) -> tuple[TaskManager, str, EventQueue, ResultAggregator, asyncio.Task]:
         context_id = params.message.context_id
 
-        # ---- FIX: A2A-INJ-01 — enforce context ownership BEFORE task lookup ----
+        # ---- FIX: A2A-INJ-01 -- enforce context ownership BEFORE task lookup ----
         # The check must happen at context_id level, not task level. An attacker
         # who sends a new task_id under an existing context_id would otherwise
-        # bypass a task-level check (get_task() returns None → check never runs).
+        # bypass a task-level check (get_task() returns None -> check never runs).
         if context_id and context_id in self._context_owners:
             self._check_context_ownership(context_id, context)
         # -----------------------------------------------------------------------
@@ -396,7 +393,11 @@ async def on_message_send(
         params: MessageSendParams,
         context: ServerCallContext | None = None,
     ) -> Message | Task:
-        """Default handler for 'message/send' (non-streaming)."""
+        """Default handler for 'message/send' interface (non-streaming).
+
+        Starts the agent execution for the message and waits for the final
+        result (Task or Message).
+        """
         (
             _task_manager,
             task_id,
@@ -461,7 +462,11 @@ async def on_message_send_stream(
         params: MessageSendParams,
         context: ServerCallContext | None = None,
     ) -> AsyncGenerator[Event]:
-        """Default handler for 'message/stream' (streaming)."""
+        """Default handler for 'message/stream' (streaming).
+
+        Starts the agent execution and yields events as they are produced
+        by the agent.
+        """
         (
             _task_manager,
             task_id,

src/a2a/utils/url_validation.py

@@ -21,27 +21,27 @@
 # Networks that must never be reachable via a resolved AgentCard URL.
 # Covers: loopback, RFC 1918 private ranges, link-local (IMDS), and other
 # IANA-reserved blocks that have no legitimate use as public agent endpoints.
-_BLOCKED_NETWORKS: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = [
+_BLOCKED_NETWORKS: tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...] = (
     # Loopback
     ipaddress.ip_network('127.0.0.0/8'),
     ipaddress.ip_network('::1/128'),
     # RFC 1918 private ranges
     ipaddress.ip_network('10.0.0.0/8'),
     ipaddress.ip_network('172.16.0.0/12'),
     ipaddress.ip_network('192.168.0.0/16'),
-    # Link-local — covers AWS/GCP/Azure/OCI IMDS (169.254.169.254)
+    # Link-local -- covers AWS/GCP/Azure/OCI IMDS (169.254.169.254)
     ipaddress.ip_network('169.254.0.0/16'),
     ipaddress.ip_network('fe80::/10'),
-    # IPv6 unique local (ULA) — equivalent of RFC 1918 for IPv6
+    # IPv6 unique local (ULA) -- equivalent of RFC 1918 for IPv6
     ipaddress.ip_network('fc00::/7'),
-    # Shared address space (RFC 6598 — carrier-grade NAT)
+    # Shared address space (RFC 6598 -- carrier-grade NAT)
     ipaddress.ip_network('100.64.0.0/10'),
     # Other IANA reserved / unroutable
     ipaddress.ip_network('0.0.0.0/8'),
     ipaddress.ip_network('192.0.0.0/24'),
     ipaddress.ip_network('198.18.0.0/15'),
     ipaddress.ip_network('240.0.0.0/4'),
-]
+)
 
 
 class A2ASSRFValidationError(ValueError):
@@ -56,7 +56,7 @@ def validate_agent_card_url(url: str) -> None:
     1. URL must be parseable and non-empty.
     2. Scheme must be ``http`` or ``https``.
     3. Hostname must be present and non-empty.
-    4. The hostname must resolve to a publicly routable IP address — it must
+    4. The hostname must resolve to a publicly routable IP address -- it must
        not resolve to a loopback, private, link-local, or otherwise reserved
        address (SSRF / IMDS protection).
 

tests/client/conftest.py

@@ -0,0 +1,24 @@
+"""conftest.py for tests/client/
+
+Patches out SSRF DNS validation so that card resolver and transport tests can
+use test hostnames (localhost, testserver, example.com) without real DNS
+lookups. The validate_agent_card_url function is tested directly in
+tests/utils/test_url_validation.py.
+
+Target: tests/client/conftest.py
+"""
+
+import pytest
+from unittest.mock import patch
+
+
+@pytest.fixture(autouse=True)
+def bypass_ssrf_url_validation():
+    """Bypass DNS-based SSRF validation for all tests in tests/client/.
+
+    Tests here mock HTTP transports and use synthetic hostnames that do not
+    resolve to real IP addresses. SSRF URL validation is exercised by its own
+    dedicated test suite in tests/utils/test_url_validation.py.
+    """
+    with patch('a2a.client.card_resolver.validate_agent_card_url'):
+        yield

tests/integration/conftest.py

@@ -0,0 +1,25 @@
+"""conftest.py for tests/integration/
+
+Patches out SSRF DNS validation so that integration tests can use httpx
+TestClient's synthetic 'testserver' hostname in AgentCard.url without
+triggering real DNS resolution. The validate_agent_card_url function is
+tested directly in tests/utils/test_url_validation.py.
+
+Target: tests/integration/conftest.py
+"""
+
+import pytest
+from unittest.mock import patch
+
+
+@pytest.fixture(autouse=True)
+def bypass_ssrf_url_validation():
+    """Bypass DNS-based SSRF validation for all tests in tests/integration/.
+
+    Integration tests use httpx's TestClient which binds to the synthetic
+    'testserver' hostname. This hostname cannot be resolved via DNS.
+    SSRF URL validation is exercised by its own dedicated test suite in
+    tests/utils/test_url_validation.py.
+    """
+    with patch('a2a.client.card_resolver.validate_agent_card_url'):
+        yield

tests/utils/test_url_validation.py

@@ -0,0 +1,95 @@
+"""Tests for a2a.utils.url_validation (A2A-SSRF-01 fix).
+
+Target: tests/utils/test_url_validation.py
+"""
+
+import pytest
+
+from a2a.utils.url_validation import A2ASSRFValidationError, validate_agent_card_url
+
+
+class TestValidateAgentCardUrlScheme:
+    """URL scheme validation."""
+
+    @pytest.mark.parametrize('url', [
+        'file:///etc/passwd',
+        'gopher://internal/1',
+        'ftp://files.example.com/secret',
+        'dict://internal/',
+        'ldap://ldap.example.com/',
+        '',
+    ])
+    def test_non_http_schemes_are_blocked(self, url):
+        with pytest.raises(A2ASSRFValidationError):
+            validate_agent_card_url(url)
+
+    @pytest.mark.parametrize('url', [
+        'http://example.com/rpc',
+        'https://example.com/rpc',
+        'HTTP://EXAMPLE.COM/RPC',
+        'HTTPS://EXAMPLE.COM/RPC',
+    ])
+    def test_http_and_https_are_allowed(self, url):
+        # Should not raise — only scheme + hostname check, DNS may vary
+        # We only verify scheme acceptance here; real DNS tested separately.
+        try:
+            validate_agent_card_url(url)
+        except A2ASSRFValidationError as exc:
+            # Accept DNS resolution failure — scheme was accepted
+            assert 'could not be resolved' in str(exc) or 'blocked network' in str(exc)
+
+
+class TestValidateAgentCardUrlPrivateIPs:
+    """Private / reserved IP range blocking."""
+
+    @pytest.mark.parametrize('url,label', [
+        ('http://127.0.0.1/rpc',       'loopback IPv4'),
+        ('http://127.1.2.3/rpc',       'loopback IPv4 (non-zero host)'),
+        ('http://[::1]/rpc',           'loopback IPv6'),
+        ('http://10.0.0.1/rpc',        'RFC 1918 10/8'),
+        ('http://10.255.255.255/rpc',  'RFC 1918 10/8 broadcast'),
+        ('http://172.16.0.1/rpc',      'RFC 1918 172.16/12'),
+        ('http://172.31.255.255/rpc',  'RFC 1918 172.31 (last in range)'),
+        ('http://192.168.1.1/rpc',     'RFC 1918 192.168/16'),
+        ('http://169.254.169.254/latest/meta-data/', 'AWS IMDS'),
+        ('http://169.254.0.1/rpc',     'link-local'),
+        ('http://100.64.0.1/rpc',      'shared address space RFC 6598'),
+    ])
+    def test_private_addresses_are_blocked(self, url, label):
+        with pytest.raises(A2ASSRFValidationError, match='blocked network'):
+            validate_agent_card_url(url)
+
+    def test_public_ip_is_allowed(self):
+        """A routable public IP should not be blocked."""
+        # 93.184.216.34 is example.com — guaranteed public
+        try:
+            validate_agent_card_url('http://93.184.216.34/rpc')
+        except A2ASSRFValidationError as exc:
+            # Only acceptable failure is DNS (not a blocked-network error)
+            assert 'could not be resolved' in str(exc)
+            pytest.skip('DNS not available in this environment')
+
+
+class TestValidateAgentCardUrlHostname:
+    """Hostname-level checks."""
+
+    def test_missing_hostname_is_blocked(self):
+        with pytest.raises(A2ASSRFValidationError, match='no hostname'):
+            validate_agent_card_url('http:///path')
+
+    def test_empty_url_is_blocked(self):
+        with pytest.raises(A2ASSRFValidationError, match='must not be empty'):
+            validate_agent_card_url('')
+
+
+class TestA2ASSRFValidationError:
+    """Exception type tests."""
+
+    def test_is_subclass_of_value_error(self):
+        assert issubclass(A2ASSRFValidationError, ValueError)
+
+    def test_raises_with_descriptive_message(self):
+        with pytest.raises(A2ASSRFValidationError) as exc_info:
+            validate_agent_card_url('http://127.0.0.1/rpc')
+        assert '127.0.0.1' in str(exc_info.value)
+        assert 'CWE-918' in str(exc_info.value)