scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

Merge pull request #2448 from BerriAI/litellm_delete_key_security_fix

Browse source 
fix(proxy_server.py): prevent user from deleting non-user owned keys
This commit is contained in:
Krish Dholakia 2024-03-11 15:47:34 -07:00 • committed by GitHub
commit 5647d945ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 103 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
litellm/proxy/proxy_server.py
View file

@ -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,33 @@ 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
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"
):
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(
f"/key/delete - deleted_keys={number_deleted_keys['deleted_keys']}"
)
try:
assert len(keys) == number_deleted_keys["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 +6556,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(

20
litellm/proxy/utils.py
View file

@ -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,25 @@ 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}
verbose_proxy_logger.debug(f"deleted_tokens: {deleted_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 +1400,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}}
)

57
litellm/tests/test_key_generate_prisma.py
View file

@ -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():
@ -737,8 +738,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]}
@ -776,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]}
@ -791,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")
@ -835,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:
@ -916,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: