scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

[Fix] Remove duplicate get_team_object call in _validate_update_key_data

Browse source 
Move the non-admin team validation into the existing get_team_object call
site to avoid an extra DB round-trip. The existing call already fetches
the team for limits checking — we now add the LIT-1884 guard there when
team_obj is None for non-admin callers.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang 2026-03-16 17:40:42 -07:00
parent bc752fb109
commit 208740a87c
2 changed files with 16 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1865,31 +1865,6 @@ async def _validate_update_key_data(
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,
@ -1929,6 +1904,13 @@ async def _validate_update_key_data(
check_db_only=True,
)
# Validate team exists when non-admin sets a new team_id (LIT-1884)
if team_obj is None and data.team_id is not None and not _is_proxy_admin:
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.",
)
if team_obj is not None:
await _check_team_key_limits(
team_table=team_obj,

12
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -7768,12 +7768,15 @@ class TestLIT1884KeyUpdateValidation:
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.
get_team_object raises HTTPException(404) when team doesn't exist in DB.
"""
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
existing_key_row.organization_id = None
existing_key_row.project_id = None
user_api_key_dict = UserAPIKeyAuth(
user_id="internal-user-123",
@ -7782,7 +7785,10 @@ class TestLIT1884KeyUpdateValidation:
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
AsyncMock(side_effect=Exception("Team not found")),
AsyncMock(side_effect=HTTPException(
status_code=404,
detail="Team doesn't exist in db. Team=nonexistent-team.",
)),
):
with pytest.raises(HTTPException) as exc_info:
await _validate_update_key_data(
@ -7794,8 +7800,8 @@ class TestLIT1884KeyUpdateValidation:
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)
assert exc_info.value.status_code == 404
assert "Team doesn't exist" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_admin_can_remove_user_id(self):