scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

[Fix] Proxy/Key Management: Align Key-Org Membership Checks On Generate And Regenerate

Browse source 
Mirrors the membership rule on /key/update so that /key/generate and
/key/{key}/regenerate apply the same `_validate_caller_can_assign_key_org`
gate when the caller specifies an `organization_id`. Proxy admins bypass.
The check no-ops when `organization_id` is not being set.
This commit is contained in:
Yuneng Jiang 2026-05-01 18:19:24 -07:00
parent e78d87ee00
commit 92d3bdbb27
No known key found for this signature in database
1 changed files with 30 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -809,6 +809,20 @@ async def _common_key_generation_helper( # noqa: PLR0915
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client:
# Mirror the membership rule applied to /key/update: when the
# caller specifies an organization_id, require that they are a
# member of (or proxy admin over) the target organization.
_is_proxy_admin = (
user_api_key_dict.user_role is not None
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
if not _is_proxy_admin:
await _validate_caller_can_assign_key_org(
user_api_key_dict=user_api_key_dict,
organization_id=data.organization_id,
prisma_client=prisma_client,
)
org_table = await get_org_object(
org_id=data.organization_id,
user_api_key_cache=user_api_key_cache,
@ -3920,6 +3934,22 @@ async def _execute_virtual_key_regeneration(
"""Generate new token, update DB, invalidate cache, and return response."""
from litellm.proxy.proxy_server import hash_token
# Apply the same membership rule used on /key/update: when the caller
# asks to point the regenerated key at a different organization_id,
# require they are a member of (or proxy admin over) the target org.
if data is not None and data.organization_id is not None:
_existing_org_id = getattr(key_in_db, "organization_id", None)
_is_proxy_admin = (
user_api_key_dict.user_role is not None
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
if data.organization_id != _existing_org_id and not _is_proxy_admin:
await _validate_caller_can_assign_key_org(
user_api_key_dict=user_api_key_dict,
organization_id=data.organization_id,
prisma_client=prisma_client,
)
new_token = await get_new_token(data=data)
new_token_hash = hash_token(new_token)
new_token_key_name = f"sk-...{new_token[-4:]}"