fix(users): validate custom IDs and reject concurrent duplicate creation

This commit is contained in:
Mrinal Chanshetty 2026-11-02 16:07:54 -08:00
parent 6c4bad4604
commit 5da3813d63
8 changed files with 149 additions and 6 deletions

View file

@ -5482,7 +5482,9 @@ async def generate_key_helper_fn(
if query_type == "insert_data":
user_row = cast( # cast-ok: table_name="user" is the insert_data branch returning the user row
"prisma_models.LiteLLM_UserTable | None",
await prisma_client.insert_data(data=user_data, table_name="user"),
await prisma_client.insert_data(
data=user_data, table_name="user", allow_existing_user=request_type != "user"
),
)
if user_row is None:

View file

@ -6452,6 +6452,8 @@ class PrismaClient:
self,
data: Mapping[str, object],
table_name: Literal["user", "key", "config", "spend", "team", "user_notification"],
*,
allow_existing_user: bool = True,
):
"""
Add a key to the database. If it already exists, do nothing.
@ -6487,14 +6489,23 @@ class PrismaClient:
verbose_proxy_logger.info("Data Inserted into Keys Table")
return new_verification_token
elif table_name == "user":
from prisma.errors import UniqueViolationError
db_data = self.jsonify_object(data=data)
try:
new_user_row: Final = await UserRepository(self).upsert_user(db_data)
new_user_row: Final = await UserRepository(self).upsert_user(
db_data, allow_existing=allow_existing_user
)
if new_user_row is None:
raise HTTPException(
status_code=403,
detail="SSO for more than 5 users requires a LiteLLM Enterprise license. Existing users can still sign in.",
)
except UniqueViolationError as e:
raise HTTPException(
status_code=409,
detail={"error": f"User with id {db_data['user_id']} already exists"},
) from e
except Exception as e:
if (
"Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`"

View file

@ -168,7 +168,9 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
"""
return await _count_billable_users(self.table)
async def upsert_user(self, data: Mapping[str, object]) -> "prisma_models.LiteLLM_UserTable | None":
async def upsert_user(
self, data: Mapping[str, object], *, allow_existing: bool = True
) -> "prisma_models.LiteLLM_UserTable | None":
"""Return the user row, or None when a new free SSO user would exceed capacity."""
scope: Final = _FREE_SSO_CREATION.get()
if (
@ -176,6 +178,8 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
or (scope[0] is not None and scope[0] != data["user_id"])
or scope[1] is not asyncio.current_task()
):
if not allow_existing:
return await self.table.create(data=dict(data))
return await self.table.upsert(
where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}}
)

View file

@ -120,6 +120,7 @@ class BulkNewUserItem(NewUserRequest):
model_config = ConfigDict(extra="forbid", protected_namespaces=())
user_id: str | None = None
auto_create_key: bool = False
@field_validator("send_invite_email")

View file

@ -1,7 +1,7 @@
from collections.abc import Mapping
from typing import Literal
from typing import Annotated, Final, Literal
from pydantic import Field, field_validator, model_validator
from pydantic import AfterValidator, Field, field_validator, model_validator
from litellm.models.user import LiteLLM_UserTable
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
@ -15,7 +15,17 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase):
user_role: Literal["user", "admin"] = "user"
def _validate_custom_user_id(value: str) -> str:
normalized: Final = value.strip()
if not normalized:
raise ValueError("user_id must not be empty or contain only whitespace")
if "\x00" in normalized:
raise ValueError("user_id must not contain NUL characters")
return normalized
class NewUserRequest(GenerateRequestBase):
user_id: Annotated[str, AfterValidator(_validate_custom_user_id)] | None = None
tpm_limit: int | None = Field(default=None, ge=0)
rpm_limit: int | None = Field(default=None, ge=0)
max_budget: float | None = None

View file

