From 2f15686ea2cca8faeb2f52ad64b82a1c83314dde Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Thu, 5 Mar 2026 03:46:03 +0530 Subject: [PATCH] fix: address greptile feedback - redact hashed tokens, proper error codes, add tests - Remove token field from JWTKeyMappingResponse to prevent hashed key exposure - Use _to_response() helper on all CRUD endpoints to control returned fields - Return 409 for unique constraint violations, 400 for FK violations, 404 for not found - Add response_model to endpoint decorators - Add 8 new unit tests covering error handling and token redaction Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_types.py | 1 - .../jwt_key_mapping_endpoints.py | 123 +++++++-- .../proxy_unit_tests/test_jwt_key_mapping.py | 238 +++++++++++++++++- 3 files changed, 333 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6b51df709f..e7aba56f00 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3691,7 +3691,6 @@ class JWTKeyMappingResponse(LiteLLMPydanticObjectBase): id: str jwt_claim_name: str jwt_claim_value: str - token: str description: Optional[str] = None is_active: bool created_at: datetime diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index c5d91d3699..779700caf5 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,12 +1,41 @@ +from typing import List + from fastapi import APIRouter, Depends, HTTPException, Query -from litellm.proxy._types import * -from litellm.proxy._types import hash_token + +from litellm.proxy._types import ( + CreateJWTKeyMappingRequest, + DeleteJWTKeyMappingRequest, + JWTKeyMappingResponse, + LitellmUserRoles, + UpdateJWTKeyMappingRequest, + UserAPIKeyAuth, + hash_token, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth router = APIRouter() -@router.post("/jwt/key/mapping/new", tags=["JWT Key Mapping"]) +def _to_response(mapping) -> JWTKeyMappingResponse: + """Convert a Prisma mapping object to a safe response (no hashed token).""" + return JWTKeyMappingResponse( + id=mapping.id, + jwt_claim_name=mapping.jwt_claim_name, + jwt_claim_value=mapping.jwt_claim_value, + description=mapping.description, + is_active=mapping.is_active, + created_at=mapping.created_at, + updated_at=mapping.updated_at, + created_by=mapping.created_by, + updated_by=mapping.updated_by, + ) + + +@router.post( + "/jwt/key/mapping/new", + tags=["JWT Key Mapping"], + response_model=JWTKeyMappingResponse, +) async def create_jwt_key_mapping( data: CreateJWTKeyMappingRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -41,12 +70,29 @@ async def create_jwt_key_mapping( cache_key = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - return new_mapping + return _to_response(new_mapping) + except HTTPException: + raise except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + error_str = str(e).lower() + if "unique" in error_str or "p2002" in error_str: + raise HTTPException( + status_code=409, + detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.", + ) + if "foreign" in error_str or "p2003" in error_str: + raise HTTPException( + status_code=400, + detail="The provided key does not match an existing virtual key.", + ) + raise HTTPException(status_code=500, detail="Failed to create JWT key mapping.") -@router.post("/jwt/key/mapping/update", tags=["JWT Key Mapping"]) +@router.post( + "/jwt/key/mapping/update", + tags=["JWT Key Mapping"], + response_model=JWTKeyMappingResponse, +) async def update_jwt_key_mapping( data: UpdateJWTKeyMappingRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -72,9 +118,11 @@ async def update_jwt_key_mapping( where={"id": data.id} ) - if old_mapping: - cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" - await user_api_key_cache.async_delete_cache(cache_key) + if old_mapping is None: + raise HTTPException(status_code=404, detail="Mapping not found") + + cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" + await user_api_key_cache.async_delete_cache(cache_key) updated_mapping = await prisma_client.db.litellm_jwtkeymapping.update( where={"id": data.id}, data=update_data @@ -84,9 +132,17 @@ async def update_jwt_key_mapping( cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - return updated_mapping + return _to_response(updated_mapping) + except HTTPException: + raise except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + error_str = str(e).lower() + if "unique" in error_str or "p2002" in error_str: + raise HTTPException( + status_code=409, + detail="A mapping with those claim values already exists.", + ) + raise HTTPException(status_code=500, detail="Failed to update JWT key mapping.") @router.post("/jwt/key/mapping/delete", tags=["JWT Key Mapping"]) @@ -110,19 +166,24 @@ async def delete_jwt_key_mapping( where={"id": data.id} ) - if old_mapping: - cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" - await user_api_key_cache.async_delete_cache(cache_key) + if old_mapping is None: + raise HTTPException(status_code=404, detail="Mapping not found") - await prisma_client.db.litellm_jwtkeymapping.delete( - where={"id": data.id} - ) + cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" + await user_api_key_cache.async_delete_cache(cache_key) + + await prisma_client.db.litellm_jwtkeymapping.delete(where={"id": data.id}) return {"status": "success"} - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + except HTTPException: + raise + except Exception: + raise HTTPException(status_code=500, detail="Failed to delete JWT key mapping.") -@router.get("/jwt/key/mapping/list", tags=["JWT Key Mapping"]) +@router.get( + "/jwt/key/mapping/list", + tags=["JWT Key Mapping"], +) async def list_jwt_key_mappings( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), page: int = Query(1, description="Page number", ge=1), @@ -147,16 +208,22 @@ async def list_jwt_key_mappings( ) total_count = await prisma_client.db.litellm_jwtkeymapping.count() return { - "mappings": mappings, + "mappings": [_to_response(m) for m in mappings], "total_count": total_count, "current_page": page, "total_pages": -(-total_count // size), # ceiling division } - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + except HTTPException: + raise + except Exception: + raise HTTPException(status_code=500, detail="Failed to list JWT key mappings.") -@router.get("/jwt/key/mapping/info", tags=["JWT Key Mapping"]) +@router.get( + "/jwt/key/mapping/info", + tags=["JWT Key Mapping"], + response_model=JWTKeyMappingResponse, +) async def info_jwt_key_mapping( id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -177,8 +244,10 @@ async def info_jwt_key_mapping( ) if mapping is None: raise HTTPException(status_code=404, detail="Mapping not found") - return mapping + return _to_response(mapping) except HTTPException: raise - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) + except Exception: + raise HTTPException( + status_code=500, detail="Failed to get JWT key mapping info." + ) diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 7d3e7371b1..b67dd2792f 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -1,6 +1,7 @@ import pytest import sys import os +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch # Add project root to sys.path @@ -10,8 +11,26 @@ from litellm.proxy.auth.user_api_key_auth import ( _resolve_jwt_to_virtual_key, ) from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy._types import LiteLLM_JWTAuth, UserAPIKeyAuth +from litellm.proxy._types import ( + JWTKeyMappingResponse, + LiteLLM_JWTAuth, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( + _to_response, + create_jwt_key_mapping, + delete_jwt_key_mapping, + info_jwt_key_mapping, + update_jwt_key_mapping, +) from litellm.caching.caching import DualCache +from fastapi import HTTPException + + +# ────────────────────────────────────────────── +# Tests: _resolve_jwt_to_virtual_key +# ────────────────────────────────────────────── @pytest.mark.asyncio @@ -114,3 +133,220 @@ async def test_jwt_to_virtual_key_mapping_no_mapping(): ) assert result_cached is None prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called() + + +# ────────────────────────────────────────────── +# Tests: _to_response redacts hashed token +# ────────────────────────────────────────────── + + +def test_to_response_excludes_token(): + """_to_response should not expose the hashed token field.""" + now = datetime.now(timezone.utc) + mock_mapping = MagicMock() + mock_mapping.id = "mapping-1" + mock_mapping.jwt_claim_name = "email" + mock_mapping.jwt_claim_value = "user@example.com" + mock_mapping.token = "hashed_secret_value" + mock_mapping.description = "test" + mock_mapping.is_active = True + mock_mapping.created_at = now + mock_mapping.updated_at = now + mock_mapping.created_by = "admin" + mock_mapping.updated_by = "admin" + + resp = _to_response(mock_mapping) + + assert isinstance(resp, JWTKeyMappingResponse) + assert resp.id == "mapping-1" + assert resp.jwt_claim_name == "email" + assert "token" not in resp.model_fields + + +# ────────────────────────────────────────────── +# Helpers +# ────────────────────────────────────────────── + + +def _make_admin_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + +def _make_non_admin_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="sk-user", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + +def _mock_prisma(): + prisma = MagicMock() + prisma.db.litellm_jwtkeymapping.create = AsyncMock() + prisma.db.litellm_jwtkeymapping.find_unique = AsyncMock() + prisma.db.litellm_jwtkeymapping.find_many = AsyncMock() + prisma.db.litellm_jwtkeymapping.update = AsyncMock() + prisma.db.litellm_jwtkeymapping.delete = AsyncMock() + prisma.db.litellm_jwtkeymapping.count = AsyncMock(return_value=0) + return prisma + + +def _mock_mapping( + id="mapping-1", + claim_name="email", + claim_value="user@example.com", +): + now = datetime.now(timezone.utc) + m = MagicMock() + m.id = id + m.jwt_claim_name = claim_name + m.jwt_claim_value = claim_value + m.token = "hashed_token" + m.description = None + m.is_active = True + m.created_at = now + m.updated_at = now + m.created_by = "admin" + m.updated_by = "admin" + return m + + +# ────────────────────────────────────────────── +# Tests: CRUD endpoint error handling +# ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_create_returns_409_on_unique_violation(): + """Duplicate mapping should return 409, not 500.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception( + "Unique constraint failed (P2002)" + ) + mock_cache = AsyncMock() + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + with pytest.raises(HTTPException) as exc_info: + await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 409 + assert "already exists" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_create_returns_400_on_foreign_key_violation(): + """Non-existent key should return 400, not 500.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception( + "Foreign key constraint failed on field: `token` (P2003)" + ) + mock_cache = AsyncMock() + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="sub", jwt_claim_value="user-999", key="sk-nonexistent", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + with pytest.raises(HTTPException) as exc_info: + await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 400 + assert "does not match" in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_create_non_admin_returns_403(): + """Non-admin users should get 403.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test", + ) + + with pytest.raises(HTTPException) as exc_info: + await create_jwt_key_mapping(data=data, user_api_key_dict=_make_non_admin_auth()) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_delete_returns_404_when_not_found(): + """Deleting non-existent mapping should return 404.""" + from litellm.proxy._types import DeleteJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None + mock_cache = AsyncMock() + + data = DeleteJWTKeyMappingRequest(id="nonexistent-id") + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + with pytest.raises(HTTPException) as exc_info: + await delete_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_update_returns_404_when_not_found(): + """Updating non-existent mapping should return 404.""" + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None + mock_cache = AsyncMock() + + data = UpdateJWTKeyMappingRequest(id="nonexistent-id", description="test") + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + with pytest.raises(HTTPException) as exc_info: + await update_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_info_returns_404_when_not_found(): + """Getting info for non-existent mapping should return 404.""" + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with pytest.raises(HTTPException) as exc_info: + await info_jwt_key_mapping(id="nonexistent-id", user_api_key_dict=_make_admin_auth()) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_create_success_returns_response_without_token(): + """Successful create should return JWTKeyMappingResponse without hashed token.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + mock_prisma = _mock_prisma() + mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping() + mock_cache = AsyncMock() + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_cache + ): + result = await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth()) + assert isinstance(result, JWTKeyMappingResponse) + assert "token" not in result.model_fields + assert result.jwt_claim_name == "email"