mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix: tighten budget field validation and authorization checks (#27897)
* fix: block NaN/Inf budget bypass and add missing non-admin guards
Addresses three security issues:
GHSA-wvg4-6222-3q4r: /user/update exposes max_budget, soft_budget, spend
to self-editing non-admin users with no server-side guard. Non-admin callers
now receive HTTP 403 if any of those fields appear in the update payload.
GHSA-q775-qw9r-2r4g: _enforce_upperbound_key_params returned early (no-op)
when upperbound_key_generate_params was absent from config, letting any
authenticated user generate a key with unlimited max_budget. Fix adds a
delegated-authority ceiling in _common_key_generation_helper: non-admins
cannot grant a key more budget than their own key carries.
GHSA-2rv4-xv66-fpjg: float('nan') passes every `value < 0` guard because
nan < 0 is False in Python, and spend >= nan is always False, permanently
disabling budget enforcement for any entity carrying a NaN max_budget.
All write-time budget guards now use `not math.isfinite(v) or v < 0`.
_enforce_upperbound_key_params validates finiteness unconditionally (before
the early-return). All spend-enforcement comparisons in auth_checks.py are
now guarded with math.isfinite(max_budget) as defense-in-depth.
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
* fix: close budget ceiling bypass for callers with no max_budget (GHSA-q775)
Non-admin callers whose API key has no explicit max_budget (None) could
bypass the delegated-authority ceiling and create keys with arbitrary
budgets. Now blocks budget assignment when caller has no budget configured.
Also removes redundant inline import of LitellmUserRoles.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: only apply budget ceiling to explicitly requested max_budget
Capture the caller-supplied max_budget before _enforce_upperbound_key_params
can fill it with a default, so auto-filled defaults don't trigger the
ceiling guard for non-admin users with no budget on their own key.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: capture requested max_budget before any defaults are applied
Move _requested_max_budget capture before both default_key_generate_params
and upperbound_key_generate_params mutations, so auto-filled values don't
trigger the ceiling check for non-admin users.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: allow unlimited-budget callers to delegate any budget
Callers with max_budget=None (unlimited) can legitimately create
budget-capped keys. Only block when caller has an explicit budget
and the requested budget exceeds it.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---------
Co-authored-by: yuneng-jiang <yuneng@berri.ai>
Co-authored-by: ryan-crabbe-berri <ryan@berri.ai>
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
de1747dca8
commit
410ce761dc
9 changed files with 565 additions and 82 deletions
|
|
@ -10,6 +10,7 @@ Run checks for:
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import math
|
||||||
import re
|
import re
|
||||||
import time
|
import time
|
||||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast
|
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast
|
||||||
|
|
@ -328,7 +329,10 @@ def _global_proxy_budget_check(
|
||||||
and route != "/v1/models"
|
and route != "/v1/models"
|
||||||
and route != "/models"
|
and route != "/models"
|
||||||
):
|
):
|
||||||
if global_proxy_spend > litellm.max_budget:
|
if (
|
||||||
|
math.isfinite(litellm.max_budget)
|
||||||
|
and global_proxy_spend > litellm.max_budget
|
||||||
|
):
|
||||||
raise litellm.BudgetExceededError(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=global_proxy_spend, max_budget=litellm.max_budget
|
current_cost=global_proxy_spend, max_budget=litellm.max_budget
|
||||||
)
|
)
|
||||||
|
|
@ -645,7 +649,7 @@ async def common_checks( # noqa: PLR0915
|
||||||
counter_key=f"spend:user:{user_object.user_id}",
|
counter_key=f"spend:user:{user_object.user_id}",
|
||||||
fallback_spend=user_object.spend or 0.0,
|
fallback_spend=user_object.spend or 0.0,
|
||||||
)
|
)
|
||||||
if user_spend >= user_budget:
|
if math.isfinite(user_budget) and user_spend >= user_budget:
|
||||||
raise litellm.BudgetExceededError(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=user_spend,
|
current_cost=user_spend,
|
||||||
max_budget=user_budget,
|
max_budget=user_budget,
|
||||||
|
|
@ -3280,7 +3284,10 @@ async def _virtual_key_max_budget_check(
|
||||||
# collect information for alerting #
|
# collect information for alerting #
|
||||||
####################################
|
####################################
|
||||||
|
|
||||||
if spend >= valid_token.max_budget:
|
# Defense-in-depth (GHSA-2rv4-xv66-fpjg): spend >= NaN is always False,
|
||||||
|
# so a NaN max_budget would silently disable enforcement. Treat a
|
||||||
|
# non-finite max_budget as "no configured limit" rather than as a bypass.
|
||||||
|
if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget:
|
||||||
raise litellm.BudgetExceededError(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=spend,
|
current_cost=spend,
|
||||||
max_budget=valid_token.max_budget,
|
max_budget=valid_token.max_budget,
|
||||||
|
|
@ -3313,7 +3320,7 @@ async def _virtual_key_multi_budget_check(
|
||||||
counter_key=counter_key,
|
counter_key=counter_key,
|
||||||
fallback_spend=0.0,
|
fallback_spend=0.0,
|
||||||
)
|
)
|
||||||
if window_spend >= w["max_budget"]:
|
if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:
|
||||||
raise litellm.BudgetExceededError(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=window_spend,
|
current_cost=window_spend,
|
||||||
max_budget=w["max_budget"],
|
max_budget=w["max_budget"],
|
||||||
|
|
@ -3568,7 +3575,10 @@ async def _check_team_member_budget(
|
||||||
fallback_spend=team_member_spend,
|
fallback_spend=team_member_spend,
|
||||||
)
|
)
|
||||||
|
|
||||||
if team_member_spend >= team_member_budget:
|
if (
|
||||||
|
math.isfinite(team_member_budget)
|
||||||
|
and team_member_spend >= team_member_budget
|
||||||
|
):
|
||||||
raise litellm.BudgetExceededError(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=team_member_spend,
|
current_cost=team_member_spend,
|
||||||
max_budget=team_member_budget,
|
max_budget=team_member_budget,
|
||||||
|
|
@ -3650,7 +3660,7 @@ async def _team_max_budget_check(
|
||||||
fallback_spend=team_object.spend or 0.0,
|
fallback_spend=team_object.spend or 0.0,
|
||||||
)
|
)
|
||||||
|
|
||||||
if spend > team_object.max_budget:
|
if math.isfinite(team_object.max_budget) and spend > team_object.max_budget:
|
||||||
if valid_token:
|
if valid_token:
|
||||||
call_info = CallInfo(
|
call_info = CallInfo(
|
||||||
token=valid_token.token,
|
token=valid_token.token,
|
||||||
|
|
@ -3698,7 +3708,7 @@ async def _team_multi_budget_check(
|
||||||
counter_key=counter_key,
|
counter_key=counter_key,
|
||||||
fallback_spend=0.0,
|
fallback_spend=0.0,
|
||||||
)
|
)
|
||||||
if window_spend >= w["max_budget"]:
|
if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:
|
||||||
raise litellm.BudgetExceededError(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=window_spend,
|
current_cost=window_spend,
|
||||||
max_budget=w["max_budget"],
|
max_budget=w["max_budget"],
|
||||||
|
|
@ -3812,6 +3822,7 @@ async def _project_max_budget_check(
|
||||||
if (
|
if (
|
||||||
max_budget is not None
|
max_budget is not None
|
||||||
and project_object.spend is not None
|
and project_object.spend is not None
|
||||||
|
and math.isfinite(max_budget)
|
||||||
and project_object.spend > max_budget
|
and project_object.spend > max_budget
|
||||||
):
|
):
|
||||||
if valid_token:
|
if valid_token:
|
||||||
|
|
@ -4004,7 +4015,7 @@ async def _organization_max_budget_check(
|
||||||
)
|
)
|
||||||
|
|
||||||
# Check if organization spend exceeds max budget
|
# Check if organization spend exceeds max budget
|
||||||
if org_spend >= org_max_budget:
|
if math.isfinite(org_max_budget) and org_spend >= org_max_budget:
|
||||||
# Trigger budget alert
|
# Trigger budget alert
|
||||||
call_info = CallInfo(
|
call_info = CallInfo(
|
||||||
token=valid_token.token,
|
token=valid_token.token,
|
||||||
|
|
|
||||||
|
|
@ -12,6 +12,8 @@ All /budget management endpoints
|
||||||
"""
|
"""
|
||||||
|
|
||||||
#### BUDGET TABLE MANAGEMENT ####
|
#### BUDGET TABLE MANAGEMENT ####
|
||||||
|
import math
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends, HTTPException
|
from fastapi import APIRouter, Depends, HTTPException
|
||||||
|
|
||||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||||
|
|
@ -57,18 +59,22 @@ async def new_budget(
|
||||||
)
|
)
|
||||||
|
|
||||||
# Validate budget values are not negative
|
# Validate budget values are not negative
|
||||||
if budget_obj.max_budget is not None and budget_obj.max_budget < 0:
|
if budget_obj.max_budget is not None and (
|
||||||
|
not math.isfinite(budget_obj.max_budget) or budget_obj.max_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"
|
"error": f"max_budget must be a non-negative finite number. Received: {budget_obj.max_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0:
|
if budget_obj.soft_budget is not None and (
|
||||||
|
not math.isfinite(budget_obj.soft_budget) or budget_obj.soft_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"
|
"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -146,18 +152,22 @@ async def update_budget(
|
||||||
raise HTTPException(status_code=400, detail={"error": "budget_id is required"})
|
raise HTTPException(status_code=400, detail={"error": "budget_id is required"})
|
||||||
|
|
||||||
# Validate budget values are not negative
|
# Validate budget values are not negative
|
||||||
if budget_obj.max_budget is not None and budget_obj.max_budget < 0:
|
if budget_obj.max_budget is not None and (
|
||||||
|
not math.isfinite(budget_obj.max_budget) or budget_obj.max_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"
|
"error": f"max_budget must be a non-negative finite number. Received: {budget_obj.max_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0:
|
if budget_obj.soft_budget is not None and (
|
||||||
|
not math.isfinite(budget_obj.soft_budget) or budget_obj.soft_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"
|
"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1152,6 +1152,46 @@ def _update_internal_user_params(
|
||||||
return non_default_values
|
return non_default_values
|
||||||
|
|
||||||
|
|
||||||
|
async def _schedule_user_update_audit_log(
|
||||||
|
response: Dict[str, Any],
|
||||||
|
existing_user_row: Optional[BaseModel],
|
||||||
|
litellm_changed_by: Optional[str],
|
||||||
|
user_api_key_dict: UserAPIKeyAuth,
|
||||||
|
litellm_proxy_admin_name: Optional[str],
|
||||||
|
) -> None:
|
||||||
|
from litellm.proxy.proxy_server import prisma_client
|
||||||
|
|
||||||
|
if prisma_client is None:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
updated_user_row = await prisma_client.db.litellm_usertable.find_first(
|
||||||
|
where={"user_id": response["user_id"]}
|
||||||
|
)
|
||||||
|
if updated_user_row:
|
||||||
|
user_row_typed = LiteLLM_UserTable(
|
||||||
|
**updated_user_row.model_dump(exclude_none=True)
|
||||||
|
)
|
||||||
|
asyncio.create_task(
|
||||||
|
UserManagementEventHooks.create_internal_user_audit_log(
|
||||||
|
user_id=user_row_typed.user_id,
|
||||||
|
action="updated",
|
||||||
|
litellm_changed_by=litellm_changed_by or user_api_key_dict.user_id,
|
||||||
|
user_api_key_dict=user_api_key_dict,
|
||||||
|
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||||
|
before_value=(
|
||||||
|
existing_user_row.model_dump_json(exclude_none=True)
|
||||||
|
if existing_user_row
|
||||||
|
else None
|
||||||
|
),
|
||||||
|
after_value=user_row_typed.model_dump_json(exclude_none=True),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
except Exception as audit_error:
|
||||||
|
verbose_proxy_logger.warning(
|
||||||
|
f"Failed to create audit log for user {response.get('user_id')}: {audit_error}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _check_user_update_authz(
|
def _check_user_update_authz(
|
||||||
user_request: UpdateUserRequest,
|
user_request: UpdateUserRequest,
|
||||||
user_api_key_dict: UserAPIKeyAuth,
|
user_api_key_dict: UserAPIKeyAuth,
|
||||||
|
|
@ -1229,6 +1269,32 @@ async def _update_single_user_helper(
|
||||||
**existing_user_row.model_dump(exclude_none=True)
|
**existing_user_row.model_dump(exclude_none=True)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Prevent budget self-escalation (GHSA-wvg4-6222-3q4r): non-admin callers
|
||||||
|
# must not be able to raise their own budget/spend fields.
|
||||||
|
# can_user_call_user_update() already restricts non-admins to self-updates,
|
||||||
|
# so this guard only fires for self-escalation attempts.
|
||||||
|
_target_user_id = user_request.user_id or (
|
||||||
|
getattr(existing_user_row, "user_id", None)
|
||||||
|
if existing_user_row is not None
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
_is_self_update = (
|
||||||
|
_target_user_id is not None and user_api_key_dict.user_id == _target_user_id
|
||||||
|
)
|
||||||
|
if (
|
||||||
|
_is_self_update
|
||||||
|
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||||
|
):
|
||||||
|
_protected_fields = ("max_budget", "soft_budget", "spend")
|
||||||
|
for _field in _protected_fields:
|
||||||
|
if _field in non_default_values:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=403,
|
||||||
|
detail={
|
||||||
|
"error": f"Non-admin users cannot modify '{_field}' on their own record. Contact your proxy admin."
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
existing_metadata = (
|
existing_metadata = (
|
||||||
cast(Dict, getattr(existing_user_row, "metadata", {}) or {})
|
cast(Dict, getattr(existing_user_row, "metadata", {}) or {})
|
||||||
if existing_user_row is not None
|
if existing_user_row is not None
|
||||||
|
|
@ -1280,39 +1346,14 @@ async def _update_single_user_helper(
|
||||||
data=non_default_values, table_name="user"
|
data=non_default_values, table_name="user"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Create audit log for successful update
|
|
||||||
if response is not None:
|
if response is not None:
|
||||||
try:
|
await _schedule_user_update_audit_log(
|
||||||
updated_user_row = await prisma_client.db.litellm_usertable.find_first(
|
response=response,
|
||||||
where={"user_id": response["user_id"]}
|
existing_user_row=existing_user_row,
|
||||||
)
|
litellm_changed_by=litellm_changed_by,
|
||||||
|
user_api_key_dict=user_api_key_dict,
|
||||||
if updated_user_row:
|
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||||
user_row_typed = LiteLLM_UserTable(
|
)
|
||||||
**updated_user_row.model_dump(exclude_none=True)
|
|
||||||
)
|
|
||||||
|
|
||||||
# Create audit log asynchronously
|
|
||||||
asyncio.create_task(
|
|
||||||
UserManagementEventHooks.create_internal_user_audit_log(
|
|
||||||
user_id=user_row_typed.user_id,
|
|
||||||
action="updated",
|
|
||||||
litellm_changed_by=litellm_changed_by
|
|
||||||
or user_api_key_dict.user_id,
|
|
||||||
user_api_key_dict=user_api_key_dict,
|
|
||||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
|
||||||
before_value=(
|
|
||||||
existing_user_row.model_dump_json(exclude_none=True)
|
|
||||||
if existing_user_row
|
|
||||||
else None
|
|
||||||
),
|
|
||||||
after_value=user_row_typed.model_dump_json(exclude_none=True),
|
|
||||||
)
|
|
||||||
)
|
|
||||||
except Exception as audit_error:
|
|
||||||
verbose_proxy_logger.warning(
|
|
||||||
f"Failed to create audit log for user {response.get('user_id')}: {audit_error}"
|
|
||||||
)
|
|
||||||
|
|
||||||
if response is None:
|
if response is None:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,7 @@ import asyncio
|
||||||
import copy
|
import copy
|
||||||
import inspect
|
import inspect
|
||||||
import json
|
import json
|
||||||
|
import math
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
import secrets
|
import secrets
|
||||||
|
|
@ -608,6 +609,11 @@ async def validate_team_id_used_in_service_account_request(
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
_BUDGET_NUMERIC_KEYS = frozenset(
|
||||||
|
["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit"]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _enforce_upperbound_key_params(
|
def _enforce_upperbound_key_params(
|
||||||
data: Union[GenerateKeyRequest, UpdateKeyRequest],
|
data: Union[GenerateKeyRequest, UpdateKeyRequest],
|
||||||
fill_defaults: bool = True,
|
fill_defaults: bool = True,
|
||||||
|
|
@ -618,6 +624,21 @@ def _enforce_upperbound_key_params(
|
||||||
For key generation (fill_defaults=True): fills None values with upperbound defaults.
|
For key generation (fill_defaults=True): fills None values with upperbound defaults.
|
||||||
For key update (fill_defaults=False): only validates explicitly provided values.
|
For key update (fill_defaults=False): only validates explicitly provided values.
|
||||||
"""
|
"""
|
||||||
|
# Always reject NaN / Inf regardless of whether an upperbound config is set
|
||||||
|
# (GHSA-2rv4-xv66-fpjg): float('nan') passes every `< 0` check because
|
||||||
|
# nan < 0 is False, and spend >= nan is always False, permanently disabling
|
||||||
|
# budget enforcement for any key that carries it.
|
||||||
|
for elem in data:
|
||||||
|
key, value = elem
|
||||||
|
if key in _BUDGET_NUMERIC_KEYS and value is not None:
|
||||||
|
if not math.isfinite(value):
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=400,
|
||||||
|
detail={
|
||||||
|
"error": f"{key} must be a finite number. Received: {value}"
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
if litellm.upperbound_key_generate_params is None:
|
if litellm.upperbound_key_generate_params is None:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
|
@ -687,6 +708,11 @@ async def _common_key_generation_helper( # noqa: PLR0915
|
||||||
prisma_client=prisma_client,
|
prisma_client=prisma_client,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Capture the caller-supplied max_budget before any defaults or upperbound
|
||||||
|
# params can fill it, so the ceiling check only fires when the caller
|
||||||
|
# explicitly requested a budget.
|
||||||
|
_requested_max_budget = data.max_budget
|
||||||
|
|
||||||
# check if user set default key/generate params on config.yaml
|
# check if user set default key/generate params on config.yaml
|
||||||
if litellm.default_key_generate_params is not None:
|
if litellm.default_key_generate_params is not None:
|
||||||
for elem in data:
|
for elem in data:
|
||||||
|
|
@ -710,6 +736,25 @@ async def _common_key_generation_helper( # noqa: PLR0915
|
||||||
# check if user set upperbound key/generate params on config.yaml
|
# check if user set upperbound key/generate params on config.yaml
|
||||||
_enforce_upperbound_key_params(data, fill_defaults=True)
|
_enforce_upperbound_key_params(data, fill_defaults=True)
|
||||||
|
|
||||||
|
# Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller
|
||||||
|
# with an explicit budget cannot grant a key a higher budget than their own.
|
||||||
|
# Callers with max_budget=None (unlimited) can delegate any budget.
|
||||||
|
if (
|
||||||
|
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||||
|
and _requested_max_budget is not None
|
||||||
|
and user_api_key_dict.max_budget is not None
|
||||||
|
and _requested_max_budget > user_api_key_dict.max_budget
|
||||||
|
):
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=400,
|
||||||
|
detail={
|
||||||
|
"error": (
|
||||||
|
f"max_budget ({_requested_max_budget}) cannot exceed the caller's "
|
||||||
|
f"own max_budget ({user_api_key_dict.max_budget})."
|
||||||
|
)
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
# APPLY ENTERPRISE KEY MANAGEMENT PARAMS
|
# APPLY ENTERPRISE KEY MANAGEMENT PARAMS
|
||||||
try:
|
try:
|
||||||
from litellm_enterprise.proxy.management_endpoints.key_management_endpoints import (
|
from litellm_enterprise.proxy.management_endpoints.key_management_endpoints import (
|
||||||
|
|
@ -1416,19 +1461,24 @@ async def generate_key_fn(
|
||||||
|
|
||||||
await check_org_admin_can_generate_keys(user_api_key_dict=user_api_key_dict)
|
await check_org_admin_can_generate_keys(user_api_key_dict=user_api_key_dict)
|
||||||
|
|
||||||
# Validate budget values are not negative
|
# Validate budget values are not negative and are finite numbers
|
||||||
if data.max_budget is not None and data.max_budget < 0:
|
# (GHSA-2rv4-xv66-fpjg): float('nan') passes `< 0` because nan < 0 is False.
|
||||||
|
if data.max_budget is not None and (
|
||||||
|
not math.isfinite(data.max_budget) or data.max_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
|
"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if data.soft_budget is not None and data.soft_budget < 0:
|
if data.soft_budget is not None and (
|
||||||
|
not math.isfinite(data.soft_budget) or data.soft_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"
|
"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -1880,10 +1930,12 @@ def _validate_max_budget(max_budget: Optional[float]) -> None:
|
||||||
Raises:
|
Raises:
|
||||||
HTTPException: If max_budget is negative
|
HTTPException: If max_budget is negative
|
||||||
"""
|
"""
|
||||||
if max_budget is not None and max_budget < 0:
|
if max_budget is not None and (not math.isfinite(max_budget) or max_budget < 0):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={"error": f"max_budget cannot be negative. Received: {max_budget}"},
|
detail={
|
||||||
|
"error": f"max_budget must be a non-negative finite number. Received: {max_budget}"
|
||||||
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -2422,12 +2474,14 @@ async def update_key_fn( # noqa: PLR0915
|
||||||
)
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Validate budget values are not negative
|
# Validate budget values are not negative and are finite numbers
|
||||||
if data.max_budget is not None and data.max_budget < 0:
|
if data.max_budget is not None and (
|
||||||
|
not math.isfinite(data.max_budget) or data.max_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
|
"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,5 @@
|
||||||
|
import math
|
||||||
|
|
||||||
"""
|
"""
|
||||||
Endpoints for /organization operations
|
Endpoints for /organization operations
|
||||||
|
|
||||||
|
|
@ -220,18 +222,22 @@ async def new_organization(
|
||||||
)
|
)
|
||||||
|
|
||||||
# Validate budget values are not negative
|
# Validate budget values are not negative
|
||||||
if data.max_budget is not None and data.max_budget < 0:
|
if data.max_budget is not None and (
|
||||||
|
not math.isfinite(data.max_budget) or data.max_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
|
"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if data.soft_budget is not None and data.soft_budget < 0:
|
if data.soft_budget is not None and (
|
||||||
|
not math.isfinite(data.soft_budget) or data.soft_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"
|
"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -482,18 +488,22 @@ async def update_organization(
|
||||||
data = LiteLLM_OrganizationTableUpdate(**raw_data_with_flat_budget_fields)
|
data = LiteLLM_OrganizationTableUpdate(**raw_data_with_flat_budget_fields)
|
||||||
|
|
||||||
# Validate budget values are not negative
|
# Validate budget values are not negative
|
||||||
if data.max_budget is not None and data.max_budget < 0:
|
if data.max_budget is not None and (
|
||||||
|
not math.isfinite(data.max_budget) or data.max_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
|
"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if data.soft_budget is not None and data.soft_budget < 0:
|
if data.soft_budget is not None and (
|
||||||
|
not math.isfinite(data.soft_budget) or data.soft_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"
|
"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -10,6 +10,7 @@ All /team management endpoints
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import math
|
||||||
import json
|
import json
|
||||||
import traceback
|
import traceback
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
|
|
@ -914,25 +915,31 @@ async def new_team( # noqa: PLR0915
|
||||||
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
raise HTTPException(status_code=500, detail={"error": "No db connected"})
|
||||||
|
|
||||||
# Validate budget values are not negative
|
# Validate budget values are not negative
|
||||||
if data.max_budget is not None and data.max_budget < 0:
|
if data.max_budget is not None and (
|
||||||
|
not math.isfinite(data.max_budget) or data.max_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
|
"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if data.team_member_budget is not None and data.team_member_budget < 0:
|
if data.team_member_budget is not None and (
|
||||||
|
not math.isfinite(data.team_member_budget) or data.team_member_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"
|
"error": f"team_member_budget must be a non-negative finite number. Received: {data.team_member_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if data.soft_budget is not None and data.soft_budget < 0:
|
if data.soft_budget is not None and (
|
||||||
|
not math.isfinite(data.soft_budget) or data.soft_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"
|
"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -1595,25 +1602,31 @@ async def update_team( # noqa: PLR0915
|
||||||
verbose_proxy_logger.debug("/team/update - %s", data)
|
verbose_proxy_logger.debug("/team/update - %s", data)
|
||||||
|
|
||||||
# Validate budget values are not negative
|
# Validate budget values are not negative
|
||||||
if data.max_budget is not None and data.max_budget < 0:
|
if data.max_budget is not None and (
|
||||||
|
not math.isfinite(data.max_budget) or data.max_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
|
"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if data.team_member_budget is not None and data.team_member_budget < 0:
|
if data.team_member_budget is not None and (
|
||||||
|
not math.isfinite(data.team_member_budget) or data.team_member_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"
|
"error": f"team_member_budget must be a non-negative finite number. Received: {data.team_member_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if data.soft_budget is not None and data.soft_budget < 0:
|
if data.soft_budget is not None and (
|
||||||
|
not math.isfinite(data.soft_budget) or data.soft_budget < 0
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=400,
|
status_code=400,
|
||||||
detail={
|
detail={
|
||||||
"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"
|
"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -186,7 +186,7 @@ async def test_new_budget_negative_max_budget(client_and_mocks):
|
||||||
assert resp.status_code == 400, resp.text
|
assert resp.status_code == 400, resp.text
|
||||||
|
|
||||||
detail = resp.json()["detail"]
|
detail = resp.json()["detail"]
|
||||||
assert "max_budget cannot be negative" in str(detail)
|
assert "max_budget must be a non-negative finite number" in str(detail)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -204,7 +204,7 @@ async def test_new_budget_negative_soft_budget(client_and_mocks):
|
||||||
assert resp.status_code == 400, resp.text
|
assert resp.status_code == 400, resp.text
|
||||||
|
|
||||||
detail = resp.json()["detail"]
|
detail = resp.json()["detail"]
|
||||||
assert "soft_budget cannot be negative" in str(detail)
|
assert "soft_budget must be a non-negative finite number" in str(detail)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -222,7 +222,7 @@ async def test_update_budget_negative_max_budget(client_and_mocks):
|
||||||
assert resp.status_code == 400, resp.text
|
assert resp.status_code == 400, resp.text
|
||||||
|
|
||||||
detail = resp.json()["detail"]
|
detail = resp.json()["detail"]
|
||||||
assert "max_budget cannot be negative" in str(detail)
|
assert "max_budget must be a non-negative finite number" in str(detail)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -240,7 +240,7 @@ async def test_update_budget_negative_soft_budget(client_and_mocks):
|
||||||
assert resp.status_code == 400, resp.text
|
assert resp.status_code == 400, resp.text
|
||||||
|
|
||||||
detail = resp.json()["detail"]
|
detail = resp.json()["detail"]
|
||||||
assert "soft_budget cannot be negative" in str(detail)
|
assert "soft_budget must be a non-negative finite number" in str(detail)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
|
||||||
|
|
@ -2838,3 +2838,125 @@ def test_enforce_user_info_access_blocks_cross_user_lookup():
|
||||||
|
|
||||||
assert exc_info.value.status_code == 403
|
assert exc_info.value.status_code == 403
|
||||||
assert "key not allowed to access this user's info" in str(exc_info.value.detail)
|
assert "key not allowed to access this user's info" in str(exc_info.value.detail)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Regression tests for GHSA-wvg4-6222-3q4r: budget self-escalation via
|
||||||
|
# /user/update
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker):
|
||||||
|
"""Non-admin updating their own record must be blocked from modifying
|
||||||
|
max_budget (self-escalation)."""
|
||||||
|
from fastapi import HTTPException
|
||||||
|
|
||||||
|
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||||
|
_update_single_user_helper,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = mocker.MagicMock()
|
||||||
|
existing_user = mocker.MagicMock()
|
||||||
|
existing_user.model_dump.return_value = {
|
||||||
|
"user_id": "user-1",
|
||||||
|
"max_budget": 100,
|
||||||
|
}
|
||||||
|
existing_user.user_id = "user-1"
|
||||||
|
mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(
|
||||||
|
return_value=existing_user
|
||||||
|
)
|
||||||
|
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||||
|
|
||||||
|
user_request = UpdateUserRequest(
|
||||||
|
user_id="user-1",
|
||||||
|
max_budget=999999,
|
||||||
|
)
|
||||||
|
caller = UserAPIKeyAuth(
|
||||||
|
user_id="user-1",
|
||||||
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(HTTPException) as exc:
|
||||||
|
await _update_single_user_helper(
|
||||||
|
user_request=user_request, user_api_key_dict=caller
|
||||||
|
)
|
||||||
|
assert exc.value.status_code == 403
|
||||||
|
assert "max_budget" in str(exc.value.detail)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_wvg4_non_admin_cannot_self_escalate_spend(mocker):
|
||||||
|
"""Non-admin must not be able to reset their own spend to zero."""
|
||||||
|
from fastapi import HTTPException
|
||||||
|
|
||||||
|
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||||
|
_update_single_user_helper,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = mocker.MagicMock()
|
||||||
|
existing_user = mocker.MagicMock()
|
||||||
|
existing_user.model_dump.return_value = {
|
||||||
|
"user_id": "user-1",
|
||||||
|
"spend": 50.0,
|
||||||
|
}
|
||||||
|
existing_user.user_id = "user-1"
|
||||||
|
mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(
|
||||||
|
return_value=existing_user
|
||||||
|
)
|
||||||
|
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||||
|
|
||||||
|
user_request = UpdateUserRequest(
|
||||||
|
user_id="user-1",
|
||||||
|
spend=0,
|
||||||
|
)
|
||||||
|
caller = UserAPIKeyAuth(
|
||||||
|
user_id="user-1",
|
||||||
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(HTTPException) as exc:
|
||||||
|
await _update_single_user_helper(
|
||||||
|
user_request=user_request, user_api_key_dict=caller
|
||||||
|
)
|
||||||
|
assert exc.value.status_code == 403
|
||||||
|
assert "spend" in str(exc.value.detail)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker):
|
||||||
|
"""PROXY_ADMIN must still be able to modify another user's budget."""
|
||||||
|
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||||
|
_update_single_user_helper,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = mocker.MagicMock()
|
||||||
|
existing_user = mocker.MagicMock()
|
||||||
|
existing_user.model_dump.return_value = {
|
||||||
|
"user_id": "target-user",
|
||||||
|
"max_budget": 100,
|
||||||
|
}
|
||||||
|
existing_user.user_id = "target-user"
|
||||||
|
mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(
|
||||||
|
return_value=existing_user
|
||||||
|
)
|
||||||
|
mock_prisma_client.update_data = mocker.AsyncMock(
|
||||||
|
return_value={"user_id": "target-user", "max_budget": 500}
|
||||||
|
)
|
||||||
|
mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x)
|
||||||
|
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||||
|
mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
|
||||||
|
|
||||||
|
user_request = UpdateUserRequest(
|
||||||
|
user_id="target-user",
|
||||||
|
max_budget=500,
|
||||||
|
)
|
||||||
|
admin_caller = UserAPIKeyAuth(
|
||||||
|
user_id="admin-1",
|
||||||
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = await _update_single_user_helper(
|
||||||
|
user_request=user_request, user_api_key_dict=admin_caller
|
||||||
|
)
|
||||||
|
assert result is not None
|
||||||
|
|
|
||||||
|
|
@ -5540,6 +5540,9 @@ async def test_validate_max_budget():
|
||||||
_validate_max_budget(-10.0)
|
_validate_max_budget(-10.0)
|
||||||
|
|
||||||
assert exc_info.value.status_code == 400
|
assert exc_info.value.status_code == 400
|
||||||
|
assert "max_budget must be a non-negative finite number" in str(
|
||||||
|
exc_info.value.detail
|
||||||
|
)
|
||||||
assert "negative" in str(exc_info.value.detail)
|
assert "negative" in str(exc_info.value.detail)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -10979,3 +10982,222 @@ async def test_regenerate_premium_gate_allows_actual_master_key_holder():
|
||||||
)
|
)
|
||||||
|
|
||||||
assert result.token == "sk-new-master"
|
assert result.token == "sk-new-master"
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Regression tests for GHSA-q775-qw9r-2r4g: budget escalation via key/generate
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_q775_non_admin_unlimited_can_delegate_budget():
|
||||||
|
"""
|
||||||
|
Non-admin caller with max_budget=None (unlimited) can legitimately create
|
||||||
|
budget-capped keys. Any finite budget is within an unlimited ceiling.
|
||||||
|
"""
|
||||||
|
data = GenerateKeyRequest(max_budget=999999)
|
||||||
|
user_api_key_dict = UserAPIKeyAuth(
|
||||||
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||||
|
api_key="sk-internal",
|
||||||
|
user_id="user-1",
|
||||||
|
max_budget=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = AsyncMock()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||||
|
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
|
||||||
|
patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
|
||||||
|
patch(
|
||||||
|
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
return_value=MagicMock(),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
result = await generate_key_fn(
|
||||||
|
data=data,
|
||||||
|
user_api_key_dict=user_api_key_dict,
|
||||||
|
litellm_changed_by=None,
|
||||||
|
)
|
||||||
|
assert result is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_q775_non_admin_cannot_exceed_own_budget():
|
||||||
|
"""
|
||||||
|
Non-admin caller with max_budget=100 must not be able to create a key
|
||||||
|
with max_budget=500.
|
||||||
|
"""
|
||||||
|
data = GenerateKeyRequest(max_budget=500)
|
||||||
|
user_api_key_dict = UserAPIKeyAuth(
|
||||||
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||||
|
api_key="sk-internal",
|
||||||
|
user_id="user-1",
|
||||||
|
max_budget=100,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = AsyncMock()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||||
|
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
|
||||||
|
patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
|
||||||
|
):
|
||||||
|
with pytest.raises((HTTPException, ProxyException)) as exc_info:
|
||||||
|
await generate_key_fn(
|
||||||
|
data=data,
|
||||||
|
user_api_key_dict=user_api_key_dict,
|
||||||
|
litellm_changed_by=None,
|
||||||
|
)
|
||||||
|
err = exc_info.value
|
||||||
|
code = getattr(err, "status_code", None) or getattr(err, "code", None)
|
||||||
|
msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
|
||||||
|
assert str(code) == "400"
|
||||||
|
assert "cannot exceed" in msg.lower()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_q775_non_admin_within_budget_allowed():
|
||||||
|
"""
|
||||||
|
Non-admin caller with max_budget=100 can create a key with max_budget=50.
|
||||||
|
"""
|
||||||
|
data = GenerateKeyRequest(max_budget=50)
|
||||||
|
user_api_key_dict = UserAPIKeyAuth(
|
||||||
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||||
|
api_key="sk-internal",
|
||||||
|
user_id="user-1",
|
||||||
|
max_budget=100,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = AsyncMock()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||||
|
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
|
||||||
|
patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
|
||||||
|
patch(
|
||||||
|
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
return_value=MagicMock(),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
result = await generate_key_fn(
|
||||||
|
data=data,
|
||||||
|
user_api_key_dict=user_api_key_dict,
|
||||||
|
litellm_changed_by=None,
|
||||||
|
)
|
||||||
|
assert result is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_q775_upperbound_default_not_rejected():
|
||||||
|
"""
|
||||||
|
When upperbound_key_generate_params fills max_budget as a default, the
|
||||||
|
ceiling check must NOT fire — only explicitly requested budgets trigger it.
|
||||||
|
"""
|
||||||
|
data = GenerateKeyRequest()
|
||||||
|
assert data.max_budget is None
|
||||||
|
|
||||||
|
user_api_key_dict = UserAPIKeyAuth(
|
||||||
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||||
|
api_key="sk-internal",
|
||||||
|
user_id="user-1",
|
||||||
|
max_budget=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = AsyncMock()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||||
|
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
|
||||||
|
patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
|
||||||
|
patch(
|
||||||
|
"litellm.upperbound_key_generate_params",
|
||||||
|
MagicMock(max_budget=100.0),
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
return_value=MagicMock(),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
result = await generate_key_fn(
|
||||||
|
data=data,
|
||||||
|
user_api_key_dict=user_api_key_dict,
|
||||||
|
litellm_changed_by=None,
|
||||||
|
)
|
||||||
|
assert result is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_q775_default_key_generate_params_not_rejected():
|
||||||
|
"""
|
||||||
|
When default_key_generate_params fills max_budget, the ceiling check must
|
||||||
|
NOT fire — only caller-supplied budgets trigger it.
|
||||||
|
"""
|
||||||
|
data = GenerateKeyRequest()
|
||||||
|
assert data.max_budget is None
|
||||||
|
|
||||||
|
user_api_key_dict = UserAPIKeyAuth(
|
||||||
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||||
|
api_key="sk-internal",
|
||||||
|
user_id="user-1",
|
||||||
|
max_budget=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = AsyncMock()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||||
|
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
|
||||||
|
patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
|
||||||
|
patch(
|
||||||
|
"litellm.default_key_generate_params",
|
||||||
|
{"max_budget": 50.0},
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
return_value=MagicMock(),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
result = await generate_key_fn(
|
||||||
|
data=data,
|
||||||
|
user_api_key_dict=user_api_key_dict,
|
||||||
|
litellm_changed_by=None,
|
||||||
|
)
|
||||||
|
assert result is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_ghsa_q775_admin_bypasses_budget_ceiling():
|
||||||
|
"""
|
||||||
|
Admin caller can set any max_budget regardless of own budget.
|
||||||
|
"""
|
||||||
|
data = GenerateKeyRequest(max_budget=999999)
|
||||||
|
user_api_key_dict = UserAPIKeyAuth(
|
||||||
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||||
|
api_key="sk-admin",
|
||||||
|
user_id="admin-1",
|
||||||
|
max_budget=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_prisma_client = AsyncMock()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||||
|
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
|
||||||
|
patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
|
||||||
|
patch(
|
||||||
|
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
return_value=MagicMock(),
|
||||||
|
),
|
||||||
|
):
|
||||||
|
result = await generate_key_fn(
|
||||||
|
data=data,
|
||||||
|
user_api_key_dict=user_api_key_dict,
|
||||||
|
litellm_changed_by=None,
|
||||||
|
)
|
||||||
|
assert result is not None
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue