diff --git a/headroom/providers/vertex/runtime.py b/headroom/providers/vertex/runtime.py index 5a1931ff7..4cf6194de 100644 --- a/headroom/providers/vertex/runtime.py +++ b/headroom/providers/vertex/runtime.py @@ -2,10 +2,22 @@ from __future__ import annotations +import re from dataclasses import dataclass from headroom.providers.registry import DEFAULT_VERTEX_API_URL +# The public (multi-region) Vertex endpoint, used for ``global`` and for any +# location that is not a well-formed region. +_VERTEX_GLOBAL_API_URL = "https://aiplatform.googleapis.com" + +# A GCP region label: lowercase alphanumeric groups joined by single hyphens +# (e.g. ``us-central1``, ``europe-west4``, ``asia-northeast1``). Anchored and +# deliberately strict — no dots, colons, slashes, ``#``, ``@``, uppercase, or +# empty groups — so a user-controlled ``location`` can never carry a host, +# port, path, or URL-fragment delimiter into the interpolated hostname. +_VERTEX_REGION_RE = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$") + VERTEX_GOOGLE_PUBLISHER = "google" VERTEX_ANTHROPIC_PUBLISHER = "anthropic" VERTEX_GOOGLE_PROVIDER_NAME = "vertex:google" @@ -50,9 +62,24 @@ def vertex_anthropic_target(base_url: str, *, versionless_route: bool = False) - def vertex_target_for_location(configured_target: str, location: str) -> str: - """Return the Vertex upstream target for a request location.""" + """Return the Vertex upstream target for a request location. + + ``location`` is a user-controlled URL path segment that is interpolated into + the upstream hostname, so it must be validated against the GCP region shape + before use. Without that check a value such as ``169.254.169.254#`` (decoded + from a percent-encoded ``%23`` in the path) produces + ``https://169.254.169.254#-aiplatform.googleapis.com``, which an HTTP client + parses as host ``169.254.169.254`` with the remainder treated as a URL + fragment — a server-side request forgery to the cloud metadata endpoint + (CWE-918). Any ``location`` that is not a well-formed region (including port, + path, or fragment-delimiter payloads) falls back to the default public + endpoint, which can never resolve to an attacker-chosen host. + + An explicitly configured gateway target still wins outright; region + derivation only applies when running against the default Vertex endpoint. + """ if configured_target and configured_target != DEFAULT_VERTEX_API_URL: return configured_target - if not location or location == "global": - return "https://aiplatform.googleapis.com" + if not location or location == "global" or not _VERTEX_REGION_RE.match(location): + return _VERTEX_GLOBAL_API_URL return f"https://{location}-aiplatform.googleapis.com" diff --git a/tests/test_provider_vertex_runtime.py b/tests/test_provider_vertex_runtime.py index cf0bad667..8c3584054 100644 --- a/tests/test_provider_vertex_runtime.py +++ b/tests/test_provider_vertex_runtime.py @@ -1,5 +1,9 @@ from __future__ import annotations +from urllib.parse import urlsplit + +import pytest + from headroom.providers.registry import DEFAULT_VERTEX_API_URL from headroom.providers.vertex import ( VERTEX_ANTHROPIC_PROVIDER_NAME, @@ -71,3 +75,54 @@ def test_vertex_target_for_location_honors_explicit_gateway() -> None: assert vertex_target_for_location("https://vertex-gateway.internal", "europe-west1") == ( "https://vertex-gateway.internal" ) + + +_VERTEX_PUBLIC_ENDPOINT = "https://aiplatform.googleapis.com" + + +@pytest.mark.parametrize( + "region", + ["us-central1", "europe-west4", "asia-northeast1", "me-central1", "us-east5"], +) +def test_vertex_target_for_location_accepts_real_regions(region: str) -> None: + assert vertex_target_for_location(DEFAULT_VERTEX_API_URL, region) == ( + f"https://{region}-aiplatform.googleapis.com" + ) + + +@pytest.mark.parametrize( + "malicious", + [ + "169.254.169.254#", # fragment delimiter -> cloud metadata IP (the reported PoC) + "127.0.0.1:44919#", # host:port + fragment + "169.254.169.254/latest/meta-data/iam#", # path injection + "169.254.169.254:80", # port injection + "evil.example", # dotted host + "foo@evil.example", # userinfo delimiter + "us_central1", # underscore (not a region) + "US-CENTRAL1", # uppercase + "-leading", # leading hyphen + "trailing-", # trailing hyphen + "a--b", # empty hyphen group + ], +) +def test_vertex_target_for_location_rejects_ssrf_payloads(malicious: str) -> None: + """A non-region ``location`` must never carry an attacker-chosen host into + the interpolated Vertex hostname (CWE-918). It falls back to the public + endpoint, and the parsed host is always the legitimate Vertex host — never + a metadata IP, loopback, or injected authority. + """ + target = vertex_target_for_location(DEFAULT_VERTEX_API_URL, malicious) + assert target == _VERTEX_PUBLIC_ENDPOINT + parsed = urlsplit(target) + assert parsed.hostname == "aiplatform.googleapis.com" + assert parsed.port is None + + +def test_vertex_target_for_location_ssrf_fallback_only_on_default_target() -> None: + """A validated non-region value still cannot override an explicitly + configured gateway (that path returns the operator's target verbatim and + is not user-derived).""" + assert vertex_target_for_location("https://gw.internal", "169.254.169.254#") == ( + "https://gw.internal" + )