diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 506b8811d3a..304c707fa0b 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/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 diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index b57c20fb7c6..cf9fee6bacf 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/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" + )