From 8ccbd82bd2f190dd3bc98c2a2e5628a6743dcfe6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:13:26 -0700 Subject: [PATCH] fix: authorize custom-logger spend payload against its own owner The custom-logger detail branch reads the payload from cold storage, which is written independently of the spend-log table and can outlive its row. The DB owner pre-check then has nothing to verify for an id lookup that matches no row, so a foreign tenant's stored payload could be returned. Authorize the returned payload against the owner recorded inside it (metadata user/team id), failing closed when none is recorded. Also fold the three identical 403 raises into one helper. --- .../spend_management_endpoints.py | 52 +++++++++++--- .../test_spend_management_endpoints.py | 71 +++++++++++++++---- 2 files changed, 99 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index d0608e1572f..15489632f9d 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2897,9 +2897,10 @@ async def ui_view_request_response_for_request_id( ) if payload is not None: if not caller_is_admin and prisma_client is not None: - await _assert_user_can_view_request_id( + await _assert_user_owns_cold_storage_payload( prisma_client=prisma_client, user_api_key_dict=user_api_key_dict, + payload=cast(Mapping[str, object], payload), # cast-ok: custom-logger payload is untyped request_id=request_id, ) return payload @@ -4345,6 +4346,13 @@ async def _user_can_view_spend_log_owner( return False +def _spend_log_forbidden(request_id: str) -> HTTPException: + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": f"Not authorized to view spend log for request_id={request_id}"}, + ) + + async def _assert_user_can_view_request_id( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, @@ -4360,10 +4368,7 @@ async def _assert_user_can_view_request_id( owners: Final = await _find_spend_log_owners(prisma_client, request_id) for owner in owners: if not await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, owner["user"], owner["team_id"]): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": f"Not authorized to view spend log for request_id={request_id}"}, - ) + raise _spend_log_forbidden(request_id) def _fetched_row_owner(row: Mapping[str, object]) -> tuple[str | None, str | None]: @@ -4390,10 +4395,39 @@ async def _assert_user_owns_fetched_spend_rows( """ for user, team_id in frozenset(_fetched_row_owner(row) for row in rows): if not await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, user, team_id): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": f"Not authorized to view spend log for request_id={request_id}"}, - ) + raise _spend_log_forbidden(request_id) + + +def _cold_storage_payload_owner(payload: Mapping[str, object]) -> tuple[str | None, str | None]: + metadata: Final = payload.get("metadata") + if not isinstance(metadata, Mapping): + return (None, None) + owner: Final = cast(Mapping[str, object], metadata) # cast-ok: cold-storage JSON is untyped + user: Final = owner.get("user_api_key_user_id") + team_id: Final = owner.get("user_api_key_team_id") + return ( + user if isinstance(user, str) else None, + team_id if isinstance(team_id, str) else None, + ) + + +async def _assert_user_owns_cold_storage_payload( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + payload: Mapping[str, object], + request_id: str, +) -> None: + """ + Authorize a cold-storage payload against the owner recorded inside it. + The custom logger reads the payload straight from cold storage, written + independently of the spend-log table and able to outlive its row, so a + request_id lookup could otherwise hand back another tenant's stored payload + when no row exists for the pre-check to catch. Verifying the payload's own + owner closes that gap, and a payload that records no owner fails closed. + """ + owner_user, owner_team_id = _cold_storage_payload_owner(payload) + if not await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, owner_user, owner_team_id): + raise _spend_log_forbidden(request_id) async def _get_permitted_team_ids_for_spend_logs( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index afcb5b36b7e..2d01010cd84 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -2567,36 +2567,34 @@ async def test_ui_view_request_response_rejects_foreign_row_inserted_after_owner @pytest.mark.asyncio -async def test_ui_view_request_response_custom_logger_rechecks_after_fetch(client, monkeypatch): - """The custom-logger payload branch re-verifies ownership after fetching. A row - that appears between the pre-check and the payload read (so the pre-check saw only - owned rows) is caught on the post-fetch check, so the foreign payload is not served.""" - owner_states = iter( - [ - [{"user": "user_1", "team_id": None}], - [{"user": "user_1", "team_id": None}, {"user": "victim_user", "team_id": None}], - ] - ) +async def test_ui_view_request_response_custom_logger_denies_foreign_payload_owner(client, monkeypatch): + """The custom-logger payload comes straight from cold storage, written independently + of the spend-log table and able to outlive its row. When an id lookup matches no row, + the DB owner pre-check has nothing to verify, so the payload is authorized against the + owner recorded inside it. A foreign tenant's stored payload is denied even though no + spend-log row exists for the pre-check to catch.""" class MockDB: async def query_raw(self, sql_query, *params): - if 'SELECT DISTINCT "user", team_id' in sql_query: - return next(owner_states) return [] class MockPrisma: def __init__(self): self.db = MockDB() - class LeakyLogger: + class ColdStorageLogger: async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc): - return {"messages": [{"role": "user", "content": "victim prompt"}], "response": {"id": "r"}} + return { + "messages": [{"role": "user", "content": "victim prompt"}], + "response": {"id": "r"}, + "metadata": {"user_api_key_user_id": "victim_user", "user_api_key_team_id": None}, + } monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrisma()) monkeypatch.setattr( litellm.logging_callback_manager, "get_active_additional_logging_utils_from_custom_logger", - lambda: [LeakyLogger()], + lambda: [ColdStorageLogger()], ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" @@ -2613,6 +2611,49 @@ async def test_ui_view_request_response_custom_logger_rechecks_after_fetch(clien app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_request_response_custom_logger_allows_own_payload_without_db_row(client, monkeypatch): + """The payload-owner authorization must not false-deny a legitimate owner whose + spend-log row is already gone from the DB. An empty owner lookup with a cold-storage + payload the caller owns still serves the payload.""" + + class MockDB: + async def query_raw(self, sql_query, *params): + return [] + + class MockPrisma: + def __init__(self): + self.db = MockDB() + + class ColdStorageLogger: + async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc): + return { + "messages": [{"role": "user", "content": "my own prompt"}], + "response": {"id": "r"}, + "metadata": {"user_api_key_user_id": "user_1", "user_api_key_team_id": None}, + } + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrisma()) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_additional_logging_utils_from_custom_logger", + lambda: [ColdStorageLogger()], + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" + ) + try: + response = client.get( + "/spend/logs/ui/shared-id", + params={"start_date": "2026-01-01 00:00:00"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + assert "my own prompt" in response.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( client, monkeypatch