diff --git a/headroom/providers/proxy_routes.py b/headroom/providers/proxy_routes.py index 8fe2a0fb4..5cf11ea47 100644 --- a/headroom/providers/proxy_routes.py +++ b/headroom/providers/proxy_routes.py @@ -67,7 +67,7 @@ from headroom.proxy.passthrough import ( custom_base_passthrough_telemetry as _custom_base_passthrough_telemetry, ) from headroom.proxy.request_scope import normalize_request_path -from headroom.proxy.upstream_guard import is_safe_upstream_url +from headroom.proxy.upstream_guard import is_safe_upstream_url_async logger = logging.getLogger("headroom.proxy.routes") @@ -267,7 +267,7 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None: # OpenAI-compatible and generic passthrough routes. custom_base = request.headers.get("x-headroom-base-url", "").strip() if custom_base: - if not is_safe_upstream_url(custom_base): + if not await is_safe_upstream_url_async(custom_base): logger.warning("rejecting unsafe x-headroom-base-url: %r", custom_base) raise HTTPException(status_code=400, detail="Rejected unsafe upstream base URL") return await proxy.handle_anthropic_messages( @@ -500,7 +500,7 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None: # proxy at loopback/RFC1918/cloud-metadata and read the response back # (CVE-2026-77775). custom_base = request.headers.get("x-headroom-base-url", "").strip() - if custom_base and not is_safe_upstream_url(custom_base): + if custom_base and not await is_safe_upstream_url_async(custom_base): logger.warning("rejecting unsafe x-headroom-base-url: %r", custom_base) raise HTTPException(status_code=400, detail="Rejected unsafe upstream base URL") return await proxy.handle_passthrough( @@ -518,7 +518,7 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None: async def passthrough(request: Request, path: str): custom_base = request.headers.get("x-headroom-base-url") if custom_base: - if not is_safe_upstream_url(custom_base): + if not await is_safe_upstream_url_async(custom_base): logger.warning("rejecting unsafe x-headroom-base-url: %r", custom_base) raise HTTPException(status_code=400, detail="Rejected unsafe upstream base URL") base_url = custom_base.rstrip("/") diff --git a/headroom/proxy/upstream_guard.py b/headroom/proxy/upstream_guard.py index 764ef944a..bfe172696 100644 --- a/headroom/proxy/upstream_guard.py +++ b/headroom/proxy/upstream_guard.py @@ -20,13 +20,38 @@ import from any handler without risking an import cycle. from __future__ import annotations +import asyncio import ipaddress import os import socket +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import TimeoutError as _FutureTimeout from urllib.parse import urlparse ALLOWED_BASE_URLS_ENV = "HEADROOM_ALLOWED_BASE_URLS" +# `socket.getaddrinfo` has no timeout parameter and runs on whatever thread +# calls it -- which, for the proxy, is the event loop. A caller-supplied host +# that resolves slowly therefore stalls every other in-flight request, so the +# lookup is bounded here and fails closed when it overruns. Callers already in +# async context should prefer `is_safe_upstream_url_async`, which keeps the +# wait off the loop entirely. +RESOLVE_TIMEOUT_ENV = "HEADROOM_UPSTREAM_RESOLVE_TIMEOUT_S" +_DEFAULT_RESOLVE_TIMEOUT_S = 3.0 +_RESOLVER_POOL = ThreadPoolExecutor(max_workers=8, thread_name_prefix="hr-upstream-dns") + + +def _resolve_timeout_seconds() -> float: + raw = (os.environ.get(RESOLVE_TIMEOUT_ENV) or "").strip() + if not raw: + return _DEFAULT_RESOLVE_TIMEOUT_S + try: + value = float(raw) + except ValueError: + return _DEFAULT_RESOLVE_TIMEOUT_S + return value if value > 0 else _DEFAULT_RESOLVE_TIMEOUT_S + + _SAFE_SCHEMES = {"http", "https", "ws", "wss"} @@ -135,10 +160,22 @@ def is_safe_upstream_url(url: str) -> bool: return (parsed.scheme.lower(), host.lower(), port) in origins try: - infos = socket.getaddrinfo(host, None, proto=socket.IPPROTO_TCP) - except OSError: + infos = _RESOLVER_POOL.submit( + socket.getaddrinfo, host, None, 0, 0, socket.IPPROTO_TCP + ).result(timeout=_resolve_timeout_seconds()) + except (OSError, _FutureTimeout): # Resolution and connection are separate operations, so allowing a DNS # miss here would fail open if the name resolves on the later lookup. + # A lookup that overruns the budget is treated the same way. # Operators can explicitly allowlist split-horizon/internal endpoints. return False return all(not _is_internal_address(str(info[4][0])) for info in infos) + + +async def is_safe_upstream_url_async(url: str) -> bool: + """Async form of :func:`is_safe_upstream_url` for event-loop callers. + + Same policy; the blocking resolution runs off the loop so a hostile or + slow-resolving hostname cannot stall unrelated in-flight requests. + """ + return await asyncio.to_thread(is_safe_upstream_url, url) diff --git a/tests/test_provider_proxy_targets.py b/tests/test_provider_proxy_targets.py index b8f4412c5..ed1385745 100644 --- a/tests/test_provider_proxy_targets.py +++ b/tests/test_provider_proxy_targets.py @@ -1,11 +1,14 @@ from __future__ import annotations +from unittest.mock import patch + from headroom.providers.proxy_targets import ( api_target, select_passthrough_base_url, vertex_target_for_location, ) from headroom.providers.registry import DEFAULT_VERTEX_API_URL +from headroom.proxy import upstream_guard def _proxy(**legacy_targets: str): @@ -56,13 +59,21 @@ def test_select_passthrough_base_url_handles_special_auth_modes() -> None: assert select_passthrough_base_url(proxy, {"x-goog-api-key": "test"}) == ( "https://legacy.gemini.test" ) - assert ( - select_passthrough_base_url( - proxy, - {"api-key": "azure", "x-headroom-base-url": "https://azure.example/base/"}, + # The Azure branch honours the override only after the SSRF guard clears + # the destination (CVE-2026-77775), and `azure.example` does not resolve. + # Pin a public answer so this stays a test of target *precedence*. + with patch.object( + upstream_guard.socket, + "getaddrinfo", + return_value=[(None, None, None, None, ("20.10.10.10", 443))], + ): + assert ( + select_passthrough_base_url( + proxy, + {"api-key": "azure", "x-headroom-base-url": "https://azure.example/base/"}, + ) + == "https://azure.example/base" ) - == "https://azure.example/base" - ) assert select_passthrough_base_url(proxy, {"x-api-key": "anthropic"}) == ( "https://legacy.anthropic.test" ) diff --git a/tests/test_provider_registry_extended.py b/tests/test_provider_registry_extended.py index c3ef4559a..e70e14175 100644 --- a/tests/test_provider_registry_extended.py +++ b/tests/test_provider_registry_extended.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging from types import SimpleNamespace from typing import Any +from unittest.mock import patch import pytest @@ -13,6 +14,7 @@ from headroom.providers.registry import ( create_proxy_backend, format_backend_status, ) +from headroom.proxy import upstream_guard class DummyStorage: @@ -85,12 +87,20 @@ def test_proxy_provider_runtime_selects_targets_and_providers() -> None: assert runtime.select_passthrough_base_url({"x-goog-api-key": "test"}) == ( "https://gemini.example" ) - assert ( - runtime.select_passthrough_base_url( - {"api-key": "azure-key", "x-headroom-base-url": "https://azure.example/openai/"} + # The Azure branch honours the override only after the SSRF guard clears + # the destination (CVE-2026-77775), and `azure.example` does not resolve. + # Pin a public answer so this stays a test of target *precedence*. + with patch.object( + upstream_guard.socket, + "getaddrinfo", + return_value=[(None, None, None, None, ("20.10.10.10", 443))], + ): + assert ( + runtime.select_passthrough_base_url( + {"api-key": "azure-key", "x-headroom-base-url": "https://azure.example/openai/"} + ) + == "https://azure.example/openai" ) - == "https://azure.example/openai" - ) assert runtime.select_passthrough_base_url({}) == "https://openai.example" diff --git a/tests/test_upstream_guard.py b/tests/test_upstream_guard.py index 1bf0e6829..e38693324 100644 --- a/tests/test_upstream_guard.py +++ b/tests/test_upstream_guard.py @@ -265,3 +265,36 @@ def test_operator_allowlist_still_permits_an_internal_azure_endpoint( headers = {"api-key": "x", "x-headroom-base-url": "https://gateway.internal/v1"} assert select_passthrough_base_url(_StubProxy(), headers) == "https://gateway.internal/v1" + + +def test_slow_resolution_is_bounded_and_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + """A hostile hostname must not hold the caller for the resolver's timeout. + + `socket.getaddrinfo` takes no timeout and runs on the calling thread, which + for the proxy is the event loop -- so an unbounded lookup is an + unauthenticated stall of every in-flight request. + """ + import time as _time + + def slow_resolution(*args: object, **kwargs: object) -> list[object]: + _time.sleep(5.0) + return [(None, None, None, None, ("8.8.8.8", 443))] + + monkeypatch.setenv("HEADROOM_UPSTREAM_RESOLVE_TIMEOUT_S", "0.25") + monkeypatch.setattr(socket, "getaddrinfo", slow_resolution) + + started = _time.perf_counter() + result = is_safe_upstream_url("https://slow.example/v1") + elapsed = _time.perf_counter() - started + + assert result is False, "a lookup that overruns its budget must fail closed" + assert elapsed < 2.0, f"resolution was not bounded (took {elapsed:.2f}s)" + + +async def test_async_guard_matches_the_sync_policy() -> None: + """The off-loop wrapper must not diverge from the blocking form.""" + from headroom.proxy.upstream_guard import is_safe_upstream_url_async + + assert await is_safe_upstream_url_async("http://127.0.0.1/") is False + assert await is_safe_upstream_url_async("http://169.254.169.254/") is False + assert await is_safe_upstream_url_async("https://8.8.8.8/v1") is True