Merge 0c361389e7 into 795e622c63

litellm/proxy/auth/user_api_key_auth.py

@@ -1122,6 +1122,8 @@ def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) ->
 def get_bearer_token(
     api_key: str,
 ) -> str:
+    if api_key.startswith("Token "):  # Deepgram SDKs send their key as `Token <key>`
+        return api_key[len("Token ") :]
     if api_key.startswith("Bearer "):  # ensure Bearer token passed in
         api_key = api_key.replace("Bearer ", "")  # extract the token
     elif api_key.startswith("Basic "):
@@ -1246,13 +1248,11 @@ async def user_api_key_auth_websocket_for_model(
             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 "):
+        api_key = get_bearer_token(authorization).strip()
+        if not api_key:
             await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
             raise HTTPException(status_code=403, detail="Invalid Authorization header format")
 
-        api_key = authorization[len("Bearer ") :].strip()
-
     # 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:

tests/unit/proxy/auth/test_user_api_key_auth.py

@@ -2282,3 +2282,95 @@ async def test_non_object_metadata_is_served_like_unset_mappings_whatever_the_ba
     )
     assert mapped == unset
     assert mapped[1] == "usecase-deployment"
+
+
+def _websocket_authorized_by(sent: list[Message], authorization: str) -> WebSocket:
+    async def receive() -> Message:
+        return {"type": "websocket.connect"}
+
+    async def send(message: Message) -> None:
+        sent.append(message)
+
+    return WebSocket(
+        {
+            "type": "websocket",
+            "path": "/deepgram/v1/listen",
+            "query_string": b"model=nova-3",
+            "headers": [(b"authorization", authorization.encode())],
+        },
+        receive,
+        send,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "authorization",
+    [
+        pytest.param("Bearer {key}", id="bearer"),
+        pytest.param("bearer {key}", id="deepgram sdk access_token mode"),
+        pytest.param("Token {key}", id="deepgram sdk api_key mode"),
+        pytest.param("Token {key} ", id="token with trailing space"),
+    ],
+)
+async def test_websocket_auth_finds_the_virtual_key_under_each_scheme_http_auth_accepts(
+    monkeypatch: pytest.MonkeyPatch, authorization: str
+) -> None:
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket_for_model
+    from litellm.proxy.proxy_server import hash_token
+
+    user_key: Final = "sk-websocket-key-for-nova"
+    _serve_virtual_key(monkeypatch, user_key, UserAPIKeyAuth(token=hash_token(user_key), models=["nova-3"]))
+    sent: Final[list[Message]] = []
+
+    authorized: Final = await user_api_key_auth_websocket_for_model(
+        _websocket_authorized_by(sent, authorization.format(key=user_key)), "nova-3"
+    )
+
+    assert authorized.token == hash_token(user_key)
+    assert sent == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "authorization",
+    [
+        pytest.param("Token{key}", id="no space after the scheme"),
+        pytest.param("Digest {key}", id="unknown scheme"),
+        pytest.param("Token ", id="scheme without a key"),
+    ],
+)
+async def test_websocket_auth_refuses_an_unparseable_authorization_with_a_policy_close(
+    monkeypatch: pytest.MonkeyPatch, authorization: str
+) -> None:
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket_for_model
+    from litellm.proxy.proxy_server import hash_token
+
+    user_key: Final = "sk-websocket-key-for-nova"
+    _serve_virtual_key(monkeypatch, user_key, UserAPIKeyAuth(token=hash_token(user_key), models=["nova-3"]))
+    sent: Final[list[Message]] = []
+
+    with pytest.raises(HTTPException) as refusal:
+        await user_api_key_auth_websocket_for_model(
+            _websocket_authorized_by(sent, authorization.format(key=user_key)), "nova-3"
+        )
+
+    assert refusal.value.status_code == 403
+    assert refusal.value.detail == "Invalid Authorization header format"
+    assert sent == [{"type": "websocket.close", "code": status.WS_1008_POLICY_VIOLATION, "reason": ""}]
+
+
+@pytest.mark.asyncio
+async def test_http_auth_finds_the_virtual_key_in_a_token_authorization(monkeypatch: pytest.MonkeyPatch) -> None:
+    from fastapi import Request as HTTPRequest
+
+    from litellm.proxy.proxy_server import hash_token
+
+    user_key: Final = "sk-http-key-for-nova"
+    _serve_virtual_key(monkeypatch, user_key, UserAPIKeyAuth(token=hash_token(user_key)))
+    request: Final = HTTPRequest(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []})
+    request._url = URL(url="/chat/completions")  # pyright: ignore[reportPrivateUsage]  # Starlette caches the URL here
+
+    authorized: Final = await user_api_key_auth(request=request, api_key=f"Token {user_key}")
+
+    assert authorized.token == hash_token(user_key)

tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py

@@ -472,3 +472,35 @@ def test_deepgram_listen_echoes_the_browser_subprotocol_that_carries_the_litellm
         MappingProxyType({"Authorization": "Token dg-provider-key"})
     ]
     assert [call.forward_headers for call in relay.calls] == [False]
+
+
+def test_deepgram_listen_accepts_the_litellm_key_in_a_deepgram_sdk_token_header(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Deepgram SDKs send ``Authorization: Token <key>``, so pointing one at the proxy with a LiteLLM key as its API
+    key must authenticate through the real key auth path and still send only the server's key upstream"""
+    monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
+    monkeypatch.setattr(litellm, "max_budget", 0.0)
+    cache: Final = asyncio.run(_cache_restricted_key("sk-deepgram-sdk", ["nova-3"]))
+    relay: Final = _FakeRelay()
+    client: Final = TestClient(_app_with_relay(relay))
+
+    with (
+        patch(GET_CREDENTIALS, return_value="dg-provider-key"),
+        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(),
+            user_api_key_cache=cache,
+            llm_model_list=None,
+            llm_router=None,
+        ),
+    ):
+        with client.websocket_connect(
+            "/deepgram/listen?model=nova-3", headers={"Authorization": "Token sk-deepgram-sdk"}
+        ):
+            pass
+
+    assert [(call.target, call.custom_headers) for call in relay.calls] == [
+        ("wss://api.deepgram.com/v1/listen?model=nova-3", MappingProxyType({"Authorization": "Token dg-provider-key"}))
+    ]

tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py

@@ -4648,7 +4648,7 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
         assert self.VKEY not in " ".join(f"{name}:{value}" for name, value in forwarded.items())
 
     @pytest.mark.asyncio
-    @pytest.mark.parametrize("scheme", ["Bearer", "bearer", "Basic"])
+    @pytest.mark.parametrize("scheme", ["Bearer", "bearer", "Basic", "Token"])
     async def test_virtual_key_echoed_in_authorization_with_any_scheme_is_stripped(self, monkeypatch, scheme):
         raised, forwarded = await self._run(
             monkeypatch,