mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
feat(auth): accept a Token Authorization header on the Deepgram WebSocket route
Deepgram SDKs send their API key as `Authorization: Token <key>`, 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
This commit is contained in:
parent
c905cd065d
commit
75c063124a
4 changed files with 149 additions and 5 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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": ""}]
|
||||
|
|
|
|||
|
|
@ -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 <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"}))
|
||||
]
|
||||
|
||||
|
||||
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 == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue