From 74c025c50ce94f9632bc167f939223e440a5e514 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 11:47:21 -0800 Subject: [PATCH 1/5] add enforcement for unique key aliases --- .../key_management_endpoints.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ba6df9a2e..b44e233d9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -420,6 +420,11 @@ async def generate_key_fn( # noqa: PLR0915 data_json.pop("tags") + await _enforce_unique_key_alias( + key_alias=data_json.get("key_alias", None), + prisma_client=prisma_client, + ) + response = await generate_key_helper_fn( request_type="key", **data_json, table_name="key" ) @@ -585,6 +590,11 @@ async def update_key_fn( data=data, existing_key_row=existing_key_row ) + await _enforce_unique_key_alias( + key_alias=non_default_values.get("key_alias", None), + prisma_client=prisma_client, + ) + response = await prisma_client.update_data( token=key, data={**non_default_values, "token": key} ) @@ -1883,3 +1893,30 @@ async def test_key_logging( status="healthy", details=f"No logger exceptions triggered, system is healthy. Manually check if logs were sent to {logging_callbacks} ", ) + + +async def _enforce_unique_key_alias( + key_alias: Optional[str], + prisma_client: Any, +) -> None: + """ + Helper to enforce unique key aliases across all keys. + + Args: + key_alias (Optional[str]): The key alias to check + prisma_client (Any): Prisma client instance + + Raises: + HTTPException: If key alias already exists + """ + if key_alias is not None and prisma_client is not None: + existing_key = await prisma_client.db.litellm_verificationtoken.find_first( + where={"key_alias": key_alias} + ) + if existing_key is not None: + raise HTTPException( + status_code=400, + detail={ + "error": f"Key alias '{key_alias}' already exists. Unique key aliases across all keys are required." + }, + ) From 3e50a06fb4006838e6854e20ba5995ce79f486fc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 16:22:04 -0800 Subject: [PATCH 2/5] fix _enforce_unique_key_alias --- .../management_endpoints/key_management_endpoints.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b44e233d9..d0670d062 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -593,6 +593,7 @@ async def update_key_fn( await _enforce_unique_key_alias( key_alias=non_default_values.get("key_alias", None), prisma_client=prisma_client, + existing_key_id=key, ) response = await prisma_client.update_data( @@ -1898,6 +1899,7 @@ async def test_key_logging( async def _enforce_unique_key_alias( key_alias: Optional[str], prisma_client: Any, + existing_key_id: Optional[str] = None, ) -> None: """ Helper to enforce unique key aliases across all keys. @@ -1905,13 +1907,19 @@ async def _enforce_unique_key_alias( Args: key_alias (Optional[str]): The key alias to check prisma_client (Any): Prisma client instance + existing_key_id (Optional[str]): ID of existing key being updated, to exclude from uniqueness check Raises: - HTTPException: If key alias already exists + HTTPException: If key alias already exists on a different key """ if key_alias is not None and prisma_client is not None: + where_clause: dict[str, Any] = {"key_alias": key_alias} + if existing_key_id: + # Exclude the current key from the uniqueness check + where_clause["NOT"] = {"token": existing_key_id} + existing_key = await prisma_client.db.litellm_verificationtoken.find_first( - where={"key_alias": key_alias} + where=where_clause ) if existing_key is not None: raise HTTPException( From de6ecdbc829c835800f531cedc680240ce02be68 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 17:17:59 -0800 Subject: [PATCH 3/5] fix _enforce_unique_key_alias --- .../management_endpoints/key_management_endpoints.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d0670d062..edaaf824f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1922,9 +1922,9 @@ async def _enforce_unique_key_alias( where=where_clause ) if existing_key is not None: - raise HTTPException( - status_code=400, - detail={ - "error": f"Key alias '{key_alias}' already exists. Unique key aliases across all keys are required." - }, + raise ProxyException( + message=f"Key alias '{key_alias}' already exists. Unique key aliases across all keys are required.", + type=ProxyErrorTypes.bad_request_error, + param="key_alias", + code=status.HTTP_400_BAD_REQUEST, ) From 9c3acfbb5dcd2e6e3379f095eb64f977ce0d6cac Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 17:46:47 -0800 Subject: [PATCH 4/5] fix _enforce_unique_key_alias --- .../key_management_endpoints.py | 4 +- .../test_key_generate_prisma.py | 73 +++++++++++++++++++ 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index edaaf824f..ef5b7b30d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -593,7 +593,7 @@ async def update_key_fn( await _enforce_unique_key_alias( key_alias=non_default_values.get("key_alias", None), prisma_client=prisma_client, - existing_key_id=key, + existing_key_id=existing_key_row.token, ) response = await prisma_client.update_data( @@ -1923,7 +1923,7 @@ async def _enforce_unique_key_alias( ) if existing_key is not None: raise ProxyException( - message=f"Key alias '{key_alias}' already exists. Unique key aliases across all keys are required.", + message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.", type=ProxyErrorTypes.bad_request_error, param="key_alias", code=status.HTTP_400_BAD_REQUEST, diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index e6f8ca541..935e4c352 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3550,3 +3550,76 @@ async def test_key_generate_with_secret_manager_call(prisma_client): ################################################################################ + + +@pytest.mark.asyncio +async def test_key_alias_uniqueness(prisma_client): + """ + Test that: + 1. We cannot create two keys with the same alias + 2. We cannot update a key to use an alias that's already taken + 3. We can update a key while keeping its existing alias + """ + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + try: + # Create first key with an alias + unique_alias = f"test-alias-{uuid.uuid4()}" + key1 = await generate_key_fn( + data=GenerateKeyRequest(key_alias=unique_alias), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + # Try to create second key with same alias - should fail + try: + key2 = await generate_key_fn( + data=GenerateKeyRequest(key_alias=unique_alias), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + pytest.fail("Should not be able to create a second key with the same alias") + except Exception as e: + print("vars(e)=", vars(e)) + assert "Unique key aliases across all keys are required" in str(e.message) + + # Create another key with different alias + another_alias = f"test-alias-{uuid.uuid4()}" + key3 = await generate_key_fn( + data=GenerateKeyRequest(key_alias=another_alias), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + # Try to update key3 to use key1's alias - should fail + try: + await update_key_fn( + data=UpdateKeyRequest(key=key3.key, key_alias=unique_alias), + request=Request(scope={"type": "http"}), + ) + pytest.fail("Should not be able to update a key to use an existing alias") + except Exception as e: + assert "Unique key aliases across all keys are required" in str(e.message) + + # Update key1 with its own existing alias - should succeed + updated_key = await update_key_fn( + data=UpdateKeyRequest(key=key1.key, key_alias=unique_alias), + request=Request(scope={"type": "http"}), + ) + assert updated_key is not None + + except Exception as e: + print("got exceptions, e=", e) + print("vars(e)=", vars(e)) + pytest.fail(f"An unexpected error occurred: {str(e)}") From 818a118bda6ea7abd02f0b424c0d5c0f51917246 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 18:02:33 -0800 Subject: [PATCH 5/5] test_enforce_unique_key_alias --- .../key_management_endpoints.py | 11 +-- .../test_key_generate_prisma.py | 76 +++++++++++++++++++ 2 files changed, 82 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ef5b7b30d..1856028c6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -593,7 +593,7 @@ async def update_key_fn( await _enforce_unique_key_alias( key_alias=non_default_values.get("key_alias", None), prisma_client=prisma_client, - existing_key_id=existing_key_row.token, + existing_key_token=existing_key_row.token, ) response = await prisma_client.update_data( @@ -1899,7 +1899,7 @@ async def test_key_logging( async def _enforce_unique_key_alias( key_alias: Optional[str], prisma_client: Any, - existing_key_id: Optional[str] = None, + existing_key_token: Optional[str] = None, ) -> None: """ Helper to enforce unique key aliases across all keys. @@ -1907,16 +1907,17 @@ async def _enforce_unique_key_alias( Args: key_alias (Optional[str]): The key alias to check prisma_client (Any): Prisma client instance - existing_key_id (Optional[str]): ID of existing key being updated, to exclude from uniqueness check + existing_key_token (Optional[str]): ID of existing key being updated, to exclude from uniqueness check + (The Admin UI passes key_alias, in all Edit key requests. So we need to be sure that if we find a key with the same alias, it's not the same key we're updating) Raises: HTTPException: If key alias already exists on a different key """ if key_alias is not None and prisma_client is not None: where_clause: dict[str, Any] = {"key_alias": key_alias} - if existing_key_id: + if existing_key_token: # Exclude the current key from the uniqueness check - where_clause["NOT"] = {"token": existing_key_id} + where_clause["NOT"] = {"token": existing_key_token} existing_key = await prisma_client.db.litellm_verificationtoken.find_first( where=where_clause diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 935e4c352..87bbe76d3 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3623,3 +3623,79 @@ async def test_key_alias_uniqueness(prisma_client): print("got exceptions, e=", e) print("vars(e)=", vars(e)) pytest.fail(f"An unexpected error occurred: {str(e)}") + + +@pytest.mark.asyncio +async def test_enforce_unique_key_alias(prisma_client): + """ + Unit test the _enforce_unique_key_alias function: + 1. Test it allows unique aliases + 2. Test it blocks duplicate aliases for new keys + 3. Test it allows updating a key with its own existing alias + 4. Test it blocks updating a key with another key's alias + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _enforce_unique_key_alias, + ) + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + await litellm.proxy.proxy_server.prisma_client.connect() + + try: + # Test 1: Allow unique alias + unique_alias = f"test-alias-{uuid.uuid4()}" + await _enforce_unique_key_alias( + key_alias=unique_alias, + prisma_client=prisma_client, + ) # Should pass + + # Create a key with this alias in the database + key1 = await generate_key_fn( + data=GenerateKeyRequest(key_alias=unique_alias), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + # Test 2: Block duplicate alias for new key + try: + await _enforce_unique_key_alias( + key_alias=unique_alias, + prisma_client=prisma_client, + ) + pytest.fail("Should not allow duplicate alias") + except Exception as e: + assert "Unique key aliases across all keys are required" in str(e.message) + + # Test 3: Allow updating key with its own alias + await _enforce_unique_key_alias( + key_alias=unique_alias, + existing_key_token=hash_token(key1.key), + prisma_client=prisma_client, + ) # Should pass + + # Test 4: Block updating with another key's alias + another_key = await generate_key_fn( + data=GenerateKeyRequest(key_alias=f"test-alias-{uuid.uuid4()}"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + try: + await _enforce_unique_key_alias( + key_alias=unique_alias, + existing_key_token=another_key.key, + prisma_client=prisma_client, + ) + pytest.fail("Should not allow using another key's alias") + except Exception as e: + assert "Unique key aliases across all keys are required" in str(e.message) + + except Exception as e: + print("Unexpected error:", e) + pytest.fail(f"An unexpected error occurred: {str(e)}")