diff --git a/src/openai/lib/_azure_websocket.py b/src/openai/lib/_azure_websocket.py new file mode 100644 index 0000000000..36620cd21f --- /dev/null +++ b/src/openai/lib/_azure_websocket.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from typing_extensions import override + +from websockets.uri import parse_uri +from websockets.exceptions import SecurityError +from websockets.asyncio.client import connect + +__all__ = ["_AzureWebSocketConnect"] + + +class _AzureWebSocketConnect(connect): + """Keep Azure's WebSocket authentication on the original origin.""" + + @override + def process_redirect(self, exc: Exception) -> Exception | str: + uri_or_exc = super().process_redirect(exc) + if isinstance(uri_or_exc, str): + current = parse_uri(self.uri) + target = parse_uri(uri_or_exc) + if (current.secure, current.host, current.port) != (target.secure, target.host, target.port): + return SecurityError("Cross-origin Azure WebSocket redirects are not allowed") + return uri_or_exc diff --git a/src/openai/lib/azure.py b/src/openai/lib/azure.py index 531496a0a9..aeae62168a 100644 --- a/src/openai/lib/azure.py +++ b/src/openai/lib/azure.py @@ -56,6 +56,26 @@ def _has_auth_header(headers: Headers) -> bool: return _has_header(headers, "Authorization") or _has_header(headers, "api-key") +_AZURE_AUTH_ORIGIN = "openai.azure_auth_origin" + + +def _origin(url: httpx2.URL) -> tuple[str, str, int | None]: + port = url.port + if port is None: + port = {"http": 80, "https": 443}.get(url.scheme) + return url.scheme, url.host, port + + +def _strip_azure_api_key_on_redirect(request: httpx2.Request) -> None: + origin = request.extensions.get(_AZURE_AUTH_ORIGIN) + if origin is not None and origin != _origin(request.url): + request.headers.pop("api-key", None) + + +async def _async_strip_azure_api_key_on_redirect(request: httpx2.Request) -> None: + _strip_azure_api_key_on_redirect(request) + + class MutuallyExclusiveAuthError(OpenAIError): def __init__(self) -> None: super().__init__( @@ -129,7 +149,11 @@ def _build_request( if model is not None and "/deployments" not in str(self.base_url.path): options.url = path_template("/deployments/{model}", model=model) + options.url - return super()._build_request(options, retries_taken=retries_taken) + request = super()._build_request(options, retries_taken=retries_taken) + # HTTPX preserves request extensions through redirects. Scope the hook + # to this Azure request, including when its HTTP client is shared. + request.extensions[_AZURE_AUTH_ORIGIN] = _origin(request.url) + return request @override def _prepare_url(self, url: str) -> httpx2.URL: @@ -335,6 +359,10 @@ def __init__( self._azure_deployment = azure_deployment if azure_endpoint else None self._azure_endpoint = httpx2.URL(azure_endpoint) if azure_endpoint else None + hooks = self._client.event_hooks["request"] + if _strip_azure_api_key_on_redirect not in hooks: + hooks.append(_strip_azure_api_key_on_redirect) + @override def copy( self, @@ -682,6 +710,10 @@ def __init__( self._azure_deployment = azure_deployment if azure_endpoint else None self._azure_endpoint = httpx2.URL(azure_endpoint) if azure_endpoint else None + hooks = self._client.event_hooks["request"] + if _async_strip_azure_api_key_on_redirect not in hooks: + hooks.append(_async_strip_azure_api_key_on_redirect) + @override def copy( self, diff --git a/src/openai/resources/beta/realtime/realtime.py b/src/openai/resources/beta/realtime/realtime.py index bf99da8abc..5365001707 100644 --- a/src/openai/resources/beta/realtime/realtime.py +++ b/src/openai/resources/beta/realtime/realtime.py @@ -360,6 +360,8 @@ async def __aenter__(self) -> AsyncRealtimeConnection: await self.__client._refresh_api_key() auth_headers = self.__client.auth_headers if is_async_azure_client(self.__client): + from ....lib._azure_websocket import _AzureWebSocketConnect as connect + url, auth_headers = await self.__client._configure_realtime(self.__model, extra_query) else: url = self._prepare_url().copy_with( diff --git a/src/openai/resources/realtime/realtime.py b/src/openai/resources/realtime/realtime.py index 929aab6096..62ba3e4e54 100644 --- a/src/openai/resources/realtime/realtime.py +++ b/src/openai/resources/realtime/realtime.py @@ -691,6 +691,8 @@ async def _connect_ws(self, extra_query: Query, extra_headers: Headers) -> Async if self.__call_id is not omit: extra_query = {**extra_query, "call_id": self.__call_id} if is_async_azure_client(self.__client): + from ...lib._azure_websocket import _AzureWebSocketConnect as connect + model = self.__model if not model: raise OpenAIError("`model` is required for Azure Realtime API") diff --git a/tests/lib/test_azure_auth.py b/tests/lib/test_azure_auth.py index c649555ae1..cbeefafe23 100644 --- a/tests/lib/test_azure_auth.py +++ b/tests/lib/test_azure_auth.py @@ -1,7 +1,7 @@ from __future__ import annotations import inspect -from typing import Any +from typing import Any, NoReturn import httpx2 import pytest @@ -204,6 +204,11 @@ async def test_internal_sentinel_is_not_a_conflicting_credential(asynchronous: b @pytest.mark.parametrize("asynchronous", [False, True]) async def test_callable_api_key_refresh_and_copy(asynchronous: bool, monkeypatch: pytest.MonkeyPatch) -> None: + def unexpected_connection(*_args: Any, **_kwargs: Any) -> NoReturn: + pytest.fail("The callable API-key test must not open a network connection") + + monkeypatch.setattr("socket.socket.connect", unexpected_connection) + monkeypatch.setattr("socket.socket.connect_ex", unexpected_connection) requests: list[httpx2.Request] = [] calls = 0 @@ -249,7 +254,9 @@ async def async_connect(*args: Any, **kwargs: Any) -> Any: return connect(*args, **kwargs) monkeypatch.setattr("websockets.sync.client.connect", connect) - monkeypatch.setattr("websockets.asyncio.client.connect", async_connect) + # Azure uses its own async connector. Patching its base class misses + # already-imported subclasses and makes the test depend on import order. + monkeypatch.setattr("openai.lib._azure_websocket._AzureWebSocketConnect", async_connect) for resource in (copied.realtime, copied.beta.realtime): await resolve(resource.connect(model="test-model").enter()) assert websocket_headers == [expected_auth("api_key", f"fake-key-{i}") for i in (4, 5)] diff --git a/tests/lib/test_azure_redirects.py b/tests/lib/test_azure_redirects.py new file mode 100644 index 0000000000..ecacc380c2 --- /dev/null +++ b/tests/lib/test_azure_redirects.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +from typing import Any, Callable +from importlib import import_module +from unittest.mock import AsyncMock, MagicMock + +import httpx2 +import pytest + +from openai import AzureOpenAI, AsyncAzureOpenAI, DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient +from openai._models import FinalRequestOptions +from openai.lib.azure import API_KEY_SENTINEL + +ORIGIN = "https://origin.test" +FAKE_KEY = "fake-azure-redirect-key" +SYNC_BACKENDS = ["httpx", "httpx2", "default"] +ASYNC_BACKENDS = [*SYNC_BACKENDS, "aiohttp"] +TARGETS = [ + ("/final", True), + ("https://ORIGIN.test:443/final", True), + ("https://other.test/final", False), + ("//other.test/final", False), + ("https://origin.test:444/final", False), + ("http://origin.test/final", False), + ("https://origin.test@other.test/final", False), +] + + +def http_module(backend: str) -> Any: + return pytest.importorskip("httpx") if backend == "httpx" else httpx2 + + +def azure_options() -> dict[str, Any]: + return dict(api_key=FAKE_KEY, azure_endpoint=ORIGIN, api_version="2024-01-01", max_retries=0) + + +def mock_aiohttp(monkeypatch: pytest.MonkeyPatch, handler: Callable[[Any], Any]) -> None: + # Exercise the real vendored adapter, but never open a network connection. + aiohttp = import_module("aiohttp") + + def request(_session: Any, method: str, url: str, **kwargs: Any) -> Any: + assert kwargs["allow_redirects"] is False + result = handler(httpx2.Request(method, url, headers=kwargs["headers"], content=kwargs["data"])) + + async def chunks(_size: int) -> Any: + yield result.content + + response = MagicMock(spec=aiohttp.ClientResponse) + response.status = result.status_code + response.reason = result.reason_phrase + response.raw_headers = result.headers.raw + response.content.iter_chunked.side_effect = chunks + response.__aexit__ = AsyncMock() + context = MagicMock() + context.__aenter__ = AsyncMock(return_value=response) + return context + + monkeypatch.setattr(aiohttp.ClientSession, "request", request) + + +def make_http_client( + backend: str, + handler: Callable[[Any], Any], + *, + asynchronous: bool = False, + monkeypatch: pytest.MonkeyPatch | None = None, + **kwargs: Any, +) -> Any: + module = http_module(backend) + if backend == "aiohttp": + assert monkeypatch is not None + mock_aiohttp(monkeypatch, handler) + return DefaultAioHttpClient(trust_env=False, **kwargs) + if backend == "default": + cls = DefaultAsyncHttpxClient if asynchronous else DefaultHttpxClient + else: + cls = module.AsyncClient if asynchronous else module.Client + kwargs.setdefault("follow_redirects", True) + return cls(transport=module.MockTransport(handler), trust_env=False, **kwargs) + + +def redirect_handler(backend: str, targets: list[str], seen: list[Any], status: int = 307) -> Callable[[Any], Any]: + module = http_module(backend) + + def handler(request: Any) -> Any: + seen.append(request) + if len(seen) <= len(targets): + return module.Response(status, headers={"location": targets[len(seen) - 1]}, content=b"redirect") + return module.Response(200, json={"object": "list", "data": []}) + + return handler + + +@pytest.mark.parametrize("backend", SYNC_BACKENDS) +@pytest.mark.parametrize("status", [301, 302, 303, 307, 308]) +@pytest.mark.parametrize("target,retained", TARGETS) +def test_sync_redirect_origin(backend: str, status: int, target: str, retained: bool) -> None: + seen: list[Any] = [] + transport = make_http_client(backend, redirect_handler(backend, [target], seen, status)) + with AzureOpenAI(**azure_options(), http_client=transport) as client: + result = client.models.with_raw_response.list(extra_headers={"Authorization": "Bearer fake-token"}) + assert len(result.http_response.history) == 1 + assert len(seen) == 2 + assert seen[0].headers["api-key"] == FAKE_KEY + assert seen[1].headers.get("api-key") == (FAKE_KEY if retained else None) + assert ("authorization" in seen[1].headers) == retained + + +@pytest.mark.parametrize("backend", ASYNC_BACKENDS) +@pytest.mark.parametrize("status", [301, 302, 303, 307, 308]) +@pytest.mark.parametrize("target,retained", TARGETS) +async def test_async_redirect_origin( + backend: str, status: int, target: str, retained: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + seen: list[Any] = [] + transport = make_http_client( + backend, redirect_handler(backend, [target], seen, status), asynchronous=True, monkeypatch=monkeypatch + ) + async with AsyncAzureOpenAI(**azure_options(), http_client=transport) as client: + result = await client.models.with_raw_response.list(extra_headers={"Authorization": "Bearer fake-token"}) + assert len(result.http_response.history) == 1 + assert len(seen) == 2 + assert seen[0].headers["api-key"] == FAKE_KEY + assert seen[1].headers.get("api-key") == (FAKE_KEY if retained else None) + assert ("authorization" in seen[1].headers) == retained + + +@pytest.mark.parametrize("backend", SYNC_BACKENDS) +@pytest.mark.parametrize("follow", [False, True]) +def test_sync_redirect_options_and_chain(backend: str, follow: bool) -> None: + seen: list[Any] = [] + targets = ["/same", "https://other.test/final", ORIGIN + "/return"] + transport = make_http_client(backend, redirect_handler(backend, targets, seen), follow_redirects=not follow) + with AzureOpenAI(**azure_options(), http_client=transport) as client: + response = client._send_request( + client._build_request(FinalRequestOptions(method="get", url="/start")), + stream=True, + follow_redirects=follow, + ) + assert response.status_code == (200 if follow else 307) + assert len(response.history) == (3 if follow else 0) + response.close() + assert [r.headers.get("api-key") for r in seen] == ([FAKE_KEY, FAKE_KEY, None, None] if follow else [FAKE_KEY]) + + +@pytest.mark.parametrize("backend", ASYNC_BACKENDS) +@pytest.mark.parametrize("follow", [False, True]) +async def test_async_redirect_options_and_chain(backend: str, follow: bool, monkeypatch: pytest.MonkeyPatch) -> None: + seen: list[Any] = [] + targets = ["/same", "https://other.test/final", ORIGIN + "/return"] + transport = make_http_client( + backend, + redirect_handler(backend, targets, seen), + asynchronous=True, + monkeypatch=monkeypatch, + follow_redirects=not follow, + ) + async with AsyncAzureOpenAI(**azure_options(), http_client=transport) as client: + response = await client._send_request( + client._build_request(FinalRequestOptions(method="get", url="/start")), + stream=True, + follow_redirects=follow, + ) + assert response.status_code == (200 if follow else 307) + assert len(response.history) == (3 if follow else 0) + await response.aclose() + assert [r.headers.get("api-key") for r in seen] == ([FAKE_KEY, FAKE_KEY, None, None] if follow else [FAKE_KEY]) + + +@pytest.mark.parametrize("backend", SYNC_BACKENDS) +def test_sync_shared_client_and_explicit_credentials(backend: str) -> None: + seen: list[Any] = [] + transport = make_http_client(backend, redirect_handler(backend, ["https://other.test/final"], seen)) + with AzureOpenAI(**azure_options(), http_client=transport) as client: + copy = client.copy() + assert len(transport.event_hooks["request"]) == 1 + copy.models.list(extra_headers={"API-KEY": "fake-explicit-key"}) + assert "fake-explicit-key" in seen[0].headers.get_list("api-key") + assert "api-key" not in seen[1].headers + transport.get("https://other.test/direct", headers={"api-key": "fake-unrelated-key"}) + assert seen[-1].headers["api-key"] == "fake-unrelated-key" + + +@pytest.mark.parametrize("backend", ASYNC_BACKENDS) +async def test_async_shared_client_and_explicit_credentials(backend: str, monkeypatch: pytest.MonkeyPatch) -> None: + seen: list[Any] = [] + transport = make_http_client( + backend, + redirect_handler(backend, ["https://other.test/final"], seen), + asynchronous=True, + monkeypatch=monkeypatch, + ) + async with AsyncAzureOpenAI(**azure_options(), http_client=transport) as client: + copy = client.copy() + assert len(transport.event_hooks["request"]) == 1 + await copy.models.list(extra_headers={"API-KEY": "fake-explicit-key"}) + assert "fake-explicit-key" in seen[0].headers.get_list("api-key") + assert "api-key" not in seen[1].headers + await transport.get("https://other.test/direct", headers={"api-key": "fake-unrelated-key"}) + assert seen[-1].headers["api-key"] == "fake-unrelated-key" + + +@pytest.mark.parametrize("backend", SYNC_BACKENDS) +@pytest.mark.parametrize("target,retained", TARGETS[:3]) +def test_sync_bearer_authentication(backend: str, target: str, retained: bool) -> None: + seen: list[Any] = [] + transport = make_http_client(backend, redirect_handler(backend, [target], seen)) + options = azure_options() + options["api_key"] = API_KEY_SENTINEL + with AzureOpenAI(**options, azure_ad_token_provider=lambda: "fake-ad-token", http_client=transport) as client: + client.models.list() + assert seen[0].headers["authorization"] == "Bearer fake-ad-token" + assert seen[1].headers.get("authorization") == ("Bearer fake-ad-token" if retained else None) + assert all("api-key" not in request.headers for request in seen) + + +@pytest.mark.parametrize("backend", ASYNC_BACKENDS) +@pytest.mark.parametrize("target,retained", TARGETS[:3]) +async def test_async_bearer_authentication( + backend: str, target: str, retained: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + seen: list[Any] = [] + transport = make_http_client( + backend, redirect_handler(backend, [target], seen), asynchronous=True, monkeypatch=monkeypatch + ) + options = azure_options() + options["api_key"] = API_KEY_SENTINEL + + async def token() -> str: + return "fake-ad-token" + + async with AsyncAzureOpenAI(**options, azure_ad_token_provider=token, http_client=transport) as client: + await client.models.list() + assert seen[0].headers["authorization"] == "Bearer fake-ad-token" + assert seen[1].headers.get("authorization") == ("Bearer fake-ad-token" if retained else None) + assert all("api-key" not in request.headers for request in seen) diff --git a/tests/lib/test_azure_websocket_redirects.py b/tests/lib/test_azure_websocket_redirects.py new file mode 100644 index 0000000000..16a54cdcb5 --- /dev/null +++ b/tests/lib/test_azure_websocket_redirects.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import threading +from typing import Any +from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler +from typing_extensions import override + +import httpx2 +import pytest +from websockets.http11 import Request, Response +from websockets.exceptions import InvalidStatus, SecurityError +from websockets.asyncio.client import connect +from websockets.asyncio.server import ServerConnection, serve +from websockets.datastructures import Headers + +from openai import AzureOpenAI, AsyncAzureOpenAI +from openai.lib.azure import API_KEY_SENTINEL +from openai.lib._azure_websocket import _AzureWebSocketConnect + + +def azure_options(base_url: str, bearer: bool) -> dict[str, Any]: + return { + "api_key": API_KEY_SENTINEL if bearer else "fake-websocket-key", + "azure_ad_token": "fake-websocket-token" if bearer else None, + "azure_endpoint": "https://origin.test", + "websocket_base_url": base_url, + "api_version": "2024-01-01", + "max_retries": 0, + } + + +@pytest.fixture(autouse=True) +def no_proxies(monkeypatch: pytest.MonkeyPatch) -> None: + for scheme in ("all", "http", "https", "ws", "wss", "socks"): + monkeypatch.delenv(f"{scheme}_proxy", raising=False) + monkeypatch.delenv(f"{scheme.upper()}_PROXY", raising=False) + # Also bypass proxies discovered from platform settings, not just the env. + monkeypatch.setenv("no_proxy", "*") + monkeypatch.setenv("NO_PROXY", "*") + + +@pytest.mark.parametrize("beta", [False, True], ids=["stable", "beta"]) +@pytest.mark.parametrize("bearer", [False, True], ids=["api-key", "bearer"]) +@pytest.mark.parametrize("redirect", ["none", "same-origin", "cross-origin"]) +async def test_async_azure_websocket_redirects(beta: bool, bearer: bool, redirect: str) -> None: + if redirect == "same-origin" and not hasattr(connect, "process_redirect"): + pytest.skip("This websockets version does not follow handshake redirects") + + source_headers: list[Headers] = [] + target_headers: list[Headers] = [] + + async def connected(connection: ServerConnection) -> None: + await connection.wait_closed() + + def record_target(_connection: ServerConnection, request: Request) -> None: + target_headers.append(request.headers) + + async with serve(connected, "127.0.0.1", 0, process_request=record_target) as target: + target_url = f"ws://127.0.0.1:{next(iter(target.sockets)).getsockname()[1]}/final" + + def process_source(_connection: ServerConnection, request: Request) -> Response | None: + source_headers.append(request.headers) + if redirect != "none" and request.path.startswith("/realtime?"): + location = "/final" if redirect == "same-origin" else target_url + return Response(302, "Found", Headers({"Location": location})) + return None + + async with serve(connected, "127.0.0.1", 0, process_request=process_source) as source: + base_url = f"ws://127.0.0.1:{next(iter(source.sockets)).getsockname()[1]}" + async with AsyncAzureOpenAI( + **azure_options(base_url, bearer), http_client=httpx2.AsyncClient(trust_env=False) + ) as client: + resource = client.beta.realtime if beta else client.realtime + if redirect == "cross-origin": + expected = ( + SecurityError if hasattr(connect, "process_redirect") else (InvalidStatus, ConnectionError) + ) + with pytest.raises(expected): + async with resource.connect(model="fake-model"): + pytest.fail("Cross-origin redirect must not connect") + else: + async with resource.connect(model="fake-model"): + pass + + assert len(source_headers) == (2 if redirect == "same-origin" else 1) + assert target_headers == [] + for headers in source_headers: + assert headers.get("api-key") == (None if bearer else "fake-websocket-key") + assert headers.get("Authorization") == ("Bearer fake-websocket-token" if bearer else None) + + +@pytest.mark.parametrize("beta", [False, True], ids=["stable", "beta"]) +@pytest.mark.parametrize("bearer", [False, True], ids=["api-key", "bearer"]) +def test_sync_azure_websocket_does_not_follow_redirects(beta: bool, bearer: bool) -> None: + received: list[tuple[str, str | None, str | None]] = [] + + class RedirectHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self) -> None: + received.append((self.path, self.headers.get("api-key"), self.headers.get("Authorization"))) + self.send_response(302) + self.send_header("Location", "/final") + self.send_header("Content-Length", "0") + self.end_headers() + + @override + def log_message(self, *_args: object, **_kwargs: object) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), RedirectHandler) + thread = threading.Thread(target=server.serve_forever) + thread.start() + try: + with AzureOpenAI( + **azure_options(f"ws://127.0.0.1:{server.server_port}", bearer), + http_client=httpx2.Client(trust_env=False), + ) as client: + resource = client.beta.realtime if beta else client.realtime + with pytest.raises(InvalidStatus): + with resource.connect(model="fake-model"): + pytest.fail("Synchronous WebSocket redirects must not connect") + finally: + server.shutdown() + thread.join() + server.server_close() + + assert len(received) == 1 + assert received[0][0].startswith("/realtime?") + assert received[0][1:] == ( + None if bearer else "fake-websocket-key", + "Bearer fake-websocket-token" if bearer else None, + ) + + +@pytest.mark.skipif(not hasattr(connect, "process_redirect"), reason="No automatic handshake redirects") +@pytest.mark.parametrize( + "target,allowed", + [ + ("/final", True), + ("wss://ORIGIN.test:443/final", True), + ("wss://other.test/final", False), + ("wss://origin.test:444/final", False), + ("ws://origin.test/final", False), + ], +) +def test_websocket_redirect_origin(target: str, allowed: bool) -> None: + connection = _AzureWebSocketConnect("wss://origin.test/realtime") + result = connection.process_redirect(InvalidStatus(Response(302, "Found", Headers({"Location": target})))) + assert isinstance(result, str) is allowed + if not allowed: + assert isinstance(result, SecurityError) + + +@pytest.mark.skipif(not hasattr(connect, "process_redirect"), reason="No automatic handshake redirects") +def test_websocket_non_redirect_error_is_preserved() -> None: + error = InvalidStatus(Response(401, "Unauthorized", Headers())) + assert _AzureWebSocketConnect("wss://origin.test/realtime").process_redirect(error) is error