From 88d6a61d7428646493d9ab47608254cedbaa266f Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Mon, 26 Oct 2026 15:46:43 -0700 Subject: [PATCH 1/3] feat(auth): accept a Token Authorization header on the Deepgram WebSocket route Deepgram SDKs send their API key as `Authorization: Token `, so pointing one at the proxy with a LiteLLM key as its API key was refused at the handshake. WebSocket key auth now takes a per-route tuple of accepted schemes, Bearer only by default, and the Deepgram listen route opts in to Bearer and Token. Every other WebSocket route still accepts Bearer only Refs LIT-10070 --- litellm/proxy/auth/user_api_key_auth.py | 17 +++- .../llm_passthrough_endpoints.py | 7 +- .../unit/proxy/auth/test_user_api_key_auth.py | 84 +++++++++++++++++++ .../test_deepgram_ws_passthrough_routes.py | 46 ++++++++++ 4 files changed, 149 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b1e8bdbe787..4d04f15eb82 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1202,8 +1202,14 @@ 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, + accepted_authorization_schemes: tuple[str, ...] = ("Bearer",), ) -> UserAPIKeyAuth: + """``accepted_authorization_schemes`` are the ``Authorization`` schemes that may carry the LiteLLM key. A route + adds one only when a provider's own SDK sends it to the same path""" ws_scope: Final = websocket.scope or {} scope_headers: Final = list(ws_scope.get("headers") or []) # ``get_request_route`` falls back to ``request.url.path`` when @@ -1246,12 +1252,15 @@ 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 "): + scheme_prefix: Final = next( + (f"{scheme} " for scheme in accepted_authorization_schemes if authorization.startswith(f"{scheme} ")), + None, + ) + if scheme_prefix is None: 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() + api_key = authorization[len(scheme_prefix) :].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 diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9eef4f0ef9d..63b3c300e7f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3450,9 +3450,14 @@ _DEEPGRAM_WS_UNPRICED_REASON: Final = ( ) +_DEEPGRAM_WS_AUTHORIZATION_SCHEMES: Final = ("Bearer", "Token") + + async def deepgram_listen_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth: return await user_api_key_auth_websocket_for_model( - websocket, model=deepgram_listen_requested_model(websocket.url.query) + websocket, + model=deepgram_listen_requested_model(websocket.url.query), + accepted_authorization_schemes=_DEEPGRAM_WS_AUTHORIZATION_SCHEMES, ) 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..cef7fd17382 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -2282,3 +2282,87 @@ 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, + ) + + +DEEPGRAM_SDK_SCHEMES: Final = ("Bearer", "Token") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("authorization", "accepted_schemes"), + [ + pytest.param("Bearer {key}", ("Bearer",), id="bearer by default"), + pytest.param("Bearer {key}", DEEPGRAM_SDK_SCHEMES, id="bearer where token is also accepted"), + pytest.param("Token {key}", DEEPGRAM_SDK_SCHEMES, id="token where the route accepts it"), + pytest.param("Token {key} ", DEEPGRAM_SDK_SCHEMES, id="token with trailing space"), + ], +) +async def test_websocket_auth_finds_the_virtual_key_under_each_accepted_scheme( + monkeypatch: pytest.MonkeyPatch, authorization: str, accepted_schemes: tuple[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", + accepted_authorization_schemes=accepted_schemes, + ) + + assert authorized.token == hash_token(user_key) + assert sent == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("authorization", "accepted_schemes"), + [ + pytest.param("Token {key}", ("Bearer",), id="token where the route did not opt in"), + pytest.param("token {key}", DEEPGRAM_SDK_SCHEMES, id="scheme is case sensitive like bearer"), + pytest.param("Token{key}", DEEPGRAM_SDK_SCHEMES, id="no space after the scheme"), + pytest.param("Basic {key}", DEEPGRAM_SDK_SCHEMES, id="scheme not in the list"), + ], +) +async def test_websocket_auth_refuses_other_schemes_with_a_policy_close( + monkeypatch: pytest.MonkeyPatch, authorization: str, accepted_schemes: tuple[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", + accepted_authorization_schemes=accepted_schemes, + ) + + 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": ""}] 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..19418c4c349 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 @@ -472,3 +472,49 @@ 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 ``, 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"})) + ] + + +def test_a_non_deepgram_websocket_route_still_refuses_a_token_header() -> None: + relay: Final = _FakeRelay() + client: Final = TestClient(_app_with_relay(relay)) + + with pytest.raises(WebSocketDisconnect) as disconnect: + with client.websocket_connect( + "/openai/v1/realtime?model=gpt-4o-realtime-preview", headers={"Authorization": "Token sk-litellm"} + ): + pass + + assert disconnect.value.code == 1008 + assert relay.calls == [] From 0ce6fabe13465a0a0e089c7dbe0c47721cc85b9d Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Tue, 27 Oct 2026 12:37:50 -0700 Subject: [PATCH 2/3] refactor(auth): read WebSocket keys with the shared Authorization parser WebSocket auth had its own Bearer-only parser, so a Deepgram SDK sending Token, or bearer in its access_token mode, was refused even though HTTP auth accepted bearer. The socket path now uses get_bearer_token, which also learns Token, and the per-route scheme list is gone Every route now accepts Token, and the passthrough leak check strips a caller key echoed under it because it reuses the same parser --- litellm/proxy/auth/user_api_key_auth.py | 19 ++----- .../llm_passthrough_endpoints.py | 7 +-- .../unit/proxy/auth/test_user_api_key_auth.py | 54 +++++++++++-------- .../test_deepgram_ws_passthrough_routes.py | 14 ----- .../test_llm_pass_through_endpoints.py | 2 +- 5 files changed, 38 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4d04f15eb82..a4997f7f223 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/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 "): + 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 "): @@ -1202,14 +1204,8 @@ 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, - accepted_authorization_schemes: tuple[str, ...] = ("Bearer",), + websocket: WebSocket, model: str | None, *, close_on_auth_error: bool = True ) -> UserAPIKeyAuth: - """``accepted_authorization_schemes`` are the ``Authorization`` schemes that may carry the LiteLLM key. A route - adds one only when a provider's own SDK sends it to the same path""" ws_scope: Final = websocket.scope or {} scope_headers: Final = list(ws_scope.get("headers") or []) # ``get_request_route`` falls back to ``request.url.path`` when @@ -1252,16 +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: - scheme_prefix: Final = next( - (f"{scheme} " for scheme in accepted_authorization_schemes if authorization.startswith(f"{scheme} ")), - None, - ) - if scheme_prefix is None: + 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(scheme_prefix) :].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: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 63b3c300e7f..9eef4f0ef9d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3450,14 +3450,9 @@ _DEEPGRAM_WS_UNPRICED_REASON: Final = ( ) -_DEEPGRAM_WS_AUTHORIZATION_SCHEMES: Final = ("Bearer", "Token") - - async def deepgram_listen_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth: return await user_api_key_auth_websocket_for_model( - websocket, - model=deepgram_listen_requested_model(websocket.url.query), - accepted_authorization_schemes=_DEEPGRAM_WS_AUTHORIZATION_SCHEMES, + websocket, model=deepgram_listen_requested_model(websocket.url.query) ) 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 cef7fd17382..b30ac874136 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -2303,21 +2303,18 @@ def _websocket_authorized_by(sent: list[Message], authorization: str) -> WebSock ) -DEEPGRAM_SDK_SCHEMES: Final = ("Bearer", "Token") - - @pytest.mark.asyncio @pytest.mark.parametrize( - ("authorization", "accepted_schemes"), + "authorization", [ - pytest.param("Bearer {key}", ("Bearer",), id="bearer by default"), - pytest.param("Bearer {key}", DEEPGRAM_SDK_SCHEMES, id="bearer where token is also accepted"), - pytest.param("Token {key}", DEEPGRAM_SDK_SCHEMES, id="token where the route accepts it"), - pytest.param("Token {key} ", DEEPGRAM_SDK_SCHEMES, id="token with trailing space"), + 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_accepted_scheme( - monkeypatch: pytest.MonkeyPatch, authorization: str, accepted_schemes: tuple[str, ...] +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 @@ -2327,9 +2324,7 @@ async def test_websocket_auth_finds_the_virtual_key_under_each_accepted_scheme( 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", - accepted_authorization_schemes=accepted_schemes, + _websocket_authorized_by(sent, authorization.format(key=user_key)), "nova-3" ) assert authorized.token == hash_token(user_key) @@ -2338,16 +2333,15 @@ async def test_websocket_auth_finds_the_virtual_key_under_each_accepted_scheme( @pytest.mark.asyncio @pytest.mark.parametrize( - ("authorization", "accepted_schemes"), + "authorization", [ - pytest.param("Token {key}", ("Bearer",), id="token where the route did not opt in"), - pytest.param("token {key}", DEEPGRAM_SDK_SCHEMES, id="scheme is case sensitive like bearer"), - pytest.param("Token{key}", DEEPGRAM_SDK_SCHEMES, id="no space after the scheme"), - pytest.param("Basic {key}", DEEPGRAM_SDK_SCHEMES, id="scheme not in the list"), + 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_other_schemes_with_a_policy_close( - monkeypatch: pytest.MonkeyPatch, authorization: str, accepted_schemes: tuple[str, ...] +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 @@ -2358,11 +2352,25 @@ async def test_websocket_auth_refuses_other_schemes_with_a_policy_close( 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", - accepted_authorization_schemes=accepted_schemes, + _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) 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 19418c4c349..6f984523cda 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 @@ -504,17 +504,3 @@ def test_deepgram_listen_accepts_the_litellm_key_in_a_deepgram_sdk_token_header( 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"})) ] - - -def test_a_non_deepgram_websocket_route_still_refuses_a_token_header() -> None: - relay: Final = _FakeRelay() - client: Final = TestClient(_app_with_relay(relay)) - - with pytest.raises(WebSocketDisconnect) as disconnect: - with client.websocket_connect( - "/openai/v1/realtime?model=gpt-4o-realtime-preview", headers={"Authorization": "Token sk-litellm"} - ): - pass - - assert disconnect.value.code == 1008 - assert relay.calls == [] diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 60043daa402..d3ceabfa31e 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/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, From 0c361389e755b34f7d21426df9f05098d02474d6 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Tue, 27 Oct 2026 14:02:41 -0700 Subject: [PATCH 3/3] docs(auth): note that the Token Authorization scheme comes from Deepgram SDKs --- litellm/proxy/auth/user_api_key_auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a4997f7f223..b7261ad9d7f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1122,7 +1122,7 @@ 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 "): + if api_key.startswith("Token "): # Deepgram SDKs send their key as `Token ` return api_key[len("Token ") :] if api_key.startswith("Bearer "): # ensure Bearer token passed in api_key = api_key.replace("Bearer ", "") # extract the token