mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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:
parent
658fe8118c
commit
64e6b0ce60
2 changed files with 242 additions and 0 deletions
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue