mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
chore(proxy): close /key/regenerate ownership-rebind + premium-gate bypass
A non-admin caller could rebind their own key's ``user_id`` via ``/key/regenerate``. ``_execute_virtual_key_regeneration`` had org/team guards but no ``user_id`` guard, and ``prepare_key_update_data`` did not strip the field — it survived ``model_dump(exclude_unset=True)`` into the Prisma update. On the next request, ``_return_user_api_key_auth_obj`` resolved the rebound ``user_id`` against ``litellm_usertable`` and returned ``PROXY_ADMIN`` whenever the target row's ``user_role`` was admin (e.g. the default ``user_id="default_user_id"`` created on first password-UI login). ``/key/update`` had the equivalent guard inline at ``_validate_update_key_data``; extract it to a shared helper ``_validate_caller_can_change_key_ownership`` and call from both ``/key/update`` and ``_execute_virtual_key_regeneration``. Future regenerate-style endpoints inherit the guard for free. Also tighten the premium gate that allowed the master-key rotation branch to skip the enterprise check. The previous predicate was ``data.new_master_key is not None`` — a field-presence test, not an identity check. Any non-premium caller could send any value in that field and the premium check would no-op. Verify the caller actually holds the master key via ``_is_master_key`` before allowing the non-premium path. Tests: - ``test_regenerate_user_id_rebind_guard`` — parametrized table over cross-user rebind (blocked), empty-string removal (blocked), and same-user no-op rebind (allowed). - ``test_regenerate_premium_gate_requires_actual_master_key`` / ``test_regenerate_premium_gate_allows_actual_master_key_holder`` — ensure the premium check requires the caller actually present the master key, and that legitimate master-key rotation still works.
This commit is contained in:
user
parent
63a2d1ddc9
commit
7738b03501
2 changed files with 208 additions and 18 deletions
|
|
@ -463,6 +463,45 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
|
|||
_NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS = frozenset({"llm_api_routes", "info_routes"})
|
||||
|
||||
|
||||
def _validate_caller_can_change_key_ownership(
|
||||
data: Optional[BaseModel],
|
||||
existing_key_row: Any,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""
|
||||
Non-admin callers must not rebind a key's ``user_id`` to a different
|
||||
user. The ``user_id`` on a verification token is what
|
||||
``_return_user_api_key_auth_obj`` resolves against ``litellm_usertable``
|
||||
to derive the request's role; a non-admin rebinding their own key's
|
||||
``user_id`` to a ``PROXY_ADMIN`` row promotes themselves.
|
||||
|
||||
``/key/update`` already enforces this inline; ``/key/regenerate`` did
|
||||
not. Sharing the check keeps both endpoints — and any future
|
||||
regenerate-style endpoint — consistent.
|
||||
"""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
return
|
||||
if data is None:
|
||||
return
|
||||
incoming_user_id = getattr(data, "user_id", None)
|
||||
if incoming_user_id is None:
|
||||
return
|
||||
if incoming_user_id == "":
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Non-admin users cannot remove the user_id from a key.",
|
||||
)
|
||||
existing_user_id = getattr(existing_key_row, "user_id", None)
|
||||
if incoming_user_id != existing_user_id:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(
|
||||
f"User={incoming_user_id} is not allowed to update the key "
|
||||
f"to belong to user={existing_user_id}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _check_allowed_routes_caller_permission(
|
||||
allowed_routes: Optional[list],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -2088,23 +2127,11 @@ async def _validate_update_key_data(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Prevent non-admin from removing user_id (setting to empty string) (LIT-1884)
|
||||
if data.user_id is not None and data.user_id == "" and not _is_proxy_admin:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Non-admin users cannot remove the user_id from a key.",
|
||||
)
|
||||
|
||||
# sanity check - prevent non-proxy admin user from updating key to belong to a different user
|
||||
if (
|
||||
data.user_id is not None
|
||||
and data.user_id != existing_key_row.user_id
|
||||
and not _is_proxy_admin
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"User={data.user_id} is not allowed to update key={data.key} to belong to user={existing_key_row.user_id}",
|
||||
)
|
||||
_validate_caller_can_change_key_ownership(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
common_key_access_checks(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -4152,6 +4179,13 @@ async def _execute_virtual_key_regeneration(
|
|||
"""Generate new token, update DB, invalidate cache, and return response."""
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
# Mirror the /key/update ownership rebind guard. See helper docstring.
|
||||
_validate_caller_can_change_key_ownership(
|
||||
data=data,
|
||||
existing_key_row=key_in_db,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# 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.
|
||||
|
|
@ -4328,7 +4362,17 @@ async def regenerate_key_fn( # noqa: PLR0915
|
|||
allow_safe_presets=True,
|
||||
)
|
||||
|
||||
is_master_key_regeneration = data and data.new_master_key is not None
|
||||
# Premium-gate bypass for master-key rotation must verify the
|
||||
# caller actually holds the master key, not just that the request
|
||||
# body has a ``new_master_key`` field. A presence-only check let
|
||||
# any non-premium caller skip the enterprise gate by sending any
|
||||
# value in that field.
|
||||
regenerate_target_key = data.key if data and data.key else key
|
||||
is_master_key_regeneration = (
|
||||
data is not None
|
||||
and data.new_master_key is not None
|
||||
and _is_master_key(api_key=regenerate_target_key, _master_key=master_key)
|
||||
)
|
||||
|
||||
if (
|
||||
premium_user is not True and not is_master_key_regeneration
|
||||
|
|
|
|||
|
|
@ -10596,3 +10596,149 @@ async def test_bulk_update_team_keys_blocks_metadata_allowed_passthrough_routes(
|
|||
assert exc.value.status_code == 403
|
||||
assert "allowed_passthrough_routes" in str(exc.value.detail)
|
||||
mock.update_data.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /key/regenerate ownership-rebind guard + premium-gate identity check
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
import contextlib # noqa: E402
|
||||
|
||||
|
||||
def _non_admin_user_api_key_dict():
|
||||
return UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
api_key="sk-internal",
|
||||
user_id="user-1",
|
||||
)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _patch_regenerate_side_effects():
|
||||
"""Mock out token creation + DB write + cache + rotation hook so
|
||||
``_execute_virtual_key_regeneration`` runs to completion under test."""
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value="sk-newtoken1234ab12",
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"incoming_user_id,expected_status,expected_substring",
|
||||
[
|
||||
# Cross-user rebind: the privesc primitive.
|
||||
("default_user_id", 403, "not allowed to update the key"),
|
||||
# Empty-string removal: companion guard.
|
||||
("", 403, "remove the user_id"),
|
||||
# No-op rebind (caller sends their own user_id): must succeed.
|
||||
("user-1", None, None),
|
||||
],
|
||||
ids=["rebind_blocked", "empty_blocked", "same_user_id_allowed"],
|
||||
)
|
||||
async def test_regenerate_user_id_rebind_guard(
|
||||
incoming_user_id, expected_status, expected_substring
|
||||
):
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_execute_virtual_key_regeneration,
|
||||
)
|
||||
|
||||
existing_key = _make_regenerate_existing_key()
|
||||
data = RegenerateKeyRequest(user_id=incoming_user_id)
|
||||
|
||||
async def _run():
|
||||
await _execute_virtual_key_regeneration(
|
||||
prisma_client=_make_regenerate_mock_prisma(),
|
||||
key_in_db=existing_key,
|
||||
hashed_api_key="abc123",
|
||||
key="abc123",
|
||||
data=data,
|
||||
user_api_key_dict=_non_admin_user_api_key_dict(),
|
||||
litellm_changed_by=None,
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
if expected_status is None:
|
||||
with _patch_regenerate_side_effects():
|
||||
await _run()
|
||||
return
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _run()
|
||||
assert exc.value.status_code == expected_status
|
||||
assert expected_substring in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_premium_gate_requires_actual_master_key():
|
||||
# ``regenerate_key_fn``'s decorator wraps the underlying ValueError
|
||||
# into a ProxyException with empty ``message``. The exception type
|
||||
# alone confirms the premium gate fired.
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
regenerate_key_fn,
|
||||
)
|
||||
|
||||
data = RegenerateKeyRequest(key="sk-not-master", new_master_key="anything")
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-the-real-master-key"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
|
||||
pytest.raises((ValueError, HTTPException, ProxyException)),
|
||||
):
|
||||
await regenerate_key_fn(
|
||||
key="sk-not-master",
|
||||
data=data,
|
||||
user_api_key_dict=_non_admin_user_api_key_dict(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_premium_gate_allows_actual_master_key_holder():
|
||||
from litellm.proxy._types import RegenerateKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
regenerate_key_fn,
|
||||
)
|
||||
|
||||
master = "sk-the-real-master-key"
|
||||
data = RegenerateKeyRequest(key=master, new_master_key="sk-new-master")
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||||
patch("litellm.proxy.proxy_server.master_key", master),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._rotate_master_key",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
):
|
||||
result = await regenerate_key_fn(
|
||||
key=master,
|
||||
data=data,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key=master,
|
||||
user_id="admin",
|
||||
),
|
||||
)
|
||||
|
||||
assert result.token == "sk-new-master"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue