fix(tasks): bound imported MCP server DNS validation

posthog/security/test/test_url_validation.py

@@ -1,4 +1,6 @@
 import ipaddress
+from concurrent.futures import Future, ThreadPoolExecutor
+from threading import Barrier, BoundedSemaphore
 
 import pytest
 
@@ -12,6 +14,76 @@ def force_prod(monkeypatch):
 
 
 class TestUrlValidation:
+    def test_resolve_host_ips_uses_bounded_lifetime(self, monkeypatch):
+        class Answers:
+            def addresses(self):
+                return iter(["93.184.216.34"])
+
+        class Resolver:
+            def resolve_name(self, host, *, lifetime):
+                assert host == "example.com"
+                assert lifetime == uv.DNS_RESOLUTION_LIFETIME_SECONDS
+                return Answers()
+
+        monkeypatch.setattr(uv.dns.resolver, "Resolver", Resolver)
+
+        assert uv.resolve_host_ips("example.com") == {ipaddress.ip_address("93.184.216.34")}
+
+    def test_resolve_url_hosts_ips_deduplicates_hosts(self, monkeypatch):
+        def fake_resolve_hosts_ips(hosts):
+            assert hosts == {"shared.example.com"}
+            return {"shared.example.com": {ipaddress.ip_address("93.184.216.34")}}
+
+        monkeypatch.setattr(uv, "resolve_hosts_ips", fake_resolve_hosts_ips)
+
+        assert uv.resolve_url_hosts_ips(["https://shared.example.com/first", "https://shared.example.com/second"]) == {
+            "shared.example.com": {ipaddress.ip_address("93.184.216.34")}
+        }
+
+    def test_resolve_hosts_ips_stops_at_batch_deadline(self, monkeypatch):
+        pending_future: Future[uv.ResolvedIPs] = Future()
+
+        class Executor:
+            def submit(self, _function, _host):
+                return pending_future
+
+        monkeypatch.setattr(uv, "_dns_resolution_executor", Executor())
+        monkeypatch.setattr(uv, "DNS_RESOLUTION_BATCH_TIMEOUT_SECONDS", 0)
+
+        assert uv.resolve_hosts_ips({"slow.example.com"}) == {"slow.example.com": set()}
+        assert pending_future.cancelled()
+
+    def test_resolve_hosts_ips_starts_all_allowed_hosts_within_the_batch_deadline(self, monkeypatch):
+        hosts = {f"host-{index}.example.com" for index in range(20)}
+        all_workers_started = Barrier(len(hosts))
+        public_ip = ipaddress.ip_address("93.184.216.34")
+
+        def resolve_after_all_workers_start(_host):
+            all_workers_started.wait(timeout=1)
+            return {public_ip}
+
+        executor = ThreadPoolExecutor(max_workers=uv.DNS_RESOLUTION_MAX_WORKERS)
+        monkeypatch.setattr(uv, "_dns_resolution_executor", executor)
+        monkeypatch.setattr(uv, "resolve_host_ips", resolve_after_all_workers_start)
+
+        try:
+            assert uv.resolve_hosts_ips(hosts) == {host: {public_ip} for host in hosts}
+        finally:
+            executor.shutdown(wait=True, cancel_futures=True)
+
+    def test_resolve_hosts_ips_fails_closed_when_global_capacity_is_exhausted(self, monkeypatch):
+        capacity = BoundedSemaphore(1)
+        capacity.acquire()
+
+        class Executor:
+            def submit(self, _function, _host):
+                raise AssertionError("capacity exhaustion must prevent queueing")
+
+        monkeypatch.setattr(uv, "_dns_resolution_capacity", capacity)
+        monkeypatch.setattr(uv, "_dns_resolution_executor", Executor())
+
+        assert uv.resolve_hosts_ips({"busy.example.com"}) == {"busy.example.com": set()}
+
     def test_is_url_allowed_disallowed_scheme(self):
         ok, err = uv.is_url_allowed("javascript:alert(1)")
         assert not ok and "scheme" in (err or "")

posthog/security/url_validation.py

@@ -1,15 +1,30 @@
-import socket
 import ipaddress
 import urllib.parse as urlparse
