mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-25 08:19:03 +00:00
Internal User Endpoint - vulnerability fix + response type fix (#8228)
* fix(key_management_endpoints.py): fix vulnerability where a user could update another user's keys Resolves https://github.com/BerriAI/litellm/issues/8031 * test(key_management_endpoints.py): return consistent 403 forbidden error when modifying key that doesn't belong to user * fix(internal_user_endpoints.py): return model max budget in internal user create response Fixes https://github.com/BerriAI/litellm/issues/7047 * test: fix test * test: update test to handle gemini token counter change * fix(factory.py): fix bedrock http:// handling * docs: fix typo in lm_studio.md (#8222) * test: fix testing * test: fix test --------- Co-authored-by: foreign-sub <51928805+foreign-sub@users.noreply.github.com>
This commit is contained in:
co-authored by
foreign-sub
parent
f6bd48a1c5
commit
df93debbc7
@@ -1314,6 +1314,11 @@ def test_generate_and_update_key(prisma_client):
|
||||
budget_duration="1mo",
|
||||
max_budget=100,
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-1234",
|
||||
user_id="1234",
|
||||
),
|
||||
)
|
||||
|
||||
print("response1=", response1)
|
||||
@@ -1322,6 +1327,11 @@ def test_generate_and_update_key(prisma_client):
|
||||
response2 = await update_key_fn(
|
||||
request=Request,
|
||||
data=UpdateKeyRequest(key=generated_key, team_id=_team_2),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-1234",
|
||||
user_id="1234",
|
||||
),
|
||||
)
|
||||
print("response2=", response2)
|
||||
|
||||
@@ -2956,7 +2966,11 @@ async def test_generate_key_with_model_tpm_limit(prisma_client):
|
||||
_request = Request(scope={"type": "http"})
|
||||
_request._url = URL(url="/update/key")
|
||||
|
||||
await update_key_fn(data=request, request=_request)
|
||||
await update_key_fn(
|
||||
data=request,
|
||||
request=_request,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
result = await info_key_fn(
|
||||
key=generated_key,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
@@ -3017,7 +3031,11 @@ async def test_generate_key_with_guardrails(prisma_client):
|
||||
_request = Request(scope={"type": "http"})
|
||||
_request._url = URL(url="/update/key")
|
||||
|
||||
await update_key_fn(data=request, request=_request)
|
||||
await update_key_fn(
|
||||
data=request,
|
||||
request=_request,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
result = await info_key_fn(
|
||||
key=generated_key,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
@@ -3710,6 +3728,11 @@ async def test_key_alias_uniqueness(prisma_client):
|
||||
await update_key_fn(
|
||||
data=UpdateKeyRequest(key=key3.key, key_alias=unique_alias),
|
||||
request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-1234",
|
||||
user_id="1234",
|
||||
),
|
||||
)
|
||||
pytest.fail("Should not be able to update a key to use an existing alias")
|
||||
except Exception as e:
|
||||
@@ -3719,6 +3742,11 @@ async def test_key_alias_uniqueness(prisma_client):
|
||||
updated_key = await update_key_fn(
|
||||
data=UpdateKeyRequest(key=key1.key, key_alias=unique_alias),
|
||||
request=Request(scope={"type": "http"}),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
api_key="sk-1234",
|
||||
user_id="1234",
|
||||
),
|
||||
)
|
||||
assert updated_key is not None
|
||||
|
||||
|
||||
@@ -1216,14 +1216,14 @@ def test_litellm_verification_token_view_response_with_budget_table(
|
||||
)
|
||||
|
||||
|
||||
def test_is_allowed_to_create_key():
|
||||
def test_is_allowed_to_make_key_request():
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_is_allowed_to_create_key,
|
||||
_is_allowed_to_make_key_request,
|
||||
)
|
||||
|
||||
assert (
|
||||
_is_allowed_to_create_key(
|
||||
_is_allowed_to_make_key_request(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="test_user_id", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
),
|
||||
@@ -1234,7 +1234,7 @@ def test_is_allowed_to_create_key():
|
||||
)
|
||||
|
||||
assert (
|
||||
_is_allowed_to_create_key(
|
||||
_is_allowed_to_make_key_request(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="test_user_id",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
@@ -1553,6 +1553,7 @@ async def test_spend_logs_cleanup_after_error():
|
||||
mock_client.spend_log_transactions == original_logs[100:]
|
||||
), "Should remove processed logs even after error"
|
||||
|
||||
|
||||
def test_provider_specific_header():
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
add_provider_specific_headers_to_request,
|
||||
|
||||
Reference in New Issue
Block a user