From bd5e0464643666e20a5bc2a2be5964f2e0430942 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 23:03:09 +0300 Subject: [PATCH] fix(bedrock): surface web identity token aud/iss on InvalidIdentityToken (#31412) When STS rejects a web identity token with InvalidIdentityToken (the "Incorrect token audience" case), litellm propagated the raw botocore error, which never names the aud LiteLLM actually sent. Diagnosing an audience mismatch then required enabling LITELLM_LOG=DEBUG on the prod instance, which degrades performance. _auth_with_web_identity_token now catches InvalidIdentityTokenException, decodes the public aud/iss claims of the resolved JWT without verifying its signature (no secret is read), and raises an AwsAuthError that preserves the STS reason and names the token audience and issuer, so the mismatch is visible from the error alone. Resolves LIT-4026 --- litellm/llms/bedrock/base_aws_llm.py | 43 ++++++++++- .../test_web_identity_session_policy.py | 73 +++++++++++++++++++ 2 files changed, 114 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index c31462a735b..b71f37023e8 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,3 +1,4 @@ +import base64 import hashlib import json import os @@ -19,7 +20,7 @@ from typing import ( ) import httpx -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from litellm._logging import verbose_logger from litellm.caching.caching import DualCache @@ -56,6 +57,11 @@ class Boto3CredentialsInfo(BaseModel): aws_bedrock_runtime_endpoint: Optional[str] +class _WebIdentityTokenClaims(BaseModel): + aud: Optional[Union[str, list[str]]] = None + iss: Optional[str] = None + + class AwsAuthError(Exception): def __init__(self, status_code, message): self.status_code = status_code @@ -817,6 +823,25 @@ class BaseAWSLLM: return False + @staticmethod + def _unverified_web_identity_audience(oidc_token: str) -> Optional[str]: + """Return the public ``aud``/``iss`` claims of a web identity JWT + without verifying its signature, so a rejected-token error can name + the audience LiteLLM actually sent. The signature is never read, so no + secret is exposed.""" + segments = oidc_token.split(".") + if len(segments) != 3: + return None + payload = segments[1] + try: + decoded = base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)) + claims = _WebIdentityTokenClaims.model_validate_json(decoded) + except (ValueError, ValidationError): + return None + if claims.aud is None and claims.iss is None: + return None + return f"aud={claims.aud!r}, iss={claims.iss!r}" + @tracer.wrap() def _auth_with_web_identity_token( self, @@ -925,7 +950,21 @@ class BaseAWSLLM: if aws_external_id is not None: assume_role_params["ExternalId"] = aws_external_id - sts_response = sts_client.assume_role_with_web_identity(**assume_role_params) + try: + sts_response = sts_client.assume_role_with_web_identity( + **assume_role_params + ) + except sts_client.exceptions.InvalidIdentityTokenException as e: + audience = ( + self._unverified_web_identity_audience(oidc_token) + if isinstance(oidc_token, str) + else None + ) + detail = f" Token {audience}" if audience else "" + raise AwsAuthError( + status_code=401, + message=f"AWS STS rejected the web identity token: {e}.{detail}", + ) from e iam_creds_dict = { "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 2cf1fa16e91..7e9c8a273ae 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -29,6 +29,7 @@ claude_platform statement are present and cover every documented action. """ +import base64 import json from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch @@ -157,6 +158,78 @@ class TestClaudePlatformActionsCovered: ) +def _make_jwt(payload: dict) -> str: + def _segment(data: dict) -> str: + return base64.urlsafe_b64encode(json.dumps(data).encode()).rstrip(b"=").decode() + + return f"{_segment({'alg': 'RS256', 'typ': 'JWT'})}.{_segment(payload)}.signature" + + +class TestInvalidIdentityTokenSurfacesAudience: + """LIT-4026: when STS rejects the web identity token with + ``InvalidIdentityToken`` (the "Incorrect token audience" case), the raised + error must name the ``aud``/``iss`` the token actually carries so an + operator can diagnose the mismatch without enabling LITELLM_LOG=DEBUG on a + prod instance.""" + + _AUD = "https://guidepoint.litellm-prod.ai" + _ISS = "https://accounts.google.com" + _STS_MESSAGE = ( + "An error occurred (InvalidIdentityToken) when calling the " + "AssumeRoleWithWebIdentity operation: Incorrect token audience" + ) + + def _raise_invalid_identity_token(self) -> Exception: + from litellm.llms.bedrock.base_aws_llm import AwsAuthError, BaseAWSLLM + + token = _make_jwt({"aud": self._AUD, "iss": self._ISS, "sub": "svc-account"}) + + mock_sts = MagicMock() + + class _InvalidIdentityTokenException(Exception): + pass + + mock_sts.exceptions.InvalidIdentityTokenException = ( + _InvalidIdentityTokenException + ) + mock_sts.assume_role_with_web_identity.side_effect = ( + _InvalidIdentityTokenException(self._STS_MESSAGE) + ) + + with ( + patch("boto3.client", return_value=mock_sts), + patch( + "litellm.llms.bedrock.base_aws_llm.get_secret", + return_value=token, + ), + pytest.raises(AwsAuthError) as exc_info, + ): + BaseAWSLLM()._auth_with_web_identity_token( + aws_web_identity_token="oidc/google/" + self._AUD, + aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role", + aws_session_name="test-session", + aws_region_name="us-east-1", + aws_sts_endpoint=None, + ) + return exc_info.value + + def test_error_names_token_audience(self): + err = self._raise_invalid_identity_token() + assert self._AUD in str(err) + + def test_error_names_token_issuer(self): + err = self._raise_invalid_identity_token() + assert self._ISS in str(err) + + def test_error_preserves_original_sts_reason(self): + err = self._raise_invalid_identity_token() + assert "Incorrect token audience" in str(err) + + def test_error_is_401(self): + err = self._raise_invalid_identity_token() + assert err.status_code == 401 + + class TestPolicyTransportConditions: def test_bedrock_statement_keeps_secure_transport_condition(self): policy = _captured_policy()