mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
chore(proxy): carry ASGI path into WebSocket auth synthetic Request
``user_api_key_auth_websocket`` built a synthetic ``Request`` with a
two-key scope (``type`` + ``headers``) and set ``request._url =
websocket.url``. ``get_request_route`` reads ``scope.get("path", ...)``
and falls back to ``request.url.path`` only when ``path`` is absent.
For the WebSocket flow that fallback fires and resolves to the
Host-header-derived value (Starlette reconstructs ``websocket.url``
from the Host header), so a malformed Host collapses the resolved
route and lets the auth gate compare against the wrong value.
Carry the ASGI scope's ``path``, ``root_path``, and ``app_root_path``
into the synthetic scope so the lookup never reaches the fallback on
the legitimate path.
Regression test pins that the request handed to ``user_api_key_auth``
has ``scope["path"]`` equal to the ASGI scope's path.
This commit is contained in:
parent
27de6547dc
commit
5ee2a5b088
2 changed files with 47 additions and 3 deletions
|
|
@ -12,7 +12,7 @@ import fnmatch
|
|||
import re
|
||||
import secrets
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Iterator, List, Optional, Tuple, Union, cast
|
||||
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast
|
||||
|
||||
import fastapi
|
||||
from fastapi import HTTPException, Request, WebSocket, status
|
||||
|
|
@ -332,8 +332,22 @@ def _apply_budget_limits_to_end_user_params(
|
|||
async def user_api_key_auth_websocket(websocket: WebSocket):
|
||||
# Accept the WebSocket connection
|
||||
|
||||
scope_headers = list(websocket.scope.get("headers") or [])
|
||||
request = Request(scope={"type": "http", "headers": scope_headers})
|
||||
ws_scope = websocket.scope or {}
|
||||
scope_headers = list(ws_scope.get("headers") or [])
|
||||
# ``get_request_route`` falls back to ``request.url.path`` when
|
||||
# ``scope["path"]`` is absent. On WebSockets that fallback reads
|
||||
# ``websocket.url``, which Starlette reconstructs from the (poisonable)
|
||||
# Host header. Carry the ASGI scope's path / root_path so the lookup
|
||||
# never reaches the fallback.
|
||||
synthetic_scope: Dict[str, Any] = {
|
||||
"type": "http",
|
||||
"headers": scope_headers,
|
||||
"path": ws_scope.get("path", ""),
|
||||
}
|
||||
for key in ("root_path", "app_root_path"):
|
||||
if key in ws_scope:
|
||||
synthetic_scope[key] = ws_scope[key]
|
||||
request = Request(scope=synthetic_scope)
|
||||
|
||||
request._url = websocket.url
|
||||
|
||||
|
|
|
|||
|
|
@ -915,6 +915,36 @@ async def test_user_api_key_auth_websocket():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_api_key_auth_websocket_carries_asgi_path():
|
||||
"""
|
||||
The synthetic Request must carry the ASGI scope's ``path`` so
|
||||
``get_request_route`` returns the real WebSocket path, not a value
|
||||
reconstructed from the (Host-poisonable) ``websocket.url``.
|
||||
"""
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket
|
||||
|
||||
mock_websocket = MagicMock(spec=WebSocket)
|
||||
mock_websocket.query_params = {"model": "some_model"}
|
||||
mock_websocket.headers = {"authorization": "Bearer some_api_key"}
|
||||
mock_websocket.scope = {
|
||||
"type": "websocket",
|
||||
"path": "/v1/realtime",
|
||||
"root_path": "",
|
||||
"headers": [(b"authorization", b"Bearer some_api_key")],
|
||||
}
|
||||
mock_websocket.url = URL(url="/v1/realtime")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True
|
||||
) as mock_user_api_key_auth:
|
||||
await user_api_key_auth_websocket(mock_websocket)
|
||||
|
||||
request_arg = mock_user_api_key_auth.call_args.kwargs["request"]
|
||||
assert request_arg.scope.get("path") == "/v1/realtime"
|
||||
assert request_arg.scope.get("root_path") == ""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enforce_rbac", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_jwt_user_api_key_auth_builder_enforce_rbac(enforce_rbac, monkeypatch):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue