From a857e4b8859e4b2366f92ce3c0bbe9caf191d385 Mon Sep 17 00:00:00 2001 From: Aarkin Karnik <56022539+Aarkin7@users.noreply.github.com> Date: Mon, 1 Jun 2026 15:35:24 +0530 Subject: [PATCH] fix(proxy): scope /spend/keys and /spend/users to the calling user (#28873) * fix(proxy): scope /spend/keys and /spend/users to the calling user Both endpoints returned every row in the table to any authenticated caller, leaking keys and user records across tenants. Non-admin callers are now scoped to their own user_id, matching the existing pattern in /spend/logs. Admins keep the full-table view. Fixes #28864 * fix(proxy): route non-admin /spend/keys through get_data and 403 on cross-user /spend/users Addresses review feedback on the original fix: both paths now share the get_data helper (same response shape, no direct Prisma access), and a non-admin caller passing another user's id to /spend/users now gets a 403 instead of a silent rewrite, making the attempt show up in logs. --- .../spend_management_endpoints.py | 51 ++- .../test_spend_management_endpoints.py | 355 ++++++++++++++++++ 2 files changed, 401 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index c7ecd64c0fb..6c57fd95b34 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -36,9 +36,18 @@ router = APIRouter() dependencies=[Depends(user_api_key_auth)], include_in_schema=False, ) -async def spend_key_fn(): +async def spend_key_fn( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ - View all keys created, ordered by spend + View keys created, ordered by spend. + + - Admin callers (PROXY_ADMIN / PROXY_ADMIN_VIEW_ONLY) see every key in + the database. + - All other callers (INTERNAL_USER / INTERNAL_USER_VIEW_ONLY, etc.) are + scoped to keys they own (``user_id == caller``). A caller with no + ``user_id`` has no scope and receives an empty list rather than the + full table. Example Request: ``` @@ -55,8 +64,17 @@ async def spend_key_fn(): "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - key_info = await prisma_client.get_data(table_name="key", query_type="find_all") - return key_info + if _is_admin_view_safe(user_api_key_dict=user_api_key_dict): + return await prisma_client.get_data(table_name="key", query_type="find_all") + + caller_user_id = user_api_key_dict.user_id + if not caller_user_id: + return [] + return await prisma_client.get_data( + table_name="key", + query_type="find_all", + user_id=caller_user_id, + ) except Exception as e: raise HTTPException( @@ -85,9 +103,19 @@ async def spend_user_fn( default=None, description="Get User Table row for user_id", ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - View all users created, ordered by spend + View users created, ordered by spend. + + - Admin callers (PROXY_ADMIN / PROXY_ADMIN_VIEW_ONLY) see every user, or + a specific user when ``user_id`` is supplied. + - All other callers may only read their own row. If they supply a + ``user_id`` query parameter that does not match their authenticated + ``user_id`` the request is rejected with HTTP 403; supplying their + own id (or none at all) returns just their row. A caller with no + ``user_id`` on their key has no scope and receives an empty list + rather than the full table. Example Request: ``` @@ -109,6 +137,17 @@ async def spend_user_fn( "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) + if not _is_admin_view_safe(user_api_key_dict=user_api_key_dict): + caller_user_id = user_api_key_dict.user_id + if not caller_user_id: + return [] + if user_id is not None and user_id != caller_user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": "Not authorized to view spend for another user."}, + ) + user_id = caller_user_id + if user_id is not None: user_info = await prisma_client.get_data( table_name="user", query_type="find_unique", user_id=user_id @@ -123,6 +162,8 @@ async def spend_user_fn( _strip_password_from_users(result) return result + except HTTPException: + raise except Exception as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, 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 4bcabfe853a..3559008642f 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 @@ -3185,3 +3185,358 @@ async def test_view_spend_logs_date_range_hashes_sk_api_key(client, monkeypatch) assert where["api_key"] == "hashed::sk-raw-admin-token" finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +class _SpendScopeMockPrismaClient: + + def __init__(self, get_data_returns=None, find_many_returns=None): + self._get_data_returns = ( + get_data_returns if get_data_returns is not None else [] + ) + self._find_many_returns = ( + find_many_returns if find_many_returns is not None else [] + ) + self.get_data_calls = [] + self.find_many_calls = [] + + client = self + + class _VerificationTokenTable: + async def find_many(self, where=None, order=None, include=None): + client.find_many_calls.append( + {"where": where, "order": order, "include": include} + ) + return client._find_many_returns + + class _DB: + def __init__(self): + self.litellm_verificationtoken = _VerificationTokenTable() + + self.db = _DB() + + async def get_data(self, table_name=None, query_type=None, **kwargs): + self.get_data_calls.append( + {"table_name": table_name, "query_type": query_type, **kwargs} + ) + if query_type == "find_unique": + return self._get_data_returns[0] if self._get_data_returns else None + return self._get_data_returns + + +@pytest.mark.asyncio +async def test_spend_key_fn_proxy_admin_returns_all_keys(client, monkeypatch): + """Admins keep their existing full-table view of /spend/keys.""" + mock_keys = [ + {"token": "hashed-a", "user_id": "alice", "spend": 10.0}, + {"token": "hashed-b", "user_id": "bob", "spend": 5.0}, + ] + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=mock_keys) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + try: + response = client.get( + "/spend/keys", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 200 + # Admin path: goes through get_data (full table), never the scoped find_many + assert len(mock_prisma.get_data_calls) == 1 + assert mock_prisma.get_data_calls[0]["table_name"] == "key" + assert mock_prisma.get_data_calls[0]["query_type"] == "find_all" + assert mock_prisma.find_many_calls == [] + assert response.json() == mock_keys + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_key_fn_proxy_admin_view_only_returns_all_keys(client, monkeypatch): + """View-only admins are still admins for this endpoint.""" + mock_keys = [{"token": "hashed-a", "user_id": "alice"}] + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=mock_keys) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, user_id="admin_viewer" + ) + try: + response = client.get( + "/spend/keys", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 200 + assert mock_prisma.find_many_calls == [] + assert len(mock_prisma.get_data_calls) == 1 + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "role", + [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY], +) +async def test_spend_key_fn_internal_user_scoped_to_own_keys(client, monkeypatch, role): + """Both internal-user roles must only see keys they own.""" + caller_owned_keys = [ + {"token": "hashed-mine-1", "user_id": "alice", "spend": 2.0}, + {"token": "hashed-mine-2", "user_id": "alice", "spend": 1.0}, + ] + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=caller_owned_keys) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=role, user_id="alice" + ) + try: + response = client.get( + "/spend/keys", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 200 + # Non-admin path goes through the same get_data helper as admin, + # but with a user_id scope so only the caller's rows come back. + assert mock_prisma.find_many_calls == [] + assert len(mock_prisma.get_data_calls) == 1 + call = mock_prisma.get_data_calls[0] + assert call["table_name"] == "key" + assert call["query_type"] == "find_all" + assert call["user_id"] == "alice" + assert response.json() == caller_owned_keys + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_key_fn_internal_user_without_user_id_returns_empty( + client, monkeypatch +): + """ + A non-admin key with no user_id has no tenant scope. Returning the full + table would re-introduce the leak; return an empty list instead. + """ + mock_prisma = _SpendScopeMockPrismaClient( + get_data_returns=[{"token": "do-not-leak"}], + find_many_returns=[{"token": "do-not-leak"}], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id=None + ) + try: + response = client.get( + "/spend/keys", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 200 + assert response.json() == [] + assert mock_prisma.get_data_calls == [] + assert mock_prisma.find_many_calls == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_user_fn_proxy_admin_returns_all_users_without_user_id( + client, monkeypatch +): + """Admins keep their existing full-table view of /spend/users.""" + mock_users = [ + {"user_id": "alice", "user_email": "alice@example.com", "spend": 1.0}, + {"user_id": "bob", "user_email": "bob@example.com", "spend": 2.0}, + ] + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=mock_users) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + try: + response = client.get( + "/spend/users", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 200 + assert len(mock_prisma.get_data_calls) == 1 + assert mock_prisma.get_data_calls[0]["table_name"] == "user" + assert mock_prisma.get_data_calls[0]["query_type"] == "find_all" + assert response.json() == mock_users + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_user_fn_proxy_admin_can_query_specific_user_id( + client, monkeypatch +): + """Admins can still target a specific user_id.""" + mock_user = { + "user_id": "carol", + "user_email": "carol@example.com", + "spend": 7.0, + } + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=[mock_user]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + try: + response = client.get( + "/spend/users", + params={"user_id": "carol"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + assert len(mock_prisma.get_data_calls) == 1 + assert mock_prisma.get_data_calls[0]["query_type"] == "find_unique" + assert mock_prisma.get_data_calls[0]["user_id"] == "carol" + assert response.json() == [mock_user] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "role", + [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY], +) +async def test_spend_user_fn_internal_user_scoped_without_user_id( + client, monkeypatch, role +): + """No user_id supplied -> must query the caller's own row, not the table.""" + own_row = {"user_id": "alice", "user_email": "alice@example.com", "spend": 3.0} + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=[own_row]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=role, user_id="alice" + ) + try: + response = client.get( + "/spend/users", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 200 + assert len(mock_prisma.get_data_calls) == 1 + assert mock_prisma.get_data_calls[0]["query_type"] == "find_unique" + assert mock_prisma.get_data_calls[0]["user_id"] == "alice" + assert response.json() == [own_row] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_user_fn_internal_user_supplying_other_user_id_returns_403( + client, monkeypatch +): + """ + An internal user passing user_id=victim must be rejected outright, not + silently rewritten. A 403 makes the attempt observable in logs. + """ + leaked_victim_row = { + "user_id": "victim", + "user_email": "victim@example.com", + "spend": 999.0, + } + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=[leaked_victim_row]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice" + ) + try: + response = client.get( + "/spend/users", + params={"user_id": "victim"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + assert mock_prisma.get_data_calls == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_user_fn_internal_user_supplying_own_user_id_is_allowed( + client, monkeypatch +): + """ + Passing your own user_id explicitly is fine — the 403 only fires when + the supplied id differs from the caller's. + """ + own_row = {"user_id": "alice", "user_email": "alice@example.com", "spend": 3.0} + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=[own_row]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice" + ) + try: + response = client.get( + "/spend/users", + params={"user_id": "alice"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + assert len(mock_prisma.get_data_calls) == 1 + assert mock_prisma.get_data_calls[0]["query_type"] == "find_unique" + assert mock_prisma.get_data_calls[0]["user_id"] == "alice" + assert response.json() == [own_row] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_user_fn_internal_user_without_user_id_returns_empty( + client, monkeypatch +): + """ + A non-admin key with no user_id has no tenant scope -> return empty, + never the full table. Same defensive contract as /spend/keys. + """ + mock_prisma = _SpendScopeMockPrismaClient( + get_data_returns=[{"user_id": "do-not-leak"}] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, user_id=None + ) + try: + response = client.get( + "/spend/users", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 200 + assert response.json() == [] + assert mock_prisma.get_data_calls == [] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_user_fn_strips_password_field(client, monkeypatch): + """ + Existing password-redaction behavior must be preserved on the scoped + path so we don't regress a separate disclosure when adding the fix. + """ + own_row = { + "user_id": "alice", + "user_email": "alice@example.com", + "password": "hashed-password-must-not-leak", + "spend": 1.0, + } + mock_prisma = _SpendScopeMockPrismaClient(get_data_returns=[own_row]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice" + ) + try: + response = client.get( + "/spend/users", headers={"Authorization": "Bearer sk-test"} + ) + assert response.status_code == 200 + body = response.json() + assert len(body) == 1 + assert "password" not in body[0] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None)