From 071cb49d320c9b31fc1d359fca368b8bc5564523 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 02:22:00 -0700 Subject: [PATCH] fix(jwt): accept a team alias in x-litellm-team-id (#42445) * fix(jwt): accept a team alias in x-litellm-team-id The header only matched canonical team ids, so a JWT caller selecting one of their teams by its alias got a 403 even though they belonged to it. The header value is now resolved through the existing alias lookup before the JWT allowed-team check and the DB membership fallback, while a value that is already a team id never costs an alias lookup and denials keep naming the value the caller sent Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(jwt): only alias a header team id the database provably lacks Under fallback_to_db_teams a header value whose team row read fails for any reason other than TeamNotFoundError now keeps the membership denial instead of falling through to the alias lookup, so a degraded read cannot select a different team that carries the value as an alias. Drops the HeaderTeam docstring that only restated its fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: ryan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/handle_jwt.py | 164 ++++++-- tests/e2e/coverage_registry/other.yaml | 2 + tests/e2e/other/other_client.py | 24 +- tests/e2e/other/test_jwt_auth_e2e.py | 81 ++++ .../proxy/auth/test_handle_jwt.py | 392 ++++++++++++++++-- 5 files changed, 575 insertions(+), 88 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 07d8d00d202..63bbaae862e 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -70,6 +70,7 @@ from litellm.types.agents import AgentResponse from litellm.types.proxy.auth.auth_checks import UserNotFoundError from .auth_checks import ( + TeamNotFoundError, _allowed_routes_check, allowed_routes_check, get_actual_routes, @@ -147,6 +148,12 @@ class _JWTProvisioning: team_id_upsert: bool +@dataclass(frozen=True, slots=True) +class HeaderTeam: + header_value: str + team_id: str + + class AgentLookup(Protocol): """The registered-agent lookups a JWT agent claim is matched against.""" @@ -1871,48 +1878,104 @@ class JWTAuthManager: return True @staticmethod - def get_team_id_from_header( - request_headers: Mapping[str, str] | None, - allowed_team_ids: set[str], - fallback_to_db_teams: bool = False, - ) -> str | None: - """ - Extract team_id from x-litellm-team-id header if present. - Validates that the team is in the user's allowed teams from JWT. - - Args: - request_headers: Dictionary of request headers - allowed_team_ids: Set of team IDs the user is allowed to access (from JWT) - fallback_to_db_teams: When True and the JWT carries no team claims - (allowed_team_ids is empty), the header value is returned - provisionally and validated against DB memberships later in - auth_builder instead of being rejected here. - - Returns: - The team_id from header if valid, None otherwise - - Raises: - HTTPException: If team_id is provided but not in allowed_team_ids - """ + def _team_header_value(request_headers: Mapping[str, str] | None) -> str | None: if not request_headers: return None - - # Normalize headers to lowercase for case-insensitive lookup normalized_headers: Final = {k.lower(): v for k, v in request_headers.items()} - header_team_id: Final = normalized_headers.get("x-litellm-team-id") + return normalized_headers.get("x-litellm-team-id") - if not header_team_id: + @staticmethod + def _raise_header_team_not_allowed(header_value: str, allowed_team_ids: set[str]) -> NoReturn: + raise HTTPException( + status_code=403, + detail=f"Team '{header_value}' from x-litellm-team-id header is not in your JWT's allowed teams. Allowed teams: {list(allowed_team_ids)}", + ) + + @staticmethod + async def _team_id_by_alias( + team_alias: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + ) -> str | None: + if prisma_client is None: + return None + try: + team: Final = await get_team_object_by_alias( + team_alias=team_alias, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException as exc: + if exc.status_code >= 500: + raise + return None + return team.team_id + + @staticmethod + async def resolve_team_from_header( + request_headers: Mapping[str, str] | None, + allowed_team_ids: set[str], + fallback_to_db_teams: bool, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + ) -> HeaderTeam | None: + """ + The team named by x-litellm-team-id, which may carry a team id or a team + alias. A value that is already an allowed team id (or, under the DB + fallback, an existing team id) never costs an alias lookup; an alias is + accepted only when the team it names would have been accepted by id. + Under the DB fallback only a team row that is provably absent falls + through to the alias lookup; a read that failed for any other reason + keeps the membership denial the id path already gives. + + Raises: + HTTPException: 403 when neither the value nor the team it aliases is + an allowed team, or the DB fallback's membership denial when the + value names no team at all; a 5xx from the alias lookup itself + is re-raised rather than reported as a denial + """ + header_value: Final = JWTAuthManager._team_header_value(request_headers) + if not header_value: return None - defer_to_db_membership: Final = fallback_to_db_teams and not allowed_team_ids - if not defer_to_db_membership and header_team_id not in allowed_team_ids: - raise HTTPException( - status_code=403, - detail=f"Team '{header_team_id}' from x-litellm-team-id header is not in your JWT's allowed teams. Allowed teams: {list(allowed_team_ids)}", - ) + if fallback_to_db_teams and not allowed_team_ids: + try: + await get_team_object( + team_id=header_value, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + team_id_upsert=False, + ) + except TeamNotFoundError: + aliased_team_id: Final = await JWTAuthManager._team_id_by_alias( + header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj + ) + if aliased_team_id is None: + JWTAuthManager._raise_header_team_membership_denial(header_value) + return HeaderTeam(header_value=header_value, team_id=aliased_team_id) + except HTTPException: + JWTAuthManager._raise_header_team_membership_denial(header_value) + return HeaderTeam(header_value=header_value, team_id=header_value) - verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_team_id) - return header_team_id + if header_value in allowed_team_ids: + verbose_proxy_logger.debug("Using team_id from x-litellm-team-id header: %s", header_value) + return HeaderTeam(header_value=header_value, team_id=header_value) + + team_id_by_alias: Final = await JWTAuthManager._team_id_by_alias( + header_value, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj + ) + if team_id_by_alias is None or team_id_by_alias not in allowed_team_ids: + JWTAuthManager._raise_header_team_not_allowed(header_value, allowed_team_ids) + verbose_proxy_logger.debug("Using team_id %s for x-litellm-team-id alias: %s", team_id_by_alias, header_value) + return HeaderTeam(header_value=header_value, team_id=team_id_by_alias) @staticmethod async def map_user_to_teams( @@ -2264,31 +2327,34 @@ class JWTAuthManager: ) @staticmethod - def _raise_header_team_membership_denial(team_id: str) -> NoReturn: + def _raise_header_team_membership_denial(header_value: str) -> NoReturn: """ The single denial shape for a provisional x-litellm-team-id header, raised identically for nonexistent teams and for teams the user is not - a member of, so the response does not reveal whether a team id exists. + a member of, and naming only the value the caller sent, so the response + reveals neither whether a team exists nor which id an alias maps to. """ raise HTTPException( status_code=403, - detail=(f"Team '{team_id}' (from x-litellm-team-id header) is not in your team memberships."), + detail=(f"Team '{header_value}' (from x-litellm-team-id header) is not in your team memberships."), ) @staticmethod def _validate_header_team_in_db_membership( team_id: str, user_object: LiteLLM_UserTable | None, + header_value: str, ) -> None: """ A provisional team_id from the x-litellm-team-id header (accepted without JWT-team validation when the JWT carries no team claims) must exist in the - user's DB team memberships before it becomes request context. + user's DB team memberships before it becomes request context. The denial + names `header_value`, the id or alias the caller sent, not `team_id`. """ user_team_ids: Final = user_object.teams if user_object else [] if team_id in user_team_ids: return - JWTAuthManager._raise_header_team_membership_denial(team_id) + JWTAuthManager._raise_header_team_membership_denial(header_value) @staticmethod async def auth_builder( @@ -2514,13 +2580,17 @@ class JWTAuthManager: if specific_team_id and not db_team_fallback: all_team_ids.add(specific_team_id) - header_team_id: Final = JWTAuthManager.get_team_id_from_header( + header_team: Final = await JWTAuthManager.resolve_team_from_header( request_headers=request_headers, allowed_team_ids=all_team_ids, fallback_to_db_teams=db_team_fallback, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - if header_team_id: - team_id = header_team_id + if header_team: + team_id = header_team.team_id # A provisional header team (accepted only because the JWT carries no # team claims) is validated against DB membership further down; never # upsert it here or an attacker-supplied x-litellm-team-id would create @@ -2538,7 +2608,7 @@ class JWTAuthManager: except HTTPException: if not db_team_fallback: raise - JWTAuthManager._raise_header_team_membership_denial(team_id) + JWTAuthManager._raise_header_team_membership_denial(header_team.header_value) elif not team_id and not db_team_fallback: ## SPECIFIC TEAM ID ( @@ -2679,10 +2749,11 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=team_id_upsert, ) - elif db_team_fallback and team_id == header_team_id: + elif db_team_fallback and header_team is not None and team_id == header_team.team_id: JWTAuthManager._validate_header_team_in_db_membership( team_id=team_id, user_object=user_object, + header_value=header_team.header_value, ) if not JWTAuthManager._is_team_route_allowed( route=route, @@ -2692,7 +2763,8 @@ class JWTAuthManager: raise HTTPException( status_code=403, detail=( - f"Team '{team_id}' (from x-litellm-team-id header) is not allowed to access route '{route}'." + f"Team '{header_team.header_value}' (from x-litellm-team-id header) " + f"is not allowed to access route '{route}'." ), ) diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 9292e5f07db..1706b2f8f25 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -15,6 +15,8 @@ - {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:1158-1166 _decode_jwt_with_public_key", rationale: "A genuine token whose signature bytes were altered fails verification with 401"} - {id: other.auth.jwt.unknown_team_denied, module: other, tier: P0, area: auth, assertions: [unknown_team_denied], source: "handle_jwt.py:1549-1632 find_team_with_model_access", rationale: "A verified JWT whose groups claim resolves to no existing team is denied with 403 naming the unresolved team, never silently admitted without a team. The proxy words it as a model-access denial, the same body an existing team without model access gets"} - {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"} +- {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"} +- {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"} - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 4313bbe4068..057a735e5e5 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -15,15 +15,25 @@ from __future__ import annotations from dataclasses import dataclass -from e2e_http import NoBody, ProbeResult, Result +from e2e_http import AuthHeaders, NoBody, ProbeResult, Result from idp import Keycloak, keycloak_from_env from models import ( + ChatBody, + ChatResponse, ReadinessDetailsResponse, ReadinessResponse, UserListParams, UserListResponse, ) from proxy_client import ProxyClient +from pydantic import Field + + +class TeamHeaders(AuthHeaders): + """Bearer auth plus ``x-litellm-team-id``, the header a JWT caller sends to + pick one of the teams it belongs to.""" + + x_litellm_team_id: str = Field(serialization_alias="x-litellm-team-id") @dataclass(frozen=True, slots=True) @@ -66,6 +76,18 @@ class OtherClient: response_type=ReadinessDetailsResponse, ) + def chat_as_team(self, token: str, team: str, body: ChatBody) -> Result[ChatResponse]: + """POST /chat/completions under `token` with `x-litellm-team-id: team`.""" + return self.proxy.transport.post( + "/chat/completions", + headers=TeamHeaders( + authorization=self.proxy.transport.bearer(token).authorization, + x_litellm_team_id=team, + ), + json=body, + response_type=ChatResponse, + ) + def list_users_as(self, key: str) -> Result[UserListResponse]: """GET /user/list under `key`. Admin-only, so it doubles as the master key's authorization proof: the master key (proxy admin) reads it, a diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py index 2bed40f4d69..de06c376f76 100644 --- a/tests/e2e/other/test_jwt_auth_e2e.py +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -4,6 +4,7 @@ from __future__ import annotations import base64 import time +from dataclasses import dataclass from typing import Final import pytest @@ -52,6 +53,31 @@ def identity(client: OtherClient, resources: ResourceManager) -> Identity: 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, @@ -155,3 +181,58 @@ class TestJwtAuth: 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, "") == for_unknown.replace(unknown, ""), ( + "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}" + ) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index cd14f630130..208af23f1c5 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -28,10 +28,12 @@ from litellm.proxy._types import ( ) from litellm.caching.dual_cache import DualCache from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry +from litellm.proxy.auth.auth_checks import TeamNotFoundError from litellm.proxy.auth.handle_jwt import ( JWKS_FETCH_ATTEMPTS, STALE_CACHE_KEY_PREFIX, STALE_WRITTEN_AT_CACHE_KEY_PREFIX, + HeaderTeam, JWKSUnreachableError, JWTAuthManager, JWTHandler, @@ -1993,29 +1995,41 @@ async def test_auth_builder_oidc_enabled_falls_back_to_jwt_auth_for_jwt_tokens() assert result["user_object"] == user_object -def test_get_team_id_from_header(): - """Test get_team_id_from_header returns team when valid, None when missing, raises on invalid.""" - from fastapi import HTTPException - - # Valid team in allowed list - result = JWTAuthManager.get_team_id_from_header( +@pytest.mark.asyncio +async def test_resolve_team_from_header_returns_allowed_id_none_without_header_and_403_on_invalid(): + """Without a DB, x-litellm-team-id resolves to the team when it names an allowed + team id, to None when the header is absent, and to a 403 for any other value.""" + allowed = await JWTAuthManager.resolve_team_from_header( request_headers={"x-litellm-team-id": "team-1"}, allowed_team_ids={"team-1", "team-2"}, + fallback_to_db_teams=False, + prisma_client=None, + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), ) - assert result == "team-1" + assert allowed == HeaderTeam(header_value="team-1", team_id="team-1") - # No header returns None - result = JWTAuthManager.get_team_id_from_header( + absent = await JWTAuthManager.resolve_team_from_header( request_headers={"authorization": "Bearer token"}, allowed_team_ids={"team-1"}, + fallback_to_db_teams=False, + prisma_client=None, + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), ) - assert result is None + assert absent is None - # Invalid team raises 403 with pytest.raises(HTTPException) as exc_info: - JWTAuthManager.get_team_id_from_header( + await JWTAuthManager.resolve_team_from_header( request_headers={"x-litellm-team-id": "invalid-team"}, allowed_team_ids={"team-1", "team-2"}, + fallback_to_db_teams=False, + prisma_client=None, + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), ) assert exc_info.value.status_code == 403 @@ -5309,32 +5323,23 @@ def test_build_decode_kwargs_warns_for_unscoped_global_fallback_in_mixed_deploym # --------------------------------------------------------------------------- -def test_get_team_id_from_header_defers_to_db_membership_only_without_jwt_claims(): - """With fallback_to_db_teams=True, an x-litellm-team-id header is accepted - provisionally only when the JWT carries no team claims (allowed set empty). - When the JWT does carry team claims, the header must still be validated +@pytest.mark.asyncio +async def test_resolve_team_from_header_defers_to_db_membership_only_without_jwt_claims(): + """With fallback_to_db_teams=True, an x-litellm-team-id header naming an existing + team is accepted provisionally only when the JWT carries no team claims (allowed + set empty). When the JWT does carry team claims, the header must still be validated against them, and the flag-off behavior must keep rejecting unknown teams.""" - deferred = JWTAuthManager.get_team_id_from_header( - request_headers={"x-litellm-team-id": "team-from-db"}, - allowed_team_ids=set(), - fallback_to_db_teams=True, - ) - assert deferred == "team-from-db" + known_ids = frozenset({"team-from-db"}) + + deferred, _, _ = await _resolve_header("team-from-db", set(), True, _teams_by_id(known_ids), _team_alias_lookup_404) + assert deferred == HeaderTeam(header_value="team-from-db", team_id="team-from-db") with pytest.raises(HTTPException) as exc_info: - JWTAuthManager.get_team_id_from_header( - request_headers={"x-litellm-team-id": "team-x"}, - allowed_team_ids={"team-1", "team-2"}, - fallback_to_db_teams=True, - ) + await _resolve_header("team-x", {"team-1", "team-2"}, True, _teams_by_id(known_ids), _team_alias_lookup_404) assert exc_info.value.status_code == 403 with pytest.raises(HTTPException): - JWTAuthManager.get_team_id_from_header( - request_headers={"x-litellm-team-id": "team-from-db"}, - allowed_team_ids=set(), - fallback_to_db_teams=False, - ) + await _resolve_header("team-from-db", set(), False, _teams_by_id(known_ids), _team_alias_lookup_404) @pytest.mark.asyncio @@ -5788,6 +5793,7 @@ def test_validate_header_team_in_db_membership_does_not_leak_team_ids(): JWTAuthManager._validate_header_team_in_db_membership( team_id="outsider_team", user_object=user_object, + header_value="outsider_team", ) detail = exc_info.value.detail @@ -5797,6 +5803,43 @@ def test_validate_header_team_in_db_membership_does_not_leak_team_ids(): assert "outsider_team" in detail +async def _team_lookup_404(team_id, **kwargs): + raise HTTPException( + status_code=404, + detail=f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call.", + ) + + +async def _team_alias_lookup_404(team_alias, **kwargs): + raise HTTPException( + status_code=404, + detail={"error": f"Team with alias '{team_alias}' doesn't exist in db. Create team via `/team/new` call."}, + ) + + +def _teams_by_alias(aliases: Mapping[str, str]): + """An alias lookup over `aliases` (alias -> team_id) that 404s like the real one otherwise.""" + + async def lookup(team_alias, **kwargs): + if team_alias not in aliases: + return await _team_alias_lookup_404(team_alias) + return LiteLLM_TeamTable(team_id=aliases[team_alias], team_alias=team_alias) + + return lookup + + +def _teams_by_id(team_ids: frozenset[str]): + """A team lookup that knows exactly `team_ids` and, like the real one, reports + any other id as provably absent.""" + + async def lookup(team_id, **kwargs): + if team_id not in team_ids: + raise TeamNotFoundError(team_id=team_id) + return LiteLLM_TeamTable(team_id=team_id) + + return lookup + + async def _run_auth_builder_with_header_team( jwt_auth_config: LiteLLM_JWTAuth, token: dict, @@ -5804,6 +5847,8 @@ async def _run_auth_builder_with_header_team( user_object: LiteLLM_UserTable, fake_get_team, allowed_team_ids: set, + fake_get_team_by_alias=_team_alias_lookup_404, + route: str = "/chat/completions", ): jwt_handler = JWTHandler() jwt_handler.litellm_jwtauth = jwt_auth_config @@ -5848,14 +5893,19 @@ async def _run_auth_builder_with_header_team( new_callable=AsyncMock, side_effect=fake_get_team, ), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", + new_callable=AsyncMock, + side_effect=fake_get_team_by_alias, + ), ): return await JWTAuthManager.auth_builder( api_key="test_jwt_token", jwt_handler=jwt_handler, request_data={"model": "gpt-4"}, general_settings={"enforce_rbac": False}, - route="/chat/completions", - prisma_client=None, + route=route, + prisma_client=MagicMock(), user_api_key_cache=None, parent_otel_span=None, proxy_logging_obj=None, @@ -5863,13 +5913,6 @@ async def _run_auth_builder_with_header_team( ) -async def _team_lookup_404(team_id, **kwargs): - raise HTTPException( - status_code=404, - detail=f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call.", - ) - - @pytest.mark.asyncio async def test_auth_builder_header_team_not_found_matches_non_membership_denial() -> ( None @@ -6942,6 +6985,273 @@ async def test_auth_builder_header_team_enforces_team_allowed_routes_under_db_fa assert result["team_id"] == header_team +async def _resolve_header( + header_value: str, + allowed_team_ids: set[str], + fallback_to_db_teams: bool, + fake_get_team, + fake_get_team_by_alias, +) -> tuple[HeaderTeam | None, AsyncMock, AsyncMock]: + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + side_effect=fake_get_team, + ) as by_id, + patch( + "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", + new_callable=AsyncMock, + side_effect=fake_get_team_by_alias, + ) as by_alias, + ): + resolved = await JWTAuthManager.resolve_team_from_header( + request_headers={"X-LiteLLM-Team-Id": header_value}, + allowed_team_ids=allowed_team_ids, + fallback_to_db_teams=fallback_to_db_teams, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + return resolved, by_id, by_alias + + +@pytest.mark.asyncio +async def test_resolve_team_from_header_accepts_the_alias_of_an_allowed_team(): + """x-litellm-team-id may carry the team alias instead of the team id (LIT-7181). + The alias resolves to its team id before the allowed-teams check, so a + caller whose JWT grants team_a gets team_a whether it sends the id or the + alias, and the id path never pays for an alias lookup.""" + aliases = {"alias_a": "team_a", "alias_b": "team_b"} + + by_alias_value, lookups_by_id, lookups_by_alias = await _resolve_header( + "alias_a", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases) + ) + assert by_alias_value == HeaderTeam(header_value="alias_a", team_id="team_a") + lookups_by_alias.assert_awaited_once() + assert lookups_by_alias.await_args.kwargs["team_alias"] == "alias_a" + lookups_by_id.assert_not_awaited() + + by_id_value, lookups_by_id, lookups_by_alias = await _resolve_header( + "team_a", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases) + ) + assert by_id_value == HeaderTeam(header_value="team_a", team_id="team_a") + lookups_by_alias.assert_not_awaited() + lookups_by_id.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_team_from_header_denies_aliases_of_teams_the_jwt_does_not_grant(): + """An alias that exists but names a team outside the JWT's allowed teams is + refused with the same 403 as an unknown value, and the detail names only + what the caller sent, so the response reveals neither that the alias exists + nor which team id it maps to.""" + aliases = {"alias_a": "team_a", "alias_b": "team_b"} + + with pytest.raises(HTTPException) as other_team: + await _resolve_header("alias_b", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases)) + with pytest.raises(HTTPException) as unknown: + await _resolve_header("no_such", {"team_a"}, False, _team_lookup_404, _teams_by_alias(aliases)) + + assert other_team.value.status_code == 403 + assert unknown.value.status_code == 403 + assert "team_b" not in other_team.value.detail + assert other_team.value.detail.replace("alias_b", "") == unknown.value.detail.replace("no_such", "") + + +@pytest.mark.asyncio +async def test_resolve_team_from_header_under_db_fallback_tries_the_id_before_the_alias(): + """Under fallback_to_db_teams a claimless JWT's header is provisional: a value + that is an existing team id resolves to itself without an alias lookup, a + value that is only an alias resolves to that team's id, and a value that is + neither gets the membership denial the id path already uses.""" + known_ids = frozenset({"team_a"}) + aliases = {"alias_a": "team_a"} + + as_id, _, lookups_by_alias = await _resolve_header( + "team_a", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases) + ) + assert as_id == HeaderTeam(header_value="team_a", team_id="team_a") + lookups_by_alias.assert_not_awaited() + + as_alias, _, _ = await _resolve_header("alias_a", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases)) + assert as_alias == HeaderTeam(header_value="alias_a", team_id="team_a") + + with pytest.raises(HTTPException) as neither: + await _resolve_header("ghost", set(), True, _teams_by_id(known_ids), _teams_by_alias(aliases)) + assert neither.value.status_code == 403 + assert neither.value.detail == ("Team 'ghost' (from x-litellm-team-id header) is not in your team memberships.") + + +@pytest.mark.asyncio +async def test_resolve_team_from_header_under_db_fallback_never_aliases_a_team_id_it_could_not_read(): + """Only a team row the database provably lacks falls through to the alias + lookup. When the id read fails for any other reason (the generic 404 the + team lookup uses for an unreadable database) the value keeps the id path's + membership denial, so an outage can never turn a team id into the team + that happens to carry it as an alias.""" + aliases = {"team_a": "team_b"} + + with ( + patch("litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock, side_effect=_team_lookup_404), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", + new_callable=AsyncMock, + side_effect=_teams_by_alias(aliases), + ) as lookups_by_alias, + pytest.raises(HTTPException) as unreadable, + ): + await JWTAuthManager.resolve_team_from_header( + request_headers={"x-litellm-team-id": "team_a"}, + allowed_team_ids=set(), + fallback_to_db_teams=True, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert unreadable.value.status_code == 403 + assert unreadable.value.detail == ("Team 'team_a' (from x-litellm-team-id header) is not in your team memberships.") + lookups_by_alias.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_team_from_header_treats_a_duplicate_alias_as_no_match_but_surfaces_lookup_errors(): + """An alias two teams share cannot name one team, so it is refused like an + unknown value (a 4xx from the lookup is a miss), while a lookup failure + (5xx) is not disguised as a denial and propagates as is.""" + + async def duplicate_alias(team_alias, **kwargs): + raise HTTPException(status_code=400, detail={"error": f"Multiple teams found with alias '{team_alias}'."}) + + async def db_down(team_alias, **kwargs): + raise HTTPException(status_code=500, detail={"error": f"Error looking up team by alias '{team_alias}'"}) + + with pytest.raises(HTTPException) as duplicate: + await _resolve_header("shared_alias", {"team_a"}, False, _team_lookup_404, duplicate_alias) + assert duplicate.value.status_code == 403 + assert "Multiple teams" not in str(duplicate.value.detail) + + with pytest.raises(HTTPException) as failure: + await _resolve_header("alias_a", {"team_a"}, False, _team_lookup_404, db_down) + assert failure.value.status_code == 500 + + +@pytest.mark.asyncio +async def test_auth_builder_header_alias_binds_the_aliased_team_under_claims_and_db_fallback(): + """End to end through auth_builder, x-litellm-team-id carrying a team alias + binds the request to the aliased team (result team_id is the canonical id) + both when the JWT grants that team by claim and when a claimless JWT relies + on fallback_to_db_teams and DB membership.""" + aliases = {"alias_member": "team_member"} + user_object = LiteLLM_UserTable( + user_id="u_alias", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + + claims_config = LiteLLM_JWTAuth(team_ids_jwt_field="team_ids") + by_claim = await _run_auth_builder_with_header_team( + claims_config, + {"sub": "u_alias", "scope": "", "team_ids": ["team_member"]}, + "alias_member", + user_object, + _teams_by_id(frozenset({"team_member"})), + {"team_member"}, + _teams_by_alias(aliases), + ) + assert by_claim["team_id"] == "team_member" + assert by_claim["team_object"].team_id == "team_member" + + fallback_config = LiteLLM_JWTAuth(fallback_to_db_teams=True) + by_membership = await _run_auth_builder_with_header_team( + fallback_config, + {"sub": "u_alias", "scope": ""}, + "alias_member", + user_object, + _teams_by_id(frozenset({"team_member"})), + set(), + _teams_by_alias(aliases), + ) + assert by_membership["team_id"] == "team_member" + assert by_membership["team_object"].team_id == "team_member" + + +@pytest.mark.asyncio +async def test_auth_builder_header_alias_of_a_non_member_team_is_denied_like_an_unknown_value_under_db_fallback(): + """Under fallback_to_db_teams, an alias naming a team the user is not a member + of is denied with the exact same 403 as an unknown value, naming the alias + the caller sent rather than the team id it resolved to.""" + aliases = {"alias_other": "team_other"} + known_ids = frozenset({"team_member", "team_other"}) + user_object = LiteLLM_UserTable( + user_id="u_alias_outsider", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True) + token = {"sub": "u_alias_outsider", "scope": ""} + + with pytest.raises(HTTPException) as outsider_alias: + await _run_auth_builder_with_header_team( + config, token, "alias_other", user_object, _teams_by_id(known_ids), set(), _teams_by_alias(aliases) + ) + with pytest.raises(HTTPException) as unknown: + await _run_auth_builder_with_header_team( + config, token, "alias_ghost", user_object, _teams_by_id(known_ids), set(), _teams_by_alias(aliases) + ) + + assert outsider_alias.value.status_code == 403 + assert unknown.value.status_code == 403 + assert "team_other" not in outsider_alias.value.detail + assert outsider_alias.value.detail.replace("alias_other", "") == unknown.value.detail.replace( + "alias_ghost", "" + ) + + +@pytest.mark.asyncio +async def test_auth_builder_header_alias_under_db_fallback_keeps_the_team_allowed_routes_gate(): + """Under fallback_to_db_teams, a member team selected by alias is still held + to team_allowed_routes, and the denial names the alias the caller sent.""" + aliases = {"alias_member": "team_member"} + user_object = LiteLLM_UserTable( + user_id="u_alias_routes", + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_member"], + ) + config = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_allowed_routes=["openai_routes"]) + token = {"sub": "u_alias_routes", "scope": ""} + + with pytest.raises(HTTPException) as exc_info: + await _run_auth_builder_with_header_team( + config, + token, + "alias_member", + user_object, + _teams_by_id(frozenset({"team_member"})), + set(), + _teams_by_alias(aliases), + route="/key/info", + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == ( + "Team 'alias_member' (from x-litellm-team-id header) is not allowed to access route '/key/info'." + ) + + allowed = await _run_auth_builder_with_header_team( + config, + token, + "alias_member", + user_object, + _teams_by_id(frozenset({"team_member"})), + set(), + _teams_by_alias(aliases), + route="/chat/completions", + ) + assert allowed["team_id"] == "team_member" + + @pytest.mark.asyncio async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_flag(): """Reading the singular team claim during sync is scoped to fallback_to_db_teams.