From b0714dec82656536b38e0df839c3ba73eb54e5b6 Mon Sep 17 00:00:00 2001 From: Diwak4r Date: Mon, 10 Aug 2026 01:51:39 +0545 Subject: [PATCH] fix(key_management): cap /key/update budgets at caller delegation ceiling A non-admin caller who passes the admin-identity check (key owner, team admin, or org admin) could still raise a key's max_budget, budget_limits, or temp_budget_increase above their own delegation ceiling. Adds the same value-vs-ceiling guards the generate path already has in _common_key_generation_helper (GHSA-q775-qw9r-2r4g). Closes #35796 --- .../key_management_endpoints.py | 63 ++++++++ .../test_key_management_endpoints.py | 148 ++++++++++++++++++ 2 files changed, 211 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2de1d177b33..0b8baf92217 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2524,6 +2524,69 @@ async def _validate_update_key_data( prisma_client=prisma_client, ) + # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller + # must not be able to raise a key's budget above their own authority on + # /key/update. _check_key_admin_access above proves identity (proxy admin + # / key owner / team admin / org admin) but never compares the requested + # values against the caller's delegation ceiling, which is what let a team + # admin grant a key a higher max_budget / budget_limits / + # temp_budget_increase than they are allowed to delegate. Mirrors the + # generate-path guards in _common_key_generation_helper. + _is_ui_session_team_key: Final = ( + user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _team_id_to_check is not None + ) + _delegation_ceiling: Final = ( + user_api_key_dict.max_budget + if user_api_key_dict.max_budget is not None + else (team_obj.max_budget if user_api_key_dict.is_session_token and team_obj is not None else None) + ) + if not _is_proxy_admin and not _is_ui_session_team_key: + _requested_max_budget: Final = data.max_budget + if ( + _requested_max_budget is not None + and _requested_max_budget != existing_key_row.max_budget + and _delegation_ceiling is not None + and _requested_max_budget > _delegation_ceiling + ): + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"max_budget ({_requested_max_budget}) cannot exceed the caller's " + f"own max_budget ({_delegation_ceiling})." + ) + }, + ) + if data.temp_budget_increase is not None and _delegation_ceiling is not None: + # temp_budget_increase is applied on top of the key's max_budget at + # request time (user_api_key_auth._update_key_budget_with_temp_budget_increase), + # so the effective budget must stay under the caller's ceiling too. + _effective_max_budget: Final = ( + data.max_budget if data.max_budget is not None else existing_key_row.max_budget + ) + if ( + _effective_max_budget is not None + and _effective_max_budget + data.temp_budget_increase > _delegation_ceiling + ): + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"max_budget plus temp_budget_increase " + f"({_effective_max_budget} + {data.temp_budget_increase} = " + f"{_effective_max_budget + data.temp_budget_increase}) cannot exceed " + f"the caller's own max_budget ({_delegation_ceiling})." + ) + }, + ) + _check_budget_limits_delegation_ceiling( + budget_limits=data.budget_limits, + delegation_ceiling=_delegation_ceiling, + user_api_key_dict=user_api_key_dict, + is_ui_session_team_key=_is_ui_session_team_key, + team_table=team_obj, + ) + TeamMemberPermissionChecks.enforce_member_can_assign_access_groups( user_api_key_dict=user_api_key_dict, team_table=team_obj, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index e8709f3af34..310fe99fd6a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -781,6 +781,154 @@ async def test_update_key_personal_non_admin_denied_access_groups( assert "Access groups" in str(exc.value.detail) +# --------------------------------------------------------------------------- +# Regression tests for GHSA-q775-qw9r-2r4g on the UPDATE path: /key/update must +# not let a non-admin caller raise a key's max_budget / budget_limits / +# temp_budget_increase above their own delegation ceiling. +# _check_key_admin_access proves identity (proxy admin / key owner / team admin +# / org admin) but never compares the requested values against the caller's +# ceiling; the fix adds that value-vs-ceiling check to _validate_update_key_data. +# prisma_client=None plus a personal key (team_id=None) skips the admin-identity +# and team-member paths so these tests exercise only the ceiling gate. +# --------------------------------------------------------------------------- + + +def _update_key_ceiling_caller(max_budget): + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + max_budget=max_budget, + ) + + +def _update_key_ceiling_existing_row(max_budget): + return MagicMock( + token="hashed_alice_personal_key", + user_id="alice", + team_id=None, + created_by="alice", + max_budget=max_budget, + organization_id=None, + project_id=None, + ) + + +@pytest.mark.asyncio +async def test_update_key_non_admin_cannot_raise_max_budget_above_ceiling(): + """Caller with max_budget=100 must not raise their key from 50 to 200.""" + from litellm.proxy._types import UpdateKeyRequest + + data = UpdateKeyRequest(key="sk-alice-personal", max_budget=200) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=data, + existing_key_row=_update_key_ceiling_existing_row(50), + user_api_key_dict=_update_key_ceiling_caller(100), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + assert exc.value.status_code == 400 + assert "cannot exceed the caller's own max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_update_key_non_admin_within_ceiling_max_budget_allowed(): + """Caller with max_budget=100 may raise their key to 80 (within ceiling).""" + from litellm.proxy._types import UpdateKeyRequest + + data = UpdateKeyRequest(key="sk-alice-personal", max_budget=80) + await _validate_update_key_data( + data=data, + existing_key_row=_update_key_ceiling_existing_row(50), + user_api_key_dict=_update_key_ceiling_caller(100), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_update_key_non_admin_cannot_raise_budget_limits_above_ceiling(): + """budget_limits window of 150 must be rejected for a caller capped at 100.""" + from litellm.models.team import BudgetLimitEntry + from litellm.proxy._types import UpdateKeyRequest + + data = UpdateKeyRequest( + key="sk-alice-personal", + budget_limits=[BudgetLimitEntry(budget_duration="30d", max_budget=150)], + ) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=data, + existing_key_row=_update_key_ceiling_existing_row(50), + user_api_key_dict=_update_key_ceiling_caller(100), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + assert exc.value.status_code == 400 + assert "cannot exceed the caller's own max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_update_key_non_admin_cannot_raise_temp_budget_increase_above_ceiling(): + """temp_budget_increase is applied on top of max_budget at request time, so + 50 + 100 = 150 must be rejected for a caller capped at 100.""" + from datetime import datetime, timedelta, timezone + + from litellm.proxy._types import UpdateKeyRequest + + data = UpdateKeyRequest( + key="sk-alice-personal", + max_budget=50, + temp_budget_increase=100, + temp_budget_expiry=datetime.now(timezone.utc) + timedelta(days=1), + ) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=data, + existing_key_row=_update_key_ceiling_existing_row(50), + user_api_key_dict=_update_key_ceiling_caller(100), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + assert exc.value.status_code == 400 + assert "temp_budget_increase" in str(exc.value.detail) + assert "cannot exceed the caller's own max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_update_key_proxy_admin_can_raise_max_budget_above_ceiling(): + """PROXY_ADMIN is not bound by the delegation ceiling.""" + from litellm.proxy._types import UpdateKeyRequest + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + + data = UpdateKeyRequest(key="sk-alice-personal", max_budget=200) + await _validate_update_key_data( + data=data, + existing_key_row=_update_key_ceiling_existing_row(50), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin", + max_budget=50, + ), + llm_router=None, + premium_user=True, + prisma_client=None, + user_api_key_cache=MagicMock(), + ) + + @pytest.mark.asyncio async def test_generate_key_helper_fn_with_access_group_ids(monkeypatch): """Ensure generate_key_helper_fn passes access_group_ids into the key insert payload."""