mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
feat: add opaque code exchange for /v3/login
/v3/login now returns a single-use opaque code (60s TTL) instead of the JWT directly. A new /v3/login/exchange endpoint redeems the code for the actual JWT. This prevents the JWT from appearing in the initial login response, reducing exposure if the response is logged or intercepted.
This commit is contained in:
parent
f8d1be2e9b
commit
d3d7699fdf
2 changed files with 180 additions and 9 deletions
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue