Merge 91aeff9300 into bc18e10f6b

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"

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])

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