diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 27fa4731baf..aa9d33d445b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1200,7 +1200,11 @@ def _has_required_baggage_tag_mapping() -> bool: async def user_api_key_auth_websocket_for_model( - websocket: WebSocket, model: str | None, *, close_on_auth_error: bool = True + websocket: WebSocket, + model: str | None, + *, + close_on_auth_error: bool = True, + api_key_headers: tuple[str, ...] = ("api-key",), ) -> UserAPIKeyAuth: ws_scope: Final = websocket.scope or {} scope_headers: Final = list(ws_scope.get("headers") or []) @@ -1232,7 +1236,7 @@ async def user_api_key_auth_websocket_for_model( authorization: Final = websocket.headers.get("authorization") # If no Authorization header, try the api-key header if not authorization: - api_key = websocket.headers.get("api-key") + api_key = next((value for header in api_key_headers if (value := websocket.headers.get(header))), None) if not api_key: # Try extracting from WebSocket subprotocol (browser clients) for protocol in websocket.headers.get("sec-websocket-protocol", "").split(","): diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 60a8dd73a34..c5da44bcc2f 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -5,6 +5,7 @@ import litellm.proxy import litellm.proxy.proxy_server +from types import SimpleNamespace from typing import Dict, Final, List, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -931,6 +932,57 @@ async def test_user_api_key_auth_live_websocket_ignores_query_model(): mock_user_api_key_auth.assert_called_once() +def _header_websocket(headers: dict[str, str]) -> MagicMock: + websocket = MagicMock(spec=WebSocket) + websocket.headers = headers + websocket.scope = {"headers": [(name.encode(), value.encode()) for name, value in headers.items()]} + websocket.url = URL(url="/elevenlabs/v1/text-to-speech/voice1/stream-input") + websocket.close = AsyncMock() + return websocket + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("headers", "api_key_headers", "expected"), + [ + pytest.param({"api-key": "sk-a"}, ("api-key",), "Bearer sk-a", id="default header"), + pytest.param({"xi-api-key": "sk-x"}, ("api-key", "xi-api-key"), "Bearer sk-x", id="route header"), + pytest.param( + {"api-key": "sk-a", "xi-api-key": "sk-x"}, ("api-key", "xi-api-key"), "Bearer sk-a", id="first header wins" + ), + pytest.param( + {"api-key": "", "xi-api-key": "sk-x"}, ("api-key", "xi-api-key"), "Bearer sk-x", id="blank skipped" + ), + ], +) +async def test_websocket_auth_reads_the_key_from_the_headers_the_route_names(headers, api_key_headers, expected): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket_for_model + + async def authenticate(request, api_key): + return SimpleNamespace(presented_key=api_key) + + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", side_effect=authenticate): + authorized = await user_api_key_auth_websocket_for_model( + _header_websocket(headers), "eleven_flash_v2_5", api_key_headers=api_key_headers + ) + + assert authorized.presented_key == expected + + +@pytest.mark.asyncio +async def test_websocket_auth_refuses_a_key_in_a_header_the_route_did_not_name(): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket_for_model + + websocket = _header_websocket({"xi-api-key": "sk-x"}) + with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth: + with pytest.raises(HTTPException) as refusal: + await user_api_key_auth_websocket_for_model(websocket, "gpt-realtime") + + assert refusal.value.status_code == 403 + websocket.close.assert_awaited_once_with(code=status.WS_1008_POLICY_VIOLATION) + mock_user_api_key_auth.assert_not_called() + + @pytest.mark.asyncio async def test_accepted_websocket_reauthorizes_the_frame_model_without_closing(): from litellm.proxy._types import ProxyException