diff --git a/packages/aws-credentials-sts/.changes/next-release/aws-credentials-sts-bugfix-f3c1259f231e4a9fa3c0a647c6ea29e3.json b/packages/aws-credentials-sts/.changes/next-release/aws-credentials-sts-bugfix-f3c1259f231e4a9fa3c0a647c6ea29e3.json new file mode 100644 index 00000000..67fb050b --- /dev/null +++ b/packages/aws-credentials-sts/.changes/next-release/aws-credentials-sts-bugfix-f3c1259f231e4a9fa3c0a647c6ea29e3.json @@ -0,0 +1,4 @@ +{ + "type": "bugfix", + "description": "The AssumeRole chain provider no longer claims a profile that sets `role_arn` and `web_identity_token_file` without `source_profile` or `credential_source`, so the profile resolves through web identity instead of failing." +} \ No newline at end of file diff --git a/packages/aws-credentials-sts/.changes/next-release/aws-credentials-sts-feature-bf5893f95aab4f85b16b61a090bcaf26.json b/packages/aws-credentials-sts/.changes/next-release/aws-credentials-sts-feature-bf5893f95aab4f85b16b61a090bcaf26.json new file mode 100644 index 00000000..29f2bbe5 --- /dev/null +++ b/packages/aws-credentials-sts/.changes/next-release/aws-credentials-sts-feature-bf5893f95aab4f85b16b61a090bcaf26.json @@ -0,0 +1,4 @@ +{ + "type": "feature", + "description": "Added `WebIdentityCredentialsResolver` and `AsyncWebIdentityCredentialsResolver`, which resolve credentials with STS `AssumeRoleWithWebIdentity` from an OIDC token file, and the `WebIdentityTokenEnvProvider` and `ProfileWebIdentityProvider` chain providers (with `Async*` equivalents), so the default credential chain uses web identity configured by the `AWS_WEB_IDENTITY_TOKEN_FILE` and `AWS_ROLE_ARN` environment variables or a profile's `web_identity_token_file` and `role_arn`." +} \ No newline at end of file diff --git a/packages/aws-credentials-sts/README.md b/packages/aws-credentials-sts/README.md index ccfb0d50..58da3766 100644 --- a/packages/aws-credentials-sts/README.md +++ b/packages/aws-credentials-sts/README.md @@ -1,11 +1,15 @@ # aws-credentials-sts -This package provides STS-based credential resolvers and a chain provider: +This package provides STS-based credential resolvers and chain providers: - `AssumeRoleCredentialsResolver` - assumes an explicit `role_arn` using credentials from a `source_resolver`. - `ProfileAssumeRoleCredentialsResolver` - assumes a role configured in a named profile in the shared config/credentials files. +- `WebIdentityCredentialsResolver` - assumes `role_arn` with an OIDC token read + from `web_identity_token_file`, using STS `AssumeRoleWithWebIdentity`. + +Each resolver and provider has an async equivalent in `aws_credentials_sts.aio`. ## Installation @@ -13,11 +17,18 @@ This package provides STS-based credential resolvers and a chain provider: uv pip install aws-credentials-sts ``` -Once installed, the provider registers itself with the SDK's modular credential -chain. When a client resolves credentials through the default chain, it -will attempt to assume the role configured in the active profile -(`role_arn`/`source_profile`/`credential_source`), unless a higher-precedence -source resolves credentials first. +Once installed, the providers register themselves with the SDK's modular +credential chain. When a client resolves credentials through the default chain, +unless a higher-precedence source resolves credentials first, it will: + +- assume a role with web identity when the `AWS_WEB_IDENTITY_TOKEN_FILE` and + `AWS_ROLE_ARN` environment variables are set, and `AWS_ROLE_SESSION_NAME` + optionally names the session. This is skipped when a profile is selected + explicitly. +- assume the role configured in the active profile + (`role_arn`/`source_profile`/`credential_source`). +- assume a role with web identity configured in the active profile + (`web_identity_token_file`/`role_arn`/`role_session_name`). ## Client Configuration @@ -30,7 +41,7 @@ from aws_credentials_sts import AssumeRoleCredentialsResolver from smithy_aws_core.identity import EnvironmentCredentialsResolver service_client = ServiceClient( - config=ServiceClientConfig( + config=ServiceClientConfig.resolve( aws_credentials_identity_resolver=AssumeRoleCredentialsResolver( source_resolver=EnvironmentCredentialsResolver(), role_arn="arn:aws:iam::123456789012:role/example-role", @@ -50,7 +61,7 @@ from smithy_aws_core.config import load_config def build_client() -> ServiceClient: return ServiceClient( - config=ServiceClientConfig( + config=ServiceClientConfig.resolve( aws_credentials_identity_resolver=ProfileAssumeRoleCredentialsResolver( profile_name="my-profile", config_file=load_config(), @@ -59,9 +70,26 @@ def build_client() -> ServiceClient: ) ``` +To assume a role with web identity explicitly, use +`WebIdentityCredentialsResolver`. The token file is read each time the +credentials are refreshed, so a token rotated on disk is picked up: + +```python +from aws_credentials_sts import WebIdentityCredentialsResolver + +service_client = ServiceClient( + config=ServiceClientConfig.resolve( + aws_credentials_identity_resolver=WebIdentityCredentialsResolver( + role_arn="arn:aws:iam::123456789012:role/example-role", + web_identity_token_file="/var/run/secrets/token", + ), + ) +) +``` + ## Standalone -Either resolver can also be used on its own to fetch credentials directly: +Any resolver can also be used on its own to fetch credentials directly: ```python import asyncio diff --git a/packages/aws-credentials-sts/pyproject.toml b/packages/aws-credentials-sts/pyproject.toml index f15a8fbd..c756b4da 100644 --- a/packages/aws-credentials-sts/pyproject.toml +++ b/packages/aws-credentials-sts/pyproject.toml @@ -39,9 +39,13 @@ dependencies = [ [project.entry-points."smithy_aws_core.identity.chain_providers"] ProfileAssumeRole = "aws_credentials_sts.providers:ProfileAssumeRoleProvider" +ProfileWebIdentity = "aws_credentials_sts.providers:ProfileWebIdentityProvider" +WebIdentityTokenEnv = "aws_credentials_sts.providers:WebIdentityTokenEnvProvider" [project.entry-points."smithy_aws_core.aio.identity.chain_providers"] ProfileAssumeRole = "aws_credentials_sts.aio.providers:AsyncProfileAssumeRoleProvider" +ProfileWebIdentity = "aws_credentials_sts.aio.providers:AsyncProfileWebIdentityProvider" +WebIdentityTokenEnv = "aws_credentials_sts.aio.providers:AsyncWebIdentityTokenEnvProvider" [dependency-groups] test = [ diff --git a/packages/aws-credentials-sts/src/aws_credentials_sts/__init__.py b/packages/aws-credentials-sts/src/aws_credentials_sts/__init__.py index 05b156a8..c88cdbf9 100644 --- a/packages/aws-credentials-sts/src/aws_credentials_sts/__init__.py +++ b/packages/aws-credentials-sts/src/aws_credentials_sts/__init__.py @@ -2,11 +2,17 @@ # SPDX-License-Identifier: Apache-2.0 __version__ = "0.4.0" -from .providers import ProfileAssumeRoleProvider +from .providers import ( + ProfileAssumeRoleProvider, + ProfileWebIdentityProvider, + WebIdentityTokenEnvProvider, +) from .resolvers import ( AssumeRoleConfigurationError, AssumeRoleCredentialsResolver, ProfileAssumeRoleCredentialsResolver, + WebIdentityConfigurationError, + WebIdentityCredentialsResolver, ) __all__ = ( @@ -14,4 +20,8 @@ "AssumeRoleCredentialsResolver", "ProfileAssumeRoleCredentialsResolver", "ProfileAssumeRoleProvider", + "ProfileWebIdentityProvider", + "WebIdentityConfigurationError", + "WebIdentityCredentialsResolver", + "WebIdentityTokenEnvProvider", ) diff --git a/packages/aws-credentials-sts/src/aws_credentials_sts/_assume_role.py b/packages/aws-credentials-sts/src/aws_credentials_sts/_assume_role.py index 262e1a47..33880eac 100644 --- a/packages/aws-credentials-sts/src/aws_credentials_sts/_assume_role.py +++ b/packages/aws-credentials-sts/src/aws_credentials_sts/_assume_role.py @@ -31,6 +31,7 @@ _DURATION_SECONDS = "duration_seconds" _SOURCE_PROFILE = "source_profile" _CREDENTIAL_SOURCE = "credential_source" +_WEB_IDENTITY_TOKEN_FILE = "web_identity_token_file" _REGION = "region" _CREDENTIAL_SOURCE_SLOTS = { @@ -64,11 +65,22 @@ class AssumeRolePlan: source: StaticSource | CredentialSource | AssumeRolePlan -def has_role_arn(config_file: MergedConfig | None, profile_name: str | None) -> bool: +def is_assume_role_profile( + config_file: MergedConfig | None, profile_name: str | None +) -> bool: + """Return whether the profile is configured for AssumeRole. + + A profile with a role ARN and a web identity token file, but no source, is + configured for web identity instead. + """ + if profile_name is None or config_file is None: + return False + if config_file.get(profile_name, _ROLE_ARN) is None: + return False return ( - profile_name is not None - and config_file is not None - and config_file.get(profile_name, _ROLE_ARN) is not None + config_file.get(profile_name, _SOURCE_PROFILE) is not None + or config_file.get(profile_name, _CREDENTIAL_SOURCE) is not None + or not config_file.get(profile_name, _WEB_IDENTITY_TOKEN_FILE) ) @@ -226,7 +238,7 @@ def _plan_source_profile( "complete static credentials." ) - if config_file.get(source_profile_name, _ROLE_ARN) is not None: + if is_assume_role_profile(config_file, source_profile_name): return _plan_assume_role( config_file, source_profile_name, (*visited, source_profile_name) ) diff --git a/packages/aws-credentials-sts/src/aws_credentials_sts/_web_identity.py b/packages/aws-credentials-sts/src/aws_credentials_sts/_web_identity.py new file mode 100644 index 00000000..c1f77ca6 --- /dev/null +++ b/packages/aws-credentials-sts/src/aws_credentials_sts/_web_identity.py @@ -0,0 +1,243 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Web identity logic shared by the sync and async resolvers and providers.""" + +from __future__ import annotations + +import os +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from smithy_aws_core.identity import AWSCredentialsIdentity, AWSIdentityProperties +from smithy_core.exceptions import SmithyError, SmithyIdentityError +from smithy_core.interfaces.identity import IdentityResolver +from smithy_core.program import Program + +from ._assume_role import _account_id_from_arn, resolve_sts_region # pyright: ignore[reportPrivateUsage] + +if TYPE_CHECKING: + from aws_sdk_sts.models import ( + AssumeRoleWithWebIdentityInput, + AssumeRoleWithWebIdentityOutput, + ) + from smithy_aws_core.aio.identity.chain.provider import AsyncChainSetup + from smithy_aws_core.identity.chain.provider import ChainSetup + from smithy_core.interfaces.retries import RetryBackoffStrategy + +MAX_RETRIES = 3 +"""The number of retries after the first attempt to fetch credentials.""" + +DEFAULT_ROLE_SESSION_NAME = f"aws-sdk-python-{uuid.uuid4().hex[:16]}" +"""The role session name used when none is configured. + +It's generated once per process and shared by every web identity resolver. +""" + +_WEB_IDENTITY_TOKEN_FILE_ENV = "AWS_WEB_IDENTITY_TOKEN_FILE" +_ROLE_ARN_ENV = "AWS_ROLE_ARN" +_ROLE_SESSION_NAME_ENV = "AWS_ROLE_SESSION_NAME" +_WEB_IDENTITY_TOKEN_FILE = "web_identity_token_file" +_ROLE_ARN = "role_arn" +_ROLE_SESSION_NAME = "role_session_name" + + +class WebIdentityConfigurationError(SmithyError): + """Raised when web identity credential configuration is invalid.""" + + +UNUSED_CREDENTIALS_MESSAGE = ( + "AssumeRoleWithWebIdentity is sent unsigned, so its STS client has no credentials." +) + + +class UnusedCredentialsResolver( + IdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties] +): + """The credentials resolver for the STS client that calls AssumeRoleWithWebIdentity. + + The call is never signed, so this is never asked for credentials. Setting it + keeps the STS client from building its own default credential chain, which + would include the web identity providers again. + """ + + def get_identity( + self, *, properties: AWSIdentityProperties + ) -> AWSCredentialsIdentity: + raise SmithyIdentityError(UNUSED_CREDENTIALS_MESSAGE) + + +@dataclass(frozen=True, kw_only=True) +class WebIdentityConfig: + """The resolver settings a web identity provider found during chain setup.""" + + role_arn: str + web_identity_token_file: str + role_session_name: str | None + region: str + + +def env_config(setup: ChainSetup | AsyncChainSetup) -> WebIdentityConfig | None: + """Read web identity settings from the environment, if they're complete. + + Returns None when either required variable is missing, so the chain moves on to + the next source, or when a profile was selected explicitly, since that profile + takes precedence over environment credentials. + """ + if setup.profile_name is not None: + return None + + web_identity_token_file = os.getenv(_WEB_IDENTITY_TOKEN_FILE_ENV) + role_arn = os.getenv(_ROLE_ARN_ENV) + if not web_identity_token_file or not role_arn: + return None + + return WebIdentityConfig( + role_arn=role_arn, + web_identity_token_file=web_identity_token_file, + role_session_name=os.getenv(_ROLE_SESSION_NAME_ENV) or None, + region=setup.region_override or resolve_sts_region(), + ) + + +def profile_config(setup: ChainSetup | AsyncChainSetup) -> WebIdentityConfig | None: + """Read web identity settings from the active profile, if it sets a token file. + + :raises WebIdentityConfigurationError: If the profile sets a token file but no + role ARN. + """ + config_file = setup.config_file + profile_name = setup.profile_name + if config_file is None or profile_name is None: + return None + + web_identity_token_file = config_file.get(profile_name, _WEB_IDENTITY_TOKEN_FILE) + # A blank value, such as a leftover 'web_identity_token_file =', counts as unset. + if not web_identity_token_file: + return None + + role_arn = config_file.get(profile_name, _ROLE_ARN) + if role_arn is None: + raise WebIdentityConfigurationError( + f"Profile '{profile_name}' sets '{_WEB_IDENTITY_TOKEN_FILE}' but not " + f"'{_ROLE_ARN}'." + ) + + return WebIdentityConfig( + role_arn=role_arn, + web_identity_token_file=web_identity_token_file, + role_session_name=config_file.get(profile_name, _ROLE_SESSION_NAME), + region=setup.region_override + or resolve_sts_region(config_file=config_file, profile_name=profile_name), + ) + + +def read_token(path: Path, *, is_last_attempt: bool) -> str | None: + """Read the OIDC token, stripped of surrounding whitespace. + + Returns None when the file couldn't be read, isn't ASCII, or is empty, and the + read should be retried. On the last attempt, the failure is raised instead. + """ + try: + token = path.read_text(encoding="ascii").strip() + except (OSError, UnicodeDecodeError) as e: + if is_last_attempt: + raise SmithyIdentityError( + f"Failed to read the web identity token file '{path}': {e}" + ) from e + return None + if not token: + if is_last_attempt: + raise SmithyIdentityError(f"The web identity token file '{path}' is empty.") + return None + return token + + +def fetch_web_identity_credentials( + *, + read_token: Callable[..., Any], + call_sts: Callable[[AssumeRoleWithWebIdentityInput], Any], + sleep: Callable[[float], Any], + backoff: RetryBackoffStrategy, + web_identity_token_file: Path, + role_arn: str, + role_session_name: str, +) -> Program[AWSCredentialsIdentity]: + """Call STS, retrying the failures the web identity SEP requires. + + A failure to read the token file and an ``ExpiredTokenException`` both read the + token again before retrying. An ``IDPCommunicationError`` retries the STS call + with the same token. Other errors are raised immediately. + + The sync and async resolvers run this with ``run_program`` and + ``run_program_async``, passing in their own token read, STS call, and sleep. + """ + from aws_sdk_sts.models import ( + ExpiredTokenException, + IDPCommunicationErrorException, + ) + + token: str | None = None + for attempt in range(MAX_RETRIES + 1): + is_last_attempt = attempt == MAX_RETRIES + if token is None: + token = yield lambda: read_token( + web_identity_token_file, is_last_attempt=is_last_attempt + ) + if token is not None: + request = web_identity_input( + role_arn=role_arn, + role_session_name=role_session_name, + web_identity_token=token, + ) + try: + response = yield lambda: call_sts(request) + except ExpiredTokenException: + if is_last_attempt: + raise + token = None # Read a fresh token before retrying. + except IDPCommunicationErrorException: + if is_last_attempt: + raise + else: + return identity_from_response(response) + delay = backoff.compute_next_backoff_delay(attempt + 1) + yield lambda: sleep(delay) + + raise AssertionError("unreachable") + + +def web_identity_input( + *, role_arn: str, role_session_name: str, web_identity_token: str +) -> AssumeRoleWithWebIdentityInput: + from aws_sdk_sts.models import AssumeRoleWithWebIdentityInput + + return AssumeRoleWithWebIdentityInput( + role_arn=role_arn, + role_session_name=role_session_name, + web_identity_token=web_identity_token, + ) + + +def identity_from_response( + response: AssumeRoleWithWebIdentityOutput, +) -> AWSCredentialsIdentity: + credentials = response.credentials + if credentials is None: + raise SmithyIdentityError( + "STS AssumeRoleWithWebIdentity response did not contain credentials." + ) + assumed_role_arn = ( + response.assumed_role_user.arn + if response.assumed_role_user is not None + else None + ) + return AWSCredentialsIdentity( + access_key_id=credentials.access_key_id, + secret_access_key=credentials.secret_access_key, + session_token=credentials.session_token, + expiration=credentials.expiration, + account_id=_account_id_from_arn(assumed_role_arn), + ) diff --git a/packages/aws-credentials-sts/src/aws_credentials_sts/aio/__init__.py b/packages/aws-credentials-sts/src/aws_credentials_sts/aio/__init__.py index d3f95897..c257031a 100644 --- a/packages/aws-credentials-sts/src/aws_credentials_sts/aio/__init__.py +++ b/packages/aws-credentials-sts/src/aws_credentials_sts/aio/__init__.py @@ -1,13 +1,21 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from .providers import AsyncProfileAssumeRoleProvider +from .providers import ( + AsyncProfileAssumeRoleProvider, + AsyncProfileWebIdentityProvider, + AsyncWebIdentityTokenEnvProvider, +) from .resolvers import ( AsyncAssumeRoleCredentialsResolver, AsyncProfileAssumeRoleCredentialsResolver, + AsyncWebIdentityCredentialsResolver, ) __all__ = ( "AsyncAssumeRoleCredentialsResolver", "AsyncProfileAssumeRoleCredentialsResolver", "AsyncProfileAssumeRoleProvider", + "AsyncProfileWebIdentityProvider", + "AsyncWebIdentityCredentialsResolver", + "AsyncWebIdentityTokenEnvProvider", ) diff --git a/packages/aws-credentials-sts/src/aws_credentials_sts/aio/providers.py b/packages/aws-credentials-sts/src/aws_credentials_sts/aio/providers.py index 0667a5b9..3a95c14a 100644 --- a/packages/aws-credentials-sts/src/aws_credentials_sts/aio/providers.py +++ b/packages/aws-credentials-sts/src/aws_credentials_sts/aio/providers.py @@ -5,8 +5,12 @@ from smithy_aws_core.identity.chain import Standard, StandardProvider from smithy_core.interfaces.identity import Identity -from .._assume_role import has_role_arn -from .resolvers import AsyncProfileAssumeRoleCredentialsResolver +from .._assume_role import is_assume_role_profile +from .._web_identity import WebIdentityConfig, env_config, profile_config +from .resolvers import ( + AsyncProfileAssumeRoleCredentialsResolver, + AsyncWebIdentityCredentialsResolver, +) class AsyncProfileAssumeRoleProvider: @@ -25,14 +29,14 @@ def ordering(self) -> Standard: async def setup( self, identity_type: type[Identity], setup: AsyncChainSetup ) -> None: - """Add a terminal resolver when the active profile declares a role ARN.""" + """Add a terminal resolver when the active profile assumes a role.""" if identity_type is not AWSCredentialsIdentity: return profile_name = setup.profile_name config_file = setup.config_file if profile_name is None or config_file is None: return - if has_role_arn(config_file, profile_name): + if is_assume_role_profile(config_file, profile_name): setup.add_terminal_resolver( AsyncProfileAssumeRoleCredentialsResolver( profile_name=profile_name, @@ -41,3 +45,61 @@ async def setup( http_client=setup.http_client, ) ) + + +class AsyncWebIdentityTokenEnvProvider: + """Adds a web identity resolver configured by environment variables.""" + + @property + def name(self) -> str: + """Return the canonical provider name.""" + return StandardProvider.WEB_IDENTITY_TOKEN_ENV.canonical_name + + @property + def ordering(self) -> Standard: + """Return the provider's standard chain position.""" + return Standard(slot=StandardProvider.WEB_IDENTITY_TOKEN_ENV) + + async def setup( + self, identity_type: type[Identity], setup: AsyncChainSetup + ) -> None: + """Add a terminal resolver when the web identity variables are set.""" + if identity_type is not AWSCredentialsIdentity: + return + if (config := env_config(setup)) is not None: + setup.add_terminal_resolver(_web_identity_resolver(config, setup)) + + +class AsyncProfileWebIdentityProvider: + """Adds a web identity resolver configured by the active profile.""" + + @property + def name(self) -> str: + """Return the canonical provider name.""" + return StandardProvider.PROFILE_WEB_IDENTITY.canonical_name + + @property + def ordering(self) -> Standard: + """Return the provider's standard chain position.""" + return Standard(slot=StandardProvider.PROFILE_WEB_IDENTITY) + + async def setup( + self, identity_type: type[Identity], setup: AsyncChainSetup + ) -> None: + """Add a terminal resolver when the active profile sets a token file.""" + if identity_type is not AWSCredentialsIdentity: + return + if (config := profile_config(setup)) is not None: + setup.add_terminal_resolver(_web_identity_resolver(config, setup)) + + +def _web_identity_resolver( + config: WebIdentityConfig, setup: AsyncChainSetup +) -> AsyncWebIdentityCredentialsResolver: + return AsyncWebIdentityCredentialsResolver( + role_arn=config.role_arn, + web_identity_token_file=config.web_identity_token_file, + role_session_name=config.role_session_name, + region=config.region, + http_client=setup.http_client, + ) diff --git a/packages/aws-credentials-sts/src/aws_credentials_sts/aio/resolvers.py b/packages/aws-credentials-sts/src/aws_credentials_sts/aio/resolvers.py index 151066e8..3f2e6ff5 100644 --- a/packages/aws-credentials-sts/src/aws_credentials_sts/aio/resolvers.py +++ b/packages/aws-credentials-sts/src/aws_credentials_sts/aio/resolvers.py @@ -4,6 +4,8 @@ from __future__ import annotations import asyncio +import os +from pathlib import Path from typing import TYPE_CHECKING from smithy_aws_core.aio.config import AsyncAwsConfigOverrides @@ -19,6 +21,10 @@ from smithy_aws_core.identity import AWSCredentialsIdentity, AWSIdentityProperties from smithy_aws_core.identity.chain import StandardProvider from smithy_core.aio.interfaces.identity import AsyncIdentityResolver +from smithy_core.aio.program import run_program_async +from smithy_core.auth import NoAuthResolver +from smithy_core.exceptions import SmithyIdentityError +from smithy_core.retries import StandardRetryStrategy from smithy_http.aio.interfaces import AsyncHTTPClient from .._assume_role import ( @@ -36,6 +42,12 @@ resolve_sts_region, validate_profile, ) +from .._web_identity import ( + DEFAULT_ROLE_SESSION_NAME, + UNUSED_CREDENTIALS_MESSAGE, + fetch_web_identity_credentials, + read_token, +) if TYPE_CHECKING: from aws_sdk_sts.aio.client import AsyncSTSClient @@ -192,3 +204,130 @@ async def _build_source( return first_resolver(setup.resolvers, source) case AssumeRolePlan(): return await self._build_assume_role(source) + + +class _AsyncUnusedCredentialsResolver( + AsyncIdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties] +): + """The credentials resolver for the STS client that calls AssumeRoleWithWebIdentity. + + The call is never signed, so this is never asked for credentials. Setting it + keeps the STS client from building its own default credential chain, which + would include the web identity providers again. + """ + + async def get_identity( + self, *, properties: AWSIdentityProperties + ) -> AWSCredentialsIdentity: + raise SmithyIdentityError(UNUSED_CREDENTIALS_MESSAGE) + + +class AsyncWebIdentityCredentialsResolver( + AsyncIdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties] +): + """Resolves credentials with an STS AssumeRoleWithWebIdentity call. + + The OIDC token is read from ``web_identity_token_file`` each time the credentials + are refreshed, so a token that is rotated on disk is picked up. + + :param role_arn: The ARN of the role to assume. + :param web_identity_token_file: The path of the file that contains the OIDC + token. A relative path is relative to the process's working directory. + :param role_session_name: The name of the role session. When it isn't set, a + name generated once per process is used. + :param region: The region of the STS endpoint to call. When it isn't set, it's + taken from ``AWS_REGION`` or ``AWS_DEFAULT_REGION``, then ``us-east-1``. + :param http_client: The HTTP client to call STS with. + """ + + def __init__( + self, + *, + role_arn: str, + web_identity_token_file: str | os.PathLike[str], + role_session_name: str | None = None, + region: str | None = None, + http_client: AsyncHTTPClient | None = None, + ) -> None: + self._role_arn = role_arn + self._web_identity_token_file = Path(web_identity_token_file) + self._role_session_name = role_session_name or DEFAULT_ROLE_SESSION_NAME + self._region = region or resolve_sts_region() + self._http_client = http_client + self._credentials: AWSCredentialsIdentity | None = None + self._client: AsyncSTSClient | None = None + self._refresh_lock = asyncio.Lock() + # Wait between attempts the same way the SDK's standard retry strategy does. + self._backoff = StandardRetryStrategy().backoff_strategy + + async def get_identity( + self, + *, + properties: AWSIdentityProperties, + ) -> AWSCredentialsIdentity: + """Return cached credentials if valid, otherwise call STS for new ones.""" + if self._credentials is not None and not self._credentials.is_expired: + return self._credentials + + async with self._refresh_lock: + if self._credentials is not None and not self._credentials.is_expired: + return self._credentials + self._credentials = await self._assume_role_with_web_identity() + return self._credentials + + async def invalidate(self) -> None: + """Discard the cached credentials so the next call fetches new ones.""" + async with self._refresh_lock: + self._credentials = None + + async def _assume_role_with_web_identity(self) -> AWSCredentialsIdentity: + from aws_sdk_sts.aio.client import AsyncSTSClient + from aws_sdk_sts.aio.config import AsyncSTSConfig + + overrides: AsyncAwsConfigOverrides = { + "region": self._region, + "aws_credentials_identity_resolver": _AsyncUnusedCredentialsResolver(), + } + # The call must be unsigned. The generated STS auth resolver doesn't + # recognize the operation, so it would otherwise be signed. + auth_scheme_resolver = NoAuthResolver() + if self._client is None and self._http_client is not None: + # CAUTION: never closed, since closing it would close the caller's transport. + overrides["transport"] = self._http_client + self._client = AsyncSTSClient( + config=await AsyncSTSConfig.resolve( + **overrides, auth_scheme_resolver=auth_scheme_resolver + ) + ) + if self._client is not None: + return await self._call(self._client) + config = await AsyncSTSConfig.resolve( + **overrides, auth_scheme_resolver=auth_scheme_resolver + ) + try: + async with AsyncSTSClient(config=config) as client: + return await self._call(client) + finally: + # The client closes its transport only after it sends a request. Close it + # here too, in case the token couldn't be read and nothing was sent. The + # sync resolver doesn't need this: an unused urllib3 transport holds nothing + # that needs closing, but an aiohttp session must be closed even if it was + # never used. + if (close := getattr(config.transport, "close", None)) is not None: + await close() + + async def _call(self, client: AsyncSTSClient) -> AWSCredentialsIdentity: + return await run_program_async( + fetch_web_identity_credentials( + # Read the file in a thread so a slow disk doesn't block the loop. + read_token=lambda path, **kwargs: asyncio.to_thread( + read_token, path, **kwargs + ), + call_sts=client.assume_role_with_web_identity, + sleep=asyncio.sleep, + backoff=self._backoff, + web_identity_token_file=self._web_identity_token_file, + role_arn=self._role_arn, + role_session_name=self._role_session_name, + ) + ) diff --git a/packages/aws-credentials-sts/src/aws_credentials_sts/providers.py b/packages/aws-credentials-sts/src/aws_credentials_sts/providers.py index b5fd017d..572c1e4a 100644 --- a/packages/aws-credentials-sts/src/aws_credentials_sts/providers.py +++ b/packages/aws-credentials-sts/src/aws_credentials_sts/providers.py @@ -5,8 +5,12 @@ from smithy_aws_core.identity.chain.provider import ChainSetup from smithy_core.interfaces.identity import Identity -from ._assume_role import has_role_arn -from .resolvers import ProfileAssumeRoleCredentialsResolver +from ._assume_role import is_assume_role_profile +from ._web_identity import WebIdentityConfig, env_config, profile_config +from .resolvers import ( + ProfileAssumeRoleCredentialsResolver, + WebIdentityCredentialsResolver, +) class ProfileAssumeRoleProvider: @@ -23,14 +27,14 @@ def ordering(self) -> Standard: return Standard(slot=StandardProvider.PROFILE_ASSUME_ROLE) def setup(self, identity_type: type[Identity], setup: ChainSetup) -> None: - """Add a terminal resolver when the active profile declares a role ARN.""" + """Add a terminal resolver when the active profile assumes a role.""" if identity_type is not AWSCredentialsIdentity: return profile_name = setup.profile_name config_file = setup.config_file if profile_name is None or config_file is None: return - if has_role_arn(config_file, profile_name): + if is_assume_role_profile(config_file, profile_name): setup.add_terminal_resolver( ProfileAssumeRoleCredentialsResolver( profile_name=profile_name, @@ -39,3 +43,64 @@ def setup(self, identity_type: type[Identity], setup: ChainSetup) -> None: http_client=setup.http_client, ) ) + + +# TODO: Declare feature_ids on the web identity providers once smithy-aws-core +# releases the web identity feature IDs: CREDENTIALS_ENV_VARS_STS_WEB_ID_TOKEN and +# CREDENTIALS_STS_ASSUME_ROLE_WEB_ID for the environment provider, and +# CREDENTIALS_PROFILE_STS_WEB_ID_TOKEN and CREDENTIALS_STS_ASSUME_ROLE_WEB_ID for +# the profile provider. + + +class WebIdentityTokenEnvProvider: + """Adds a web identity resolver configured by environment variables.""" + + @property + def name(self) -> str: + """Return the canonical provider name.""" + return StandardProvider.WEB_IDENTITY_TOKEN_ENV.canonical_name + + @property + def ordering(self) -> Standard: + """Return the provider's standard chain position.""" + return Standard(slot=StandardProvider.WEB_IDENTITY_TOKEN_ENV) + + def setup(self, identity_type: type[Identity], setup: ChainSetup) -> None: + """Add a terminal resolver when the web identity variables are set.""" + if identity_type is not AWSCredentialsIdentity: + return + if (config := env_config(setup)) is not None: + setup.add_terminal_resolver(_web_identity_resolver(config, setup)) + + +class ProfileWebIdentityProvider: + """Adds a web identity resolver configured by the active profile.""" + + @property + def name(self) -> str: + """Return the canonical provider name.""" + return StandardProvider.PROFILE_WEB_IDENTITY.canonical_name + + @property + def ordering(self) -> Standard: + """Return the provider's standard chain position.""" + return Standard(slot=StandardProvider.PROFILE_WEB_IDENTITY) + + def setup(self, identity_type: type[Identity], setup: ChainSetup) -> None: + """Add a terminal resolver when the active profile sets a token file.""" + if identity_type is not AWSCredentialsIdentity: + return + if (config := profile_config(setup)) is not None: + setup.add_terminal_resolver(_web_identity_resolver(config, setup)) + + +def _web_identity_resolver( + config: WebIdentityConfig, setup: ChainSetup +) -> WebIdentityCredentialsResolver: + return WebIdentityCredentialsResolver( + role_arn=config.role_arn, + web_identity_token_file=config.web_identity_token_file, + role_session_name=config.role_session_name, + region=config.region, + http_client=setup.http_client, + ) diff --git a/packages/aws-credentials-sts/src/aws_credentials_sts/resolvers.py b/packages/aws-credentials-sts/src/aws_credentials_sts/resolvers.py index ef883f79..0156b53f 100644 --- a/packages/aws-credentials-sts/src/aws_credentials_sts/resolvers.py +++ b/packages/aws-credentials-sts/src/aws_credentials_sts/resolvers.py @@ -3,7 +3,10 @@ from __future__ import annotations +import os import threading +import time +from pathlib import Path from typing import TYPE_CHECKING from smithy_aws_core.config import AwsConfigOverrides @@ -16,7 +19,10 @@ ) from smithy_aws_core.identity.chain import StandardProvider from smithy_aws_core.identity.chain.provider import ChainIdentityProvider, ChainSetup +from smithy_core.auth import NoAuthResolver from smithy_core.interfaces.identity import IdentityResolver +from smithy_core.program import run_program +from smithy_core.retries import StandardRetryStrategy from smithy_http.interfaces import HTTPClient from ._assume_role import ( @@ -35,6 +41,13 @@ resolve_sts_region, validate_profile, ) +from ._web_identity import ( + DEFAULT_ROLE_SESSION_NAME, + UnusedCredentialsResolver, + WebIdentityConfigurationError, + fetch_web_identity_credentials, + read_token, +) if TYPE_CHECKING: from aws_sdk_sts.client import STSClient @@ -43,6 +56,8 @@ "AssumeRoleConfigurationError", "AssumeRoleCredentialsResolver", "ProfileAssumeRoleCredentialsResolver", + "WebIdentityConfigurationError", + "WebIdentityCredentialsResolver", ) _CHAIN_PROVIDER_ENTRY_POINT_GROUP = "smithy_aws_core.identity.chain_providers" @@ -189,3 +204,105 @@ def _build_source( return first_resolver(setup.resolvers, source) case AssumeRolePlan(): return self._build_assume_role(source) + + +class WebIdentityCredentialsResolver( + IdentityResolver[AWSCredentialsIdentity, AWSIdentityProperties] +): + """Resolves credentials with an STS AssumeRoleWithWebIdentity call. + + The OIDC token is read from ``web_identity_token_file`` each time the credentials + are refreshed, so a token that is rotated on disk is picked up. + + :param role_arn: The ARN of the role to assume. + :param web_identity_token_file: The path of the file that contains the OIDC + token. A relative path is relative to the process's working directory. + :param role_session_name: The name of the role session. When it isn't set, a + name generated once per process is used. + :param region: The region of the STS endpoint to call. When it isn't set, it's + taken from ``AWS_REGION`` or ``AWS_DEFAULT_REGION``, then ``us-east-1``. + :param http_client: The HTTP client to call STS with. + """ + + def __init__( + self, + *, + role_arn: str, + web_identity_token_file: str | os.PathLike[str], + role_session_name: str | None = None, + region: str | None = None, + http_client: HTTPClient | None = None, + ) -> None: + self._role_arn = role_arn + self._web_identity_token_file = Path(web_identity_token_file) + self._role_session_name = role_session_name or DEFAULT_ROLE_SESSION_NAME + self._region = region or resolve_sts_region() + self._http_client = http_client + self._credentials: AWSCredentialsIdentity | None = None + self._client: STSClient | None = None + self._refresh_lock = threading.Lock() + # Wait between attempts the same way the SDK's standard retry strategy does. + self._backoff = StandardRetryStrategy().backoff_strategy + + def get_identity( + self, + *, + properties: AWSIdentityProperties, + ) -> AWSCredentialsIdentity: + """Return cached credentials if valid, otherwise call STS for new ones.""" + # Read the cache once: invalidate() can clear it from another thread. + credentials = self._credentials + if credentials is not None and not credentials.is_expired: + return credentials + + with self._refresh_lock: + if self._credentials is not None and not self._credentials.is_expired: + return self._credentials + self._credentials = self._assume_role_with_web_identity() + return self._credentials + + def invalidate(self) -> None: + """Discard the cached credentials so the next call fetches new ones.""" + with self._refresh_lock: + self._credentials = None + + def _assume_role_with_web_identity(self) -> AWSCredentialsIdentity: + from aws_sdk_sts.client import STSClient + from aws_sdk_sts.config import STSConfig + + overrides: AwsConfigOverrides = { + "region": self._region, + "aws_credentials_identity_resolver": UnusedCredentialsResolver(), + } + # The call must be unsigned. The generated STS auth resolver doesn't + # recognize the operation, so it would otherwise be signed. + auth_scheme_resolver = NoAuthResolver() + if self._client is None and self._http_client is not None: + # CAUTION: never closed, since closing it would close the caller's transport. + overrides["transport"] = self._http_client + self._client = STSClient( + config=STSConfig.resolve( + **overrides, auth_scheme_resolver=auth_scheme_resolver + ) + ) + if self._client is not None: + return self._call(self._client) + with STSClient( + config=STSConfig.resolve( + **overrides, auth_scheme_resolver=auth_scheme_resolver + ) + ) as client: + return self._call(client) + + def _call(self, client: STSClient) -> AWSCredentialsIdentity: + return run_program( + fetch_web_identity_credentials( + read_token=read_token, + call_sts=client.assume_role_with_web_identity, + sleep=time.sleep, + backoff=self._backoff, + web_identity_token_file=self._web_identity_token_file, + role_arn=self._role_arn, + role_session_name=self._role_session_name, + ) + ) diff --git a/packages/aws-credentials-sts/tests/unit/test_assume_role.py b/packages/aws-credentials-sts/tests/unit/test_assume_role.py index f76e75b9..66d5f687 100644 --- a/packages/aws-credentials-sts/tests/unit/test_assume_role.py +++ b/packages/aws-credentials-sts/tests/unit/test_assume_role.py @@ -5,11 +5,19 @@ from collections.abc import Callable, Mapping import pytest -from aws_credentials_sts._assume_role import _account_id_from_arn, resolve_sts_region +from aws_credentials_sts._assume_role import ( + AssumeRoleConfigurationError, + _account_id_from_arn, + is_assume_role_profile, + plan_profile, + resolve_sts_region, +) from smithy_aws_core.config.file_parser import Section, StandardizedOutput from smithy_aws_core.config.merged_config import MergedConfig ASSUMED_ROLE_ARN = "arn:aws:sts::123456789012:assumed-role/MyRole/session" +ROLE_ARN = "arn:aws:iam::123456789012:role/MyRole" +TOKEN_FILE = "/var/run/secrets/token" @pytest.fixture @@ -70,3 +78,54 @@ def testresolve_sts_region( ) def test_account_id_from_arn(arn: str | None, expected: str | None) -> None: assert _account_id_from_arn(arn) == expected + + +@pytest.mark.parametrize( + ("profile", "expected"), + [ + ({"role_arn": ROLE_ARN, "source_profile": "base"}, True), + ({"role_arn": ROLE_ARN, "credential_source": "Environment"}, True), + # Left to AssumeRole so it reports the missing source. + ({"role_arn": ROLE_ARN}, True), + # A web identity profile. + ({"role_arn": ROLE_ARN, "web_identity_token_file": TOKEN_FILE}, False), + ( + { + "role_arn": ROLE_ARN, + "web_identity_token_file": TOKEN_FILE, + "source_profile": "base", + }, + True, + ), + ({"web_identity_token_file": TOKEN_FILE}, False), + # A blank token file counts as unset, so this is a role without a source. + ({"role_arn": ROLE_ARN, "web_identity_token_file": ""}, True), + ], +) +def test_is_assume_role_profile( + profile: dict[str, str], + expected: bool, + merged_config: Callable[..., MergedConfig], +) -> None: + config_file = merged_config({"default": profile}) + + assert is_assume_role_profile(config_file, "default") is expected + + +def test_web_identity_source_profile_is_reported_as_unsupported( + merged_config: Callable[..., MergedConfig], +) -> None: + config_file = merged_config( + { + "top": {"role_arn": ROLE_ARN, "source_profile": "wid"}, + "wid": {"role_arn": ROLE_ARN, "web_identity_token_file": TOKEN_FILE}, + } + ) + + # The source profile is a web identity profile, not an AssumeRole profile that + # lacks a source, so the error doesn't suggest adding a source to it. + with pytest.raises( + AssumeRoleConfigurationError, + match="Source profile 'wid' has no supported credential source", + ): + plan_profile(config_file, "top") diff --git a/packages/aws-credentials-sts/tests/unit/test_web_identity.py b/packages/aws-credentials-sts/tests/unit/test_web_identity.py new file mode 100644 index 00000000..6c7ac221 --- /dev/null +++ b/packages/aws-credentials-sts/tests/unit/test_web_identity.py @@ -0,0 +1,719 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# pyright: reportPrivateUsage=false +import asyncio +import os +import re +import time +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from datetime import UTC, datetime, timedelta +from io import BytesIO +from pathlib import Path +from typing import Any, Literal +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +import pytest +from aws_credentials_sts import ( + ProfileWebIdentityProvider, + WebIdentityConfigurationError, + WebIdentityCredentialsResolver, + WebIdentityTokenEnvProvider, +) +from aws_credentials_sts._web_identity import read_token +from aws_credentials_sts.aio import ( + AsyncProfileWebIdentityProvider, + AsyncWebIdentityCredentialsResolver, + AsyncWebIdentityTokenEnvProvider, +) +from aws_sdk_sts.models import ( + AssumedRoleUser, + AssumeRoleWithWebIdentityInput, + AssumeRoleWithWebIdentityOutput, + Credentials, + ExpiredTokenException, + IDPCommunicationErrorException, + InvalidIdentityTokenException, +) +from smithy_aws_core.aio.identity import AsyncIdentityChain +from smithy_aws_core.aio.identity.chain.provider import AsyncChainSetup +from smithy_aws_core.config.file_parser import Section, StandardizedOutput +from smithy_aws_core.config.merged_config import MergedConfig +from smithy_aws_core.identity import AWSCredentialsIdentity +from smithy_aws_core.identity.chain.provider import ChainSetup +from smithy_aws_core.identity.chain.resolver import IdentityChain +from smithy_core.exceptions import SmithyIdentityError +from smithy_http.aio.aiohttp import AIOHTTPClient +from smithy_http.aio.testing import AsyncMockHTTPClient +from smithy_http.testing import MockHTTPClient + +ROLE_ARN = "arn:aws:iam::123456789012:role/WebIdentityRole" +ASSUMED_ROLE_ARN = "arn:aws:sts::123456789012:assumed-role/WebIdentityRole/session" +ACCESS_KEY_ID = "test-access-key" +SECRET_ACCESS_KEY = "test-secret-key" +SESSION_TOKEN = "test-session-token" + +type Mode = Literal["sync", "async"] +type Resolver = WebIdentityCredentialsResolver | AsyncWebIdentityCredentialsResolver +type Outcome = AssumeRoleWithWebIdentityOutput | Exception + + +def _output( + *, access_key_id: str = ACCESS_KEY_ID, expiration: datetime | None = None +) -> AssumeRoleWithWebIdentityOutput: + """An AssumeRoleWithWebIdentity response with valid credentials.""" + return AssumeRoleWithWebIdentityOutput( + credentials=Credentials( + access_key_id=access_key_id, + secret_access_key=SECRET_ACCESS_KEY, + session_token=SESSION_TOKEN, + expiration=expiration or datetime.now(UTC) + timedelta(hours=1), + ), + assumed_role_user=AssumedRoleUser(assumed_role_id="id", arn=ASSUMED_ROLE_ARN), + ) + + +class _FakeSTS: + """Plays back scripted STS outcomes and records each request. + + After every call it writes a new token to the token file, so a test can tell + from the recorded tokens whether the resolver read the file again. + """ + + def __init__(self, token_file: Path, *outcomes: Outcome) -> None: + self._token_file = token_file + self._outcomes = list(outcomes) + self.requests: list[AssumeRoleWithWebIdentityInput] = [] + + @property + def tokens(self) -> list[str | None]: + return [request.web_identity_token for request in self.requests] + + def __call__( + self, request: AssumeRoleWithWebIdentityInput + ) -> AssumeRoleWithWebIdentityOutput: + self.requests.append(request) + self._token_file.write_text(f"token-{len(self.requests) + 1}") + outcome = self._outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + +def _resolver(mode: Mode, token_file: Path, sts: Any = None, **kwargs: Any) -> Resolver: + """Create a resolver in the given mode, with ``sts`` standing in for STS.""" + resolver_class = ( + WebIdentityCredentialsResolver + if mode == "sync" + else AsyncWebIdentityCredentialsResolver + ) + resolver = resolver_class( + role_arn=ROLE_ARN, web_identity_token_file=token_file, **kwargs + ) + resolver._backoff = Mock(compute_next_backoff_delay=Mock(return_value=0)) + if sts is not None: + client = Mock() if mode == "sync" else AsyncMock() + client.assume_role_with_web_identity.side_effect = sts + resolver._client = client + return resolver + + +async def _get_identity(resolver: Resolver) -> AWSCredentialsIdentity: + if isinstance(resolver, AsyncWebIdentityCredentialsResolver): + return await resolver.get_identity(properties={}) + return resolver.get_identity(properties={}) + + +async def _invalidate(resolver: Resolver) -> None: + if isinstance(resolver, AsyncWebIdentityCredentialsResolver): + await resolver.invalidate() + else: + resolver.invalidate() + + +@contextmanager +def _patched_client_class(mode: Mode, sts: _FakeSTS) -> Iterator[MagicMock]: + """Replace the STS client class so the resolver creates a fake client. + + The fake works whether the resolver keeps the client or uses it in a ``with`` + block. + """ + target = ( + "aws_sdk_sts.client.STSClient" + if mode == "sync" + else "aws_sdk_sts.aio.client.AsyncSTSClient" + ) + with patch(target) as client_class: + client = client_class.return_value + mock_class = Mock if mode == "sync" else AsyncMock + client.assume_role_with_web_identity = mock_class(side_effect=sts) + enter = "__enter__" if mode == "sync" else "__aenter__" + getattr(client, enter).return_value = client + yield client_class + + +@pytest.fixture(autouse=True) +def isolated_aws_environment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Keep the machine's AWS environment and config files out of every test.""" + for name in [name for name in os.environ if name.startswith("AWS_")]: + monkeypatch.delenv(name) + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "missing-config")) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "missing-creds")) + + +@pytest.fixture(params=["sync", "async"]) +def mode(request: pytest.FixtureRequest) -> Mode: + return request.param + + +@pytest.fixture +def token_file(tmp_path: Path) -> Path: + path = tmp_path / "token" + path.write_text("token-1") + return path + + +@pytest.fixture +def merged_config() -> Callable[..., MergedConfig]: + def _build(profiles: Mapping[str, Mapping[str, str]]) -> MergedConfig: + sections = { + name: Section(properties=dict(properties)) + for name, properties in profiles.items() + } + return MergedConfig(StandardizedOutput(profiles=sections), StandardizedOutput()) + + return _build + + +# --------------------------------------------------------------------------- +# Resolvers tests +# --------------------------------------------------------------------------- + + +async def test_requests_credentials_for_configured_role( + mode: Mode, token_file: Path +) -> None: + token_file.write_text(" token-1\n") + sts = _FakeSTS(token_file, _output()) + resolver = _resolver(mode, token_file, sts, role_session_name="my-session") + + identity = await _get_identity(resolver) + + request = sts.requests[0] + assert request.role_arn == ROLE_ARN + assert request.role_session_name == "my-session" + # Whitespace around the token, such as a trailing newline, isn't sent. + assert request.web_identity_token == "token-1" + assert identity.access_key_id == ACCESS_KEY_ID + assert identity.secret_access_key == SECRET_ACCESS_KEY + assert identity.session_token == SESSION_TOKEN + assert identity.expiration is not None + assert identity.account_id == "123456789012" + + +def test_default_session_name_is_shared_and_valid(token_file: Path) -> None: + sync_resolver = _resolver("sync", token_file) + async_resolver = _resolver("async", token_file) + + name = sync_resolver._role_session_name + # Every resolver in the process uses the same generated name. + assert async_resolver._role_session_name == name + # The name meets STS's RoleSessionName constraints. + assert 2 <= len(name) <= 64 + assert re.fullmatch(r"[\w+=,.@-]+", name) + + +async def test_response_without_credentials_raises(token_file: Path) -> None: + sts = _FakeSTS(token_file, AssumeRoleWithWebIdentityOutput()) + resolver = _resolver("sync", token_file, sts) + + with pytest.raises(SmithyIdentityError, match="did not contain credentials"): + await _get_identity(resolver) + + +@pytest.mark.parametrize( + ("outcomes", "expected_error", "expected_tokens"), + [ + pytest.param( + [ExpiredTokenException(message="expired"), _output()], + None, + ["token-1", "token-2"], + id="expired-token-retries-with-new-token", + ), + pytest.param( + [IDPCommunicationErrorException(message="idp down"), _output()], + None, + ["token-1", "token-1"], + id="idp-error-retries-with-same-token", + ), + pytest.param( + [ExpiredTokenException(message="expired")] * 4, + ExpiredTokenException, + ["token-1", "token-2", "token-3", "token-4"], + id="expired-token-gives-up-after-retries", + ), + pytest.param( + [IDPCommunicationErrorException(message="idp down")] * 4, + IDPCommunicationErrorException, + ["token-1"] * 4, + id="idp-error-gives-up-after-retries", + ), + pytest.param( + [InvalidIdentityTokenException(message="invalid")], + InvalidIdentityTokenException, + ["token-1"], + id="other-sts-error-not-retried", + ), + pytest.param( + [ConnectionError("network down")], + ConnectionError, + ["token-1"], + id="network-error-not-retried", + ), + ], +) +async def test_retries_sts_errors( + mode: Mode, + token_file: Path, + outcomes: list[Outcome], + expected_error: type[Exception] | None, + expected_tokens: list[str], +) -> None: + sts = _FakeSTS(token_file, *outcomes) + resolver = _resolver(mode, token_file, sts) + + if expected_error is None: + await _get_identity(resolver) + else: + with pytest.raises(expected_error): + await _get_identity(resolver) + + assert sts.tokens == expected_tokens + + +@pytest.mark.parametrize("fixed_before_retry", [True, False]) +@pytest.mark.parametrize( + "initial_content", + # A missing file, a non-ASCII file (a UTF-8 byte order mark), and an empty one. + [None, b"\xef\xbb\xbftoken-1", b"\n"], + ids=["missing", "non-ascii", "empty"], +) +async def test_retries_unreadable_token_file( + mode: Mode, + tmp_path: Path, + initial_content: bytes | None, + fixed_before_retry: bool, +) -> None: + token_file = tmp_path / "token" + if initial_content is not None: + token_file.write_bytes(initial_content) + sts = _FakeSTS(token_file, _output()) + resolver = _resolver(mode, token_file, sts) + + def fix_file_before_retry(retry_attempt: int) -> float: + if fixed_before_retry: + token_file.write_text("token-1") + return 0 + + backoff = Mock(side_effect=fix_file_before_retry) + resolver._backoff = Mock(compute_next_backoff_delay=backoff) + + if fixed_before_retry: + await _get_identity(resolver) + assert sts.tokens == ["token-1"] + else: + with pytest.raises(SmithyIdentityError, match="web identity token file"): + await _get_identity(resolver) + assert sts.requests == [] + # The file was read on every attempt: once, then on each of 3 retries. + assert backoff.call_count == 3 + + +@pytest.mark.parametrize( + ("expires_in", "expected_sts_calls"), + [(timedelta(hours=1), 1), (timedelta(hours=-1), 2)], + ids=["valid-credentials-cached", "expired-credentials-refreshed"], +) +async def test_caches_credentials_until_expired( + mode: Mode, token_file: Path, expires_in: timedelta, expected_sts_calls: int +) -> None: + expiration = datetime.now(UTC) + expires_in + sts = _FakeSTS( + token_file, _output(expiration=expiration), _output(expiration=expiration) + ) + resolver = _resolver(mode, token_file, sts) + + await _get_identity(resolver) + await _get_identity(resolver) + + assert len(sts.requests) == expected_sts_calls + + +async def test_invalidate_forces_refresh(mode: Mode, token_file: Path) -> None: + sts = _FakeSTS( + token_file, + _output(access_key_id="test-access-key-1"), + _output(access_key_id="test-access-key-2"), + ) + resolver = _resolver(mode, token_file, sts) + + await _get_identity(resolver) + await _invalidate(resolver) + identity = await _get_identity(resolver) + + assert identity.access_key_id == "test-access-key-2" + assert len(sts.requests) == 2 + + +async def test_closes_sts_client_it_creates( + mode: Mode, token_file: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("AWS_REGION", "us-west-2") + sts = _FakeSTS(token_file, _output()) + resolver = _resolver(mode, token_file) + + with _patched_client_class(mode, sts) as client_class: + await _get_identity(resolver) + + exit_method = "__exit__" if mode == "sync" else "__aexit__" + getattr(client_class.return_value, exit_method).assert_called_once() + config = client_class.call_args.kwargs["config"] + # Without an explicit region, the region comes from the environment. + assert config.region == "us-west-2" + + # The fake client didn't close the real transport the config created. + if mode == "sync": + config.transport.close() + else: + await config.transport.close() + + +async def test_closes_transport_when_token_never_read( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from aws_sdk_sts.aio.config import AsyncSTSConfig + + configs: list[AsyncSTSConfig] = [] + resolve = AsyncSTSConfig.resolve + + async def recording_resolve(**kwargs: Any) -> AsyncSTSConfig: + config = await resolve(**kwargs) + configs.append(config) + return config + + monkeypatch.setattr(AsyncSTSConfig, "resolve", recording_resolve) + resolver = _resolver("async", tmp_path / "missing-token") + + with pytest.raises(SmithyIdentityError, match="web identity token file"): + await _get_identity(resolver) + + # No request was sent, so the client never closed the transport itself. + assert len(configs) == 1 + transport = configs[0].transport + assert isinstance(transport, AIOHTTPClient) + assert transport._closed + + +async def test_reuses_sts_client_for_callers_http_client( + mode: Mode, token_file: Path +) -> None: + http_client = Mock() + # The first credentials are already expired, so the second call refreshes them + # and shows the same client is reused across refreshes. + sts = _FakeSTS(token_file, _output(expiration=datetime.now(UTC)), _output()) + resolver = _resolver(mode, token_file, region="eu-west-1", http_client=http_client) + + with _patched_client_class(mode, sts) as client_class: + await _get_identity(resolver) + await _get_identity(resolver) + + # The client is created once and never closed, since that would close the + # caller's HTTP client. + client_class.assert_called_once() + exit_method = "__exit__" if mode == "sync" else "__aexit__" + getattr(client_class.return_value, exit_method).assert_not_called() + config = client_class.call_args.kwargs["config"] + assert config.region == "eu-west-1" + assert config.transport is http_client + + +def test_returns_cached_credentials_when_invalidated_concurrently( + token_file: Path, +) -> None: + resolver = _resolver("sync", token_file) + assert isinstance(resolver, WebIdentityCredentialsResolver) + + class InvalidatedWhileChecked(AWSCredentialsIdentity): + @property + def is_expired(self) -> bool: + # Another thread invalidates the cache while this thread checks it. + resolver._credentials = None + return False + + cached = InvalidatedWhileChecked( + access_key_id=ACCESS_KEY_ID, secret_access_key=SECRET_ACCESS_KEY + ) + resolver._credentials = cached + + assert resolver.get_identity(properties={}) is cached + + +def test_concurrent_threads_share_one_refresh(token_file: Path) -> None: + fake_sts = _FakeSTS(token_file, _output()) + + def slow_sts( + request: AssumeRoleWithWebIdentityInput, + ) -> AssumeRoleWithWebIdentityOutput: + time.sleep(0.05) # Keep the refresh in flight while other threads arrive. + return fake_sts(request) + + resolver = _resolver("sync", token_file, slow_sts) + assert isinstance(resolver, WebIdentityCredentialsResolver) + with ThreadPoolExecutor(max_workers=10) as pool: + futures = [pool.submit(resolver.get_identity, properties={}) for _ in range(10)] + # result() re-raises any error a thread hit, so a failing thread fails the test. + identities = [future.result() for future in futures] + + assert len(fake_sts.requests) == 1 + assert all(identity is identities[0] for identity in identities) + + +async def test_concurrent_tasks_share_one_refresh(token_file: Path) -> None: + fake_sts = _FakeSTS(token_file, _output()) + + async def slow_sts( + request: AssumeRoleWithWebIdentityInput, + ) -> AssumeRoleWithWebIdentityOutput: + await asyncio.sleep(0.05) # Keep the refresh in flight while others arrive. + return fake_sts(request) + + resolver = _resolver("async", token_file, slow_sts) + + identities = await asyncio.gather(*(_get_identity(resolver) for _ in range(10))) + + assert len(fake_sts.requests) == 1 + assert all(identity is identities[0] for identity in identities) + + +async def test_token_read_does_not_block_event_loop( + token_file: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + def slow_read_token(path: Path, *, is_last_attempt: bool) -> str | None: + time.sleep(0.3) # A slow disk. + return read_token(path, is_last_attempt=is_last_attempt) + + monkeypatch.setattr("aws_credentials_sts.aio.resolvers.read_token", slow_read_token) + resolver = _resolver("async", token_file, _FakeSTS(token_file, _output())) + ticks = 0 + + async def tick() -> None: + nonlocal ticks + while True: + await asyncio.sleep(0.01) + ticks += 1 + + ticker = asyncio.create_task(tick()) + await _get_identity(resolver) + ticker.cancel() + + # If the read blocked the loop, the ticker couldn't run while it waited. + assert ticks >= 10 + + +_STS_RESPONSE = f"""\ + + + + {ACCESS_KEY_ID} + {SECRET_ACCESS_KEY} + {SESSION_TOKEN} + 2099-01-01T00:00:00Z + + + id + {ASSUMED_ROLE_ARN} + + +""".encode() + + +async def test_default_chain_sends_unsigned_request_with_only_web_identity( + mode: Mode, token_file: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # Only web identity is configured: no other credentials or config files. + monkeypatch.setenv("AWS_WEB_IDENTITY_TOKEN_FILE", str(token_file)) + monkeypatch.setenv("AWS_ROLE_ARN", ROLE_ARN) + monkeypatch.setenv("AWS_REGION", "us-west-2") + + # The real generated STS client sends the request through a mock transport. + if mode == "sync": + http_client = MockHTTPClient() + http_client.add_response(body=_STS_RESPONSE) + identity = IdentityChain.create( + AWSCredentialsIdentity, http_client=http_client + ).get_identity(properties={}) + else: + http_client = AsyncMockHTTPClient() + http_client.add_response(body=_STS_RESPONSE) + chain = await AsyncIdentityChain.create( + AWSCredentialsIdentity, http_client=http_client + ) + identity = await chain.get_identity(properties={}) + + assert identity.access_key_id == ACCESS_KEY_ID + assert identity.account_id == "123456789012" + assert http_client.call_count == 1 + request = http_client.captured_requests[0] + # AssumeRoleWithWebIdentity authenticates with the token, not a signature. + assert "Authorization" not in request.fields + assert isinstance(request.body, BytesIO) + assert b"Action=AssumeRoleWithWebIdentity" in request.body.getvalue() + + +# --------------------------------------------------------------------------- +# Providers tests +# --------------------------------------------------------------------------- + + +async def _setup( + mode: Mode, + provider_kind: Literal["env", "profile"], + *, + config_file: MergedConfig | None = None, + profile_name: str | None = None, +) -> ChainSetup | AsyncChainSetup: + """Run a web identity provider's setup and return the chain setup.""" + if mode == "sync": + sync_setup = ChainSetup(config_file=config_file, profile_name=profile_name) + provider = ( + WebIdentityTokenEnvProvider() + if provider_kind == "env" + else ProfileWebIdentityProvider() + ) + sync_setup.set_current_provider(provider) + provider.setup(AWSCredentialsIdentity, sync_setup) + return sync_setup + + async_setup = AsyncChainSetup(config_file=config_file, profile_name=profile_name) + async_provider = ( + AsyncWebIdentityTokenEnvProvider() + if provider_kind == "env" + else AsyncProfileWebIdentityProvider() + ) + async_setup.set_current_provider(async_provider) + await async_provider.setup(AWSCredentialsIdentity, async_setup) + return async_setup + + +def _added_resolver(mode: Mode, setup: ChainSetup | AsyncChainSetup) -> Resolver: + """Return the single terminal resolver the provider added.""" + assert setup.terminal + assert len(setup.resolvers) == 1 + resolver = setup.resolvers[0].resolver + expected_class = ( + WebIdentityCredentialsResolver + if mode == "sync" + else AsyncWebIdentityCredentialsResolver + ) + assert isinstance(resolver, expected_class) + return resolver + + +@pytest.fixture +def web_identity_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "/var/run/secrets/token") + monkeypatch.setenv("AWS_ROLE_ARN", ROLE_ARN) + monkeypatch.setenv("AWS_ROLE_SESSION_NAME", "env-session") + monkeypatch.setenv("AWS_REGION", "us-west-2") + + +@pytest.mark.usefixtures("web_identity_env") +async def test_env_provider_adds_resolver(mode: Mode) -> None: + resolver = _added_resolver(mode, await _setup(mode, "env")) + + assert resolver._role_arn == ROLE_ARN + assert resolver._web_identity_token_file == Path("/var/run/secrets/token") + assert resolver._role_session_name == "env-session" + assert resolver._region == "us-west-2" + + +@pytest.mark.usefixtures("web_identity_env") +@pytest.mark.parametrize( + ("unset_variable", "profile_name"), + [ + pytest.param("AWS_WEB_IDENTITY_TOKEN_FILE", None, id="missing-token-file"), + pytest.param("AWS_ROLE_ARN", None, id="missing-role-arn"), + # An explicitly selected profile takes precedence over the environment. + pytest.param(None, "dev", id="explicit-profile"), + ], +) +async def test_env_provider_skips_incomplete_or_overridden_config( + mode: Mode, + monkeypatch: pytest.MonkeyPatch, + unset_variable: str | None, + profile_name: str | None, +) -> None: + if unset_variable is not None: + monkeypatch.delenv(unset_variable) + + setup = await _setup(mode, "env", profile_name=profile_name) + + assert setup.resolvers == () + assert not setup.terminal + + +async def test_profile_provider_adds_resolver( + mode: Mode, merged_config: Callable[..., MergedConfig] +) -> None: + config_file = merged_config( + { + "ci": { + "role_arn": ROLE_ARN, + "web_identity_token_file": "/var/run/secrets/token", + "role_session_name": "profile-session", + "region": "eu-west-1", + } + } + ) + + setup = await _setup(mode, "profile", config_file=config_file, profile_name="ci") + + resolver = _added_resolver(mode, setup) + assert resolver._role_arn == ROLE_ARN + assert resolver._web_identity_token_file == Path("/var/run/secrets/token") + assert resolver._role_session_name == "profile-session" + assert resolver._region == "eu-west-1" + + +@pytest.mark.parametrize( + "profile", + [ + pytest.param({"role_arn": ROLE_ARN}, id="no-token-file"), + # A blank leftover key must not break a profile that resolves another way. + pytest.param( + {"credential_process": "get-creds", "web_identity_token_file": ""}, + id="blank-token-file", + ), + ], +) +async def test_profile_provider_skips_profile_without_token_file( + mode: Mode, merged_config: Callable[..., MergedConfig], profile: dict[str, str] +) -> None: + config_file = merged_config({"ci": profile}) + + setup = await _setup(mode, "profile", config_file=config_file, profile_name="ci") + + assert setup.resolvers == () + + +async def test_profile_provider_requires_role_arn( + mode: Mode, merged_config: Callable[..., MergedConfig] +) -> None: + config_file = merged_config( + {"ci": {"web_identity_token_file": "/var/run/secrets/token"}} + ) + + with pytest.raises(WebIdentityConfigurationError, match="role_arn"): + await _setup(mode, "profile", config_file=config_file, profile_name="ci")