mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(proxy): resolve router_settings.model_group_alias before key/team model auth
Key and team router_settings.model_group_alias aliases were resolved only after the key/team model allowlist checks ran, so a key allowed the alias target was denied when it requested the alias. Resolve the alias during auth and rewrite the request body to the target before the allowlist checks. The alias the client sent is kept in the request scope so the response model still echoes it. Resolves LIT-3054 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
226b1e1bb9
commit
bae2bf003e
6 changed files with 212 additions and 9 deletions
|
|
@ -317,6 +317,7 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
|||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
|
||||
CLIENT_REQUESTED_MODEL_SCOPE_KEY: Final = "litellm.client_requested_model"
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
|
|||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.constants import (
|
||||
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
|
||||
GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
INVALID_VIRTUAL_KEY_ERROR_MARKER,
|
||||
INVALID_VIRTUAL_KEY_ERROR_MESSAGE,
|
||||
|
|
@ -124,6 +125,7 @@ from litellm.proxy.utils import (
|
|||
normalize_route_for_root_path,
|
||||
)
|
||||
from litellm.repositories.table_repositories import TeamMembershipRepository
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
|
|
@ -235,13 +237,56 @@ async def _normalize_claude_model(
|
|||
request.scope[_CLAUDE_MODEL_NORMALIZED] = True
|
||||
if source is None:
|
||||
return
|
||||
request_data["model"] = source
|
||||
_rewrite_request_model(request_data, request, source)
|
||||
|
||||
|
||||
def _rewrite_request_model(
|
||||
request_data: dict, # mutable-ok: the request body is rewritten in place for every downstream reader
|
||||
request: Request | None,
|
||||
model: str,
|
||||
) -> None:
|
||||
request_data["model"] = model
|
||||
_safe_set_request_parsed_body(request=request, parsed_body=request_data)
|
||||
if request is not None:
|
||||
request._json = request_data
|
||||
request._body = orjson.dumps(request_data)
|
||||
|
||||
|
||||
_MODEL_GROUP_ALIAS_RESOLVED: Final = "litellm.model_group_alias_resolved"
|
||||
|
||||
|
||||
async def _resolve_router_settings_model_group_alias(
|
||||
request_data: dict, # mutable-ok: the request body is rewritten in place for every downstream reader
|
||||
valid_token: UserAPIKeyAuth,
|
||||
request: Request | None,
|
||||
route: str,
|
||||
) -> None:
|
||||
"""Rewrite the requested model through the key's or team's ``router_settings.model_group_alias``
|
||||
before the allowlist checks, so they authorize the model group the request is routed to.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj
|
||||
|
||||
if request is None or llm_router is None or not RouteChecks.is_llm_api_route(route=route):
|
||||
return
|
||||
if request.scope.get(_MODEL_GROUP_ALIAS_RESOLVED) is True:
|
||||
return
|
||||
request.scope[_MODEL_GROUP_ALIAS_RESOLVED] = True
|
||||
requested: Final = request_data.get("model")
|
||||
if not isinstance(requested, str) or await read_raw_json_body(request=request) is None:
|
||||
return
|
||||
settings: Final = await proxy_config.get_hierarchical_router_settings(
|
||||
user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
|
||||
)
|
||||
if not isinstance(settings, Mapping):
|
||||
return
|
||||
target: Final = resolve_model_group_alias(settings.get("model_group_alias"), requested)
|
||||
if target is None or target == requested:
|
||||
return
|
||||
verbose_proxy_logger.debug("router_settings.model_group_alias resolved %s -> %s before auth", requested, target)
|
||||
request.scope.setdefault(CLIENT_REQUESTED_MODEL_SCOPE_KEY, requested)
|
||||
_rewrite_request_model(request_data, request, target)
|
||||
|
||||
|
||||
def _get_model_names_for_budget_checks(
|
||||
model: str | list[str] | None,
|
||||
) -> list[str]:
|
||||
|
|
@ -2926,6 +2971,7 @@ async def _authorize_authenticated_request(
|
|||
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
|
||||
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)
|
||||
await _normalize_claude_model(request_data, user_api_key_auth_obj, request, route)
|
||||
await _resolve_router_settings_model_group_alias(request_data, user_api_key_auth_obj, request, route)
|
||||
|
||||
# Single authorization point. Builder paths MUST NOT call common_checks.
|
||||
# Route through the same exception handler the builder uses so
|
||||
|
|
@ -3312,6 +3358,7 @@ async def _enforce_key_and_fallback_model_access(
|
|||
Not included in common_checks — common_checks enforces team/user/project model access only.
|
||||
"""
|
||||
await _normalize_claude_model(request_data, valid_token, request, route)
|
||||
await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route)
|
||||
config: Final = valid_token.config
|
||||
|
||||
if config != {}:
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm.proxy.common_utils.callback_utils import (
|
|||
get_logging_caching_headers,
|
||||
get_remaining_tokens_and_requests_from_request_data,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
attribute_of,
|
||||
error_status_code,
|
||||
|
|
@ -622,9 +623,9 @@ async def _resolve_per_request_model_group_alias(
|
|||
holds the global config map and is shared across requests, so a per-request
|
||||
map has to be applied here instead of being forwarded to the Router.
|
||||
|
||||
Model access was authorized against the requested group, so the target is
|
||||
authorized in its own right before the rewrite; a key that may not call the
|
||||
target gets the usual 403 rather than being quietly served it.
|
||||
Auth already rewrote the body through this map for LLM API routes, so this is
|
||||
a fallback for callers that skipped it; the target is authorized in its own
|
||||
right before the rewrite, so a key that may not call it gets the usual 403.
|
||||
|
||||
Returns the target model group, or None when no alias applies.
|
||||
"""
|
||||
|
|
@ -2338,9 +2339,8 @@ class ProxyBaseLLMRequestProcessing:
|
|||
"""
|
||||
Common request processing logic for both chat completions and responses API endpoints
|
||||
"""
|
||||
requested_model_from_client: Final[str | None] = (
|
||||
self.data.get("model") if isinstance(self.data.get("model"), str) else None
|
||||
)
|
||||
client_model: Final = get_client_requested_model(request) or self.data.get("model")
|
||||
requested_model_from_client: Final[str | None] = client_model if isinstance(client_model, str) else None
|
||||
self._debug_log_request_payload()
|
||||
|
||||
if skip_pre_call_logic:
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from fastapi import Request, UploadFile, status
|
|||
from typing_extensions import NotRequired, ReadOnly, Required
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB
|
||||
from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
|
|
@ -235,6 +235,13 @@ def _safe_get_request_parsed_body(request: Request | None) -> dict | None:
|
|||
return None
|
||||
|
||||
|
||||
def get_client_requested_model(request: Request | None) -> str | None:
|
||||
if request is None or not hasattr(request, "scope"):
|
||||
return None
|
||||
model: Final = request.scope.get(CLIENT_REQUESTED_MODEL_SCOPE_KEY)
|
||||
return model if isinstance(model, str) else None
|
||||
|
||||
|
||||
def _safe_get_request_query_params(request: Request | None) -> dict:
|
||||
if request is None:
|
||||
return {}
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.proxy._types import (
|
|||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_check_key_model_budget_with_fallback,
|
||||
_ensure_litellm_received_at_on_request_state,
|
||||
|
|
@ -8137,3 +8138,99 @@ async def test_auth_flow_enters_virtual_key_mapping_when_only_an_issuer_configur
|
|||
assert resolve_mock.await_args.kwargs["jwt_claims"][JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == ISSUER_TWO
|
||||
assert result.api_key == "hashed-mapped-key"
|
||||
assert result.team_id == "svc-team"
|
||||
|
||||
|
||||
def _alias_router() -> litellm.Router:
|
||||
return litellm.Router(model_list=[{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in ("claude-haiku", "claude-sonnet")])
|
||||
|
||||
|
||||
def _alias_request(route: str, data: dict, content_type: str = "application/json"):
|
||||
"""A request as auth sees it: the body already read once and cached alongside its parsed form."""
|
||||
from starlette.requests import Request
|
||||
|
||||
headers = [(b"content-type", content_type.encode())]
|
||||
request = Request({"type": "http", "method": "POST", "path": route, "headers": headers, "query_string": b"", "parsed_body": (tuple(data), data)})
|
||||
request._body = json.dumps(data).encode()
|
||||
return request
|
||||
|
||||
|
||||
def _alias_token(monkeypatch, level: str, alias: dict, models: list) -> UserAPIKeyAuth:
|
||||
"""A key whose ``router_settings.model_group_alias`` lives on the key itself or on its cached team row."""
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
if level == "key":
|
||||
return UserAPIKeyAuth(models=models, router_settings={"model_group_alias": alias})
|
||||
cache = UserApiKeyCache()
|
||||
cache.set_cache(key="team_id:team-alias", value=LiteLLM_TeamTableCachedObj(team_id="team-alias", models=models, router_settings={"model_group_alias": alias}))
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", MagicMock())
|
||||
return UserAPIKeyAuth(team_id="team-alias", models=models)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("level", ["key", "team"])
|
||||
@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/messages", "/v1/embeddings"])
|
||||
async def test_router_settings_model_group_alias_authorizes_target_for_key(monkeypatch, level, route):
|
||||
"""LIT-3054: a key allowed only the alias target must be able to call the alias, and a key not
|
||||
allowed the target must still be denied even when the alias itself is what it requested."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access
|
||||
|
||||
router = _alias_router()
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
|
||||
data = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]}
|
||||
request = _alias_request(route, data)
|
||||
token = _alias_token(monkeypatch, level, {"AgentX-LLM": "claude-haiku"}, ["claude-haiku"])
|
||||
await _enforce_key_and_fallback_model_access(valid_token=token, request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router)
|
||||
assert data["model"] == "claude-haiku"
|
||||
assert (await request.json())["model"] == "claude-haiku"
|
||||
assert json.loads(await request.body())["model"] == "claude-haiku"
|
||||
assert request.scope["parsed_body"][1]["model"] == "claude-haiku"
|
||||
assert get_client_requested_model(request) == "AgentX-LLM"
|
||||
|
||||
denied = _alias_token(monkeypatch, level, {"AgentX-LLM": "claude-sonnet"}, ["claude-haiku"])
|
||||
denied_data = {"model": "AgentX-LLM"}
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await _enforce_key_and_fallback_model_access(valid_token=denied, request_data=denied_data, route=route, request=_alias_request(route, denied_data), llm_model_list=router.model_list, llm_router=router)
|
||||
assert "claude-sonnet" in exc.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_settings_model_group_alias_leaves_form_bodies_alone(monkeypatch):
|
||||
"""LIT-3054: a multipart body cannot be re-serialized as JSON, so auth must not rewrite it."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access
|
||||
|
||||
router = _alias_router()
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
|
||||
data = {"model": "AgentX-LLM"}
|
||||
request = _alias_request("/v1/audio/transcriptions", data, content_type="multipart/form-data; boundary=x")
|
||||
token = _alias_token(monkeypatch, "key", {"AgentX-LLM": "claude-haiku"}, ["claude-haiku", "AgentX-LLM"])
|
||||
await _enforce_key_and_fallback_model_access(valid_token=token, request_data=data, route="/v1/audio/transcriptions", request=request, llm_model_list=router.model_list, llm_router=router)
|
||||
assert data["model"] == "AgentX-LLM"
|
||||
assert get_client_requested_model(request) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("target, expect_denied", [("claude-haiku", False), ("claude-sonnet", True)])
|
||||
async def test_router_settings_model_group_alias_authorizes_target_for_team(monkeypatch, target, expect_denied):
|
||||
"""LIT-3054: the team allowlist check in common_checks must judge the alias target, not the alias."""
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from litellm.proxy.auth.user_api_key_auth import _authorize_authenticated_request
|
||||
|
||||
router = _alias_router()
|
||||
logging_obj = MagicMock(post_call_failure_hook=AsyncMock(return_value=None))
|
||||
attrs = {**_proxy_attrs_for_centralized_checks(), "llm_router": router, "proxy_logging_obj": logging_obj}
|
||||
for k, v in attrs.items():
|
||||
monkeypatch.setattr(_proxy_server_mod, k, v)
|
||||
token = _alias_token(monkeypatch, "team", {"AgentX-LLM": target}, ["claude-haiku"])
|
||||
token.team_models = ["claude-haiku"]
|
||||
data = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]}
|
||||
request = _alias_request("/v1/chat/completions", data)
|
||||
if expect_denied:
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await _authorize_authenticated_request(user_api_key_auth_obj=token, request=request, request_data=data, route="/v1/chat/completions", api_key="sk-test")
|
||||
assert exc.value.type == ProxyErrorTypes.team_model_access_denied
|
||||
assert target in exc.value.message
|
||||
return
|
||||
await _authorize_authenticated_request(user_api_key_auth_obj=token, request=request, request_data=data, route="/v1/chat/completions", api_key="sk-test")
|
||||
assert (await request.json())["model"] == target
|
||||
assert get_client_requested_model(request) == "AgentX-LLM"
|
||||
|
|
|
|||
|
|
@ -13,7 +13,11 @@ from fastapi.responses import JSONResponse, StreamingResponse
|
|||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import MAX_LITELLM_CALL_ID_LENGTH, RETURN_RAW_MODEL_NAME_METADATA_KEY
|
||||
from litellm.constants import (
|
||||
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
|
||||
MAX_LITELLM_CALL_ID_LENGTH,
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.opentelemetry import UserAPIKeyAuth
|
||||
from litellm.proxy.common_request_processing import (
|
||||
|
|
@ -4395,6 +4399,53 @@ class TestDisconnectGatherCleanup:
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("client_model, expected", [("AgentX-LLM", "AgentX-LLM"), (None, "gpt-mini")])
|
||||
async def test_response_model_echoes_the_name_the_client_sent_before_auth_rewrote_it(
|
||||
monkeypatch, client_model, expected
|
||||
):
|
||||
"""LIT-3054: auth resolves router_settings.model_group_alias in the body, so the alias the
|
||||
client sent only survives in the request scope. The response must still echo it."""
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
|
||||
async def llm():
|
||||
return litellm.ModelResponse(
|
||||
model="gpt-4o-mini", choices=[{"message": {"role": "assistant", "content": "pong"}}]
|
||||
)
|
||||
|
||||
async def fake_route_request(**_kwargs):
|
||||
return llm()
|
||||
|
||||
logging_obj = MagicMock(litellm_call_id="call-id", _defer_async_logging=False)
|
||||
proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging.during_call_hook = AsyncMock(return_value=None)
|
||||
proxy_logging.post_call_success_hook = AsyncMock(side_effect=lambda data, user_api_key_dict, response: response)
|
||||
proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
proxy_logging._callback_capabilities_cache = {}
|
||||
monkeypatch.setattr(cpr, "route_request", fake_route_request)
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-mini", "messages": []})
|
||||
monkeypatch.setattr(
|
||||
processor, "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gpt-mini"}, logging_obj))
|
||||
)
|
||||
monkeypatch.setattr(processor, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
scope = {"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": [], "query_string": b""}
|
||||
request = Request({**scope, CLIENT_REQUESTED_MODEL_SCOPE_KEY: client_model} if client_model else scope)
|
||||
|
||||
response = await processor.base_process_llm_request(
|
||||
request=request,
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(),
|
||||
proxy_logging_obj=proxy_logging,
|
||||
general_settings={},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
route_type="acompletion",
|
||||
version=None,
|
||||
)
|
||||
|
||||
assert response.model == expected
|
||||
|
||||
|
||||
class TestStreamingClientDisconnectLogging:
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_streaming_client_disconnect_sets_error_information(self):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue