scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

Merge 4245e6fa30 into 7bc80994b7

Browse source
This commit is contained in:
Diwakar Ray Yadav 2026-08-26 14:36:15 -04:00 • committed by GitHub
commit 97f43da805
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 302 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

95
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -2624,6 +2624,101 @@ 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})."
)
},
)
# 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.
_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 + _effective_temp_increase > _delegation_ceiling
):
raise HTTPException(
status_code=400,
detail={
"error": (
f"max_budget plus temp_budget_increase "
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})."
)
},
)
_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,

207
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -779,6 +779,213 @@ 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_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."""
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."""