From ee2cf0e6e89d3b29aa013292d8dd689822cf306b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 17 Apr 2026 15:11:45 -0700 Subject: [PATCH] fix: address three CI failures from recent security PR merges MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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). --- litellm/litellm_core_utils/url_utils.py | 23 +- .../key_management_endpoints.py | 67 +++--- .../test_logging_redaction_e2e_test.py | 48 ++--- .../test_key_management_endpoints.py | 203 ++++++++++++++++++ 4 files changed, 283 insertions(+), 58 deletions(-) diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index b55882819de..a65d0892aa2 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -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 diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 32227b3dd6d..98e41c409b5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 [] diff --git a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py index 0391a5a8957..3047578dbb3 100644 --- a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py +++ b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 41ab60fcf67..811fd8d5cfa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -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 # ============================================================================