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
This commit is contained in:
Diwak4r 2026-08-10 01:51:39 +05:45
parent f6b9518ddb
commit b0714dec82
2 changed files with 211 additions and 0 deletions

View file

@ -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,

View file

@ -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."""