mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(proxy): reject non-finite budget_limits windows on /key/generate (#31630)
Enforce that every `budget_limits[*].max_budget` is a finite number; applies to every caller including proxy admin and runs before the role / ceiling checks. Six parametrized regression tests cover NaN / +inf / -inf for both non-admin and admin callers.
This commit is contained in:
parent
829bfebe0f
commit
be6b28f25e
2 changed files with 68 additions and 5 deletions
|
|
@ -576,15 +576,20 @@ def _check_budget_limits_delegation_ceiling(
|
|||
is_ui_session_team_key: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Enforce the delegation ceiling on every per-window budget entry.
|
||||
Enforce two invariants on `budget_limits`:
|
||||
|
||||
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.
|
||||
- Every `budget_limits[*].max_budget` must be a finite number; applies
|
||||
to every caller including proxy admin.
|
||||
- Non-admin callers may not set a window above their delegation ceiling.
|
||||
"""
|
||||
if not budget_limits:
|
||||
return
|
||||
non_finite = next((w for w in budget_limits if not math.isfinite(w.max_budget)), None)
|
||||
if non_finite is not None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": (f"budget_limits entry max_budget ({non_finite.max_budget}) must be a finite number.")},
|
||||
)
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
if is_ui_session_team_key:
|
||||
|
|
|
|||
|
|
@ -12854,6 +12854,64 @@ async def test_budget_limits_admin_unrestricted(monkeypatch):
|
|||
assert result is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("non_finite", [float("nan"), float("inf"), float("-inf")])
|
||||
async def test_budget_limits_window_non_finite_rejected_for_non_admin(monkeypatch, non_finite):
|
||||
"""A non-admin caller submitting a non-finite `budget_limits` window
|
||||
gets 400. The finite-number invariant applies before role / ceiling
|
||||
checks."""
|
||||
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": non_finite}],
|
||||
)
|
||||
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 "finite" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("non_finite", [float("nan"), float("inf"), float("-inf")])
|
||||
async def test_budget_limits_window_non_finite_rejected_for_admin(monkeypatch, non_finite):
|
||||
"""The finite-number invariant applies to every caller including
|
||||
proxy admin."""
|
||||
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(
|
||||
budget_limits=[{"budget_duration": "1d", "max_budget": non_finite}],
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _common_key_generation_helper(
|
||||
data=request,
|
||||
user_api_key_dict=admin,
|
||||
litellm_changed_by=None,
|
||||
team_table=None,
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "finite" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue