From 4eb244c3caeeaac1c8c60daf37a07a2ddc86aef7 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 11 Mar 2024 12:13:30 -0700 Subject: [PATCH 1/5] fix(proxy_server.py): prevent user from deleting non-user owned keys when they use ui --- litellm/proxy/proxy_server.py | 38 ++++++++++++++++++++++++++--------- litellm/proxy/utils.py | 18 ++++++++++++++--- 2 files changed, 44 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8510b350166..7917f14c3f6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2103,12 +2103,14 @@ async def generate_key_helper_fn( return key_data -async def delete_verification_token(tokens: List): +async def delete_verification_token(tokens: List, user_id: Optional[str] = None): global prisma_client try: if prisma_client: # Assuming 'db' is your Prisma Client instance - deleted_tokens = await prisma_client.delete_data(tokens=tokens) + deleted_tokens = await prisma_client.delete_data( + tokens=tokens, user_id=user_id + ) else: raise Exception except Exception as e: @@ -3744,7 +3746,10 @@ async def update_key_fn(request: Request, data: UpdateKeyRequest): @router.post( "/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)] ) -async def delete_key_fn(data: KeyRequest): +async def delete_key_fn( + data: KeyRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Delete a key from the key management system. @@ -3769,11 +3774,28 @@ async def delete_key_fn(data: KeyRequest): code=status.HTTP_400_BAD_REQUEST, ) - result = await delete_verification_token(tokens=keys) - verbose_proxy_logger.debug("/key/delete - deleted_keys=", result) + ## only allow user to delete keys they own + user_id = user_api_key_dict.user_id + if ( + user_api_key_dict.user_role is not None + and user_api_key_dict.user_role == "proxy_admin" + ): + user_id = None # unless they're admin - number_deleted_keys = len(result["deleted_keys"]) - assert len(keys) == number_deleted_keys + number_deleted_keys = await delete_verification_token( + tokens=keys, user_id=user_id + ) + verbose_proxy_logger.debug("/key/delete - deleted_keys=", number_deleted_keys) + + try: + assert len(keys) == number_deleted_keys + except Exception as e: + raise HTTPException( + status_code=400, + detail={ + "error": "Not all keys passed in were deleted. This probably means you don't have access to delete all the keys passed in." + }, + ) for key in keys: user_api_key_cache.delete_cache(key) @@ -6529,8 +6551,6 @@ async def login(request: Request): algorithm="HS256", ) litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token - # if a user has logged in they should be allowed to create keys - this ensures that it's set to True - general_settings["allow_user_auth"] = True return RedirectResponse(url=litellm_dashboard_ui, status_code=303) else: raise ProxyException( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 270b53647b3..d95f5a55002 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1356,9 +1356,12 @@ class PrismaClient: tokens: Optional[List] = None, team_id_list: Optional[List] = None, table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None, + user_id: Optional[str] = None, ): """ Allow user to delete a key(s) + + Ensure user owns that key, unless admin. """ try: if tokens is not None and isinstance(tokens, List): @@ -1369,15 +1372,23 @@ class PrismaClient: else: hashed_token = token hashed_tokens.append(hashed_token) - await self.db.litellm_verificationtoken.delete_many( - where={"token": {"in": hashed_tokens}} + filter_query: dict = {} + if user_id is not None: + filter_query = { + "AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}] + } + else: + filter_query = {"token": {"in": hashed_tokens}} + deleted_tokens = await self.db.litellm_verificationtoken.delete_many( + where=filter_query # type: ignore ) - return {"deleted_keys": tokens} + return {"deleted_keys": deleted_tokens} elif ( table_name == "team" and team_id_list is not None and isinstance(team_id_list, List) ): + # admin only endpoint -> `/team/delete` await self.db.litellm_teamtable.delete_many( where={"team_id": {"in": team_id_list}} ) @@ -1387,6 +1398,7 @@ class PrismaClient: and team_id_list is not None and isinstance(team_id_list, List) ): + # admin only endpoint -> `/team/delete` await self.db.litellm_verificationtoken.delete_many( where={"team_id": {"in": team_id_list}} ) From 2addd663939a88b02eab5ccb217cfc70f269fcb6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 11 Mar 2024 13:54:58 -0700 Subject: [PATCH 2/5] fix(proxy_server.py): bug fix --- litellm/proxy/proxy_server.py | 6 ++++-- litellm/proxy/utils.py | 2 ++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7917f14c3f6..86fd69fa6ff 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3785,10 +3785,12 @@ async def delete_key_fn( number_deleted_keys = await delete_verification_token( tokens=keys, user_id=user_id ) - verbose_proxy_logger.debug("/key/delete - deleted_keys=", number_deleted_keys) + verbose_proxy_logger.debug( + f"/key/delete - deleted_keys={number_deleted_keys['deleted_keys']}" + ) try: - assert len(keys) == number_deleted_keys + assert len(keys) == number_deleted_keys["deleted_keys"] except Exception as e: raise HTTPException( status_code=400, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d95f5a55002..4bfb87058bf 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1379,9 +1379,11 @@ class PrismaClient: } else: filter_query = {"token": {"in": hashed_tokens}} + deleted_tokens = await self.db.litellm_verificationtoken.delete_many( where=filter_query # type: ignore ) + verbose_proxy_logger.debug(f"deleted_tokens: {deleted_tokens}") return {"deleted_keys": deleted_tokens} elif ( table_name == "team" From d1644db8ce0e94541729ec35619f85c492a3f400 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 11 Mar 2024 14:18:01 -0700 Subject: [PATCH 3/5] test(test_key_generate_prisma.py): fix test to only let admin delete a key --- litellm/proxy/proxy_server.py | 3 +++ litellm/tests/test_key_generate_prisma.py | 13 ++++++++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 86fd69fa6ff..45f432f9df5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3776,6 +3776,9 @@ async def delete_key_fn( ## only allow user to delete keys they own user_id = user_api_key_dict.user_id + verbose_proxy_logger.debug( + f"user_api_key_dict.user_role: {user_api_key_dict.user_role}" + ) if ( user_api_key_dict.user_role is not None and user_api_key_dict.user_role == "proxy_admin" diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 524eee6f29b..74ff61abd13 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -737,8 +737,19 @@ def test_delete_key(prisma_client): delete_key_request = KeyRequest(keys=[generated_key]) + bearer_token = "Bearer sk-1234" + + request = Request(scope={"type": "http"}) + request._url = URL(url="/key/delete") + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print(f"result: {result}") + result.user_role = "proxy_admin" # delete the key - result_delete_key = await delete_key_fn(data=delete_key_request) + result_delete_key = await delete_key_fn( + data=delete_key_request, user_api_key_dict=result + ) print("result from delete key", result_delete_key) assert result_delete_key == {"deleted_keys": [generated_key]} From 64aeb088d9f9d4d0472735392c6b41a5920348e5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 11 Mar 2024 14:59:11 -0700 Subject: [PATCH 4/5] test(test_key_generate_prisma.py): fix test --- litellm/tests/test_key_generate_prisma.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 74ff61abd13..91c6d24dbd0 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -722,6 +722,7 @@ def test_delete_key(prisma_client): setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + setattr(litellm.proxy.proxy_server, "user_custom_auth", None) try: async def test(): From 917f92800de03322682a00ce68dc8e91aee7ad86 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 11 Mar 2024 15:24:42 -0700 Subject: [PATCH 5/5] test(test_key_generate_prisma.py): fix tests --- litellm/tests/test_key_generate_prisma.py | 43 ++++++++++++++++++++--- 1 file changed, 39 insertions(+), 4 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 91c6d24dbd0..62f6c38a955 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -788,7 +788,19 @@ def test_delete_key_auth(prisma_client): delete_key_request = KeyRequest(keys=[generated_key]) # delete the key - result_delete_key = await delete_key_fn(data=delete_key_request) + bearer_token = "Bearer sk-1234" + + request = Request(scope={"type": "http"}) + request._url = URL(url="/key/delete") + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print(f"result: {result}") + result.user_role = "proxy_admin" + + result_delete_key = await delete_key_fn( + data=delete_key_request, user_api_key_dict=result + ) print("result from delete key", result_delete_key) assert result_delete_key == {"deleted_keys": [generated_key]} @@ -803,6 +815,7 @@ def test_delete_key_auth(prisma_client): ) # use generated key to auth in + bearer_token = "Bearer " + generated_key result = await user_api_key_auth(request=request, api_key=bearer_token) print("got result", result) pytest.fail(f"This should have failed!. IT's an invalid key") @@ -847,9 +860,19 @@ def test_generate_and_call_key_info(prisma_client): # cleanup - delete key delete_key_request = KeyRequest(keys=[generated_key]) + bearer_token = "Bearer sk-1234" - # delete the key - await delete_key_fn(data=delete_key_request) + request = Request(scope={"type": "http"}) + request._url = URL(url="/key/delete") + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print(f"result: {result}") + result.user_role = "proxy_admin" + + result_delete_key = await delete_key_fn( + data=delete_key_request, user_api_key_dict=result + ) asyncio.run(test()) except Exception as e: @@ -928,7 +951,19 @@ def test_generate_and_update_key(prisma_client): delete_key_request = KeyRequest(keys=[generated_key]) # delete the key - await delete_key_fn(data=delete_key_request) + bearer_token = "Bearer sk-1234" + + request = Request(scope={"type": "http"}) + request._url = URL(url="/key/delete") + + # use generated key to auth in + result = await user_api_key_auth(request=request, api_key=bearer_token) + print(f"result: {result}") + result.user_role = "proxy_admin" + + result_delete_key = await delete_key_fn( + data=delete_key_request, user_api_key_dict=result + ) asyncio.run(test()) except Exception as e: