diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 73b224d4bba..10e05ecdf1e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11120,12 +11120,23 @@ async def login_v3(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "?login=success" - json_response = JSONResponse( - content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, + # Store JWT behind a single-use opaque code (60s TTL) + code = secrets.token_urlsafe(32) + cache_key = f"login_code:{code}" + cache_value = {"token": jwt_token, "redirect_url": litellm_dashboard_ui} + if redis_usage_cache is not None: + await redis_usage_cache.async_set_cache( + key=cache_key, value=cache_value, ttl=60 + ) + else: + await user_api_key_cache.async_set_cache( + key=cache_key, value=cache_value, ttl=60 + ) + + return JSONResponse( + content={"code": code, "expires_in": 60}, status_code=status.HTTP_200_OK, ) - json_response.set_cookie(key="token", value=jwt_token) - return json_response except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.login_v3(): Exception occurred - {}".format( @@ -11151,6 +11162,66 @@ async def login_v3(request: Request): # noqa: PLR0915 ) +@router.post( + "/v3/login/exchange", include_in_schema=False +) # exchange single-use opaque code for JWT +async def login_v3_exchange(request: Request): + try: + body = await request.json() + code = body.get("code") + if not code: + raise ProxyException( + message="Missing 'code' parameter", + type=ProxyErrorTypes.auth_error, + param="code", + code=status.HTTP_400_BAD_REQUEST, + ) + + cache_key = f"login_code:{code}" + if redis_usage_cache is not None: + cached_data = await redis_usage_cache.async_get_cache(key=cache_key) + else: + cached_data = await user_api_key_cache.async_get_cache(key=cache_key) + + if not cached_data or not isinstance(cached_data, dict): + raise ProxyException( + message="Invalid or expired login code", + type=ProxyErrorTypes.auth_error, + param="code", + code=status.HTTP_401_UNAUTHORIZED, + ) + + # Single-use: delete immediately + if redis_usage_cache is not None: + await redis_usage_cache.async_delete_cache(key=cache_key) + else: + await user_api_key_cache.async_delete_cache(key=cache_key) + + json_response = JSONResponse( + content={ + "token": cached_data["token"], + "redirect_url": cached_data["redirect_url"], + }, + status_code=status.HTTP_200_OK, + ) + json_response.set_cookie(key="token", value=cached_data["token"]) + return json_response + except ProxyException: + raise + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {}".format( + str(e) + ) + ) + raise ProxyException( + message=str(e), + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + @app.get("/onboarding/get_token", include_in_schema=False) async def onboarding(invite_link: str, request: Request): """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index f23f5a39353..0e7a0eee5e6 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -253,8 +253,8 @@ def test_login_v3_rejected_without_control_plane_url(monkeypatch): assert "control_plane_url" in response.json()["error"]["message"] -def test_login_v3_includes_token_in_body(monkeypatch): - """v3/login returns token in body when control_plane_url is configured.""" +def test_login_v3_returns_code(monkeypatch): + """v3/login returns an opaque code, not the JWT directly.""" mock_prisma_client = MagicMock() monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", @@ -285,9 +285,109 @@ def test_login_v3_includes_token_in_body(monkeypatch): ) assert response.status_code == 200 - assert response.json()["token"] == "signed-token" - assert response.json()["redirect_url"] == "http://testserver/ui/?login=success" - assert response.cookies.get("token") == "signed-token" + data = response.json() + assert "code" in data + assert data["expires_in"] == 60 + assert "token" not in data + + +def test_login_v3_exchange_happy_path(monkeypatch): + """Full flow: v3/login returns code, v3/login/exchange redeems it for JWT.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_config = MagicMock() + mock_config.worker_registry = [] + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + + client = TestClient(app) + + # Step 1: login — get code + login_response = client.post( + "/v3/login", + json={"username": "alice", "password": "secret"}, + ) + assert login_response.status_code == 200 + code = login_response.json()["code"] + + # Step 2: exchange — get JWT + exchange_response = client.post( + "/v3/login/exchange", + json={"code": code}, + ) + assert exchange_response.status_code == 200 + exchange_data = exchange_response.json() + assert exchange_data["token"] == "signed-token" + assert "redirect_url" in exchange_data + assert exchange_response.cookies.get("token") == "signed-token" + + +def test_login_v3_exchange_single_use(monkeypatch): + """Code can only be redeemed once.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + AsyncMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + MagicMock(return_value={"user_id": "test-user"}), + ) + monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token")) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"control_plane_url": "https://cp.example.com"}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_config = MagicMock() + mock_config.worker_registry = [] + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config) + monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "") + monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None) + + client = TestClient(app) + + login_response = client.post( + "/v3/login", + json={"username": "alice", "password": "secret"}, + ) + code = login_response.json()["code"] + + # First exchange succeeds + first = client.post("/v3/login/exchange", json={"code": code}) + assert first.status_code == 200 + + # Second exchange fails + second = client.post("/v3/login/exchange", json={"code": code}) + assert second.status_code == 401 + + +def test_login_v3_exchange_invalid_code(): + """Random code returns 401.""" + client = TestClient(app) + response = client.post( + "/v3/login/exchange", + json={"code": "nonexistent-code"}, + ) + assert response.status_code == 401 def test_login_v3_returns_json_on_proxy_exception(monkeypatch):