fix: address three CI failures from recent security PR merges

- url_utils.py: narrow sockaddr[0] from str|int to str via a helper with a
  fail-closed isinstance check. Fixes the two mypy errors introduced by
  the SSRF hardening without masking unexpected stdlib behavior.

- key_management_endpoints.py: restore the documented team member_permissions
  path for /key/update. The cross-key admin check added to close the
  cross-org rewrite attack was over-broad: it rejected non-admin team
  members even when can_team_member_execute_key_management_endpoint had
  already validated their team membership and /key/update grant. Now skip
  the admin check when the key has a team_id and the change is non-budget
  (membership + permission already enforced above). Budget/spend changes
  still require team/org admin. The cross-org attack remains blocked:
  an outside org admin fails the earlier team membership check.

- test_logging_redaction_e2e_test.py: rename and rewrite two parametrized
  tests to assert that request-body turn_off_message_logging has no effect.
  Reflects the intentional removal of turn_off_message_logging from
  _supported_callback_params so the caller cannot override admin logging
  policy via the request body.

- test_key_management_endpoints.py: add two tests covering the restored
  team member permission path — one positive (non-budget update succeeds
  for a team member with /key/update grant), one negative (max_budget
  change still rejected without admin role).
This commit is contained in:
Yuneng Jiang 2026-04-17 15:11:45 -07:00
parent 7c66edbf4a
commit ee2cf0e6e8
No known key found for this signature in database
4 changed files with 283 additions and 58 deletions

View file

