diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 3b93e89..55e2d17 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -4,6 +4,23 @@ All notable changes to the Python package are documented here. The format follow [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the package uses [Semantic Versioning](https://semver.org/). +## [Unreleased] + +### Added + +- `ElastiCacheIAMAuthTokenManager` — caches a token for 880 seconds (its 900-second + lifetime minus a 20-second serve margin) and refreshes it on a background daemon + thread `refresh_after` seconds after signing (default 300). Concurrent callers share + a single signing cycle and its outcome; failed refreshes are retried with jittered + exponential backoff while the current token is still served, with a warning logged + per failed cycle; `on_token_changed` reports each new token; `get_credentials()` + returns `(user_id, token)`; `refresh_token()` signs a replacement on demand; + `close()` or a `with` block stops background work. +- `TokenRefreshError` — raised by the manager when a cached token expired before a + refresh could replace it, or after `close()`; `__cause__` is the last refresh error. + Not a `ToolkitInputError`. +- `DEFAULT_REFRESH_AFTER_SECONDS` constant. + ## [1.0.0] - Unreleased Initial release. diff --git a/python/README.md b/python/README.md index 31211f1..2643197 100644 --- a/python/README.md +++ b/python/README.md @@ -88,6 +88,72 @@ username, password = auth.user_id, auth.get_token() The [`examples/`](examples/) directory has redis-py and valkey-py integrations. +### Caching and background refresh + +`ElastiCacheIAMAuthTokenProvider` is stateless: every `get_token()` call resolves +credentials and signs a new token. `ElastiCacheIAMAuthTokenManager` adds caching on +top of that: + +| | `ElastiCacheIAMAuthTokenProvider` | `ElastiCacheIAMAuthTokenManager` | +| ---------------------- | --------------------------------- | ----------------------------------------------- | +| Token per `get_token()`| A freshly signed token every call | The cached token while it is valid | +| Refresh | Caller-driven | Background, `refresh_after` seconds after signing | +| Refresh failures | Raised to the caller | Retried, with the current token still served | +| Change notification | None | `on_token_changed` | +| Cleanup | None needed | `close()`, or use as a context manager | + +Use the manager when a client reconnects often or asks for a token more often than +the 900-second token lifetime; use the provider when each call should sign a new +token. + +```python +from developer_toolkit_elasticache import ElastiCacheIAMAuthTokenManager + +with ElastiCacheIAMAuthTokenManager( + serverless_cache_name="my-cache", + user_id="iam-user", + region="us-east-1", + # Optional; defaults to 300. Must be greater than 0 and less than 880. + refresh_after=300, + on_token_changed=lambda token: print("installed a new IAM auth token"), +) as auth: + username, password = auth.get_credentials() +``` + +`get_credentials()` returns `(user_id, token)` for clients that take a no-argument +credentials callable; `get_token()` remains available when the client takes the +username and password separately. `refresh_token()` signs a replacement immediately +and returns it, for a client whose `AUTH` was rejected with a token the manager still +considers valid (for example after the signing credentials were revoked); it raises +`TokenRefreshError` if no replacement could be signed. + +The region is resolved when the manager is created, but no credentials are resolved +and no token is signed until the first `get_token()` or `get_credentials()` call. +After that, `get_token()` returns the cached token immediately while it is valid. A +token is served for 880 seconds after signing — its 900-second lifetime minus a +20-second margin for clock skew and the `AUTH` round trip — and `refresh_after` must +be below that. A background thread signs a replacement `refresh_after` seconds after +each token is issued; concurrent callers with no valid token share a single signing +and receive the same outcome. + +A refresh makes up to 8 attempts, the first immediate and the rest with exponential +backoff from 100 ms to a nominal cap of 5 seconds, each delay jittered by ±20%. Every +error is retried, including `ConfigurationError`: on an instance role a transient +IMDS failure surfaces as "no credentials". A `get_token()` call with no valid token +blocks for that cycle, so it can take several seconds when signing keeps failing. + +While the current token is still valid a failed refresh is invisible to callers: the +cached token keeps being served, another refresh is scheduled, and a warning is +logged to the `developer_toolkit_elasticache.token_manager` logger. `get_token()` +raises `TokenRefreshError` only if the cached token expired before a refresh could +replace it, or after `close()`; its `__cause__` is the error from the last attempt. A +failure to sign the very first token is raised as the underlying error instead. + +`on_token_changed` must be a synchronous callable; if it raises, the error is logged +and the new token is kept. The refresh thread is a daemon, so it never keeps the +interpreter alive on its own. Call `close()` (or leave the `with` block) when the +manager is no longer needed. + ### Command line ```bash diff --git a/python/src/developer_toolkit_elasticache/__init__.py b/python/src/developer_toolkit_elasticache/__init__.py index 0f2b49b..aa6c80d 100644 --- a/python/src/developer_toolkit_elasticache/__init__.py +++ b/python/src/developer_toolkit_elasticache/__init__.py @@ -6,17 +6,25 @@ from developer_toolkit_elasticache.errors import ( ConfigurationError, InvalidParameterError, + TokenRefreshError, ToolkitInputError, ) from developer_toolkit_elasticache.token_generator import ( ElastiCacheIAMAuthTokenProvider, generate_iam_auth_token, ) +from developer_toolkit_elasticache.token_manager import ( + DEFAULT_REFRESH_AFTER_SECONDS, + ElastiCacheIAMAuthTokenManager, +) __all__ = [ + "DEFAULT_REFRESH_AFTER_SECONDS", "ConfigurationError", + "ElastiCacheIAMAuthTokenManager", "ElastiCacheIAMAuthTokenProvider", "InvalidParameterError", + "TokenRefreshError", "ToolkitInputError", "generate_iam_auth_token", ] diff --git a/python/src/developer_toolkit_elasticache/errors.py b/python/src/developer_toolkit_elasticache/errors.py index 72b507b..2050c91 100644 --- a/python/src/developer_toolkit_elasticache/errors.py +++ b/python/src/developer_toolkit_elasticache/errors.py @@ -33,3 +33,13 @@ class ConfigurationError(ToolkitInputError): Covers the missing-configuration cases (no region or no credentials). """ + + +class TokenRefreshError(Exception): + """A cached token could not be refreshed before it expired, or the token + manager was closed. + + Raised only by ``ElastiCacheIAMAuthTokenManager``. This is an operational + failure rather than an input error, so it does not extend ``ToolkitInputError``. + When a refresh failed, ``__cause__`` is the error from the last refresh attempt. + """ diff --git a/python/src/developer_toolkit_elasticache/token_manager.py b/python/src/developer_toolkit_elasticache/token_manager.py new file mode 100644 index 0000000..315837e --- /dev/null +++ b/python/src/developer_toolkit_elasticache/token_manager.py @@ -0,0 +1,377 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A cached, self-refreshing IAM auth token for Amazon ElastiCache. + +``ElastiCacheIAMAuthTokenManager`` wraps the stateless provider with a cache that is +refreshed in the background, so callers always receive a valid token without signing +on every call. +""" + +import logging +import math +import random +import threading +import time +from collections.abc import Callable +from concurrent.futures import Future +from dataclasses import dataclass, field +from typing import Protocol + +from botocore.session import Session + +from developer_toolkit_elasticache.errors import InvalidParameterError, TokenRefreshError +from developer_toolkit_elasticache.token_generator import ( + _TOKEN_TTL_SECONDS, + ElastiCacheIAMAuthTokenProvider, +) + +_logger = logging.getLogger(__name__) + +# Stop serving a token this long before its 900s lifetime ends, to absorb clock skew +# and the AUTH round trip. +_SERVE_MARGIN_SECONDS = 20 +_EFFECTIVE_LIFETIME_SECONDS = _TOKEN_TTL_SECONDS - _SERVE_MARGIN_SECONDS + +# A third of the token lifetime: two refresh windows remain if this one fails. +DEFAULT_REFRESH_AFTER_SECONDS = 300 + +# Retry policy for one refresh cycle +_MAX_ATTEMPTS = 8 +_BASE_DELAY_SECONDS = 0.1 +_MAX_DELAY_SECONDS = 5.0 +_JITTER_RATIO = 0.2 + +_CLOSED_MESSAGE = "The token manager is closed." +_REFRESH_FAILED_MESSAGE = ( + "Could not refresh the ElastiCache IAM authentication token before the previous " + "token expired." +) +_REPLACEMENT_FAILED_MESSAGE = ( + "Could not sign a replacement ElastiCache IAM authentication token." +) + +# Module attributes so tests can substitute the clock, scheduler, jitter, and sleep. +# Wall clock, not monotonic: the server judges the token by X-Amz-Date against its own +# wall clock, and a monotonic clock stands still while the host is suspended. +_now = time.time +_random = random.random # jitter only, not security-sensitive + + +def _wait(stop: threading.Event, seconds: float) -> None: + """Sleep between retries; returns early once ``stop`` is set by close().""" + stop.wait(seconds) + + +class _ScheduledCall(Protocol): + def cancel(self) -> None: ... + + +def _schedule(delay: float, callback: Callable[[], None]) -> _ScheduledCall: + """Run ``callback`` after ``delay`` seconds on a daemon thread.""" + timer = threading.Timer(delay, callback) + timer.daemon = True + timer.name = "elasticache-iam-token-refresh" + timer.start() + return timer + + +@dataclass(frozen=True) +class _CachedToken: + token: str = field(repr=False) # bearer secret: keep it out of dumps of locals + issued_at: float + expires_at: float + + +def _validate_refresh_after(value: object) -> float: + valid = ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + and 0 < value < _EFFECTIVE_LIFETIME_SECONDS + ) + if not valid: + raise InvalidParameterError( + "Invalid value for parameter 'refresh_after': must be greater than 0 " + f"and less than {_EFFECTIVE_LIFETIME_SECONDS}." + ) + return float(value) + + +class ElastiCacheIAMAuthTokenManager: + """Cache an ElastiCache IAM auth token and refresh it in the background. + + Takes the provider's arguments; configuration is validated at construction and + the first token is signed lazily. A replacement is signed ``refresh_after`` + seconds after each token (default 300; must be below the 880s effective lifetime, + which is the 900s token lifetime minus a 20s serve margin) while callers keep + receiving the current one. ``on_token_changed`` is called synchronously + with each new token; its failures are logged and otherwise ignored. + + ``refresh_token()`` signs a replacement immediately, for a client whose ``AUTH`` + was rejected with a token this manager still considers valid. + + Thread-safe: callers with no valid token share one signing cycle. Call + ``close()`` (or use as a context manager) when done; later calls raise + ``TokenRefreshError``. + """ + + def __init__( # noqa: PLR0913 # the provider's keyword-only arguments plus manager options + self, + *, + user_id: str, + region: str | None = None, + serverless_cache_name: str | None = None, + replication_group_id: str | None = None, + session: Session | None = None, + refresh_after: float = DEFAULT_REFRESH_AFTER_SECONDS, + on_token_changed: Callable[[str], None] | None = None, + ): + self._provider = ElastiCacheIAMAuthTokenProvider( + user_id=user_id, + region=region, + serverless_cache_name=serverless_cache_name, + replication_group_id=replication_group_id, + session=session, + ) + self._refresh_after = _validate_refresh_after(refresh_after) + self._on_token_changed = on_token_changed + # Guards the fields below; never held while signing or sleeping. + self._lock = threading.Lock() + self._cached: _CachedToken | None = None + self._closed = False + self._stop = threading.Event() + self._inflight: Future[str] | None = None # the signing cycle in progress + self._scheduled: _ScheduledCall | None = None + self._refresh_due_at = math.inf # in the past = missed timer, refresh now + self._refreshing = False + + @property + def user_id(self) -> str: + return self._provider.user_id + + def get_credentials(self) -> tuple[str, str]: + """Return ``(user_id, token)`` for clients that take a credentials callable.""" + return self.user_id, self.get_token() + + def get_token(self) -> str: + """Return a valid token, signing one only if the cache is empty or expired. + + Blocks for the signing cycle (retries included) only when no valid token is + cached. Raises ``TokenRefreshError`` if the cached token expired before a + refresh replaced it, or after ``close()``; a first-token failure raises the + underlying error. + """ + with self._lock: + self._ensure_open() + cached = self._cached + now = _now() + if cached is not None and now < cached.expires_at: + if now >= self._refresh_due_at: + self._start_background_refresh_locked(cached) + return cached.token + return self._refresh(reuse_unless=None) + + def refresh_token(self) -> str: + """Sign a replacement token now and return it. + + Use after the server rejected the current token (for example because the + signing credentials were revoked). Joins a signing cycle already in progress + instead of starting another. Raises ``TokenRefreshError`` if no replacement + could be signed, or after ``close()``. + """ + with self._lock: + self._ensure_open() + current = self._cached + token = self._refresh(reuse_unless=current) + with self._lock: + replaced = self._cached is not None and self._cached is not current + if not replaced: + raise TokenRefreshError(_REPLACEMENT_FAILED_MESSAGE) + return token + + def close(self) -> None: + """Stop background refresh and drop the cached token. Idempotent.""" + with self._lock: + self._closed = True + self._cached = None + scheduled, self._scheduled = self._scheduled, None + self._stop.set() + if scheduled is not None: + scheduled.cancel() + + def __enter__(self) -> "ElastiCacheIAMAuthTokenManager": # noqa: PYI034 # typing.Self is 3.11+; the package supports 3.10 + return self + + def __exit__(self, *exc_info: object) -> None: + self.close() + + def _ensure_open(self) -> None: + if self._closed: + raise TokenRefreshError(_CLOSED_MESSAGE) + + def _refresh(self, *, reuse_unless: "_CachedToken | None") -> str: + """Sign a new token, unless a valid one other than ``reuse_unless`` is cached. + + ``reuse_unless`` is the token the caller wants replaced (``None`` to accept any + valid token). Joins the signing cycle already in progress, if any, and shares + its outcome. + """ + with self._lock: + self._ensure_open() + cached = self._cached + if ( + cached is not None + and cached is not reuse_unless + and _now() < cached.expires_at + ): + return cached.token + flight = self._inflight + owner = flight is None + if owner: + flight = self._inflight = Future() + if not owner: + return flight.result() + try: + token, installed = self._run_refresh_cycle() + except BaseException as error: + flight.set_exception(error) + raise + else: + flight.set_result(token) + finally: + with self._lock: + if self._inflight is flight: + self._inflight = None + # Only after the flight is closed: a refresh timer that fires now starts its own + # cycle instead of joining this one and skipping the reschedule, and a slow + # callback cannot hold the flight open. + if installed: + with self._lock: + if not self._closed and self._cached is not None: + self._schedule_locked(self._refresh_after, self._cached) + self._notify(token) + return token + + def _run_refresh_cycle(self) -> tuple[str, bool]: + """Sign with retries; returns ``(token, newly_installed)``. + + Every error is retried: on an instance role a transient IMDS failure surfaces + as "no credentials", so even ``ConfigurationError`` may be transient. On + failure: return a still-valid cached token and reschedule, raise the last + error if nothing is cached, or raise ``TokenRefreshError`` if it expired. + """ + last_error: Exception | None = None + for attempt in range(1, _MAX_ATTEMPTS + 1): + self._raise_if_closed(last_error) + try: + issued_at = _now() + token = self._provider.get_token() + except Exception as error: # noqa: BLE001 # see docstring + last_error = error + if attempt < _MAX_ATTEMPTS: + _wait(self._stop, self._backoff_delay(attempt)) + continue + with self._lock: + self._ensure_open() + self._install_locked(token, issued_at) + return token, True + + remaining = 0.0 + with self._lock: + self._raise_if_closed(last_error) + cached = self._cached + if cached is not None and _now() < cached.expires_at: + self._schedule_locked(_MAX_DELAY_SECONDS, cached) + remaining = cached.expires_at - _now() + if cached is not None and remaining > 0: + _logger.warning( + "Could not refresh the ElastiCache IAM auth token after %d attempts " + "(%s: %s); serving the current token for another %.0fs and retrying " + "in %.0fs.", + _MAX_ATTEMPTS, + type(last_error).__name__, + last_error, + remaining, + _MAX_DELAY_SECONDS, + ) + # Re-check: the token may have expired while the lock was released, and a + # waiter that joined because its token expired must not get it back. + with self._lock: + self._raise_if_closed(last_error) + if _now() < cached.expires_at: + return cached.token, False + if cached is None and last_error is not None: + raise last_error + raise TokenRefreshError(_REFRESH_FAILED_MESSAGE) from last_error + + def _notify(self, token: str) -> None: + if self._on_token_changed is None: + return + try: + self._on_token_changed(token) + except Exception as error: # noqa: BLE001 # a callback failure must not discard the token + # Type only: the callback received the token, so its message might too. + _logger.warning("on_token_changed callback raised %s", type(error).__name__) + + def _raise_if_closed(self, cause: Exception | None) -> None: + if self._closed: + raise TokenRefreshError(_CLOSED_MESSAGE) from cause + + @staticmethod + def _backoff_delay(failed_attempt: int) -> float: + delay = min(_BASE_DELAY_SECONDS * 2 ** (failed_attempt - 1), _MAX_DELAY_SECONDS) + jitter = 1 - _JITTER_RATIO + 2 * _JITTER_RATIO * _random() + return delay * jitter + + def _install_locked(self, token: str, issued_at: float) -> None: + """Cache the token. Its refresh timer is started by ``_refresh`` once the + flight is closed; until then the due time is set so nobody treats the gap as + a missed timer.""" + self._cached = _CachedToken( + token=token, + issued_at=issued_at, + expires_at=issued_at + _EFFECTIVE_LIFETIME_SECONDS, + ) + self._refresh_due_at = _now() + self._refresh_after + + def _schedule_locked(self, delay: float, due: _CachedToken) -> None: + """Schedule a background refresh of ``due``; on failure, leave it to get_token().""" + previous, self._scheduled = self._scheduled, None + if previous is not None: + previous.cancel() + try: + self._scheduled = _schedule(delay, lambda: self._background_refresh(due)) + except Exception: # noqa: BLE001 # e.g. RuntimeError: can't start new thread + self._refresh_due_at = -math.inf + self._refreshing = False + return + self._refresh_due_at = _now() + delay + + def _start_background_refresh_locked(self, due: _CachedToken) -> None: + if not self._refreshing: + self._refreshing = True + self._schedule_locked(0, due) + + def _background_refresh(self, due: _CachedToken) -> None: + """Replace ``due`` with a fresh token. Runs on the scheduler's thread.""" + with self._lock: + if self._closed: + self._refreshing = False + return + self._refreshing = True + try: + # A cycle that fails while the token is valid logs, reschedules, and + # returns. It only raises once the token has expired; the next get_token() + # surfaces that to a caller, so nothing is lost by not re-raising here. + self._refresh(reuse_unless=due) + except Exception as error: # noqa: BLE001 # a timer thread has nobody to raise to + _logger.warning( + "Background refresh of the ElastiCache IAM auth token failed and the " + "cached token has expired (%s: %s).", + type(error).__name__, + error, + ) + finally: + with self._lock: + self._refreshing = False diff --git a/python/tests/conftest.py b/python/tests/conftest.py new file mode 100644 index 0000000..f86b662 --- /dev/null +++ b/python/tests/conftest.py @@ -0,0 +1,23 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from helpers import fake_session + + +@pytest.fixture +def mock_session(): + return fake_session() + + +@pytest.fixture(autouse=True) +def no_region_env(monkeypatch): + """Sever the region environment variables for every test. + + Autouse because AWS_REGION is read from the environment directly, so without this + the suite's result would depend on the developer's shell: a region exported + locally would mask the session fallback and quietly pass a broken build. Tests + that want the env var set it explicitly with monkeypatch.setenv. + """ + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False) diff --git a/python/tests/helpers.py b/python/tests/helpers.py new file mode 100644 index 0000000..12b3f28 --- /dev/null +++ b/python/tests/helpers.py @@ -0,0 +1,48 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Test doubles shared by the token generator and token manager suites.""" + +from unittest.mock import MagicMock + +from botocore.credentials import Credentials +from botocore.session import Session + +CACHE = "my-cache" +USER = "testuser" +REGION = "us-east-1" +SESSION_REGION = "eu-west-1" + +DEFAULT_CREDENTIALS = Credentials("FRAIDA1FODNN7EXAMPLE", "secret", "token") +ROTATED_CREDENTIALS = Credentials("AKIAI44QH8DHBEXAMPLE", "other-secret", "token2") + +_UNSET = object() + + +def fake_session(credentials=_UNSET, region=_UNSET): + """A botocore session stand-in that vends fixed credentials and a region. + + Every test signs with an injected session, so the suite never touches the real + AWS credential chain: no environment, config file, or IMDS lookup can leak in. + Pass ``credentials=None`` to simulate an empty credential chain, or + ``region=None`` to simulate a session with no region configured. + + The region is stubbed explicitly rather than left to MagicMock's auto-attribute, + because an auto-created mock is truthy and would let a broken fallback look like + it resolved a region. + """ + if credentials is _UNSET: + credentials = DEFAULT_CREDENTIALS + if region is _UNSET: + region = SESSION_REGION + session = MagicMock(spec=Session) + if credentials is None: + session.get_credentials.return_value = None + else: + session.get_credentials.return_value = MagicMock( + get_frozen_credentials=MagicMock(return_value=credentials) + ) + session.get_config_variable.side_effect = lambda name: ( + region if name == "region" else None + ) + return session diff --git a/python/tests/test_token_generator.py b/python/tests/test_token_generator.py index 0dd9d04..849b48e 100644 --- a/python/tests/test_token_generator.py +++ b/python/tests/test_token_generator.py @@ -6,7 +6,14 @@ import pytest from botocore.credentials import Credentials -from botocore.session import Session +from helpers import ( + CACHE, + REGION, + ROTATED_CREDENTIALS, + SESSION_REGION, + USER, + fake_session, +) from developer_toolkit_elasticache import ( ElastiCacheIAMAuthTokenProvider, @@ -17,67 +24,6 @@ InvalidParameterError, ) -CACHE = "my-cache" -USER = "testuser" -REGION = "us-east-1" -# Deliberately different from REGION so a test can tell an explicitly-passed region -# apart from one that fell back to the session. -SESSION_REGION = "eu-west-1" - -_DEFAULT_CREDENTIALS = Credentials("FRAIDA1FODNN7EXAMPLE", "secret", "token") -_ROTATED_CREDENTIALS = Credentials("AKIAI44QH8DHBEXAMPLE", "other-secret", "token2") - -# Sentinel so callers can ask for "no credentials at all" (None) and still be -# distinguishable from "just give me the default set". -_UNSET = object() - - -def _fake_session(credentials=_UNSET, region=_UNSET): - """A botocore session stand-in that vends fixed credentials and a region. - - Every test signs with an injected session, so the suite never touches the real - AWS credential chain: no environment, config file, or IMDS lookup can leak in. - Pass ``credentials=None`` to simulate an empty credential chain, or - ``region=None`` to simulate a session with no region configured. - - The region is stubbed explicitly rather than left to MagicMock's auto-attribute, - because an auto-created mock is truthy and would let a broken fallback look like - it resolved a region. - """ - if credentials is _UNSET: - credentials = _DEFAULT_CREDENTIALS - if region is _UNSET: - region = SESSION_REGION - session = MagicMock(spec=Session) - if credentials is None: - session.get_credentials.return_value = None - else: - session.get_credentials.return_value = MagicMock( - get_frozen_credentials=MagicMock(return_value=credentials) - ) - session.get_config_variable.side_effect = lambda name: ( - region if name == "region" else None - ) - return session - - -@pytest.fixture -def mock_session(): - return _fake_session() - - -@pytest.fixture(autouse=True) -def no_region_env(monkeypatch): - """Sever the region environment variables for every test. - - Autouse because AWS_REGION is read from the environment directly, so without this - the suite's result would depend on the developer's shell: a region exported - locally would mask the session fallback and quietly pass a broken build. Tests - that want the env var set it explicitly with monkeypatch.setenv. - """ - monkeypatch.delenv("AWS_REGION", raising=False) - monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False) - def _assert_valid_token(token, *, serverless=True): assert token is not None @@ -167,7 +113,7 @@ def test_rotated_credentials_change_the_signature(mock_session): first = auth.get_token() mock_session.get_credentials.return_value = MagicMock( - get_frozen_credentials=MagicMock(return_value=_ROTATED_CREDENTIALS) + get_frozen_credentials=MagicMock(return_value=ROTATED_CREDENTIALS) ) assert auth.get_token() != first @@ -189,7 +135,7 @@ def test_empty_credential_chain_raises(): serverless_cache_name=CACHE, user_id=USER, region=REGION, - session=_fake_session(credentials=None), + session=fake_session(credentials=None), ) with pytest.raises(ConfigurationError, match="credentials"): auth.get_token() @@ -398,7 +344,7 @@ def test_region_falls_back_to_the_session(mock_session): def test_explicit_region_overrides_the_session(): """An explicitly passed region wins over the session's configured one.""" - session = _fake_session(region=SESSION_REGION) + session = fake_session(region=SESSION_REGION) token = generate_iam_auth_token( serverless_cache_name=CACHE, user_id=USER, region=REGION, session=session ) @@ -413,7 +359,7 @@ def test_aws_region_env_var_is_honoured(mock_session, monkeypatch): is set only there would otherwise have a working CLI and a failing toolkit. """ monkeypatch.setenv("AWS_REGION", ENV_REGION) - session = _fake_session(region=None) + session = fake_session(region=None) token = generate_iam_auth_token( serverless_cache_name=CACHE, user_id=USER, session=session ) @@ -423,7 +369,7 @@ def test_aws_region_env_var_is_honoured(mock_session, monkeypatch): def test_aws_region_takes_priority_over_the_session(monkeypatch): """Environment beats config file, matching how botocore ranks its own sources.""" monkeypatch.setenv("AWS_REGION", ENV_REGION) - session = _fake_session(region=SESSION_REGION) + session = fake_session(region=SESSION_REGION) token = generate_iam_auth_token( serverless_cache_name=CACHE, user_id=USER, session=session ) @@ -438,7 +384,7 @@ def test_explicit_region_overrides_aws_region_env_var(monkeypatch): serverless_cache_name=CACHE, user_id=USER, region=REGION, - session=_fake_session(region=SESSION_REGION), + session=fake_session(region=SESSION_REGION), ) assert f"%2F{REGION}%2F" in token assert ENV_REGION not in token @@ -446,7 +392,7 @@ def test_explicit_region_overrides_aws_region_env_var(monkeypatch): def test_missing_region_raises(): """No region anywhere is a user-fixable configuration problem.""" - session = _fake_session(region=None) + session = fake_session(region=None) with pytest.raises(ConfigurationError, match="region"): generate_iam_auth_token( serverless_cache_name=CACHE, user_id=USER, session=session @@ -455,7 +401,7 @@ def test_missing_region_raises(): def test_missing_region_error_names_every_source(): """The message must name each place a region can come from.""" - session = _fake_session(region=None) + session = fake_session(region=None) with pytest.raises(ConfigurationError, match="AWS_REGION"): generate_iam_auth_token( serverless_cache_name=CACHE, user_id=USER, session=session @@ -475,7 +421,7 @@ def test_class_region_falls_back_to_the_session(mock_session): def test_class_missing_region_fails_at_construction(): """A missing region surfaces immediately, not at the first connection attempt.""" - session = _fake_session(region=None) + session = fake_session(region=None) with pytest.raises(ConfigurationError, match="region"): ElastiCacheIAMAuthTokenProvider( serverless_cache_name=CACHE, user_id=USER, session=session @@ -505,7 +451,7 @@ def test_class_missing_region_fails_at_construction(): def _kat_session(): - return _fake_session(credentials=_KAT_CREDENTIALS, region=_KAT_REGION) + return fake_session(credentials=_KAT_CREDENTIALS, region=_KAT_REGION) def _generate_at_fixed_time(**kwargs): diff --git a/python/tests/test_token_manager.py b/python/tests/test_token_manager.py new file mode 100644 index 0000000..4911256 --- /dev/null +++ b/python/tests/test_token_manager.py @@ -0,0 +1,976 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import threading +import time +from concurrent.futures import Future + +import pytest +from helpers import CACHE, REGION, ROTATED_CREDENTIALS, USER, fake_session + +import developer_toolkit_elasticache as package +from developer_toolkit_elasticache import token_manager +from developer_toolkit_elasticache.errors import ( + ConfigurationError, + InvalidParameterError, + TokenRefreshError, + ToolkitInputError, +) +from developer_toolkit_elasticache.token_manager import ( + DEFAULT_REFRESH_AFTER_SECONDS, + ElastiCacheIAMAuthTokenManager, +) + +TOKEN_TTL = 900 +SERVE_MARGIN = 20 +EFFECTIVE_LIFETIME = TOKEN_TTL - SERVE_MARGIN # 880: served this long after signing +MAX_ATTEMPTS = 8 +MAX_DELAY = 5.0 + +# Captured before any fixture patches them, for the tests of the real implementations. +REAL_SCHEDULE = token_manager._schedule +REAL_WAIT = token_manager._wait + + +class FakeClock: + """A wall clock the test advances by hand.""" + + def __init__(self, start=1_000.0): + self.now = start + + def __call__(self): + return self.now + + def advance(self, seconds): + self.now += seconds + + +@pytest.fixture +def clock(monkeypatch): + fake = FakeClock() + monkeypatch.setattr(token_manager, "_now", fake) + return fake + + +class FakeScheduledCall: + def __init__(self, delay, callback): + self.delay = delay + self.callback = callback + self.cancelled = False + + def cancel(self): + self.cancelled = True + + +class FakeScheduler: + """Captures scheduled callbacks so a test can fire them deterministically.""" + + def __init__(self): + self.calls = [] + self.fail = False + + def __call__(self, delay, callback): + if self.fail: + raise RuntimeError("can't start new thread") + call = FakeScheduledCall(delay, callback) + self.calls.append(call) + return call + + @property + def pending(self): + return [call for call in self.calls if not call.cancelled] + + def fire_pending(self): + """Run every live callback once, on this thread, as the timer would.""" + for call in self.pending: + call.cancel() + call.callback() + + +@pytest.fixture(autouse=True) +def scheduler(monkeypatch): + fake = FakeScheduler() + monkeypatch.setattr(token_manager, "_schedule", fake) + return fake + + +class FakeSleep: + """Records retry delays instead of sleeping. + + The real ``_wait`` honours close() through the stop event; that contract is + covered separately by ``test_real_wait_returns_early_when_stopped``. + """ + + def __init__(self): + self.delays = [] + self.on_sleep = None + + def __call__(self, stop, seconds): + self.delays.append(round(seconds, 6)) + if self.on_sleep is not None: + self.on_sleep() + + +@pytest.fixture(autouse=True) +def sleep(monkeypatch): + fake = FakeSleep() + monkeypatch.setattr(token_manager, "_wait", fake) + # Midpoint jitter (factor 1.0) so recorded delays equal the nominal backoff. + monkeypatch.setattr(token_manager, "_random", lambda: 0.5) + return fake + + +def _failing_then_ok(session, failures, error=None): + """Make credential resolution raise ``failures`` times, then succeed.""" + credentials = session.get_credentials.return_value + calls = {"n": 0} + + def side_effect(): + calls["n"] += 1 + if calls["n"] <= failures: + raise error or RuntimeError(f"transient failure {calls['n']}") + return credentials + + session.get_credentials.side_effect = side_effect + + +def _manager(session, **kwargs): + kwargs.setdefault("serverless_cache_name", CACHE) + kwargs.setdefault("user_id", USER) + kwargs.setdefault("region", REGION) + return ElastiCacheIAMAuthTokenManager(session=session, **kwargs) + + +def _run_concurrently(count, target, *, after_start=None): + """Start ``count`` daemon threads that all enter ``target`` at once. + + ``after_start`` runs on the test thread once the workers are going, e.g. to release + a flight they are waiting on. A deadlock fails the liveness assertion instead of + hanging pytest, because the workers are daemons. + """ + results = [] + results_lock = threading.Lock() + ready = threading.Barrier(count) + + def worker(): + ready.wait(timeout=5) + try: + outcome = target() + except Exception as error: # noqa: BLE001 # collected for the assertions + outcome = error + with results_lock: + results.append(outcome) + + threads = [threading.Thread(target=worker, daemon=True) for _ in range(count)] + for thread in threads: + thread.start() + if after_start is not None: + after_start() + for thread in threads: + thread.join(timeout=5) + assert not thread.is_alive(), "worker did not finish: deadlock?" + assert len(results) == count + return results + + +def _run_bounded(target, timeout=5): + """Run ``target`` on a daemon thread; fail (don't hang) if it does not finish.""" + box = {} + + def run(): + try: + box["result"] = target() + except Exception as error: # noqa: BLE001 # re-raised on the test thread + box["error"] = error + + thread = threading.Thread(target=run, daemon=True) + thread.start() + thread.join(timeout=timeout) + assert not thread.is_alive(), "target did not finish: deadlock?" + if "error" in box: + raise box["error"] + return box["result"] + + +class JoinCountingFuture(Future): + """A Future that counts callers entering result(), so a test can wait until + every waiter has actually joined the flight before releasing the owner.""" + + joins = 0 + joins_changed = threading.Condition() + + def result(self, timeout=None): + with JoinCountingFuture.joins_changed: + JoinCountingFuture.joins += 1 + JoinCountingFuture.joins_changed.notify_all() + return super().result(timeout) + + @classmethod + def reset(cls): + with cls.joins_changed: + cls.joins = 0 + + @classmethod + def wait_for_joins(cls, count, timeout=5): + with cls.joins_changed: + assert cls.joins_changed.wait_for(lambda: cls.joins >= count, timeout), ( + f"only {cls.joins} of {count} waiters joined the flight" + ) + + +@pytest.fixture +def join_counting_future(monkeypatch): + JoinCountingFuture.reset() + monkeypatch.setattr(token_manager, "Future", JoinCountingFuture) + return JoinCountingFuture + + +# --- errors and exports ------------------------------------------------------------ + + +def test_token_refresh_error_is_not_an_input_error(): + assert issubclass(TokenRefreshError, Exception) + assert not issubclass(TokenRefreshError, ToolkitInputError) + + +def test_package_root_exports_the_manager_api(): + assert package.ElastiCacheIAMAuthTokenManager is ElastiCacheIAMAuthTokenManager + assert package.TokenRefreshError is TokenRefreshError + assert package.DEFAULT_REFRESH_AFTER_SECONDS == DEFAULT_REFRESH_AFTER_SECONDS + assert "ElastiCacheIAMAuthTokenManager" in package.__all__ + assert "TokenRefreshError" in package.__all__ + + +# --- caching ----------------------------------------------------------------------- + + +def test_signs_lazily_and_caches_the_token(clock, mock_session): + manager = _manager(mock_session) + assert mock_session.get_credentials.call_count == 0 + + first = manager.get_token() + second = manager.get_token() + + assert first == second + assert f"{CACHE}/?" in first + assert mock_session.get_credentials.call_count == 1 + + +def test_serves_the_cached_token_until_the_serve_margin(clock, mock_session): + # Stop 20s before the real 900s expiry to absorb clock skew and the AUTH round trip. + manager = _manager(mock_session) + manager.get_token() + + clock.advance(EFFECTIVE_LIFETIME - 1) + manager.get_token() + assert mock_session.get_credentials.call_count == 1 + + clock.advance(1) + manager.get_token() + assert mock_session.get_credentials.call_count == 2 + + +def test_uses_the_wall_clock_for_expiry(): + # The server judges the token by X-Amz-Date against its wall clock; a monotonic + # clock would stand still through a host suspend and overstate validity. + assert token_manager._now is time.time + + +def test_get_credentials_returns_the_normalized_user_id_and_cached_token( + clock, mock_session +): + manager = _manager(mock_session, user_id="MixedCaseUser") + + user_id, token = manager.get_credentials() + + assert user_id == "mixedcaseuser" + assert token == manager.get_token() + assert manager.user_id == "mixedcaseuser" + + +def test_cached_token_repr_does_not_contain_the_token(clock, mock_session): + manager = _manager(mock_session) + token = manager.get_token() + + assert "X-Amz-Signature" not in repr(manager._cached) + assert token not in repr(manager._cached) + + +def test_first_signing_failure_is_the_underlying_error(clock, sleep): + manager = _manager(fake_session(credentials=None)) + + with pytest.raises(ConfigurationError): + manager.get_token() + assert len(sleep.delays) == MAX_ATTEMPTS - 1 + + +def test_validates_parameters_at_construction(clock, mock_session): + with pytest.raises(InvalidParameterError): + _manager(mock_session, serverless_cache_name="bad_name!") + with pytest.raises(ConfigurationError): + _manager(fake_session(region=None), region=None) + + +# --- single flight ----------------------------------------------------------------- + + +def test_concurrent_callers_share_one_signing(clock, join_counting_future): + session = fake_session() + credentials = session.get_credentials.return_value + release_owner = threading.Event() + + def blocked_get_credentials(): + # The owner holds the flight open here until every waiter has joined it. + assert release_owner.wait(timeout=5) + return credentials + + session.get_credentials.side_effect = blocked_get_credentials + manager = _manager(session) + + def release_once_all_joined(): + join_counting_future.wait_for_joins(4) # 5 callers: 1 owner + 4 waiters + release_owner.set() + + tokens = _run_concurrently(5, manager.get_token, after_start=release_once_all_joined) + + assert all(isinstance(token, str) for token in tokens) + assert len(set(tokens)) == 1 + assert session.get_credentials.call_count == 1 + + +def test_concurrent_callers_share_one_failure(clock, sleep, join_counting_future): + # A failed cycle is reported to every waiter, rather than each waiter running its + # own eight attempts in turn. + session = fake_session() + _failing_then_ok(session, failures=1_000) + release_owner = threading.Event() + # The owner's first retry sleep holds the flight open until the waiters joined. + sleep.on_sleep = lambda: release_owner.wait(timeout=5) + manager = _manager(session) + + def release_once_all_joined(): + join_counting_future.wait_for_joins(3) # 4 callers: 1 owner + 3 waiters + release_owner.set() + + outcomes = _run_concurrently( + 4, manager.get_token, after_start=release_once_all_joined + ) + + assert all(isinstance(outcome, RuntimeError) for outcome in outcomes) + assert len({id(outcome) for outcome in outcomes}) == 1 + assert session.get_credentials.call_count == MAX_ATTEMPTS + + +# --- refresh_after ----------------------------------------------------------------- + + +def test_default_refresh_after_is_five_minutes(): + assert DEFAULT_REFRESH_AFTER_SECONDS == 300 + + +@pytest.mark.parametrize( + "refresh_after", + [ + 0, + -1, + EFFECTIVE_LIFETIME, + EFFECTIVE_LIFETIME + 1, + TOKEN_TTL, + "300", + True, + float("nan"), + ], +) +def test_rejects_invalid_refresh_after(clock, mock_session, refresh_after): + with pytest.raises(InvalidParameterError) as excinfo: + _manager(mock_session, refresh_after=refresh_after) + assert "refresh_after" in str(excinfo.value) + + +@pytest.mark.parametrize("refresh_after", [0.001, 1, EFFECTIVE_LIFETIME - 0.001]) +def test_accepts_refresh_after_within_the_token_lifetime( + clock, scheduler, mock_session, refresh_after +): + manager = _manager(mock_session, refresh_after=refresh_after) + manager.get_token() + + assert [call.delay for call in scheduler.pending] == [refresh_after] + + +# --- background refresh ------------------------------------------------------------ + + +def test_schedules_a_background_refresh_after_the_first_token( + clock, scheduler, mock_session +): + manager = _manager(mock_session) + assert scheduler.pending == [] + + manager.get_token() + + assert [call.delay for call in scheduler.pending] == [DEFAULT_REFRESH_AFTER_SECONDS] + + +def test_background_refresh_installs_rotated_credentials(clock, scheduler): + session = fake_session() + manager = _manager(session) + first = manager.get_token() + session.get_credentials.return_value.get_frozen_credentials.return_value = ( + ROTATED_CREDENTIALS + ) + + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS) + scheduler.fire_pending() + + second = manager.get_token() + assert second != first + assert "AKIAI44QH8DHBEXAMPLE" in second + assert session.get_credentials.call_count == 2 + # The new token gets its own refresh timer. + assert [call.delay for call in scheduler.pending] == [DEFAULT_REFRESH_AFTER_SECONDS] + + +def test_get_token_past_the_refresh_point_does_not_block(clock, scheduler): + # Timer never fired (e.g. the process was suspended): the caller still gets the + # current token immediately, and a refresh is kicked off in the background. + session = fake_session() + manager = _manager(session) + first = manager.get_token() + + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS + 1) + assert manager.get_token() == first + assert session.get_credentials.call_count == 1 + assert [call.delay for call in scheduler.pending] == [0] + + scheduler.fire_pending() + assert session.get_credentials.call_count == 2 + + +def test_only_one_background_refresh_is_started(clock, scheduler, mock_session): + manager = _manager(mock_session) + manager.get_token() + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS + 1) + + manager.get_token() + manager.get_token() + + assert len(scheduler.pending) == 1 + + +def test_missed_timer_path_works_again_after_a_refresh_completes( + clock, scheduler, mock_session +): + manager = _manager(mock_session) + manager.get_token() + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS + 1) + manager.get_token() + scheduler.fire_pending() # first missed-timer refresh completes + assert mock_session.get_credentials.call_count == 2 + + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS + 1) + manager.get_token() + scheduler.fire_pending() + + assert mock_session.get_credentials.call_count == 3 + + +def test_stale_background_refresh_skips_signing_when_a_newer_token_exists( + clock, scheduler, mock_session +): + manager = _manager(mock_session) + manager.get_token() + (stale_timer,) = scheduler.pending + clock.advance(EFFECTIVE_LIFETIME) # expired: the next call signs synchronously ... + manager.get_token() + assert mock_session.get_credentials.call_count == 2 + + stale_timer.callback() # ... so the old timer, firing late, has nothing to do + + assert mock_session.get_credentials.call_count == 2 + + +def test_scheduling_failure_does_not_lose_the_token(clock, scheduler, mock_session): + scheduler.fail = True + manager = _manager(mock_session) + + token = manager.get_token() # signed, installed, timer could not be started + + assert f"{CACHE}/?" in token + assert manager.get_token() == token + # The due time is left in the past so the next call re-attempts scheduling. + scheduler.fail = False + manager.get_token() + assert [call.delay for call in scheduler.pending] == [0] + + +def test_default_scheduler_uses_a_daemon_timer(): + timer = REAL_SCHEDULE(60, lambda: None) + try: + assert isinstance(timer, threading.Timer) + assert timer.daemon + finally: + timer.cancel() + + +# --- retry and failure semantics --------------------------------------------------- + + +def test_retries_transient_failures_with_exponential_backoff(clock, sleep): + session = fake_session() + _failing_then_ok(session, failures=3) + manager = _manager(session) + + token = manager.get_token() + + assert f"{CACHE}/?" in token + assert session.get_credentials.call_count == 4 + assert sleep.delays == [0.1, 0.2, 0.4] + + +def test_backoff_is_capped(clock, sleep): + session = fake_session() + _failing_then_ok(session, failures=MAX_ATTEMPTS - 1) + manager = _manager(session) + + manager.get_token() + + assert session.get_credentials.call_count == MAX_ATTEMPTS + assert sleep.delays == [0.1, 0.2, 0.4, 0.8, 1.6, 3.2, MAX_DELAY] + + +# random() never returns 1.0, so 1.0 here just pins the formula's upper bound (+20%). +@pytest.mark.parametrize(("draw", "expected"), [(0.0, 0.08), (1.0, 0.12)]) +def test_backoff_is_jittered(clock, sleep, monkeypatch, draw, expected): + monkeypatch.setattr(token_manager, "_random", lambda: draw) + session = fake_session() + _failing_then_ok(session, failures=1) + manager = _manager(session) + + manager.get_token() + + assert sleep.delays == [expected] + + +def test_first_token_failure_after_all_attempts_is_the_underlying_error(clock, sleep): + session = fake_session() + _failing_then_ok(session, failures=MAX_ATTEMPTS + 1) + manager = _manager(session) + + with pytest.raises(RuntimeError, match="transient failure 8"): + manager.get_token() + assert session.get_credentials.call_count == MAX_ATTEMPTS + assert len(sleep.delays) == MAX_ATTEMPTS - 1 + + +def test_configuration_errors_are_retried_too(clock, sleep): + # On an instance role a transient IMDS failure surfaces as "no credentials" + # (botocore returns None), so ConfigurationError gets the same retries. + session = fake_session() + calls = {"n": 0} + credentials = session.get_credentials.return_value + + def flaky_chain(): + calls["n"] += 1 + return None if calls["n"] <= 2 else credentials + + session.get_credentials.side_effect = flaky_chain + manager = _manager(session) + + token = manager.get_token() + + assert f"{CACHE}/?" in token + assert sleep.delays == [0.1, 0.2] + + +def test_serves_the_valid_token_while_refresh_keeps_failing(clock, scheduler, sleep): + session = fake_session() + manager = _manager(session) + first = manager.get_token() + _failing_then_ok(session, failures=MAX_ATTEMPTS + 1) + + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS) + scheduler.fire_pending() + + assert manager.get_token() == first + assert len(sleep.delays) == MAX_ATTEMPTS - 1 + # Another cycle is scheduled after the maximum backoff, not at refresh_after. + assert [call.delay for call in scheduler.pending] == [MAX_DELAY] + + +def test_configuration_error_in_the_background_still_reschedules( + clock, scheduler, sleep, mock_session +): + # The cycle must end normally: keep the valid token, reschedule, and never let + # get_token() spawn a refresh per call. + manager = _manager(mock_session) + first = manager.get_token() + mock_session.get_credentials.return_value = None # credential chain went empty + + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS) + scheduler.fire_pending() + + assert len(sleep.delays) == MAX_ATTEMPTS - 1 + assert [call.delay for call in scheduler.pending] == [MAX_DELAY] + for _ in range(3): + assert manager.get_token() == first + assert mock_session.get_credentials.call_count == 1 + MAX_ATTEMPTS + assert [call.delay for call in scheduler.pending] == [MAX_DELAY] + + +def test_get_token_does_not_restart_a_refresh_during_the_backoff(clock, scheduler, sleep): + session = fake_session() + manager = _manager(session) + manager.get_token() + _failing_then_ok(session, failures=1_000) + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS) + scheduler.fire_pending() + attempts = session.get_credentials.call_count + + clock.advance(MAX_DELAY - 1) + manager.get_token() + assert session.get_credentials.call_count == attempts + assert [call.delay for call in scheduler.pending] == [MAX_DELAY] + + clock.advance(1) + manager.get_token() # backoff over and timer not fired: treated as missed + assert [call.delay for call in scheduler.pending] == [0] + + +def test_raises_token_refresh_error_once_the_cached_token_expired( + clock, scheduler, sleep +): + session = fake_session() + manager = _manager(session) + manager.get_token() + _failing_then_ok(session, failures=1_000) + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS) + scheduler.fire_pending() # background cycle fails; token still valid + + clock.advance(EFFECTIVE_LIFETIME) + with pytest.raises(TokenRefreshError) as excinfo: + manager.get_token() + + assert "before the previous token expired" in str(excinfo.value) + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert "transient failure" not in str(excinfo.value) + + +def test_fallback_token_is_rechecked_after_the_lock_is_released( + clock, scheduler, sleep, mock_session, monkeypatch +): + # The cycle decides "serve the still-valid token" under the lock, then logs with the + # lock released. If the token expires in that gap, a waiter that joined because its + # token had expired must not receive it. + manager = _manager(mock_session) + manager.get_token() + _failing_then_ok(mock_session, failures=1_000) + clock.advance(EFFECTIVE_LIFETIME - 1) # 1s of validity left when the cycle ends + monkeypatch.setattr( + token_manager._logger, "warning", lambda *a, **k: clock.advance(1) + ) + + # refresh_token() runs the cycle in the foreground. Without the re-check the + # cycle would hand back the now-expired token and refresh_token() would complain + # about a missing replacement; with it, the expiry itself is reported. + with pytest.raises(TokenRefreshError, match="before the previous token expired"): + manager.refresh_token() + + +def test_expired_token_with_a_configuration_error_is_a_refresh_error( + clock, scheduler, sleep, mock_session +): + manager = _manager(mock_session) + manager.get_token() + mock_session.get_credentials.return_value = None + clock.advance(EFFECTIVE_LIFETIME) + + with pytest.raises(TokenRefreshError) as excinfo: + manager.get_token() + + assert isinstance(excinfo.value.__cause__, ConfigurationError) + + +# --- refresh_token ----------------------------------------------------------------- + + +def test_refresh_token_signs_a_replacement_immediately(clock, scheduler): + session = fake_session() + manager = _manager(session) + first = manager.get_token() + session.get_credentials.return_value.get_frozen_credentials.return_value = ( + ROTATED_CREDENTIALS + ) + + replacement = manager.refresh_token() + + assert replacement != first + assert manager.get_token() == replacement + assert session.get_credentials.call_count == 2 + # The replacement gets a fresh refresh timer. + assert [call.delay for call in scheduler.pending] == [DEFAULT_REFRESH_AFTER_SECONDS] + + +def test_refresh_token_with_no_cached_token_signs_the_first_one(clock, mock_session): + manager = _manager(mock_session) + + token = manager.refresh_token() + + assert manager.get_token() == token + assert mock_session.get_credentials.call_count == 1 + + +def test_refresh_token_raises_when_no_replacement_could_be_signed(clock, sleep): + session = fake_session() + manager = _manager(session) + first = manager.get_token() + _failing_then_ok(session, failures=1_000) + + with pytest.raises(TokenRefreshError, match="replacement"): + manager.refresh_token() + # The still-valid token keeps being served to ordinary callers. + assert manager.get_token() == first + + +def test_refresh_token_after_close_raises(clock, mock_session): + manager = _manager(mock_session) + manager.close() + + with pytest.raises(TokenRefreshError): + manager.refresh_token() + + +# --- on_token_changed -------------------------------------------------------------- + + +def test_timer_firing_during_a_slow_callback_still_reschedules(clock, scheduler): + # The callback runs after the signing flight is closed, so a timer that fires while + # the callback is still running starts its own cycle instead of joining the old + # one, and a replacement refresh is always scheduled. + session = fake_session() + manager = _manager(session, refresh_after=0.01) + calls = [] + + def slow_callback(token): + calls.append(token) + if len(calls) == 1: + clock.advance(0.01) + scheduler.fire_pending() # the timer fires while we are "still in" the callback + + manager._on_token_changed = slow_callback + manager.get_token() + + assert session.get_credentials.call_count == 2 # first token + the timer's refresh + assert len(calls) == 2 + assert [call.delay for call in scheduler.pending] == [0.01] # replacement scheduled + + +def test_on_token_changed_receives_each_installed_token(clock, scheduler): + session = fake_session() + seen = [] + manager = _manager(session, on_token_changed=seen.append) + + first = manager.get_token() + assert seen == [first] + + session.get_credentials.return_value.get_frozen_credentials.return_value = ( + ROTATED_CREDENTIALS + ) + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS) + scheduler.fire_pending() + + assert seen == [first, manager.get_token()] + assert seen[0] != seen[1] + + +def test_on_token_changed_failures_are_ignored(clock, mock_session): + def explode(_token): + raise RuntimeError("consumer bug") + + manager = _manager(mock_session, on_token_changed=explode) + + token = manager.get_token() + + assert manager.get_token() == token + assert mock_session.get_credentials.call_count == 1 + + +def test_on_token_changed_runs_outside_the_lock(clock, mock_session): + # A client may re-authenticate and ask for the credentials again from inside the + # callback, and other threads must not be blocked behind it. + seen = [] + manager = _manager( + mock_session, + on_token_changed=lambda _t: seen.append( + (manager.get_credentials(), manager._lock.locked()) + ), + ) + + token = _run_bounded(manager.get_token) # a deadlock here must fail, not hang + + assert seen == [((USER, token), False)] + + +def test_on_token_changed_failure_is_logged_without_the_token( + clock, mock_session, caplog +): + def explode(_token): + raise RuntimeError("consumer bug") + + manager = _manager(mock_session, on_token_changed=explode) + + with caplog.at_level("WARNING", logger="developer_toolkit_elasticache"): + token = manager.get_token() + + assert [r.levelname for r in caplog.records] == ["WARNING"] + assert "on_token_changed" in caplog.text + assert "RuntimeError" in caplog.text + # The callback received the token, so its message might contain it: never logged. + assert "consumer bug" not in caplog.text + assert token not in caplog.text + + +def test_failed_background_refresh_is_logged_without_the_token( + clock, scheduler, sleep, caplog +): + session = fake_session() + manager = _manager(session) + token = manager.get_token() + _failing_then_ok(session, failures=1_000) + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS) + + with caplog.at_level("WARNING", logger="developer_toolkit_elasticache"): + scheduler.fire_pending() + + assert [r.levelname for r in caplog.records] == ["WARNING"] + assert "Could not refresh" in caplog.text + assert "RuntimeError" in caplog.text + assert "retrying in 5s" in caplog.text + assert token not in caplog.text + + +def test_on_token_changed_is_not_called_for_a_failed_refresh(clock, scheduler, sleep): + session = fake_session() + seen = [] + manager = _manager(session, on_token_changed=seen.append) + manager.get_token() + _failing_then_ok(session, failures=1_000) + + clock.advance(DEFAULT_REFRESH_AFTER_SECONDS) + scheduler.fire_pending() + + assert len(seen) == 1 + + +# --- close ------------------------------------------------------------------------- + + +def test_close_rejects_later_requests(clock, mock_session): + manager = _manager(mock_session) + manager.get_token() + + manager.close() + + with pytest.raises(TokenRefreshError): + manager.get_token() + with pytest.raises(TokenRefreshError): + manager.get_credentials() + + +def test_close_is_idempotent(clock, mock_session): + manager = _manager(mock_session) + manager.get_token() + + manager.close() + manager.close() + + with pytest.raises(TokenRefreshError): + manager.get_token() + + +def test_context_manager_closes(clock, mock_session): + with _manager(mock_session) as manager: + manager.get_token() + + with pytest.raises(TokenRefreshError): + manager.get_token() + + +def test_close_cancels_the_pending_refresh(clock, scheduler, mock_session): + manager = _manager(mock_session) + manager.get_token() + (pending,) = scheduler.pending + + manager.close() + + assert pending.cancelled + pending.callback() # a timer that already fired must be a no-op after close + assert mock_session.get_credentials.call_count == 1 + + +def test_close_during_a_retry_stops_the_cycle(clock, sleep): + session = fake_session() + _failing_then_ok(session, failures=1_000) + manager = _manager(session) + sleep.on_sleep = manager.close + + with pytest.raises(TokenRefreshError) as excinfo: + manager.get_token() + + assert "closed" in str(excinfo.value) + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert session.get_credentials.call_count == 1 + + +def test_close_during_signing_discards_the_token(clock): + session = fake_session() + seen = [] + manager = _manager(session, on_token_changed=seen.append) + credentials = session.get_credentials.return_value + + def close_mid_signing(): + manager.close() + return credentials + + session.get_credentials.side_effect = close_mid_signing + + with pytest.raises(TokenRefreshError): + manager.get_token() + assert seen == [] + assert manager._cached is None + + +def test_real_wait_is_interrupted_by_close(): + stop = threading.Event() + waiting = threading.Thread(target=REAL_WAIT, args=(stop, 60), daemon=True) + started = time.monotonic() + waiting.start() + time.sleep(0.1) + assert waiting.is_alive() # blocked in the real wait + + stop.set() # what close() does + + waiting.join(timeout=5) + assert not waiting.is_alive() + assert time.monotonic() - started < 2 + + +# --- real scheduler ---------------------------------------------------------------- + + +def test_real_timer_refreshes_through_a_slow_callback(monkeypatch): + # End to end with the real threading.Timer and wait: a short refresh_after and a + # callback slower than it must still produce a steady stream of refreshes. + monkeypatch.setattr(token_manager, "_schedule", REAL_SCHEDULE) + monkeypatch.setattr(token_manager, "_wait", REAL_WAIT) + signed = threading.Semaphore(0) + + def slow_callback(_token): + time.sleep(0.05) + signed.release() + + with _manager( + fake_session(), refresh_after=0.02, on_token_changed=slow_callback + ) as m: + m.get_token() + for _ in range(3): + assert signed.acquire(timeout=5), "background refresh did not happen"