fix(proxy): gate non-admin /key/generate budget_limits and permissions (VERIA-392) (#31469)

/key/generate validated the caller's delegation ceiling against
data.max_budget only. The per-window entries in data.budget_limits
bypassed the check, so a non-admin caller could mint a key whose
1-day window vastly exceeded their own max_budget. The data.permissions
dict also went unvalidated for non-admin callers, so they could
self-grant capabilities like allow_pii_controls (and on Enterprise,
get_spend_routes).

Both gates now live in _common_key_generation_helper, covering
/key/generate and /key/service-account/generate. The existing empty
{} default on permissions still passes for non-admin callers.
This commit is contained in:
yucheng-berri 2026-06-29 14:25:58 -07:00 • committed by fzowl
parent 658fe8118c
commit 64e6b0ce60
2 changed files with 242 additions and 0 deletions

View file

@ -547,6 +547,63 @@ def _check_allowed_routes_caller_permission(
)
def _check_permissions_caller_permission(
permissions: Optional[dict],
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""
Only proxy admins may set the `permissions` dict on a key.
The field grants ambient capabilities (e.g. `get_spend_routes` exposes
`/global/spend/*`), so it must follow the same admin gate as
`allowed_routes`. Without this gate a non-admin can self-grant capabilities
they do not hold, including read access to global spend.
"""
if not permissions:
return
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
raise HTTPException(
status_code=403,
detail={"error": "Only proxy admins can set `permissions` on a key."},
)
def _check_budget_limits_delegation_ceiling(
budget_limits: Optional[List[BudgetLimitEntry]],
delegation_ceiling: Optional[float],
user_api_key_dict: UserAPIKeyAuth,
is_ui_session_team_key: bool,
) -> None:
"""
Enforce the delegation ceiling on every per-window budget entry.
The single-value `max_budget` check upstream guards the all-time budget;
`budget_limits` lets a key carry independent concurrent windows and was
bypassing the ceiling entirely, so a non-admin caller could mint a key
with a window budget far above their own authority.
"""
if not budget_limits:
return
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
if is_ui_session_team_key:
return
if delegation_ceiling is None:
return
over_ceiling = next((w for w in budget_limits if w.max_budget > delegation_ceiling), None)
if over_ceiling is not None:
raise HTTPException(
status_code=400,
detail={
"error": (
f"budget_limits entry max_budget ({over_ceiling.max_budget}) "
f"cannot exceed the caller's own max_budget ({delegation_ceiling})."
)
},
)
async def validate_team_id_used_in_service_account_request(
team_id: Optional[str],
prisma_client: Optional[PrismaClient],
@ -744,6 +801,17 @@ async def _common_key_generation_helper(
},
)
_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,
)
_check_permissions_caller_permission(
permissions=data.permissions,
user_api_key_dict=user_api_key_dict,
)
# APPLY ENTERPRISE KEY MANAGEMENT PARAMS
try:
from litellm_enterprise.proxy.management_endpoints.key_management_endpoints import (

View file

@ -12761,3 +12761,177 @@ async def test_cli_session_token_personal_key_without_budget_allowed():
team_table=None, # no team in request = personal key, but no explicit budget
)
assert result is not None
@pytest.mark.asyncio
async def test_budget_limits_window_cannot_exceed_caller_max_budget(monkeypatch):
"""A non-admin caller may not set a `budget_limits` window above
their own `max_budget`."""
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params",
None,
raising=False,
)
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
max_budget=10.0,
)
request = GenerateKeyRequest(
budget_limits=[
{"budget_duration": "1d", "max_budget": 1_000_000.0},
],
)
with pytest.raises(HTTPException) as exc_info:
await _common_key_generation_helper(
data=request,
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=None,
)
assert exc_info.value.status_code == 400
assert "budget_limits" in str(exc_info.value.detail)
assert "1000000" in str(exc_info.value.detail).replace(",", "")
@pytest.mark.asyncio
async def test_budget_limits_window_within_caller_max_budget_allowed(monkeypatch):
"""Counterpart: a window within the caller's ceiling must still pass."""
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params",
None,
raising=False,
)
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
max_budget=100.0,
)
request = GenerateKeyRequest(
budget_limits=[{"budget_duration": "1d", "max_budget": 25.0}],
)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
):
result = await _common_key_generation_helper(
data=request,
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=None,
)
assert result is not None
@pytest.mark.asyncio
async def test_budget_limits_admin_unrestricted(monkeypatch):
"""Proxy admin can set any window budget regardless of their own max_budget."""
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params",
None,
raising=False,
)
admin = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-1",
max_budget=10.0,
)
request = GenerateKeyRequest(
budget_limits=[{"budget_duration": "1d", "max_budget": 1_000_000.0}],
)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"key": "sk-test", "expires": None, "user_id": "admin-1"},
):
result = await _common_key_generation_helper(
data=request,
user_api_key_dict=admin,
litellm_changed_by=None,
team_table=None,
)
assert result is not None
@pytest.mark.asyncio
async def test_permissions_field_rejected_for_non_admin(monkeypatch):
"""A non-admin caller may not set the `permissions` field on a key
they create."""
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params",
None,
raising=False,
)
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
max_budget=100.0,
)
request = GenerateKeyRequest(permissions={"get_spend_routes": True})
with pytest.raises(HTTPException) as exc_info:
await _common_key_generation_helper(
data=request,
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=None,
)
assert exc_info.value.status_code == 403
assert "permissions" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_permissions_empty_default_allowed_for_non_admin(monkeypatch):
"""
The empty `{}` default on GenerateKeyRequest.permissions must continue
to pass for non-admin callers; only a non-empty dict triggers the gate.
"""
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params",
None,
raising=False,
)
caller = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="user-1",
max_budget=100.0,
)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"key": "sk-test", "expires": None, "user_id": "user-1"},
):
result = await _common_key_generation_helper(
data=GenerateKeyRequest(),
user_api_key_dict=caller,
litellm_changed_by=None,
team_table=None,
)
assert result is not None
@pytest.mark.asyncio
async def test_permissions_admin_can_set_any(monkeypatch):
"""Proxy admin can still set `permissions` on a key."""
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.litellm.default_key_generate_params",
None,
raising=False,
)
admin = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-1",
)
request = GenerateKeyRequest(permissions={"get_spend_routes": True})
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"key": "sk-test", "expires": None, "user_id": "admin-1"},
):
result = await _common_key_generation_helper(
data=request,
user_api_key_dict=admin,
litellm_changed_by=None,
team_table=None,
)
assert result is not None