litellm/tests/e2e/other/test_jwt_auth_e2e.py

"""Real Keycloak tokens exercise verification, attribution and virtual-key coexistence."""

from __future__ import annotations

import base64
import time
from dataclasses import dataclass
from typing import Final

import pytest
from e2e_config import CHEAP_OPENAI_MODEL, unique_marker
from e2e_http import UnauthorizedError, UnknownApiError, unwrap
from idp import SHORT_LIVED_CLIENT_ID, WRONG_AUDIENCE_CLIENT_ID, Identity
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, TeamNewBody
from other_client import OtherClient
from pydantic import BaseModel

pytestmark = pytest.mark.e2e


class IssuedClaims(BaseModel):
    """Read the IdP's signed payload only to check the test precondition."""

    exp: int
    sub: str
    iss: str
    aud: str | list[str]


def _claims(token: str) -> IssuedClaims:
    payload: Final = token.split(".")[1]
    return IssuedClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)))


def _provision(client: OtherClient, resources: ResourceManager, *, marker: str) -> Identity:
    """A Keycloak group and a user in it, torn down with the test. The group name
    is what the token's `groups` claim carries, which is what the proxy resolves
    as a litellm team id."""
    identity: Final = client.idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer)
    resources.defer(lambda: client.proxy.delete_user(identity.user_id))
    return identity


@pytest.fixture
def identity(client: OtherClient, resources: ResourceManager) -> Identity:
    """An IdP identity whose group is also a real litellm team, so anything the
    proxy rejects is about the token and never about an unresolvable team."""
    marker: Final = unique_marker()
    provisioned: Final = _provision(client, resources, marker=marker)
    team_id: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=provisioned.group))
    resources.defer(lambda: client.proxy.delete_team(team_id))
    return provisioned


@dataclass(frozen=True, slots=True)
class BoundTeam:
    identity: Identity
    team_id: str
    team_alias: str


def _team(client: OtherClient, resources: ResourceManager, *, marker: str, team_id: str) -> str:
    """A litellm team whose alias differs from its id, so a header naming one
    cannot accidentally match the other."""
    team_alias: Final = f"e2e-jwt-alias-{marker}"
    created: Final = client.proxy.create_team(TeamNewBody(team_alias=team_alias, team_id=team_id))
    resources.defer(lambda: client.proxy.delete_team(created))
    return team_alias


@pytest.fixture
def bound_team(client: OtherClient, resources: ResourceManager) -> BoundTeam:
    """An identity whose single group is a real team, plus that team's alias."""
    marker: Final = unique_marker()
    provisioned: Final = _provision(client, resources, marker=marker)
    team_alias: Final = _team(client, resources, marker=marker, team_id=provisioned.group)
    return BoundTeam(identity=provisioned, team_id=provisioned.group, team_alias=team_alias)


def _ping() -> ChatBody:
    return ChatBody(
        model=CHEAP_OPENAI_MODEL,
        messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")],
        max_tokens=16,
    )


def _corrupt_signature(token: str) -> str:
    header, payload, signature = token.split(".")
    flipped: Final = "A" if signature[10] != "A" else "B"
    return f"{header}.{payload}.{signature[:10]}{flipped}{signature[11:]}"


