From b0714dec82656536b38e0df839c3ba73eb54e5b6 Mon Sep 17 00:00:00 2001 From: Diwak4r Date: Mon, 10 Aug 2026 01:51:39 +0545 Subject: [PATCH 1/2] 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.""" From 4245e6fa30fbda54f1ef6bdf1891b21aad96f835 Mon Sep 17 00:00:00 2001 From: Diwak4r Date: Sat, 22 Aug 2026 22:34:07 +0545 Subject: [PATCH 2/2] fix(proxy): close persisted-temp-increase and NaN bypasses in /key/update delegation ceiling Two review findings on the delegation-ceiling guard: 1. Persisted temp_budget_increase bypass: when a caller updates only max_budget without setting temp_budget_increase, the key row's metadata may still carry a non-expired persisted increase. The effective budget (new max_budget + retained increase) could exceed the caller's ceiling. Now reads the persisted value from metadata (checking expiry) and includes it in the check. 2. NaN temp_budget_increase bypass: any comparison with NaN is False, so _effective_max_budget + float('nan') > ceiling silently passed. Added a math.isfinite guard that rejects non-finite values with a 400 before the comparison runs. --- .../key_management_endpoints.py | 40 +++++++++++-- .../test_key_management_endpoints.py | 59 +++++++++++++++++++ 2 files changed, 95 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 0b8baf92217..9a5d9faba38 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2557,7 +2557,39 @@ async def _validate_update_key_data( ) }, ) - if data.temp_budget_increase is not None and _delegation_ceiling is not None: + # Determine the effective temp_budget_increase after this update. + # Either the caller is setting a new one, or a previously persisted + # one in the key's metadata still applies (it is added on top of + # max_budget at request time by _update_key_budget_with_temp_budget_increase). + _effective_temp_increase: float | None = data.temp_budget_increase + if ( + _effective_temp_increase is None + and existing_key_row.metadata is not None + ): + _persisted_increase: Final = existing_key_row.metadata.get( + "temp_budget_increase" + ) + if _persisted_increase is not None: + try: + _persisted_expiry: Final = datetime.fromisoformat( + existing_key_row.metadata.get("temp_budget_expiry", "") + ) + if _persisted_expiry > datetime.now(timezone.utc): + _effective_temp_increase = float(_persisted_increase) + except (ValueError, TypeError): + pass + + if _effective_temp_increase is not None and _delegation_ceiling is not None: + if not math.isfinite(_effective_temp_increase): + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"temp_budget_increase must be a finite number. " + f"Received: {_effective_temp_increase}" + ) + }, + ) # 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. @@ -2566,15 +2598,15 @@ async def _validate_update_key_data( ) if ( _effective_max_budget is not None - and _effective_max_budget + data.temp_budget_increase > _delegation_ceiling + and _effective_max_budget + _effective_temp_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"({_effective_max_budget} + {_effective_temp_increase} = " + f"{_effective_max_budget + _effective_temp_increase}) cannot exceed " f"the caller's own max_budget ({_delegation_ceiling})." ) }, 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 310fe99fd6a..17f4e98a981 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 @@ -906,6 +906,65 @@ async def test_update_key_non_admin_cannot_raise_temp_budget_increase_above_ceil assert "cannot exceed the caller's own max_budget" in str(exc.value.detail) +@pytest.mark.asyncio +async def test_update_key_persisted_temp_increase_blocks_max_budget_raise(): + """A key with a persisted (non-expired) temp_budget_increase of 60 in its + metadata must not be raised from 30 to 80 by a caller capped at 100, + because the effective budget (80 + 60 = 140) exceeds the ceiling.""" + from datetime import datetime, timedelta, timezone + + from litellm.proxy._types import UpdateKeyRequest + + row = _update_key_ceiling_existing_row(30) + row.metadata = { + "temp_budget_increase": 60, + "temp_budget_expiry": ( + datetime.now(timezone.utc) + timedelta(days=1) + ).isoformat(), + } + data = UpdateKeyRequest(key="sk-alice-personal", max_budget=80) + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + data=data, + existing_key_row=row, + 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_nan_temp_budget_increase_rejected(): + """NaN temp_budget_increase must be rejected: any comparison with NaN + is False, so the ceiling check would silently pass.""" + from datetime import datetime, timedelta, timezone + + from litellm.proxy._types import UpdateKeyRequest + + data = UpdateKeyRequest( + key="sk-alice-personal", + max_budget=50, + temp_budget_increase=float("nan"), + 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 "finite" 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."""