Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 91 additions & 2 deletions litellm/proxy/management_endpoints/key_management_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.management_endpoints.common_utils import (
_is_user_org_admin_for_team,
_is_user_team_admin,
_set_object_metadata_field,
)
Expand Down Expand Up @@ -1836,6 +1837,18 @@ async def _validate_update_key_data(
user_api_key_cache=user_api_key_cache,
)

# Admin-only: only proxy admins, team admins, or org admins can modify max_budget
if data.max_budget is not None and data.max_budget != existing_key_row.max_budget:
if prisma_client is not None:
hashed_key = existing_key_row.token
await _check_key_admin_access(
user_api_key_dict=user_api_key_dict,
hashed_token=hashed_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
route="/key/update (max_budget)",
)
Comment on lines +1841 to +1850
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Admin check silently skipped when prisma_client is None

The if prisma_client is not None: guard means that if the DB is somehow unavailable at this call site, the max_budget privilege check is bypassed entirely and a non-admin user can silently raise their own budget.

While update_key_fn raises earlier when prisma_client is None, _validate_update_key_data is a reusable helper that could be called from other call sites without that guarantee. The safer approach is to fail loudly rather than silently permit the mutation — for example, raising an HTTPException(status_code=500) when prisma_client is None and a budget change is requested, so the absence of DB connectivity is an explicit error rather than a hidden privilege bypass.


# Check team limits if key has a team_id (from request or existing key)
team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
_team_id_to_check = data.team_id or getattr(existing_key_row, "team_id", None)
Expand Down Expand Up @@ -4733,6 +4746,64 @@ def _get_condition_to_filter_out_ui_session_tokens() -> Dict[str, Any]:
}


async def _check_key_admin_access(
user_api_key_dict: UserAPIKeyAuth,
hashed_token: str,
prisma_client: Any,
user_api_key_cache: DualCache,
route: str,
) -> None:
"""
Check that the caller has admin privileges for the target key.

Allowed callers:
- Proxy admin
- Team admin for the key's team
- Org admin for the key's team's organization

Raises HTTPException(403) if the caller is not authorized.
"""

if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return

# Look up the target key to find its team
target_key_row = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
)
Comment on lines +4770 to +4773
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Redundant DB fetch in the /key/update max_budget path

When _check_key_admin_access is called from _validate_update_key_data, the target key was already fetched as existing_key_row (which carries team_id). This find_unique call adds a redundant DB round-trip on every max_budget update request.

A simple improvement would be to add an optional team_id parameter to _check_key_admin_access that callers can provide when the team is already known, skipping this lookup. For _validate_update_key_data, you could pass getattr(existing_key_row, "team_id", None) directly.

if target_key_row is None:
raise HTTPException(
status_code=404,
detail={"error": f"Key not found: {hashed_token}"},
)

# If the key belongs to a team, check team admin / org admin
if target_key_row.team_id:
team_obj = await get_team_object(
team_id=target_key_row.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
check_db_only=True,
)
Comment on lines +4782 to +4787
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

check_db_only=True bypasses cache for team lookup

Passing check_db_only=True forces a DB hit every time this function runs, even when the team object is already warm in user_api_key_cache. For the block/unblock path this means at minimum two sequential DB queries per request (key lookup + team lookup).

Consider dropping check_db_only=True to allow the cache to serve the team object when available — the same pattern used elsewhere in the file for management-endpoint team lookups.

if team_obj is not None:
if _is_user_team_admin(
user_api_key_dict=user_api_key_dict, team_obj=team_obj
):
return
if await _is_user_org_admin_for_team(
user_api_key_dict=user_api_key_dict, team_obj=team_obj
):
return
Comment on lines +4780 to +4796
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Team admin silently denied when team lookup returns None

When the target key has a team_id but get_team_object returns None (e.g., the team record was deleted after the key was created), the entire if team_obj is not None: block is skipped and execution falls through to the unconditional raise HTTPException(403). This means a legitimate team admin operating on an orphaned-team key receives a confusing 403 with no indication that the team itself is missing.

The safer approach is to check team_obj is None explicitly inside the if target_key_row.team_id: branch and raise a 404 for a missing team rather than silently falling through to the generic authorization error. This makes the failure mode observable and distinguishable from a genuine authorization denial.


raise HTTPException(
status_code=403,
detail={
"error": f"Only proxy admins, team admins, or org admins can call {route}. "
f"user_role={user_api_key_dict.user_role}, user_id={user_api_key_dict.user_id}"
},
)


@router.post(
"/key/block", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
)
Expand Down Expand Up @@ -4762,7 +4833,7 @@ async def block_key(
}'
```

Note: This is an admin-only endpoint. Only proxy admins can block keys.
Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can block keys.
"""
from litellm.proxy.proxy_server import (
create_audit_log_for_update,
Expand All @@ -4788,6 +4859,15 @@ async def block_key(
else:
hashed_token = data.key

# Admin-only: only proxy admins, team admins, or org admins can block keys
await _check_key_admin_access(
user_api_key_dict=user_api_key_dict,
hashed_token=hashed_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
route="/key/block",
)

if litellm.store_audit_logs is True:
# make an audit log for key update
record = await prisma_client.db.litellm_verificationtoken.find_unique(
Expand Down Expand Up @@ -4876,7 +4956,7 @@ async def unblock_key(
}'
```

Note: This is an admin-only endpoint. Only proxy admins can unblock keys.
Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can unblock keys.
"""
from litellm.proxy.proxy_server import (
create_audit_log_for_update,
Expand All @@ -4902,6 +4982,15 @@ async def unblock_key(
else:
hashed_token = data.key

# Admin-only: only proxy admins, team admins, or org admins can unblock keys
await _check_key_admin_access(
user_api_key_dict=user_api_key_dict,
hashed_token=hashed_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
route="/key/unblock",
)

if litellm.store_audit_logs is True:
# make an audit log for key update
record = await prisma_client.db.litellm_verificationtoken.find_unique(
Expand Down
Loading
Loading