@ -13,6 +13,7 @@ import httpx
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
from pydantic import ValidationError
from pytest_mock import MockerFixture
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@ -48,6 +49,48 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import (
client = TestClient(app)
@pytest.mark.parametrize("user_id", ["", " \t\n", "user\x00id", 123, True, [], {}])
def test_new_user_rejects_invalid_custom_id(user_id: object) -> None:
with pytest.raises(ValidationError) as exc:
NewUserRequest.model_validate({"user_id": user_id})
assert exc.value.errors()[0]["loc"] == ("user_id",)
@pytest.mark.parametrize("user_id", [None, "employee-42", "社員@example.test", " employee-42 "])
def test_new_user_normalizes_custom_id_and_preserves_auto_generation(user_id: str | None) -> None:
request: Final = NewUserRequest(user_id=user_id)
assert request.user_id == (user_id.strip() if user_id is not None else None)
@pytest.mark.asyncio
async def test_new_user_returns_conflict_when_id_is_taken_after_duplicate_check(mocker: MockerFixture) -> None:
from prisma.errors import UniqueViolationError
from litellm.proxy import proxy_server
from litellm.proxy.utils import PrismaClient
table: Final = SimpleNamespace(
find_first=AsyncMock(return_value=None),
count=AsyncMock(return_value=0),
create=AsyncMock(side_effect=UniqueViolationError({}, message="Unique constraint failed on user_id")),
upsert=AsyncMock(return_value=LiteLLM_UserTable(user_id="employee-42")),
)
mocker.patch("prisma.Prisma", return_value=SimpleNamespace(litellm_usertable=table))
database: Final = PrismaClient(
database_url="postgresql://unused:unused@localhost/unused",
proxy_logging_obj=proxy_server.proxy_logging_obj,
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", database)
with pytest.raises(ProxyException) as exc:
await new_user(
data=NewUserRequest(user_id="employee-42", auto_create_key=False, send_invite_email=False),
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert str(exc.value.code) == "409"
assert "employee-42" in exc.value.message
if TYPE_CHECKING:
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.utils import PrismaClient
@ -4862,7 +4905,9 @@ async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget,
models = []
class _FakePrisma:
async def insert_data(self, data, table_name):
async def insert_data(
self, data: Mapping[str, object], table_name: str, *, allow_existing_user: bool = True
) -> _FakeUserRow | SimpleNamespace:
if table_name == "user":
captured["user_data"] = dict(data)
return _FakeUserRow()

View file

@ -22,6 +22,7 @@ from litellm.proxy._types import (
LiteLLM_JWTAuth,
LiteLLM_UserTable,
NewUserResponse,
SSOUserDefinedValues,
ProxyException,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
@ -52,6 +53,62 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
from tests._master_key import MASTER_KEY
@pytest.mark.asyncio
@pytest.mark.parametrize(
("billable_users", "premium", "admitted"),
[(4, False, True), (5, False, False), (26, False, False), (26, True, True)],
)
async def test_sso_auto_creation_enforces_free_user_cap_through_new_user(
monkeypatch: pytest.MonkeyPatch, billable_users: int, premium: bool, admitted: bool
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.utils import PrismaClient
row: Final = LiteLLM_UserTable(user_id="new-sso-user", user_email="new-sso-user@example.test")
table: Final = SimpleNamespace(
find_first=AsyncMock(return_value=None),
find_unique=AsyncMock(return_value=None),
count=AsyncMock(side_effect=lambda where=None: billable_users if where is None else 0),
create=AsyncMock(return_value=row),
upsert=AsyncMock(return_value=row),
)
transaction: Final = MagicMock()
transaction.__aenter__ = AsyncMock(
return_value=SimpleNamespace(litellm_usertable=table, execute_raw=AsyncMock(), query_raw=AsyncMock())
)
transaction.__aexit__ = AsyncMock(return_value=None)
monkeypatch.setattr(
"prisma.Prisma",
lambda **kwargs: SimpleNamespace(litellm_usertable=table, tx=MagicMock(return_value=transaction)),
)
database: Final = PrismaClient(
database_url="postgresql://unused:unused@localhost/unused",
proxy_logging_obj=proxy_server.proxy_logging_obj,
)
monkeypatch.setattr(proxy_server, "prisma_client", database)
monkeypatch.setattr(proxy_server, "premium_user", premium)
monkeypatch.setattr("asyncio.sleep", AsyncMock())
values: Final[SSOUserDefinedValues] = {
"models": [],
"user_id": row.user_id,
"user_email": row.user_email,
"user_role": "internal_user_viewer",
"max_budget": None,
"budget_duration": None,
}
if admitted:
response: Final = await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values)
assert response.user_id == row.user_id
assert table.create.await_count + table.upsert.await_count == 1
else:
with pytest.raises(ProxyException) as error:
await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values)
assert str(error.value.code) == "403"
assert "SSO for more than 5 users" in error.value.message
table.create.assert_not_awaited()
table.upsert.assert_not_awaited()
def _sso_http_handler(route: Callable[[httpx.Request], httpx.Response] | None = None) -> AsyncHTTPHandler:
"""litellm's SSO httpx client answered in-process, by ``route`` or by the active respx mock."""
handler = AsyncHTTPHandler()

View file

@ -191,6 +191,19 @@ async def _run(prisma, users, caller=ADMIN, license=None, generate_key=_no_keys)
)
@pytest.mark.asyncio
@pytest.mark.parametrize("user_id", [None, "", " employee-42 ", " \t"])
async def test_bulk_keeps_existing_user_id_generation_and_whitespace_behavior(user_id: str | None) -> None:
prisma: Final = _FakePrisma()
response: Final = await _run(prisma, [{"user_id": user_id}])
result: Final = response.data[0]
assert result.success
assert result.user_id
assert result.user_id in prisma.db.litellm_usertable.rows
if user_id:
assert result.user_id == user_id
@pytest.mark.asyncio
async def test_creates_users_and_team_membership_in_every_store():
prisma = _FakePrisma(teams=[_team("t1", [Member(user_id="existing", role="admin")]), _team("t2")])