mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
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:
parent
f6b9518ddb
commit
b0714dec82
2 changed files with 211 additions and 0 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue