mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-15 10:20:58 +00:00
[Fix] Prevent internal users from creating invalid keys via key/generate and key/update
Internal users could exploit key/generate and key/update to create unbound keys (no user_id, no budget) or attach keys to non-existent teams. This adds validation for non-admin callers: auto-assign user_id on generate, reject invalid team_ids, and prevent removing user_id on update. Closes LIT-1884 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
31a677e357
commit
bc752fb109
@@ -354,6 +354,10 @@ def key_generation_check(
|
||||
|
||||
## check if key is for team or individual
|
||||
is_team_key = _is_team_key(data=data)
|
||||
_is_admin = (
|
||||
user_api_key_dict.user_role is not None
|
||||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
if is_team_key:
|
||||
if team_table is None and litellm.key_generation_settings is not None:
|
||||
raise HTTPException(
|
||||
@@ -361,7 +365,13 @@ def key_generation_check(
|
||||
detail=f"Unable to find team object in database. Team ID: {data.team_id}",
|
||||
)
|
||||
elif team_table is None:
|
||||
return True # assume user is assigning team_id without using the team table
|
||||
if _is_admin:
|
||||
return True # admins can assign team_id without team table
|
||||
# Non-admin callers must have a valid team (LIT-1884)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Unable to find team object in database. Team ID: {data.team_id}",
|
||||
)
|
||||
return _team_key_generation_check(
|
||||
team_table=team_table,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
@@ -1214,6 +1224,19 @@ async def generate_key_fn(
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN, detail=message
|
||||
)
|
||||
# For non-admin internal users: auto-assign caller's user_id if not provided
|
||||
# This prevents creating unbound keys with no user association (LIT-1884)
|
||||
_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 and data.user_id is None:
|
||||
data.user_id = user_api_key_dict.user_id
|
||||
verbose_proxy_logger.warning(
|
||||
"key/generate: auto-assigning user_id=%s for non-admin caller",
|
||||
user_api_key_dict.user_id,
|
||||
)
|
||||
|
||||
team_table: Optional[LiteLLM_TeamTableCachedObj] = None
|
||||
if data.team_id is not None:
|
||||
try:
|
||||
@@ -1228,6 +1251,12 @@ async def generate_key_fn(
|
||||
verbose_proxy_logger.debug(
|
||||
f"Error getting team object in `/key/generate`: {e}"
|
||||
)
|
||||
# For non-admin callers, team must exist (LIT-1884)
|
||||
if not _is_proxy_admin:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot create keys for non-existent teams.",
|
||||
)
|
||||
|
||||
key_generation_check(
|
||||
team_table=team_table,
|
||||
@@ -1810,17 +1839,57 @@ async def _validate_update_key_data(
|
||||
user_api_key_cache: Any,
|
||||
) -> None:
|
||||
"""Validate permissions and constraints for key update."""
|
||||
_is_proxy_admin = (
|
||||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
)
|
||||
|
||||
# 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 user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||
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 team exists when non-admin changes team_id (LIT-1884)
|
||||
if (
|
||||
data.team_id is not None
|
||||
and not _is_proxy_admin
|
||||
):
|
||||
try:
|
||||
_team_obj = await get_team_object(
|
||||
team_id=data.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
check_db_only=True,
|
||||
)
|
||||
if _team_obj is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.",
|
||||
)
|
||||
|
||||
common_key_access_checks(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
|
||||
@@ -41,12 +41,15 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_transform_verification_tokens_to_deleted_records,
|
||||
_validate_max_budget,
|
||||
_validate_reset_spend_value,
|
||||
_validate_update_key_data,
|
||||
can_modify_verification_token,
|
||||
check_org_key_model_specific_limits,
|
||||
check_team_key_model_specific_limits,
|
||||
delete_verification_tokens,
|
||||
generate_key_fn,
|
||||
generate_key_helper_fn,
|
||||
key_aliases,
|
||||
key_generation_check,
|
||||
list_keys,
|
||||
prepare_key_update_data,
|
||||
reset_key_spend_fn,
|
||||
@@ -7535,3 +7538,292 @@ async def test_update_key_non_budget_fields_allowed_for_internal_user(monkeypatc
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# LIT-1884: Internal users cannot create invalid keys
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class TestLIT1884KeyGenerateValidation:
|
||||
"""Tests for LIT-1884: internal users should not be able to generate invalid keys."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_generate_key_no_user_id_auto_assigns(self):
|
||||
"""
|
||||
When an internal_user calls /key/generate without user_id,
|
||||
the caller's user_id should be auto-assigned before reaching
|
||||
_common_key_generation_helper.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
data = GenerateKeyRequest(key_alias="test-alias")
|
||||
assert data.user_id is None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
# Patch _common_key_generation_helper to avoid needing full DB mocks.
|
||||
# We just want to verify user_id is set before we reach this point.
|
||||
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(),
|
||||
):
|
||||
await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
# The data object should have been mutated to include the caller's user_id
|
||||
assert data.user_id == "internal-user-123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_generate_key_invalid_team_id_rejected(self):
|
||||
"""
|
||||
When an internal_user provides a non-existent team_id,
|
||||
key/generate should raise ProxyException with status 400.
|
||||
"""
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
data = GenerateKeyRequest(
|
||||
key_alias="test-alias",
|
||||
team_id="nonexistent-team-id",
|
||||
)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
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.get_team_object",
|
||||
AsyncMock(side_effect=Exception("Team not found")),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
assert str(exc_info.value.code) == "400"
|
||||
assert "Team not found" in str(exc_info.value.message)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_generate_key_invalid_team_id_allowed(self):
|
||||
"""
|
||||
Admin callers should be allowed to create keys with any team_id,
|
||||
even if the team doesn't exist (team_table=None is OK for admins).
|
||||
"""
|
||||
data = GenerateKeyRequest(
|
||||
key_alias="admin-key",
|
||||
team_id="nonexistent-team-id",
|
||||
user_id="admin-user",
|
||||
)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
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.get_team_object",
|
||||
AsyncMock(side_effect=Exception("Team not found")),
|
||||
), \
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
|
||||
new_callable=AsyncMock,
|
||||
return_value=MagicMock(),
|
||||
):
|
||||
# Should NOT raise — admin bypasses team validation
|
||||
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_admin_generate_key_no_user_id_not_auto_assigned(self):
|
||||
"""
|
||||
Admin callers should NOT have user_id auto-assigned — they may
|
||||
intentionally create keys without a user_id.
|
||||
"""
|
||||
data = GenerateKeyRequest(key_alias="admin-unbound-key")
|
||||
assert data.user_id is None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
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(),
|
||||
):
|
||||
await generate_key_fn(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
# user_id should remain None for admin
|
||||
assert data.user_id is None
|
||||
|
||||
def test_key_generation_check_non_admin_no_team_table_raises(self):
|
||||
"""
|
||||
key_generation_check should raise 400 for non-admin when team_table is None
|
||||
and key_generation_settings is not set.
|
||||
"""
|
||||
data = GenerateKeyRequest(team_id="some-team-id")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
with patch.object(litellm, "key_generation_settings", None):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
key_generation_check(
|
||||
team_table=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
route="key_generate",
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Unable to find team object" in str(exc_info.value.detail)
|
||||
|
||||
def test_key_generation_check_admin_no_team_table_allowed(self):
|
||||
"""
|
||||
key_generation_check should allow admin to proceed even when team_table is None.
|
||||
"""
|
||||
data = GenerateKeyRequest(team_id="some-team-id")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
with patch.object(litellm, "key_generation_settings", None):
|
||||
result = key_generation_check(
|
||||
team_table=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
route="key_generate",
|
||||
)
|
||||
assert result is True
|
||||
|
||||
|
||||
class TestLIT1884KeyUpdateValidation:
|
||||
"""Tests for LIT-1884: internal users should not be able to update keys to remove user_id or set invalid team."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_cannot_remove_user_id(self):
|
||||
"""
|
||||
Non-admin users should not be able to set user_id to empty string (remove it).
|
||||
"""
|
||||
data = UpdateKeyRequest(key="sk-test-key", user_id="")
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.user_id = "internal-user-123"
|
||||
existing_key_row.token = "hashed_token"
|
||||
existing_key_row.team_id = None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _validate_update_key_data(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=None,
|
||||
premium_user=False,
|
||||
prisma_client=AsyncMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "cannot remove the user_id" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_internal_user_cannot_set_invalid_team_id(self):
|
||||
"""
|
||||
Non-admin users should not be able to update a key to a non-existent team.
|
||||
"""
|
||||
data = UpdateKeyRequest(key="sk-test-key", team_id="nonexistent-team")
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.user_id = "internal-user-123"
|
||||
existing_key_row.token = "hashed_token"
|
||||
existing_key_row.team_id = None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="internal-user-123",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
AsyncMock(side_effect=Exception("Team not found")),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _validate_update_key_data(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=None,
|
||||
premium_user=False,
|
||||
prisma_client=AsyncMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "Team not found" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_can_remove_user_id(self):
|
||||
"""
|
||||
Admin users should be allowed to set user_id to empty string.
|
||||
"""
|
||||
data = UpdateKeyRequest(key="sk-test-key", user_id="")
|
||||
existing_key_row = MagicMock()
|
||||
existing_key_row.user_id = "some-user"
|
||||
existing_key_row.token = "hashed_token"
|
||||
existing_key_row.team_id = None
|
||||
existing_key_row.organization_id = None
|
||||
existing_key_row.project_id = None
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
# Should NOT raise
|
||||
await _validate_update_key_data(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=None,
|
||||
premium_user=False,
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user