From 5b64a4127cd12d56b5b894bb2aaeb6248ab42ce2 Mon Sep 17 00:00:00 2001 From: chahatkesh Date: Mon, 10 Aug 2026 23:41:00 +0530 Subject: [PATCH] fix(proxy): reject duplicate virtual key secrets on /key/generate Key inserts used upsert with an empty update, so reusing a secret returned 200 while leaving only the first key in the database. Create with the unique constraint instead and surface duplicates as a 400. --- .../key_management_endpoints.py | 26 +++++++++++- litellm/proxy/utils.py | 34 ++++++++++----- .../test_key_management_endpoints.py | 41 +++++++++++++++++++ .../test_prisma_client_writes.py | 23 +++++++++++ 4 files changed, 111 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2de1d177b33..1c78717df03 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3929,7 +3929,21 @@ async def generate_key_helper_fn( "prisma_client: Creating Key= %s", {**key_data, "token": hash_token(token=token)}, ) - create_key_response: Final = await prisma_client.insert_data(data=key_data, table_name="key") + try: + create_key_response: Final = await prisma_client.insert_data( + data=key_data, + table_name="key", + ignore_duplicates=False, + ) + except Exception as e: + if _is_unique_constraint_failure(e): + raise ProxyException( + message="Key already exists. Use /key/update to modify an existing key.", + type=ProxyErrorTypes.bad_request_error, + param="key", + code=400, + ) from e + raise key_data["token_id"] = getattr(create_key_response, "token", None) key_data["litellm_budget_table"] = getattr(create_key_response, "litellm_budget_table", None) @@ -3947,7 +3961,7 @@ async def generate_key_helper_fn( except Exception as e: verbose_proxy_logger.error("litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - %s", e) verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): + if isinstance(e, (HTTPException, ProxyException)): raise e raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -6529,6 +6543,14 @@ def _validate_key_alias_format(key_alias: str | None) -> None: ) +def _is_unique_constraint_failure(exc: Exception) -> bool: + code: Final = getattr(exc, "code", None) + if code == "P2002": + return True + message: Final = str(exc) + return "P2002" in message or "Unique constraint failed" in message + + async def _enforce_unique_key_alias( key_alias: str | None, prisma_client: PrismaClient | None, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5f22ca021ac..f47fddb1b6f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3832,9 +3832,11 @@ class PrismaClient: self, data: dict, table_name: Literal["user", "key", "config", "spend", "team", "user_notification"], + ignore_duplicates: bool = True, ): """ - Add a key to the database. If it already exists, do nothing. + Add a key to the database. If it already exists and ignore_duplicates is True, do nothing. + When ignore_duplicates is False, raise on primary-key conflicts so callers can reject duplicates. """ start_time: Final = time.time() try: @@ -3851,16 +3853,26 @@ class PrismaClient: # Strip them so the DB stores NULL via the column's nullable constraint. if db_data.get("budget_limits") is None: db_data.pop("budget_limits", None) - print_verbose("PrismaClient: Before upsert into litellm_verificationtoken") - new_verification_token: Final = await VerificationTokenRepository(self).table.upsert( - where={ - "token": hashed_token, - }, - data={ - "create": {**db_data}, - "update": {}, # don't do anything if it already exists - }, - include={"litellm_budget_table": True}, + verification_token_table: Final = VerificationTokenRepository(self).table + print_verbose( + f"PrismaClient: Before {'create' if ignore_duplicates is False else 'upsert'} into litellm_verificationtoken" + ) + new_verification_token: Final = ( + await verification_token_table.create( + data={**db_data}, + include={"litellm_budget_table": True}, + ) + if ignore_duplicates is False + else await verification_token_table.upsert( + where={ + "token": hashed_token, + }, + data={ + "create": {**db_data}, + "update": {}, # don't do anything if it already exists + }, + include={"litellm_budget_table": True}, + ) ) verbose_proxy_logger.info("Data Inserted into Keys Table") return new_verification_token diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index e8709f3af34..7a16f6151f4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1483,6 +1483,47 @@ async def test_generate_key_fn_accepts_custom_key_at_minimum_length(monkeypatch) ) assert response.key == custom_key + mock_insert_data.assert_awaited() + assert mock_insert_data.await_args.kwargs.get("ignore_duplicates") is False + + +@pytest.mark.asyncio +async def test_generate_key_fn_rejects_duplicate_secret_key(monkeypatch): + """Regression for #20494: reusing an existing virtual-key secret must 400.""" + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles, ProxyException + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.insert_data = AsyncMock( + side_effect=Exception("Unique constraint failed on the fields: (`token`)") + ) + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(return_value=None) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached", + AsyncMock(return_value={}), + ) + + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=GenerateKeyRequest(key="sk-abcdefghijklmnop", key_alias="duplicate-secret"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234" + ), + ) + + assert str(exc_info.value.code) == "400" + assert "Key already exists" in str(exc_info.value.message) + assert mock_prisma_client.insert_data.await_args.kwargs.get("ignore_duplicates") is False @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py index dd241397e87..85102be4aa5 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py @@ -56,6 +56,29 @@ async def test_insert_data_hashes_token_and_upserts(prisma_client: PrismaClient) } +@pytest.mark.asyncio +async def test_insert_data_ignore_duplicates_false_uses_create( + prisma_client: PrismaClient, +) -> None: + token = "sk-secret-create-1" + expected_hash = hashlib.sha256(token.encode()).hexdigest() + response = SimpleNamespace(token=expected_hash, key_alias="alias", user_id="u1") + prisma_client.db.litellm_verificationtoken.create = AsyncMock(return_value=response) + prisma_client.db.litellm_verificationtoken.upsert = AsyncMock() + + result = await prisma_client.insert_data( + data={"token": token, "user_id": "u1", "metadata": {"a": 1}}, + table_name="key", + ignore_duplicates=False, + ) + + create_kwargs = prisma_client.db.litellm_verificationtoken.create.await_args.kwargs + assert result is response + assert create_kwargs["data"]["token"] == expected_hash + assert create_kwargs["include"] == {"litellm_budget_table": True} + prisma_client.db.litellm_verificationtoken.upsert.assert_not_awaited() + + @pytest.mark.asyncio async def test_insert_data_strips_null_budget_limits(prisma_client: PrismaClient) -> None: prisma_client.db.litellm_verificationtoken.upsert = AsyncMock(return_value=None)