fix(bedrock): address review - cross-account, SSL verify, narrow fallback

1. Cross-account false match: Added _parse_arn_account_and_role_name()
   helper that compares partition + account ID + role name (not just
   role name) to prevent same-name-different-account false matches.

2. SSL verify: _is_already_running_as_role() now passes ssl_verify to
   the STS client via self._get_ssl_verify(), consistent with all other
   boto3 client creation in this module.

3. Overbroad AccessDenied fallback: The catch in _auth_with_aws_role now
   only falls back to ambient credentials when _is_already_running_as_role
   positively confirms the caller is the target role. Genuine trust-policy
   or permission misconfigurations are re-raised.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

litellm/llms/bedrock/base_aws_llm.py

@@ -213,7 +213,7 @@ class BaseAWSLLM:
         elif aws_role_name is not None:
             # Check if we're already running as the target role and can skip assumption
             # This handles IRSA (EKS), ECS task roles, and EC2 instance profiles
-            if self._is_already_running_as_role(aws_role_name):
+            if self._is_already_running_as_role(aws_role_name, ssl_verify=ssl_verify):
                 verbose_logger.debug(
                     "Already running as target role %s, using ambient credentials",
                     aws_role_name,
@@ -541,7 +541,49 @@ class BaseAWSLLM:
                 aws_region_name = "us-west-2"
         return aws_region_name
 
-    def _is_already_running_as_role(self, aws_role_name: str) -> bool:
+    @staticmethod
+    def _parse_arn_account_and_role_name(
+        arn: str,
+    ) -> Optional[Tuple[str, str, str]]:
+        """
+        Parse an ARN and return (partition, account_id, role_name).
+
+        Handles:
+        - arn:aws:iam::123456789012:role/MyRole
+        - arn:aws:iam::123456789012:role/path/to/MyRole
+        - arn:aws:sts::123456789012:assumed-role/MyRole/session-name
+
+        Returns None if the ARN cannot be parsed.
+        """
+        # ARN format: arn:PARTITION:SERVICE:REGION:ACCOUNT:RESOURCE
+        parts = arn.split(":")
+        if len(parts) < 6 or parts[0] != "arn":
+            return None
+
+        partition = parts[1]  # e.g. "aws", "aws-cn", "aws-us-gov"
+        account_id = parts[4]
+        resource = ":".join(parts[5:])  # rejoin in case resource contains colons
+
+        if resource.startswith("role/"):
+            # arn:aws:iam::ACCOUNT:role/[path/]ROLE_NAME
+            role_name = resource.split("/")[-1]
+        elif resource.startswith("assumed-role/"):
+            # arn:aws:sts::ACCOUNT:assumed-role/ROLE_NAME/SESSION
+            role_parts = resource.split("/")
+            if len(role_parts) >= 2:
+                role_name = role_parts[1]
+            else:
+                return None
+        else:
+            return None
+
+        return partition, account_id, role_name
+
+    def _is_already_running_as_role(
+        self,
+        aws_role_name: str,
+        ssl_verify: Optional[Union[bool, str]] = None,
+    ) -> bool:
         """
         Check if the current environment is already running as the target IAM role.
 
@@ -550,9 +592,18 @@ class BaseAWSLLM:
         - ECS task roles: Uses sts:GetCallerIdentity to check current role ARN
         - EC2 instance profiles: Uses sts:GetCallerIdentity to check current role ARN
 
+        Compares partition, account ID, and role name to avoid cross-account
+        false matches.
+
         Returns True if the current identity matches the target role, meaning
         we can skip sts:AssumeRole and use ambient credentials directly.
         """
+        target_parsed = self._parse_arn_account_and_role_name(aws_role_name)
+        if target_parsed is None:
+            return False
+
+        target_partition, target_account, target_role = target_parsed
+
         # Fast path: IRSA environment check (no API call needed)
         current_role_arn = os.getenv("AWS_ROLE_ARN")
         web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
@@ -564,29 +615,20 @@ class BaseAWSLLM:
             import boto3
 
             with tracer.trace("boto3.client(sts).get_caller_identity"):
-                sts_client = boto3.client("sts")
+                sts_client = boto3.client(
+                    "sts", verify=self._get_ssl_verify(ssl_verify)
+                )
                 identity = sts_client.get_caller_identity()
                 caller_arn = identity.get("Arn", "")
 
-            # The caller ARN for an ECS task role looks like:
-            # arn:aws:sts::123456789012:assumed-role/MyRole/session-name
-            # The target role ARN looks like:
-            # arn:aws:iam::123456789012:role/MyRole
-            # We need to compare the role name portion
-            if ":assumed-role/" in caller_arn:
-                # Extract role name from assumed-role ARN
-                # Format: arn:aws:sts::ACCOUNT:assumed-role/ROLE_NAME/SESSION
-                caller_role_name = caller_arn.split(":assumed-role/")[1].split("/")[0]
-
-                # Extract role name from target role ARN
-                # Format: arn:aws:iam::ACCOUNT:role/ROLE_NAME or
-                # arn:aws:iam::ACCOUNT:role/path/ROLE_NAME
-                if ":role/" in aws_role_name:
-                    target_role_name = aws_role_name.split(":role/")[-1].split("/")[-1]
-                else:
-                    target_role_name = aws_role_name
-
-                if caller_role_name == target_role_name:
+            caller_parsed = self._parse_arn_account_and_role_name(caller_arn)
+            if caller_parsed is not None:
+                caller_partition, caller_account, caller_role = caller_parsed
+                if (
+                    caller_partition == target_partition
+                    and caller_account == target_account
+                    and caller_role == target_role
+                ):
                     verbose_logger.debug(
                         "Current identity already matches target role: %s",
                         aws_role_name,
@@ -918,17 +960,30 @@ class BaseAWSLLM:
             sts_response = sts_client.assume_role(**assume_role_params)
         except Exception as e:
             error_str = str(e)
-            # If AssumeRole fails because the caller already IS the role
-            # (e.g., ECS task role, root account, or same-role scenario),
-            # fall back to using ambient credentials directly
             if "AccessDenied" in error_str:
-                verbose_logger.warning(
-                    "AssumeRole failed for %s (%s). "
-                    "Falling back to ambient credentials (boto3 default chain).",
+                # Only fall back to ambient credentials if we can positively
+                # confirm the caller is already the target role (same account,
+                # partition, and role name).  This avoids silently using the
+                # wrong identity when there is a genuine trust-policy or
+                # permission misconfiguration.
+                if self._is_already_running_as_role(
+                    aws_role_name, ssl_verify=ssl_verify
+                ):
+                    verbose_logger.warning(
+                        "AssumeRole failed for %s (%s). "
+                        "Caller is already running as this role; "
+                        "falling back to ambient credentials.",
+                        aws_role_name,
+                        error_str,
+                    )
+                    return self._auth_with_env_vars()
+                # Genuine permission error — re-raise
+                verbose_logger.error(
+                    "AssumeRole AccessDenied for %s and caller is NOT "
+                    "the same role. Re-raising. Error: %s",
                     aws_role_name,
                     error_str,
                 )
-                return self._auth_with_env_vars()
             raise
 
         # Extract the credentials from the response and convert to Session Credentials

tests/test_litellm/llms/bedrock/test_base_aws_llm.py

@@ -853,11 +853,10 @@ def test_role_assumption_ttl_calculation():
         assert 3500 <= ttl <= 3600  # Allow some variance for test execution time
 
 
-def test_role_assumption_access_denied_falls_back_to_env_vars():
+def test_role_assumption_access_denied_falls_back_when_same_role():
     """
-    Test that when AssumeRole fails with AccessDenied, we fall back to ambient credentials.
-    This handles ECS task roles, root accounts, and same-role scenarios where
-    AssumeRole is unnecessary because the caller already has the role's permissions.
+    Test that when AssumeRole fails with AccessDenied AND the caller is confirmed
+    to already be running as the target role, we fall back to ambient credentials.
     """
     base_aws_llm = BaseAWSLLM()
 
@@ -877,17 +876,51 @@ def test_role_assumption_access_denied_falls_back_to_env_vars():
         with patch.object(
             base_aws_llm, "_auth_with_env_vars", return_value=(mock_creds, None)
         ) as mock_env_auth:
-            credentials, ttl = base_aws_llm._auth_with_aws_role(
-                aws_access_key_id=None,
-                aws_secret_access_key=None,
-                aws_session_token=None,
-                aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole",
-                aws_session_name="error-test-session",
-            )
+            # _is_already_running_as_role returns True => fallback allowed
+            with patch.object(
+                base_aws_llm, "_is_already_running_as_role", return_value=True
+            ):
+                credentials, ttl = base_aws_llm._auth_with_aws_role(
+                    aws_access_key_id=None,
+                    aws_secret_access_key=None,
+                    aws_session_token=None,
+                    aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole",
+                    aws_session_name="error-test-session",
+                )
 
-            # Should have fallen back to env vars
-            mock_env_auth.assert_called_once()
-            assert credentials.access_key == "fallback-access-key"
+                # Should have fallen back to env vars
+                mock_env_auth.assert_called_once()
+                assert credentials.access_key == "fallback-access-key"
+
+
+def test_role_assumption_access_denied_raises_when_different_role():
+    """
+    Test that when AssumeRole fails with AccessDenied but the caller is NOT
+    the same role, the error is re-raised (genuine permission failure).
+    """
+    base_aws_llm = BaseAWSLLM()
+
+    mock_sts_client = MagicMock()
+    mock_sts_client.assume_role.side_effect = Exception(
+        "An error occurred (AccessDenied) when calling the AssumeRole operation: "
+        "User is not authorized to perform sts:AssumeRole"
+    )
+
+    with patch("boto3.client", return_value=mock_sts_client):
+        # _is_already_running_as_role returns False => do NOT fallback
+        with patch.object(
+            base_aws_llm, "_is_already_running_as_role", return_value=False
+        ):
+            with pytest.raises(Exception) as exc_info:
+                base_aws_llm._auth_with_aws_role(
+                    aws_access_key_id=None,
+                    aws_secret_access_key=None,
+                    aws_session_token=None,
+                    aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole",
+                    aws_session_name="error-test-session",
+                )
+
+            assert "AccessDenied" in str(exc_info.value)
 
 
 def test_role_assumption_non_access_denied_error_propagated():
@@ -1367,3 +1400,116 @@ def test_get_credentials_ecs_same_role_skips_assume_role():
                 mock_env_auth.assert_called_once()
                 mock_role_auth.assert_not_called()
                 assert credentials.access_key == "ecs-access-key"
+
+
+def test_parse_arn_account_and_role_name():
+    """Test the ARN parser helper for various ARN formats."""
+    parse = BaseAWSLLM._parse_arn_account_and_role_name
+
+    # Standard IAM role ARN
+    assert parse("arn:aws:iam::123456789012:role/MyRole") == (
+        "aws", "123456789012", "MyRole"
+    )
+
+    # IAM role ARN with path
+    assert parse("arn:aws:iam::123456789012:role/service-role/MyRole") == (
+        "aws", "123456789012", "MyRole"
+    )
+
+    # Assumed-role ARN (from GetCallerIdentity)
+    assert parse("arn:aws:sts::123456789012:assumed-role/MyRole/session-id") == (
+        "aws", "123456789012", "MyRole"
+    )
+
+    # China partition
+    assert parse("arn:aws-cn:iam::123456789012:role/MyRole") == (
+        "aws-cn", "123456789012", "MyRole"
+    )
+
+    # GovCloud partition
+    assert parse("arn:aws-us-gov:iam::123456789012:role/MyRole") == (
+        "aws-us-gov", "123456789012", "MyRole"
+    )
+
+    # Invalid ARNs
+    assert parse("not-an-arn") is None
+    assert parse("arn:aws:iam::123456789012:user/MyUser") is None
+    assert parse("") is None
+
+
+def test_is_already_running_as_role_cross_account_same_name():
+    """
+    Test that same role NAME in different accounts does NOT match.
+    This is the cross-account false-match prevention.
+    """
+    base_aws_llm = BaseAWSLLM()
+
+    mock_sts_client = MagicMock()
+    # Caller is in account 111111111111
+    mock_sts_client.get_caller_identity.return_value = {
+        "Arn": "arn:aws:sts::111111111111:assumed-role/MyRole/session-id"
+    }
+
+    with patch.dict(os.environ, {}, clear=False):
+        env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")}
+        with patch.dict(os.environ, env, clear=True):
+            with patch("boto3.client", return_value=mock_sts_client):
+                # Target is same role name but in account 222222222222
+                assert base_aws_llm._is_already_running_as_role(
+                    "arn:aws:iam::222222222222:role/MyRole"
+                ) is False
+
+
+def test_is_already_running_as_role_cross_partition():
+    """
+    Test that same role name + account but different partition does NOT match.
+    """
+    base_aws_llm = BaseAWSLLM()
+
+    mock_sts_client = MagicMock()
+    mock_sts_client.get_caller_identity.return_value = {
+        "Arn": "arn:aws:sts::123456789012:assumed-role/MyRole/session-id"
+    }
+
+    with patch.dict(os.environ, {}, clear=False):
+        env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")}
+        with patch.dict(os.environ, env, clear=True):
+            with patch("boto3.client", return_value=mock_sts_client):
+                # Same account and role but aws-cn partition
+                assert base_aws_llm._is_already_running_as_role(
+                    "arn:aws-cn:iam::123456789012:role/MyRole"
+                ) is False
+
+
+def test_is_already_running_as_role_invalid_target_arn():
+    """
+    Test that an unparseable target ARN returns False immediately.
+    """
+    base_aws_llm = BaseAWSLLM()
+
+    # Should return False without making any API calls
+    assert base_aws_llm._is_already_running_as_role("not-a-valid-arn") is False
+
+
+def test_is_already_running_as_role_ssl_verify_passed():
+    """
+    Test that ssl_verify parameter is correctly passed to the STS client.
+    """
+    base_aws_llm = BaseAWSLLM()
+
+    mock_sts_client = MagicMock()
+    mock_sts_client.get_caller_identity.return_value = {
+        "Arn": "arn:aws:sts::123456789012:assumed-role/MyRole/session-id"
+    }
+
+    with patch.dict(os.environ, {}, clear=False):
+        env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")}
+        with patch.dict(os.environ, env, clear=True):
+            with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client:
+                base_aws_llm._is_already_running_as_role(
+                    "arn:aws:iam::123456789012:role/MyRole",
+                    ssl_verify="/path/to/ca-bundle.crt",
+                )
+                mock_boto3_client.assert_called_once_with(
+                    "sts", verify="/path/to/ca-bundle.crt"
+                )