+from collections.abc import Iterable, Mapping
+from concurrent.futures import Future, ThreadPoolExecutor, wait
+from threading import BoundedSemaphore
 
 from django.conf import settings
 
 import structlog
+import dns.resolver
+import dns.exception
 
 from posthog.cloud_utils import is_dev_mode
 
 logger = structlog.get_logger(__name__)
 
+ResolvedIPs = set[ipaddress.IPv4Address | ipaddress.IPv6Address]
+
+DNS_RESOLUTION_LIFETIME_SECONDS = 2.0
+DNS_RESOLUTION_BATCH_TIMEOUT_SECONDS = 2.5
+DNS_RESOLUTION_MAX_WORKERS = 20
+_dns_resolution_executor = ThreadPoolExecutor(
+    max_workers=DNS_RESOLUTION_MAX_WORKERS,
+    thread_name_prefix="url-validation-dns",
+)
+_dns_resolution_capacity = BoundedSemaphore(DNS_RESOLUTION_MAX_WORKERS)
+
 # Schemes that should never be allowed for external URLs
 DISALLOWED_SCHEMES = {"file", "ftp", "gopher", "ws", "wss", "data", "javascript"}
 
@@ -34,23 +49,67 @@
 )
 
 
-def resolve_host_ips(host: str) -> set[ipaddress.IPv4Address | ipaddress.IPv6Address]:
+def resolve_host_ips(host: str) -> ResolvedIPs:
     """Resolve a hostname to its IP addresses."""
     try:
-        infos = socket.getaddrinfo(host, None)
-    except socket.gaierror as e:
-        logger.warning("url_validation.dns_resolution_failed", host=host, errno=e.errno, strerror=e.strerror)
+        return {ipaddress.ip_address(host)}
+    except ValueError:
+        pass
+
+    try:
+        answers = dns.resolver.Resolver().resolve_name(host, lifetime=DNS_RESOLUTION_LIFETIME_SECONDS)
+    except dns.exception.DNSException as error:
+        logger.warning("url_validation.dns_resolution_failed", host=host, error=str(error))
         return set()
-    ips: set[ipaddress.IPv4Address | ipaddress.IPv6Address] = set()
-    for _fam, *_rest, sockaddr in infos:
-        ip = sockaddr[0]
+
+    ips: ResolvedIPs = set()
+    for address in answers.addresses():
         try:
-            ips.add(ipaddress.ip_address(ip))
+            ips.add(ipaddress.ip_address(address))
         except ValueError:
             pass
     return ips
 
 
