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 @@ ) 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 @@ def _is_already_running_as_role( 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 @@ def _auth_with_web_identity_token( 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 @@ action. """ +import base64 import json from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch @@ -157,6 +158,78 @@ def test_no_aws_external_anthropic_statement_collision(self): ) +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()