@ -78,6 +78,22 @@ def _format_host_header(hostname: str, port: int, default_port: int) -> str:
return f"{bracketed}:{port}"
def _sockaddr_host(sockaddr: Any) -> str:
"""Return the host element of a ``getaddrinfo`` sockaddr as ``str``.
``getaddrinfo`` with ``IPPROTO_TCP`` returns AF_INET / AF_INET6 sockaddrs
whose first element is always a host string. mypy types it as
``str | int`` (since sockaddrs for other families can hold ints), so we
narrow at the boundary. Fail closed if the stdlib ever returns something
unexpected — a non-string here would mean we have no IP to check against
the SSRF blocklist.
"""
host = sockaddr[0]
if not isinstance(host, str):
raise SSRFError(f"getaddrinfo returned non-string host: {host!r}")
return host
def _is_host_allowlisted(hostname: str, effective_port: int) -> bool:
"""Check whether a host is in the admin-configured allowlist.
@ -148,9 +164,10 @@ def validate_url(url: str) -> Tuple[str, str]:
if not is_allowlisted:
for family, type_, proto, canonname, sockaddr in addrinfo:
if _is_blocked_ip(sockaddr[0]):
resolved_ip = _sockaddr_host(sockaddr)
if _is_blocked_ip(resolved_ip):
raise SSRFError(
f"URL targets a blocked address ({sockaddr[0]}). "
f"URL targets a blocked address ({resolved_ip}). "
"If this is a legitimate internal service, add the host "
"to `user_url_allowed_hosts` in general_settings."
)
@ -166,7 +183,7 @@ def validate_url(url: str) -> Tuple[str, str]:
# For HTTP, rewrite URL to connect to the validated IP directly
# to prevent DNS rebinding (no TLS to bind the connection).
validated_ip = addrinfo[0][4][0]
validated_ip = _sockaddr_host(addrinfo[0][4])
is_ipv6 = addrinfo[0][0] == socket.AF_INET6
ip_host = f"[{validated_ip}]" if is_ipv6 else validated_ip

View file

@ -768,9 +768,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
request_type="key", **data_json, table_name="key"
)
response[
"soft_budget"
] = data.soft_budget # include the user-input soft budget in the response
response["soft_budget"] = (
data.soft_budget
) # include the user-input soft budget in the response
response = GenerateKeyResponse(**response)
@ -1970,26 +1970,37 @@ async def _validate_update_key_data(
# Policy:
# - Key owner (same user_id): may update non-budget fields on their
# own key without the admin check.
# - Anyone else (non-PROXY_ADMIN, not the owner): must pass
# _check_key_admin_access (PROXY_ADMIN / key-owner / team-admin /
# org-admin of the key).
# - Team member with /key/update grant (on a team key): may update
# non-budget fields. Team membership + permission is already
# enforced by can_team_member_execute_key_management_endpoint
# above, which raises 401 for non-members or members without the
# grant — so reaching this point on a team key means the caller
# was authorized via member_permissions. This preserves the
# documented member_permissions feature while still blocking the
# cross-org attack (an outside org admin is not a member of the
# victim team and gets rejected at the earlier check).
# - Anyone else (non-PROXY_ADMIN, not the owner, not a team member
# on a team key): must pass _check_key_admin_access (PROXY_ADMIN
# / key-owner / team-admin / org-admin of the key).
# - max_budget / spend: always require the admin check, even for the
# key owner (matches the existing admin-only budget semantics).
# key owner or a team member (matches the existing admin-only
# budget semantics).
is_key_owner = (
user_api_key_dict.user_id is not None
and existing_key_row.user_id == user_api_key_dict.user_id
)
_is_budget_change = (
data.max_budget is not None
and data.max_budget != existing_key_row.max_budget
data.max_budget is not None and data.max_budget != existing_key_row.max_budget
) or (
data.spend is not None
and data.spend != getattr(existing_key_row, "spend", None)
)
is_team_key = existing_key_row.team_id is not None
can_skip_admin_check_for_non_budget = is_key_owner or is_team_key
if (
(not _is_proxy_admin)
and prisma_client is not None
and (not is_key_owner or _is_budget_change)
and (_is_budget_change or not can_skip_admin_check_for_non_budget)
):
hashed_key = existing_key_row.token
await _check_key_admin_access(
@ -1998,9 +2009,7 @@ async def _validate_update_key_data(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
route=(
"/key/update (max_budget/spend)"
if _is_budget_change
else "/key/update"
"/key/update (max_budget/spend)" if _is_budget_change else "/key/update"
),
)
@ -3303,10 +3312,10 @@ async def delete_verification_tokens(
try:
if prisma_client:
tokens = [_hash_token_if_needed(token=key) for key in tokens]
_keys_being_deleted: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
_keys_being_deleted: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
)
)
if len(_keys_being_deleted) == 0:
@ -3506,9 +3515,9 @@ async def _rotate_master_key( # noqa: PLR0915
from litellm.proxy.proxy_server import proxy_config
try:
models: Optional[
List
] = await prisma_client.db.litellm_proxymodeltable.find_many()
models: Optional[List] = (
await prisma_client.db.litellm_proxymodeltable.find_many()
)
except Exception:
models = None
# 2. process model table
@ -4148,11 +4157,11 @@ async def validate_key_list_check(
param="user_id",
code=status.HTTP_403_FORBIDDEN,
)
complete_user_info_db_obj: Optional[
BaseModel
] = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
complete_user_info_db_obj: Optional[BaseModel] = (
await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
)
)
if complete_user_info_db_obj is None:
@ -4235,10 +4244,10 @@ async def _fetch_user_team_objects(
if complete_user_info is None or not complete_user_info.teams:
return []
teams: Optional[
List[BaseModel]
] = await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
teams: Optional[List[BaseModel]] = (
await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
)
)
if teams is None:
return []

View file

@ -56,7 +56,13 @@ async def test_global_redaction_on():
@pytest.mark.parametrize("turn_off_message_logging", [True, False])
@pytest.mark.asyncio
async def test_global_redaction_with_dynamic_params(turn_off_message_logging):
async def test_global_redaction_ignores_dynamic_param(turn_off_message_logging):
"""
Request-body `turn_off_message_logging` is no longer honored as a dynamic
callback param — global setting (or admin-configured key/team config) wins.
With global redaction ON, the caller cannot disable redaction via the
request body.
"""
litellm.turn_off_message_logging = True
test_custom_logger = TestCustomLogger()
litellm.callbacks = [test_custom_logger]
@ -75,23 +81,20 @@ async def test_global_redaction_with_dynamic_params(turn_off_message_logging):
json.dumps(standard_logging_payload, indent=2),
)
if turn_off_message_logging is True:
response = standard_logging_payload["response"]
assert response["choices"][0]["message"]["content"] == "redacted-by-litellm"
assert (
standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
)
else:
assert (
standard_logging_payload["response"]["choices"][0]["message"]["content"]
== "hello"
)
assert standard_logging_payload["messages"][0]["content"] == "hi"
response = standard_logging_payload["response"]
assert response["choices"][0]["message"]["content"] == "redacted-by-litellm"
assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
@pytest.mark.parametrize("turn_off_message_logging", [True, False])
@pytest.mark.asyncio
async def test_global_redaction_off_with_dynamic_params(turn_off_message_logging):
async def test_global_redaction_off_ignores_dynamic_param(turn_off_message_logging):
"""
Request-body `turn_off_message_logging` is no longer honored as a dynamic
callback param — global setting (or admin-configured key/team config) wins.
With global redaction OFF, the caller cannot enable redaction via the
request body.
"""
litellm.turn_off_message_logging = False
test_custom_logger = TestCustomLogger()
litellm.callbacks = [test_custom_logger]
@ -109,18 +112,11 @@ async def test_global_redaction_off_with_dynamic_params(turn_off_message_logging
"logged standard logging payload",
json.dumps(standard_logging_payload, indent=2),
)
if turn_off_message_logging is True:
response = standard_logging_payload["response"]
assert response["choices"][0]["message"]["content"] == "redacted-by-litellm"
assert (
standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
)
else:
assert (
standard_logging_payload["response"]["choices"][0]["message"]["content"]
== "hello"
)
assert standard_logging_payload["messages"][0]["content"] == "hi"
assert (
standard_logging_payload["response"]["choices"][0]["message"]["content"]
== "hello"
)
assert standard_logging_payload["messages"][0]["content"] == "hi"
@pytest.mark.asyncio

View file

@ -8098,6 +8098,209 @@ async def test_update_key_non_budget_rejects_cross_user_modification(monkeypatch
assert str(exc.value.code) == "403"
@pytest.mark.asyncio
async def test_update_key_team_member_with_permission_can_update_non_budget(
monkeypatch,
):
"""A team member whose team grants /key/update in member_permissions can
update non-budget fields on a team key even though they are not a team
admin. Regression: the cross-key admin check was over-broad and rejected
this documented path."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
test_hashed_token = "deadbeef" * 8
team_id = "team-with-update-grant"
member_user_id = "team-member-user"
mock_existing_key = MagicMock()
mock_existing_key.token = test_hashed_token
mock_existing_key.user_id = None # team-scoped key (no owning user)
mock_existing_key.team_id = team_id
mock_existing_key.project_id = None
mock_existing_key.max_budget = 10.0
mock_existing_key.key_alias = "original"
mock_existing_key.models = []
mock_existing_key.model_dump.return_value = {
"token": test_hashed_token,
"user_id": None,
"team_id": team_id,
"max_budget": 10.0,
}
team_table = LiteLLM_TeamTableCachedObj(
team_id=team_id,
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="some-team-admin", role="admin"),
Member(user_id=member_user_id, role="user"),
],
team_member_permissions=["/key/update", "/key/info"],
)
mock_updated_key = MagicMock()
mock_updated_key.token = test_hashed_token
mock_updated_key.key_alias = "renamed-by-member"
mock_prisma_client = AsyncMock()
mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
mock_prisma_client.update_data = AsyncMock(return_value=mock_updated_key)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=mock_existing_key
)
async def mock_get_team_object(*args, **kwargs):
return team_table
async def mock_enforce_unique_key_alias(**kwargs):
pass
async def mock_delete_cache_key_object(**kwargs):
pass
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
monkeypatch.setattr(
"litellm.proxy.management_helpers.team_member_permission_checks.get_team_object",
mock_get_team_object,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias",
mock_enforce_unique_key_alias,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
mock_delete_cache_key_object,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
monkeypatch.setattr("litellm.store_audit_logs", False)
monkeypatch.setattr(
"litellm.proxy.proxy_server.hash_token", lambda t: test_hashed_token
)
mock_request = MagicMock()
mock_request.query_params = {}
team_member = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-team-member",
user_id=member_user_id,
team_id=team_id,
)
# Non-budget update on a team key by a team member with /key/update
# permission should succeed.
result = await update_key_fn(
request=mock_request,
data=UpdateKeyRequest(key=test_hashed_token, key_alias="renamed-by-member"),
user_api_key_dict=team_member,
litellm_changed_by=None,
)
assert result is not None
@pytest.mark.asyncio
async def test_update_key_team_member_cannot_change_budget(monkeypatch):
"""A team member with /key/update in member_permissions still cannot
change max_budget — budget/spend changes require team/org admin. The
member_permissions bypass only applies to non-budget fields."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
test_hashed_token = "feedface" * 8
team_id = "team-with-update-grant"
member_user_id = "team-member-user"
mock_existing_key = MagicMock()
mock_existing_key.token = test_hashed_token
mock_existing_key.user_id = None # team-scoped key (no owning user)
mock_existing_key.team_id = team_id
mock_existing_key.project_id = None
mock_existing_key.max_budget = 10.0
mock_existing_key.key_alias = "original"
mock_existing_key.models = []
mock_existing_key.model_dump.return_value = {
"token": test_hashed_token,
"user_id": None,
"team_id": team_id,
"max_budget": 10.0,
}
team_table = LiteLLM_TeamTableCachedObj(
team_id=team_id,
team_alias="test-team",
tpm_limit=None,
rpm_limit=None,
max_budget=None,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[
Member(user_id="some-team-admin", role="admin"),
Member(user_id=member_user_id, role="user"),
],
team_member_permissions=["/key/update", "/key/info"],
)
mock_prisma_client = AsyncMock()
mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=mock_existing_key
)
async def mock_get_team_object(*args, **kwargs):
return team_table
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
mock_get_team_object,
)
monkeypatch.setattr(
"litellm.proxy.management_helpers.team_member_permission_checks.get_team_object",
mock_get_team_object,
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
monkeypatch.setattr(
"litellm.proxy.proxy_server.hash_token", lambda t: test_hashed_token
)
mock_request = MagicMock()
mock_request.query_params = {}
team_member = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER,
api_key="sk-team-member",
user_id=member_user_id,
team_id=team_id,
)
with pytest.raises(ProxyException) as exc:
await update_key_fn(
request=mock_request,
data=UpdateKeyRequest(key=test_hashed_token, max_budget=500.0),
user_api_key_dict=team_member,
litellm_changed_by=None,
)
assert str(exc.value.code) == "403"
# ============================================================================
# LIT-1884: Internal users cannot create invalid keys
# ============================================================================