+def _submit_dns_resolution(host: str) -> Future[ResolvedIPs] | None:
+    if not _dns_resolution_capacity.acquire(blocking=False):
+        logger.warning("url_validation.dns_resolution_capacity_exhausted", host=host)
+        return None
+    try:
+        future = _dns_resolution_executor.submit(resolve_host_ips, host)
+    except Exception as error:
+        _dns_resolution_capacity.release()
+        logger.exception("url_validation.dns_resolution_submit_failed", host=host, error=str(error))
+        return None
+    future.add_done_callback(lambda _future: _dns_resolution_capacity.release())
+    return future
+
+
+def resolve_hosts_ips(hosts: Iterable[str]) -> dict[str, ResolvedIPs]:
+    unique_hosts = set(hosts)
+    resolved: dict[str, ResolvedIPs] = {host: set() for host in unique_hosts}
+    futures = {host: future for host in unique_hosts if (future := _submit_dns_resolution(host)) is not None}
+    if not futures:
+        return resolved
+
+    completed, pending = wait(futures.values(), timeout=DNS_RESOLUTION_BATCH_TIMEOUT_SECONDS)
+
+    for future in pending:
+        future.cancel()
+
+    for host, future in futures.items():
+        if future not in completed:
+            logger.warning("url_validation.dns_resolution_timed_out", host=host)
+            resolved[host] = set()
+            continue
+        try:
+            resolved[host] = future.result()
+        except Exception as error:
+            logger.exception("url_validation.dns_resolution_failed", host=host, error=str(error))
+            resolved[host] = set()
+    return resolved
+
+
 def _is_internal_ip(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
     """Check if an IP address is internal/private and should be blocked."""
     return any(
@@ -121,7 +180,34 @@ def _dev_bypass_enabled() -> bool:
     return not settings.FORCE_URL_VALIDATION
 
 
-def is_url_allowed(raw_url: str) -> tuple[bool, str | None]:
+def resolve_url_hosts_ips(raw_urls: Iterable[str]) -> dict[str, ResolvedIPs]:
+    if _dev_bypass_enabled():
+        return {}
+    hosts: set[str] = set()
+    for raw_url in raw_urls:
+        if has_authority_bypass_chars(raw_url):
+            continue
+        try:
+            parsed_url = urlparse.urlparse(raw_url)
+            host = (parsed_url.hostname or "").lower()
+        except Exception:
+            continue
+        if (
+            parsed_url.scheme not in {"http", "https"}
+            or not parsed_url.netloc
+            or host in METADATA_HOSTS
+            or host in {"localhost", "127.0.0.1", "::1"}
+            or any(host.endswith(pattern) for pattern in INTERNAL_DOMAIN_PATTERNS)
+            or _is_private_ip_literal(host)
+        ):
+            continue
+        hosts.add(host)
+    return resolve_hosts_ips(hosts)
+
+
+def is_url_allowed(
+    raw_url: str, *, resolved_ips_by_host: Mapping[str, ResolvedIPs] | None = None
+) -> tuple[bool, str | None]:
     """
     Validate a URL for SSRF protection.
 
@@ -132,7 +218,7 @@ def is_url_allowed(raw_url: str) -> tuple[bool, str | None]:
     - Host must not be localhost, metadata service, or internal domain
     - Resolved IPs must not be private/internal
     """
-    allowed, reason, _ips = _validate_url_with_ips(raw_url)
+    allowed, reason, _ips = _validate_url_with_ips(raw_url, resolved_ips_by_host=resolved_ips_by_host)
     return allowed, reason
 
 
@@ -151,6 +237,8 @@ def validate_url_and_pin_ips(
 
 def _validate_url_with_ips(
     raw_url: str,
+    *,
+    resolved_ips_by_host: Mapping[str, ResolvedIPs] | None = None,
 ) -> tuple[bool, str | None, set[ipaddress.IPv4Address | ipaddress.IPv6Address]]:
     empty: set[ipaddress.IPv4Address | ipaddress.IPv6Address] = set()
 
@@ -188,7 +276,7 @@ def _blocked(
     if _is_private_ip_literal(host):
         return _blocked("Private IP address not allowed", host=host)
 
-    ips = resolve_host_ips(host)
+    ips = resolve_host_ips(host) if resolved_ips_by_host is None else resolved_ips_by_host.get(host, empty)
     if not ips:
         return _blocked("Could not resolve host", host=host)
     for ip in ips:

products/tasks/backend/presentation/serializers.py

@@ -19,7 +19,7 @@
 from posthog.event_usage import groups
 from posthog.models.integration import Integration
 from posthog.models.user_integration import UserIntegration
-from posthog.security.url_validation import is_url_allowed
+from posthog.security.url_validation import is_url_allowed, resolve_url_hosts_ips
 
 from products.tasks.backend.facade import api as tasks_facade
 from products.tasks.backend.facade.contracts import (
@@ -1535,6 +1535,13 @@ class ImportedMcpServerHeaderSerializer(serializers.Serializer):
     )
 
 
+class ImportedMcpServerListSerializer(serializers.ListSerializer):
+    def to_internal_value(self, data: object) -> list[dict[str, object]]:
+        if isinstance(data, list) and len(data) > MAX_IMPORTED_MCP_SERVERS:
+            raise serializers.ValidationError(f"At most {MAX_IMPORTED_MCP_SERVERS} imported MCP servers are allowed.")
+        return cast(list[dict[str, object]], super().to_internal_value(data))
+
+
 class ImportedMcpServerSerializer(serializers.Serializer):
     """One client-imported MCP server, in the agent server's --mcpServers entry shape."""
 
@@ -1543,14 +1550,8 @@ class ImportedMcpServerSerializer(serializers.Serializer):
     url = serializers.URLField(max_length=2048)
     headers = ImportedMcpServerHeaderSerializer(many=True, required=False, default=list)
 
-    def validate_url(self, value: str) -> str:
-        # The client classifies public vs private hosts for UX, but that is not a
-        # security boundary: the sandbox egresses from PostHog infrastructure, so a
-        # private URL here is a user-controlled SSRF vector. Re-check server-side.
-        allowed, reason = is_url_allowed(value)
-        if not allowed:
-            raise serializers.ValidationError(reason or "URL is not allowed.")
-        return value
+    class Meta:
+        list_serializer_class = ImportedMcpServerListSerializer
 
 
 class ImportedMcpServersFieldMixin(serializers.Serializer):
@@ -1572,11 +1573,15 @@ class ImportedMcpServersFieldMixin(serializers.Serializer):
     def validate_imported_mcp_servers(self, value):
         if not value:
             return None
-        if len(value) > MAX_IMPORTED_MCP_SERVERS:
-            raise serializers.ValidationError(f"At most {MAX_IMPORTED_MCP_SERVERS} imported MCP servers are allowed.")
         _validate_unique_unreserved_mcp_names(value)
         if len(json.dumps(value)) > MAX_IMPORTED_MCP_SERVERS_BYTES:
             raise serializers.ValidationError("Imported MCP servers payload is too large.")
+
+        resolved_ips_by_host = resolve_url_hosts_ips(server["url"] for server in value)
+        for server in value:
+            allowed, reason = is_url_allowed(server["url"], resolved_ips_by_host=resolved_ips_by_host)
+            if not allowed:
+                raise serializers.ValidationError(reason or "URL is not allowed.")
         return value
 
 

products/tasks/backend/tests/test_presentation_serializers.py

@@ -1,8 +1,15 @@
+import ipaddress
+
+from unittest.mock import patch
+
 from django.test import SimpleTestCase
 
 from parameterized import parameterized
 
-from products.tasks.backend.presentation.serializers import TaskRunLivingArtifactCreateRequestSerializer
+from products.tasks.backend.presentation.serializers import (
+    TaskRunCreateRequestSerializer,
+    TaskRunLivingArtifactCreateRequestSerializer,
+)
 
 
 class TestTaskRunLivingArtifactCreateRequestSerializer(SimpleTestCase):
@@ -15,3 +22,38 @@ class TestTaskRunLivingArtifactCreateRequestSerializer(SimpleTestCase):
     def test_content_source_exclusivity(self, _name: str, data: dict, expected_valid: bool) -> None:
         serializer = TaskRunLivingArtifactCreateRequestSerializer(data=data)
         assert serializer.is_valid() is expected_valid
+
+
+class TestTaskRunCreateRequestSerializer(SimpleTestCase):
+    @patch(
+        "posthog.security.url_validation.resolve_host_ips",
+        return_value={ipaddress.ip_address("93.184.216.34")},
+    )
+    def test_deduplicates_imported_mcp_server_host_resolution(self, mock_resolve_host_ips) -> None:
+        serializer = TaskRunCreateRequestSerializer(
+            data={
+                "environment": "cloud",
+                "imported_mcp_servers": [
+                    {"type": "http", "name": "first", "url": "https://shared.example.com/first"},
+                    {"type": "http", "name": "second", "url": "https://shared.example.com/second"},
+                ],
+            }
+        )
+
+        assert serializer.is_valid(), serializer.errors
+        mock_resolve_host_ips.assert_called_once_with("shared.example.com")
+
+    @patch("products.tasks.backend.presentation.serializers.resolve_url_hosts_ips")
+    def test_rejects_too_many_imported_mcp_servers_before_dns_resolution(self, mock_resolve_url_hosts_ips) -> None:
+        serializer = TaskRunCreateRequestSerializer(
+            data={
+                "environment": "cloud",
+                "imported_mcp_servers": [
+                    {"type": "http", "name": f"server-{index}", "url": f"https://{index}.example.com"}
+                    for index in range(21)
+                ],
+            }
+        )
+
+        assert not serializer.is_valid()
+        mock_resolve_url_hosts_ips.assert_not_called()