tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-09 16:24:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #26862 from stuxf/codex/control-field-sanitization

Browse Source 
chore(proxy): harden request control fields
This commit is contained in:
yuneng-jiang
2026-04-30 13:10:58 -07:00
committed by GitHub
parent 9637d8c17b f48dfdbdd9
commit aa76ab2df7
20 changed files with 1074 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/proxy_unit_tests/test_audit_logs_proxy.py
+118 -2
View File
@@ -45,8 +45,11 @@ verbose_proxy_logger.setLevel(level=logging.DEBUG)
from starlette.datastructures import URL
from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames
from litellm.proxy.management_helpers.audit_logs import (
create_audit_log_for_update,
get_audit_log_changed_by,
)
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth
from litellm.caching.caching import DualCache
from unittest.mock import patch, AsyncMock
@@ -54,6 +57,119 @@ proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
import json
def test_get_audit_log_changed_by_prefers_authenticated_user():
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
user_id="authenticated-user",
)
assert (
get_audit_log_changed_by(
litellm_changed_by="spoofed-user",
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name="proxy-admin",
)
== "authenticated-user"
)
def test_get_audit_log_changed_by_honors_header_with_admin_opt_in():
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
user_id="service-account",
metadata={"allow_litellm_changed_by_header": True},
)
assert (
get_audit_log_changed_by(
litellm_changed_by="delegated-user",
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name="proxy-admin",
)
== "delegated-user"
)
def test_get_audit_log_changed_by_honors_header_with_team_opt_in():
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
user_id="service-account",
team_metadata={"allow_litellm_changed_by_header": True},
)
assert (
get_audit_log_changed_by(
litellm_changed_by="delegated-user",
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name="proxy-admin",
)
== "delegated-user"
)
def test_get_audit_log_changed_by_ignores_header_without_opt_in_when_user_id_missing():
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
assert (
get_audit_log_changed_by(
litellm_changed_by="spoofed-user",
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name="proxy-admin",
)
== "proxy-admin"
)
def test_get_audit_log_changed_by_honors_header_with_opt_in_when_user_id_missing():
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
metadata={"allow_litellm_changed_by_header": True},
)
assert (
get_audit_log_changed_by(
litellm_changed_by="delegated-user",
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name="proxy-admin",
)
== "delegated-user"
)
@pytest.mark.asyncio
async def test_create_internal_user_audit_log_uses_changed_by_helper():
from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
user_api_key_dict = UserAPIKeyAuth(
api_key="test-key",
user_id="service-account",
metadata={"allow_litellm_changed_by_header": True},
)
with (
patch("litellm.store_audit_logs", True),
patch(
"litellm.proxy.hooks.user_management_event_hooks.create_audit_log_for_update",
new_callable=AsyncMock,
) as mock_create_audit_log_for_update,
):
await UserManagementEventHooks.create_internal_user_audit_log(
user_id="target-user",
action="updated",
litellm_changed_by="delegated-user",
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name="proxy-admin",
before_value='{"before": true}',
after_value='{"after": true}',
)
request_data = mock_create_audit_log_for_update.await_args.kwargs["request_data"]
assert request_data.changed_by == "delegated-user"
assert request_data.changed_by_api_key == "test-key"
assert request_data.object_id == "target-user"
assert request_data.action == "updated"
@pytest.mark.asyncio
async def test_create_audit_log_for_update_premium_user():
"""
tests/proxy_unit_tests/test_proxy_server.py
+8 -4
View File
@@ -1553,6 +1553,7 @@ async def test_add_callback_via_key(prisma_client):
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(
metadata={
"allow_client_mock_response": True,
"logging": [
{
"callback_name": "langfuse", # 'otel', 'langfuse', 'lunary'
@@ -1563,7 +1564,7 @@ async def test_add_callback_via_key(prisma_client):
"langfuse_host": "https://us.cloud.langfuse.com",
},
}
]
],
}
),
)
@@ -1657,6 +1658,7 @@ async def test_add_callback_via_key_litellm_pre_call_utils(
team_id=None,
max_parallel_requests=None,
metadata={
"allow_client_mock_response": True,
"logging": [
{
"callback_name": "langfuse",
@@ -1667,7 +1669,7 @@ async def test_add_callback_via_key_litellm_pre_call_utils(
"langfuse_host": "https://us.cloud.langfuse.com",
},
}
]
],
},
tpm_limit=None,
rpm_limit=None,
@@ -1813,6 +1815,7 @@ async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket(
team_id=None,
max_parallel_requests=None,
metadata={
"allow_client_mock_response": True,
"logging": [
{
"callback_name": "gcs_bucket",
@@ -1822,7 +1825,7 @@ async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket(
"gcs_path_service_account": "pathrise-convert-1606954137718-a956eef1a2a8.json",
},
}
]
],
},
tpm_limit=None,
rpm_limit=None,
@@ -1946,6 +1949,7 @@ async def test_add_callback_via_key_litellm_pre_call_utils_langsmith(
team_id=None,
max_parallel_requests=None,
metadata={
"allow_client_mock_response": True,
"logging": [
{
"callback_name": "langsmith",
@@ -1956,7 +1960,7 @@ async def test_add_callback_via_key_litellm_pre_call_utils_langsmith(
"langsmith_base_url": "https://api.smith.langchain.com",
},
}
]
],
},
tpm_limit=None,
rpm_limit=None,