mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
4086d5ccaa
* feat(proxy): hash UI passwords with PBKDF2 in FIPS mode A FIPS OpenSSL refuses scrypt, so every password write fails there and nobody can sign in with a password. When OpenSSL enforces FIPS, passwords now use PBKDF2-HMAC-SHA256 at 600k iterations with a salt from OpenSSL's DRBG, stored as "pbkdf2:" plus base64 of salt and key. Any other row is refused with one error log and left as it is, and the boot migration does nothing, because FIPS mode is for new deployments. PBKDF2 at 600k takes about five times as long as scrypt, so in FIPS mode the five request paths run it on one thread instead of the event loop. A single thread keeps a login burst from using up a CPU quota, which is what stalls the loop with asyncio.to_thread. A login that queued for that thread asks the throttle again before it answers, so a burst of guesses can't get past a block that the failures ahead of it just set. The breached password check is off in FIPS mode, because it runs the password through SHA-1. A standard OpenSSL reads False at every new branch, so standard mode keeps scrypt, inline, on the same rows. Resolves LIT-8425 * fix(proxy): take FIPS logins one at a time and recheck the throttle first A login that waited for the KDF thread asked the throttle again only after deriving, so a burst for one user kept that thread busy long after the user was blocked. Logins now wait for their turn on a lock and ask the throttle when it comes, so a blocked login costs no PBKDF2 and only one login at a time holds a place in the KDF queue. Password changes and invite claims no longer wait behind a whole burst of logins. * fix(proxy): keep PBKDF2 rows on a standard boot and warn when FIPS skips the breach check A standard boot treated a pbkdf2: row as plaintext and hashed its text, so the real password stopped working and the old row text signed in. The migration now leaves a row of that exact shape alone. Any other pbkdf2: text still migrates. FIPS mode now logs a warning when it skips a breach check that is turned on. The one-thread test no longer depends on timing, and the new test variables are annotated Final. * fix(proxy): keep the breached password check on in FIPS mode * fix(proxy): choose the password KDF from fips_mode instead of the host OpenSSL * test(proxy): fake the scrypt probe in the FIPS fixture now that the gate checks scrypt
716 lines
27 KiB
Python
716 lines
27 KiB
Python
"""
|
||
Tests for the invite-link onboarding endpoints.
|
||
|
||
Covers the security behavior of:
|
||
GET /onboarding/get_token – rejects already-used links and returns only a
|
||
short-lived onboarding token, not a UI session key
|
||
POST /onboarding/claim_token – requires that onboarding token; mints the UI
|
||
session key only after the password is written
|
||
"""
|
||
|
||
import hashlib
|
||
from datetime import timedelta
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import httpx
|
||
import jwt
|
||
import pytest
|
||
import respx
|
||
from fastapi import HTTPException
|
||
|
||
import litellm
|
||
from litellm.proxy._types import InvitationClaim, ProxyException
|
||
|
||
_POLICY_NO_BREACH_CHECK = {"password_policy_check_breached_passwords": False}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class _AsyncTx:
|
||
def __init__(self, db: MagicMock):
|
||
self.db = db
|
||
|
||
async def __aenter__(self) -> MagicMock:
|
||
return self.db
|
||
|
||
async def __aexit__(self, exc_type, exc, tb):
|
||
return False
|
||
|
||
|
||
def _make_invite(
|
||
*, is_accepted: bool, expired: bool = False, claimed: bool = False
|
||
) -> MagicMock:
|
||
now = litellm.utils.get_utc_datetime()
|
||
invite = MagicMock()
|
||
invite.id = "invite-abc"
|
||
invite.user_id = "user-123"
|
||
invite.is_accepted = is_accepted
|
||
invite.expires_at = now - timedelta(days=1) if expired else now + timedelta(days=6)
|
||
invite.accepted_at = now if claimed else None
|
||
return invite
|
||
|
||
|
||
def _make_user() -> MagicMock:
|
||
user = MagicMock()
|
||
user.user_id = "user-123"
|
||
user.user_email = "alice@example.com"
|
||
user.user_role = "internal_user"
|
||
return user
|
||
|
||
|
||
def _make_prisma(invite: MagicMock, user: MagicMock | None = None) -> MagicMock:
|
||
prisma = MagicMock()
|
||
prisma.db.litellm_invitationlink.find_unique = AsyncMock(return_value=invite)
|
||
prisma.db.litellm_invitationlink.update = AsyncMock()
|
||
prisma.db.litellm_invitationlink.update_many = AsyncMock(return_value=1)
|
||
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user)
|
||
prisma.db.litellm_usertable.update = AsyncMock(return_value=user)
|
||
prisma.db.tx = MagicMock(return_value=_AsyncTx(prisma.db))
|
||
return prisma
|
||
|
||
|
||
def _make_onboarding_token(
|
||
*,
|
||
invitation_link: str = "invite-abc",
|
||
user_id: str = "user-123",
|
||
token_type: str = "litellm_onboarding",
|
||
master_key: str = "sk-test",
|
||
) -> str:
|
||
return jwt.encode(
|
||
{
|
||
"token_type": token_type,
|
||
"invitation_link": invitation_link,
|
||
"user_id": user_id,
|
||
"exp": litellm.utils.get_utc_datetime() + timedelta(minutes=15),
|
||
},
|
||
master_key,
|
||
algorithm="HS256",
|
||
)
|
||
|
||
|
||
def _make_claim_request(token: str | None = None) -> MagicMock:
|
||
request = MagicMock()
|
||
request.headers = {"Authorization": f"Bearer {token}"} if token is not None else {}
|
||
request.base_url = "http://localhost:4000/"
|
||
return request
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GET /onboarding/get_token
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_token_rejects_already_used_link():
|
||
"""
|
||
If is_accepted is True the link was already claimed.
|
||
The endpoint must raise 401 *before* returning any user data.
|
||
"""
|
||
from litellm.proxy.proxy_server import onboarding
|
||
|
||
invite = _make_invite(is_accepted=True)
|
||
prisma = _make_prisma(invite)
|
||
request = MagicMock()
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await onboarding(invite_link="invite-abc", request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "already been used" in exc_info.value.detail["error"]
|
||
# The user table must never have been queried
|
||
prisma.db.litellm_usertable.find_unique.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_token_rejects_expired_link():
|
||
"""An expired link must raise 401 regardless of is_accepted."""
|
||
from litellm.proxy.proxy_server import onboarding
|
||
|
||
invite = _make_invite(is_accepted=False, expired=True)
|
||
prisma = _make_prisma(invite)
|
||
request = MagicMock()
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await onboarding(invite_link="invite-abc", request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "expired" in exc_info.value.detail["error"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_token_rejects_missing_link():
|
||
"""A link that does not exist in the DB must raise 401."""
|
||
from litellm.proxy.proxy_server import onboarding
|
||
|
||
prisma = _make_prisma(invite=None) # type: ignore[arg-type]
|
||
request = MagicMock()
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await onboarding(invite_link="nonexistent", request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "does not exist" in exc_info.value.detail["error"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_token_returns_onboarding_token_without_minting_ui_key():
|
||
"""
|
||
A valid, unused link should return a short-lived onboarding token, but
|
||
must not reserve the invite or mint a usable UI/API key on GET.
|
||
"""
|
||
from litellm.proxy.proxy_server import onboarding
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
user = _make_user()
|
||
prisma = _make_prisma(invite, user)
|
||
request = MagicMock()
|
||
request.base_url = "http://localhost:4000/"
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
) as mock_generate_key,
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_custom_url",
|
||
return_value="http://localhost:4000/",
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation",
|
||
return_value=False,
|
||
),
|
||
patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""),
|
||
):
|
||
result = await onboarding(invite_link="invite-abc", request=request)
|
||
|
||
# Endpoint succeeded
|
||
assert "token" in result
|
||
assert "login_url" in result
|
||
|
||
outer_claims = jwt.decode(result["token"], "sk-test", algorithms=["HS256"])
|
||
onboarding_token = outer_claims["key"]
|
||
onboarding_claims = jwt.decode(onboarding_token, "sk-test", algorithms=["HS256"])
|
||
assert onboarding_claims["token_type"] == "litellm_onboarding"
|
||
assert onboarding_claims["invitation_link"] == "invite-abc"
|
||
assert onboarding_claims["user_id"] == "user-123"
|
||
assert not onboarding_token.startswith("sk-")
|
||
|
||
mock_generate_key.assert_not_called()
|
||
prisma.db.litellm_invitationlink.update_many.assert_not_called()
|
||
prisma.db.litellm_invitationlink.update.assert_not_called()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# POST /onboarding/claim_token
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_already_used_link():
|
||
"""
|
||
If is_accepted is True, the password has already been set.
|
||
A second claim attempt must be rejected with 401.
|
||
"""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=True, claimed=True)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
|
||
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=_make_claim_request())
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "already been used" in exc_info.value.detail["error"]
|
||
# Password must never have been written
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_expired_link():
|
||
"""An expired link must be rejected even if is_accepted is False."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False, expired=True)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
|
||
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=_make_claim_request())
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "expired" in exc_info.value.detail["error"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_mismatched_user_id():
|
||
"""The user_id in the request must match the one on the invite."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="wrong-user",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
|
||
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=_make_claim_request())
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "does not match" in exc_info.value.detail["error"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_missing_onboarding_token():
|
||
"""The password endpoint must require the onboarding token returned by get_token."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=_make_claim_request())
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "Missing onboarding session" in exc_info.value.detail["error"]
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_wrong_onboarding_session():
|
||
"""The onboarding token must be bound to the invite and user being claimed."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
request = _make_claim_request(
|
||
_make_onboarding_token(invitation_link="other-invite")
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "Invalid onboarding session" in exc_info.value.detail["error"]
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_invalid_bearer_token():
|
||
"""A regular API key must not be accepted as an onboarding token."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
request = _make_claim_request("sk-regular-key")
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "Invalid onboarding session" in exc_info.value.detail["error"]
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_concurrent_reuse_before_password_write():
|
||
"""Only the first valid claim may reserve the invitation."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
prisma.db.litellm_invitationlink.update_many = AsyncMock(return_value=0)
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
|
||
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
) as mock_generate_key,
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "already been used" in exc_info.value.detail["error"]
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
mock_generate_key.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_sets_accepted_at_after_password_written():
|
||
"""
|
||
A valid first-time claim must:
|
||
1. Write the hashed password to the user table.
|
||
2. Set accepted_at on the invitation link after the password write succeeds.
|
||
"""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
user = _make_user()
|
||
prisma = _make_prisma(invite, user)
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
|
||
mock_token_response = {"token": "sk-generated-key", "user_id": "user-123"}
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
|
||
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
|
||
),
|
||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
return_value=mock_token_response,
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_custom_url",
|
||
return_value="http://localhost:4000/",
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation",
|
||
return_value=False,
|
||
),
|
||
patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""),
|
||
):
|
||
result = await claim_onboarding_link(data=data, request=request)
|
||
|
||
# Password was written
|
||
prisma.db.litellm_invitationlink.update_many.assert_called_once()
|
||
reserve_kwargs = prisma.db.litellm_invitationlink.update_many.call_args.kwargs
|
||
assert reserve_kwargs["where"] == {"id": "invite-abc", "is_accepted": False}
|
||
assert reserve_kwargs["data"]["is_accepted"] is True
|
||
prisma.db.litellm_usertable.update.assert_called_once()
|
||
call_kwargs = prisma.db.litellm_usertable.update.call_args
|
||
assert call_kwargs.kwargs["where"] == {"user_id": "user-123"}
|
||
assert call_kwargs.kwargs["data"]["password"].startswith(("scrypt:", "pbkdf2:"))
|
||
# A freshly claimed, policy-screened password lifts any pending forced
|
||
# reset and re-arms the login-time breach screen.
|
||
assert call_kwargs.kwargs["data"]["password_reset_required"] is False
|
||
assert call_kwargs.kwargs["data"]["last_breach_check_at"] is None
|
||
|
||
# is_accepted was flipped to True on the invitation link
|
||
prisma.db.litellm_invitationlink.update.assert_called_once()
|
||
link_update_data = prisma.db.litellm_invitationlink.update.call_args.kwargs["data"]
|
||
assert "is_accepted" not in link_update_data
|
||
assert link_update_data["accepted_at"] is not None
|
||
outer_claims = jwt.decode(result["token"], "sk-test", algorithms=["HS256"])
|
||
assert outer_claims["key"] == "sk-generated-key"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.usefixtures("fips_mode")
|
||
async def test_fips_claim_hashes_off_the_event_loop(pbkdf2_on_event_loop: list[bool]) -> None:
|
||
await test_claim_token_sets_accepted_at_after_password_written()
|
||
assert pbkdf2_on_event_loop == [False]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_revokes_existing_ui_sessions():
|
||
"""A claimed invite/reset link changes the password; any UI session minted
|
||
under the old password may be in hostile hands and must be revoked. The
|
||
sweep runs before the fresh session key is minted, so revoke-all is safe."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
user = _make_user()
|
||
prisma = _make_prisma(invite, user)
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
|
||
mock_token_response = {"token": "sk-generated-key", "user_id": "user-123"}
|
||
revoke_mock = AsyncMock(return_value=1)
|
||
mint_order: list[str] = []
|
||
|
||
async def _mint(*args, **kwargs):
|
||
mint_order.append("mint")
|
||
return mock_token_response
|
||
|
||
async def _revoke(*args, **kwargs):
|
||
mint_order.append("revoke")
|
||
return 1
|
||
|
||
revoke_mock.side_effect = _revoke
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
|
||
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
|
||
),
|
||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
side_effect=_mint,
|
||
),
|
||
patch(
|
||
"litellm.proxy.management_endpoints.session_endpoints.revoke_ui_session_keys",
|
||
revoke_mock,
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_custom_url",
|
||
return_value="http://localhost:4000/",
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation",
|
||
return_value=False,
|
||
),
|
||
patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""),
|
||
):
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
revoke_mock.assert_awaited_once()
|
||
assert revoke_mock.await_args.kwargs["user_id"] == "user-123"
|
||
# The sweep must precede the mint or it would kill the fresh session too.
|
||
assert mint_order == ["revoke", "mint"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rolls_back_invite_when_session_key_mint_fails():
|
||
"""A session key failure must not leave the invite permanently consumed."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
user = _make_user()
|
||
prisma = _make_prisma(invite, user)
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd123",
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
|
||
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
side_effect=Exception("key mint failed"),
|
||
),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.status_code == 500
|
||
assert "Failed to create onboarding session" in exc_info.value.detail["error"]
|
||
assert prisma.db.litellm_invitationlink.update_many.call_count == 2
|
||
rollback_kwargs = prisma.db.litellm_invitationlink.update_many.call_args_list[
|
||
1
|
||
].kwargs
|
||
assert rollback_kwargs["where"] == {
|
||
"id": "invite-abc",
|
||
"is_accepted": True,
|
||
}
|
||
assert rollback_kwargs["data"]["accepted_at"] is None
|
||
assert rollback_kwargs["data"]["is_accepted"] is False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# POST /onboarding/claim_token - password policy
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _hibp_url_for(password: str) -> str:
|
||
sha1 = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()
|
||
return f"https://api.pwnedpasswords.com/range/{sha1[:5]}"
|
||
|
||
|
||
def _hibp_suffix_for(password: str) -> str:
|
||
return hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()[5:]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_short_password_before_consuming_invite():
|
||
"""Default policy requires 12 characters; the invite must stay claimable."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite, _make_user())
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="Sh0rt!pw",
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above
|
||
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above
|
||
):
|
||
with pytest.raises(ProxyException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.code == "400"
|
||
assert "at least 12 characters" in exc_info.value.message
|
||
prisma.db.litellm_invitationlink.update_many.assert_not_called()
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@respx.mock
|
||
async def test_claim_token_rejects_breached_password_before_consuming_invite(httpx_transport):
|
||
"""A password found in the HIBP corpus must be rejected and never stored."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
password = "P@ssword123456"
|
||
respx.get(_hibp_url_for(password)).mock(
|
||
return_value=httpx.Response(200, text=f"{_hibp_suffix_for(password)}:1387")
|
||
)
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite, _make_user())
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password=password,
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above
|
||
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above
|
||
):
|
||
with pytest.raises(ProxyException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.code == "400"
|
||
assert "data breaches" in exc_info.value.message
|
||
prisma.db.litellm_invitationlink.update_many.assert_not_called()
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@respx.mock
|
||
async def test_claim_token_fails_open_when_hibp_unreachable(httpx_transport):
|
||
"""An HIBP outage must never block onboarding: the claim proceeds."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
password = "NewP@ssw0rd-2026"
|
||
respx.get(_hibp_url_for(password)).mock(side_effect=httpx.ConnectError("no route to host"))
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
user = _make_user()
|
||
prisma = _make_prisma(invite, user)
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password=password,
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above
|
||
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above
|
||
patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: same as above
|
||
patch( # test-quality-ok: same as above
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
return_value={"token": "sk-generated-key", "user_id": "user-123"},
|
||
),
|
||
patch( # test-quality-ok: same as above
|
||
"litellm.proxy.proxy_server.get_custom_url",
|
||
return_value="http://localhost:4000/",
|
||
),
|
||
patch( # test-quality-ok: same as above
|
||
"litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation",
|
||
return_value=False,
|
||
),
|
||
patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""), # test-quality-ok: same as above
|
||
):
|
||
result = await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert "token" in result
|
||
prisma.db.litellm_usertable.update.assert_called_once()
|