diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bc2728c2203..3390caa8731 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10824,18 +10824,12 @@ async def onboarding(invite_link: str, request: Request): status_code=401, detail={"error": "Invitation link has expired."} ) - #### INVALIDATE LINK - current_time = litellm.utils.get_utc_datetime() - - _ = await prisma_client.db.litellm_invitationlink.update( - where={"id": invite_link}, - data={ - "accepted_at": current_time, - "updated_at": current_time, - "is_accepted": True, - "updated_by": invite_obj.user_id, # type: ignore - }, - ) + #### CHECK IF ALREADY USED + if invite_obj.is_accepted is True: + raise HTTPException( + status_code=401, + detail={"error": "Invitation link has already been used."}, + ) ### GET USER OBJECT ### user_obj = await prisma_client.db.litellm_usertable.find_unique( @@ -10940,19 +10934,11 @@ async def claim_onboarding_link(data: InvitationClaim): status_code=401, detail={"error": "Invitation link has expired."} ) - #### CHECK IF CLAIMED - ##### if claimed - accept - ##### if unclaimed - reject - + #### CHECK IF ALREADY USED if invite_obj.is_accepted is True: - # this is a valid invite that was accepted - pass - else: raise HTTPException( status_code=401, - detail={ - "error": "The invitation link was never validated. Please file an issue, if this is not intended - https://github.com/BerriAI/litellm/issues." - }, + detail={"error": "Invitation link has already been used."}, ) #### CHECK IF VALID USER ID @@ -10976,6 +10962,18 @@ async def claim_onboarding_link(data: InvitationClaim): status_code=401, detail={"error": "User does not exist in db."} ) + #### MARK LINK AS USED + current_time = litellm.utils.get_utc_datetime() + await prisma_client.db.litellm_invitationlink.update( + where={"id": data.invitation_link}, + data={ + "accepted_at": current_time, + "updated_at": current_time, + "is_accepted": True, + "updated_by": invite_obj.user_id, # type: ignore + }, + ) + return user_obj diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py new file mode 100644 index 00000000000..18126e34a2c --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -0,0 +1,259 @@ +""" +Tests for the invite-link onboarding endpoints. + +Covers the security behavior of: + GET /onboarding/get_token – rejects already-used links before showing any user data + POST /onboarding/claim_token – rejects already-used links; marks is_accepted=True only + after the password is successfully written +""" + +from datetime import timedelta +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.proxy._types import InvitationClaim + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_invite(*, is_accepted: bool, expired: 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 = 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_usertable.find_unique = AsyncMock(return_value=user) + prisma.db.litellm_usertable.update = AsyncMock(return_value=user) + return prisma + + +# --------------------------------------------------------------------------- +# 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_does_not_set_is_accepted(): + """ + A valid, unused link should succeed and must NOT flip is_accepted to True. + That flag is only written after the password is claimed. + """ + 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/" + + 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("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, + 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 onboarding(invite_link="invite-abc", request=request) + + # Endpoint succeeded + assert "token" in result + assert "login_url" in result + + # is_accepted must NOT have been updated here + 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) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data) + + 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@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data) + + 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@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data) + + 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_sets_is_accepted_after_password_written(): + """ + A valid first-time claim must: + 1. Write the hashed password to the user table. + 2. Flip is_accepted to True on the invitation link — and only 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) + + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + result = await claim_onboarding_link(data=data) + + # Password was written + 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 "password" in call_kwargs.kwargs["data"] + + # 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 link_update_data["is_accepted"] is True + assert link_update_data["accepted_at"] is not None