class TestJwtAuth:
    @pytest.mark.covers("other.auth.jwt.valid_token_allows", "other.auth.jwt.spend_attributed_to_claims")
    def test_valid_token_for_an_existing_team_is_accepted_and_attributed(
        self, client: OtherClient, identity: Identity
    ) -> None:
        token: Final = client.idp.access_token(identity)

        assert _claims(token).sub == identity.user_id, "IdP must emit the provisioned user as sub"
        response: Final = unwrap(client.proxy.chat(token, _ping()))
        assert response.id is not None and response.choices, (
            f"chat under a valid JWT returned no completion: {response}"
        )

        rows: Final = client.proxy.poll_logs_for_request_id(response.id)
        assert rows, f"no spend log row for request {response.id} within the poll deadline"
        row: Final = rows[0]
        assert row.team_id == identity.group, (
            f"spend row must carry the team from the JWT groups claim {identity.group!r}, got {row.team_id!r}"
        )
        assert row.user == identity.user_id, (
            f"spend row must carry the user from the JWT sub claim {identity.user_id!r}, got {row.user!r}"
        )

    @pytest.mark.covers("other.auth.jwt.invalid_signature_denied")
    def test_tampered_signature_is_rejected(self, client: OtherClient, identity: Identity) -> None:
        tampered: Final = _corrupt_signature(client.idp.access_token(identity))

        result: Final = client.proxy.chat(tampered, _ping())
        assert isinstance(result, UnauthorizedError), (
            f"a JWT whose signature does not verify must be rejected with 401, got {result}"
        )
        assert "signature verification failed" in result.body.lower(), (
            f"the 401 must come from signature verification, not another auth failure, got {result.body[:300]}"
        )

    @pytest.mark.covers("other.auth.jwt.expired_denied")
    def test_expired_token_is_rejected(self, client: OtherClient, identity: Identity) -> None:
        expiring: Final = client.idp.access_token(identity, client_id=SHORT_LIVED_CLIENT_ID)
        delay: Final = _claims(expiring).exp - time.time() + 1
        assert delay <= 5, f"short-lived client expiry or IdP clock drifted: wait would be {delay}s"
        time.sleep(max(0, delay))

        result: Final = client.proxy.chat(expiring, _ping())
        assert isinstance(result, UnauthorizedError), (
            f"an expired JWT must be rejected with 401 even though its signature verifies, got {result}"
        )
        assert "expired" in result.body.lower(), f"the 401 must say the token expired, got {result.body[:300]}"

    @pytest.mark.covers("other.auth.jwt.wrong_issuer_denied")
    def test_signed_token_from_the_wrong_issuer_is_rejected(self, client: OtherClient, identity: Identity) -> None:
        token: Final = client.idp.access_token(identity, issuer_host="unexpected-issuer.invalid")
        claims: Final = _claims(token)
        assert claims.iss != client.idp.issuer and "litellm-e2e" in claims.aud

        result: Final = client.proxy.chat(token, _ping())
        assert isinstance(result, UnauthorizedError), f"wrong issuer must be rejected: {result}"
        assert "issuer" in result.body.lower(), f"expected issuer validation to reject the token: {result}"

    @pytest.mark.covers("other.auth.jwt.wrong_audience_denied")
    def test_signed_token_for_another_application_is_rejected(self, client: OtherClient, identity: Identity) -> None:
        token: Final = client.idp.access_token(identity, client_id=WRONG_AUDIENCE_CLIENT_ID)
        claims: Final = _claims(token)
        assert claims.iss == client.idp.issuer and "litellm-e2e" not in (
            [claims.aud] if isinstance(claims.aud, str) else claims.aud
        )

        result: Final = client.proxy.chat(token, _ping())
        assert isinstance(result, UnauthorizedError), f"wrong audience must be rejected: {result}"
        assert "audience" in result.body.lower(), f"expected audience validation to reject the token: {result}"

    @pytest.mark.covers("other.auth.jwt.unknown_team_denied")
    def test_token_naming_a_team_that_does_not_exist_is_rejected(
        self, client: OtherClient, resources: ResourceManager
    ) -> None:
        stranger: Final = _provision(client, resources, marker=unique_marker())
        token: Final = client.idp.access_token(stranger)

        result: Final = client.proxy.chat(token, _ping())
        assert isinstance(result, UnknownApiError) and result.status_code == 403, (
            f"a valid JWT whose groups name no existing team must be rejected with 403, got {result}"
        )
        assert stranger.group in result.body, (
            f"the 403 must name the team it could not resolve ({stranger.group}), got {result.body[:300]}"
        )

    @pytest.mark.covers("other.auth.jwt.virtual_key_unaffected")
    def test_plain_virtual_key_still_works_with_jwt_auth_enabled(self, client: OtherClient, scoped_key: str) -> None:
        response: Final = unwrap(client.proxy.chat(scoped_key, _ping()))
        assert response.choices, f"an sk- key must keep working on a proxy with enable_jwt_auth, got {response}"


def _team_of_request(client: OtherClient, token: str, team: str) -> str | None:
    response: Final = unwrap(client.chat_as_team(token, team, _ping()))
    assert response.id is not None and response.choices, (
        f"chat with x-litellm-team-id={team!r} returned no completion: {response}"
    )
    rows: Final = client.proxy.poll_logs_for_request_id(response.id)
    assert rows, f"no spend log row for request {response.id} within the poll deadline"
    return rows[0].team_id


def _denial(client: OtherClient, token: str, team: str) -> str:
    result: Final = client.chat_as_team(token, team, _ping())
    assert isinstance(result, UnknownApiError) and result.status_code == 403, (
        f"x-litellm-team-id={team!r} names no team the caller is in, so it must be rejected with 403, got {result}"
    )
    assert team in result.body, f"the 403 must name the header value it rejected ({team!r}), got {result.body[:300]}"
    return result.body


class TestJwtTeamHeader:
    @pytest.mark.covers("other.auth.jwt.team_header_alias_binds_team")
    def test_team_header_with_the_team_alias_binds_the_same_team_as_the_team_id(
        self, client: OtherClient, bound_team: BoundTeam
    ) -> None:
        token: Final = client.idp.access_token(bound_team.identity)
        assert bound_team.team_alias != bound_team.team_id

        by_id: Final = _team_of_request(client, token, bound_team.team_id)
        assert by_id == bound_team.team_id, (
            f"precondition: x-litellm-team-id with the team id must bind {bound_team.team_id!r}, got {by_id!r}"
        )

        by_alias: Final = _team_of_request(client, token, bound_team.team_alias)
        assert by_alias == bound_team.team_id, (
            f"x-litellm-team-id={bound_team.team_alias!r} must bind the same team as its id "
            f"{bound_team.team_id!r}, got {by_alias!r}"
        )

    @pytest.mark.covers("other.auth.jwt.team_header_non_member_alias_denied")
    def test_team_header_with_the_alias_of_a_team_the_caller_is_not_in_is_rejected_like_an_unknown_value(
        self, client: OtherClient, resources: ResourceManager, bound_team: BoundTeam
    ) -> None:
        token: Final = client.idp.access_token(bound_team.identity)
        other_marker: Final = unique_marker()
        other_alias: Final = _team(client, resources, marker=other_marker, team_id=f"e2e-jwt-other-{other_marker}")
        unknown: Final = f"e2e-jwt-unknown-{unique_marker()}"

        for_other_alias: Final = _denial(client, token, other_alias)
        for_unknown: Final = _denial(client, token, unknown)
        assert for_other_alias.replace(other_alias, "<value>") == for_unknown.replace(unknown, "<value>"), (
            "a non-member alias and an unknown value must get the same denial body, so the response does not "
            f"reveal whether the team exists; got {for_other_alias[:300]!r} vs {for_unknown[:300]!r}"
        )