diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 27fa4731baf..e9e36c30eed 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1240,9 +1240,6 @@ async def user_api_key_auth_websocket_for_model( if protocol.startswith("openai-insecure-api-key."): api_key = protocol[len("openai-insecure-api-key.") :] break - if not api_key: - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="No API key provided") else: # Extract the API key from the Bearer token if not authorization.startswith("Bearer "): @@ -1254,7 +1251,10 @@ async def user_api_key_auth_websocket_for_model( # Call user_api_key_auth with the extracted API key # Note: You'll need to modify this to work with WebSocket context if needed try: - authorized: Final = await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") + authorized: Final = await user_api_key_auth( + request=request, + api_key=f"Bearer {api_key}" if api_key else None, # pyright: ignore[reportArgumentType] # None = no key + ) if _has_required_baggage_tag_mapping(): raise HTTPException( 403, "Required baggage tag mappings need HTTP inference, WebSocket inference is disabled" diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 4cdd7fefd28..0fe3c3ff918 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -6,7 +6,7 @@ import subprocess import sys import time from collections.abc import Mapping, Sequence -from contextlib import contextmanager, nullcontext +from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import datetime, timedelta, timezone from functools import partial from io import BytesIO @@ -82,6 +82,7 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth_builder, get_api_key, user_api_key_auth, + user_api_key_auth_websocket, user_api_key_auth_websocket_for_model, ) from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata @@ -12438,6 +12439,96 @@ async def test_auto_register_mapping_insert_emits_a_postgres_insert_event_for_th ) +def _websocket_for_auth(headers: dict[str, str]) -> tuple[MagicMock, AsyncMock]: + from fastapi import WebSocket + from starlette.datastructures import URL + + close: Final = AsyncMock() + websocket: Final = MagicMock(spec=WebSocket) + websocket.query_params = {"model": "test-model"} + websocket.headers = headers + websocket.scope = { + "type": "websocket", + "path": "/v1/responses", + "headers": [(name.lower().encode(), value.encode()) for name, value in headers.items()], + } + websocket.url = URL(url="/v1/responses") + websocket.close = close + return websocket, close + + +def _proxy_state(master_key: str | None) -> AbstractContextManager[object]: + return patch.multiple( # test-quality-ok: master_key is a proxy_server module global with no injection seam + "litellm.proxy.proxy_server", + master_key=master_key, + prisma_client=None, + user_custom_auth=None, + general_settings={}, + llm_model_list=[], + llm_router=None, + jwt_handler=None, + open_telemetry_logger=None, + ) + + +@pytest.mark.parametrize( + "headers", + [ + pytest.param({}, id="no headers at all"), + pytest.param({"sec-websocket-protocol": "realtime"}, id="subprotocol carrying no key"), + ], +) +@pytest.mark.asyncio +async def test_websocket_auth_without_master_key_accepts_a_keyless_client(headers: dict[str, str]) -> None: + websocket, _ = _websocket_for_auth(headers) + + with _proxy_state(master_key=None): + result: Final = await user_api_key_auth_websocket(websocket) + + assert result.user_role == LitellmUserRoles.INTERNAL_USER + assert result.api_key is None + + +@pytest.mark.parametrize( + "headers", + [ + pytest.param({"authorization": "Bearer sk-master-key"}, id="bearer token"), + pytest.param({"authorization": "Bearer sk-master-key"}, id="bearer token with extra space"), + pytest.param({"api-key": "sk-master-key"}, id="api-key header"), + pytest.param( + {"sec-websocket-protocol": "realtime, openai-insecure-api-key.sk-master-key"}, + id="browser subprotocol", + ), + ], +) +@pytest.mark.asyncio +async def test_websocket_auth_accepts_the_master_key_from_every_key_source(headers: dict[str, str]) -> None: + websocket, _ = _websocket_for_auth(headers) + + with _proxy_state(master_key="sk-master-key"): + result: Final = await user_api_key_auth_websocket(websocket) + + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + +@pytest.mark.parametrize( + "headers", + [ + pytest.param({}, id="no key"), + pytest.param({"authorization": "Bearer sk-wrong-key"}, id="wrong key"), + ], +) +@pytest.mark.asyncio +async def test_websocket_auth_with_master_key_refuses_a_client_without_it(headers: dict[str, str]) -> None: + websocket, close = _websocket_for_auth(headers) + + with _proxy_state(master_key="sk-master-key"), pytest.raises(HTTPException) as exc_info: + await user_api_key_auth_websocket(websocket) + + assert exc_info.value.status_code == status.HTTP_403_FORBIDDEN + close.assert_awaited_once_with(code=status.WS_1008_POLICY_VIOLATION) + + @pytest.mark.asyncio @pytest.mark.parametrize("route", ["/lens", "/lens/traces", "/lens/settings"]) @pytest.mark.parametrize("jwt_override", [False, True]) diff --git a/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 404ed0f5800..a722d3b609c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -303,7 +303,12 @@ def test_deepgram_listen_rejects_connections_without_a_litellm_key(): relay = _FakeRelay() client = TestClient(_app_with_relay(relay)) - with patch(GET_CREDENTIALS, return_value="dg-provider-key") as get_credentials: + with ( + patch(GET_CREDENTIALS, return_value="dg-provider-key") as get_credentials, + patch.multiple( # test-quality-ok: the real key auth path reads these proxy_server globals and has no injection seam + "litellm.proxy.proxy_server", master_key="sk-master", prisma_client=MagicMock() + ), + ): with pytest.raises(WebSocketDisconnect) as disconnect: with client.websocket_connect("/deepgram/v1/listen?model=nova-3"): pass