From 64e6b0ce60fc31b981bfc282309d9e2185cf5724 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 29 Jun 2026 14:25:58 -0700 Subject: [PATCH] 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. --- .../key_management_endpoints.py | 68 +++++++ .../test_key_management_endpoints.py | 174 ++++++++++++++++++ 2 files changed, 242 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 15e228a1a9f..368e411526f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 ( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 5c68cdc32dc..d9fee47aaa7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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