diff --git a/sagemaker-core/src/sagemaker/core/token_generator/token_generator.py b/sagemaker-core/src/sagemaker/core/token_generator/token_generator.py index 836e297292..773a6541ae 100644 --- a/sagemaker-core/src/sagemaker/core/token_generator/token_generator.py +++ b/sagemaker-core/src/sagemaker/core/token_generator/token_generator.py @@ -38,6 +38,17 @@ def _generate_token(credentials: Credentials, region: str, expires: int) -> str: params={"Action": "CallWithBearerToken"}, ) + # Take an atomic snapshot of the credentials before signing. botocore's + # SigV4 signer reads access_key, token, and secret_key as three separate, + # non-atomic property accesses. With refreshable credentials, a refresh that + # lands mid-signature would mix key material from two credential versions and + # produce a token that fails verification with InvalidSignature. Freezing + # once (under botocore's refresh lock) guarantees all three reads come from a + # single credential version. Static credentials also expose this method and + # return a coherent snapshot, so this is safe for every credential type. + if hasattr(credentials, "get_frozen_credentials"): + credentials = credentials.get_frozen_credentials() + auth = SigV4QueryAuth(credentials, SERVICE_NAME, region, expires=expires) auth.add_auth(request) diff --git a/sagemaker-core/tests/unit/test_token_generator.py b/sagemaker-core/tests/unit/test_token_generator.py index 74f0fc9976..426bd12749 100644 --- a/sagemaker-core/tests/unit/test_token_generator.py +++ b/sagemaker-core/tests/unit/test_token_generator.py @@ -20,10 +20,24 @@ import pytest from unittest.mock import Mock -from botocore.credentials import Credentials +import threading +import urllib.parse +from datetime import datetime, timezone + +import botocore.auth as botocore_auth +from botocore.auth import SigV4QueryAuth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials, ReadOnlyCredentials, RefreshableCredentials from sagemaker.core.token_generator import generate_token, SageMakerTokenGenerator -from sagemaker.core.token_generator.token_generator import AUTH_PREFIX +from sagemaker.core.token_generator import token_generator as token_generator_module +from sagemaker.core.token_generator.token_generator import ( + AUTH_PREFIX, + DEFAULT_HOST, + DEFAULT_URL, + SERVICE_NAME, + _generate_token, +) class TestSageMakerTokenGenerator: @@ -168,3 +182,205 @@ def test_generate_token_with_custom_expiry_produces_different_token(self): assert token_default != token_custom assert token_default.startswith(AUTH_PREFIX) assert token_custom.startswith(AUTH_PREFIX) + + +# Two internally-consistent credential versions. A generated token is corrupt if +# it mixes fields (access key / secret / session token) across these versions. +_KEYPAIRS = [ + ("AKIDVERSION0000000001", "secretVersion1zzzzzzzzzzzzzzzzzzzzzzzzzz", "sessionTokenV1AAAA"), + ("AKIDVERSION0000000002", "secretVersion2yyyyyyyyyyyyyyyyyyyyyyyyyy", "sessionTokenV2BBBB"), +] +_AKID_TO_VERSION = {kp[0]: i for i, kp in enumerate(_KEYPAIRS)} +_TOKEN_TO_VERSION = {kp[2]: i for i, kp in enumerate(_KEYPAIRS)} +_AKID_TO_SECRET = {kp[0]: kp[1] for kp in _KEYPAIRS} +_REGION = "us-west-2" + + +def _make_rotate_on_read_credentials(): + """Build RefreshableCredentials that flip to the next version on every read. + + The expiry is always kept just inside botocore's advisory refresh window, so + ``refresh_needed()`` is perpetually True and every access to ``access_key``, + ``secret_key``, or ``token`` triggers a refresh. A single unfrozen SigV4 + signing pass therefore reads its three credential fields from three different + versions -- deterministically corrupting the signature unless the signer + snapshots the credentials first. This is the deterministic (100%-reliable) + analogue of a concurrent credential-refresh race, kept fast for CI. + """ + state = {"i": 0} + lock = threading.Lock() + + def _refresh(): + with lock: + state["i"] = (state["i"] + 1) % len(_KEYPAIRS) + access_key, secret_key, session_token = _KEYPAIRS[state["i"]] + # Expiry inside the advisory window keeps refresh_needed() True without + # tripping the still-expired RuntimeError in botocore's refresh path. + expiry = (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat() + return { + "access_key": access_key, + "secret_key": secret_key, + "token": session_token, + "expiry_time": expiry, + } + + access_key, secret_key, session_token = _KEYPAIRS[0] + return RefreshableCredentials( + access_key=access_key, + secret_key=secret_key, + token=session_token, + expiry_time=datetime.now(timezone.utc) + timedelta(seconds=60), + refresh_using=_refresh, + method="test-rotate-on-read", + ) + + +def _decode_query(token: str) -> dict: + """Decode a bearer token back into its presigned-URL query parameters.""" + presigned = base64.b64decode(token[len(AUTH_PREFIX) :]).decode("utf-8") + query = presigned.split("?", 1)[1] + return { + key: value[0] for key, value in urllib.parse.parse_qs(query, keep_blank_values=True).items() + } + + +def _recompute_signature(akid: str, embedded_token: str, amzdate: str, expires: str, region: str): + """Recompute the SigV4 signature for the exact request the token embeds. + + Signs an identical canonical request with a coherent static credential whose + secret is the one that pairs with the embedded access key, pinning the clock + to the token's own X-Amz-Date so the canonical request is byte-identical. If + the token was signed with a secret from a different credential version, the + recomputed signature will not match the embedded one. + """ + static = Credentials(akid, _AKID_TO_SECRET[akid], embedded_token) + fixed_dt = datetime.strptime(amzdate, "%Y%m%dT%H%M%SZ").replace(tzinfo=timezone.utc) + + original_now = botocore_auth.get_current_datetime + botocore_auth.get_current_datetime = lambda: fixed_dt + try: + request = AWSRequest( + method="POST", + url=DEFAULT_URL, + headers={"host": DEFAULT_HOST}, + params={"Action": "CallWithBearerToken"}, + ) + SigV4QueryAuth(static, SERVICE_NAME, region, expires=int(expires)).add_auth(request) + finally: + botocore_auth.get_current_datetime = original_now + + recomputed = { + key: value[0] + for key, value in urllib.parse.parse_qs( + urllib.parse.urlsplit(request.url).query, keep_blank_values=True + ).items() + } + return recomputed["X-Amz-Signature"] + + +def _token_is_corrupt(token: str) -> bool: + """Return True if a generated token mixes credential versions. + + Two independent checks catch the non-atomic-read race: (1) the embedded + access key and session token must belong to the same credential version, and + (2) the embedded signature must match one recomputed with the secret paired + to the embedded access key. A single credential refresh landing mid-signature + trips one or both checks. + """ + query = _decode_query(token) + embedded_akid = query["X-Amz-Credential"].split("/", 1)[0] + embedded_token = query.get("X-Amz-Security-Token") + + akid_version = _AKID_TO_VERSION.get(embedded_akid) + token_version = _TOKEN_TO_VERSION.get(embedded_token) + if akid_version is None or token_version is None or akid_version != token_version: + return True + if embedded_akid not in _AKID_TO_SECRET: + return True + + scope_parts = query["X-Amz-Credential"].split("/") + region = scope_parts[2] if len(scope_parts) > 2 else _REGION + expected_signature = _recompute_signature( + embedded_akid, + embedded_token, + query["X-Amz-Date"], + query["X-Amz-Expires"], + region, + ) + return expected_signature != query["X-Amz-Signature"] + + +class TestGenerateTokenFreezesCredentials: + """Regression tests for the credential-snapshot fix in ``_generate_token``. + + Without the ``get_frozen_credentials()`` snapshot, refreshable credentials + that rotate mid-signature yield tokens that fail server-side verification + with InvalidSignature. These tests reproduce that corruption deterministically + and assert the fix eliminates it. + """ + + def test_static_credentials_are_never_flagged_corrupt(self): + """Sanity check: coherent static credentials must produce coherent tokens.""" + access_key, secret_key, session_token = _KEYPAIRS[0] + token = _generate_token(Credentials(access_key, secret_key, session_token), _REGION, 43200) + assert not _token_is_corrupt(token) + + def test_rotating_credentials_produce_only_coherent_tokens(self): + """Rotate-on-read credentials must still yield internally consistent tokens. + + With rotate-on-read credentials, an unfrozen signer deterministically mixes + credential versions (100% corrupt). The frozen-snapshot fix must bring the + corrupt count to zero. + """ + corrupt = 0 + for _ in range(200): + token = _generate_token(_make_rotate_on_read_credentials(), _REGION, 43200) + if _token_is_corrupt(token): + corrupt += 1 + assert corrupt == 0 + + def test_rotating_credentials_are_coherent_under_concurrency(self): + """A shared rotating credential hammered by many threads yields no corruption.""" + shared = _make_rotate_on_read_credentials() + tokens = [] + tokens_lock = threading.Lock() + + def _worker(): + local = [_generate_token(shared, _REGION, 43200) for _ in range(40)] + with tokens_lock: + tokens.extend(local) + + threads = [threading.Thread(target=_worker) for _ in range(32)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + corrupt = sum(1 for token in tokens if _token_is_corrupt(token)) + assert corrupt == 0 + + def test_signer_receives_single_consistent_credential_triple(self, monkeypatch): + """The SigV4 signer must be constructed with one immutable credential snapshot. + + Captures the credentials object handed to ``SigV4QueryAuth`` and asserts it + is a frozen ``ReadOnlyCredentials`` whose three fields form a single coherent + version -- i.e. the signer never sees the live, rotating credential object. + """ + captured = {} + real_signer = token_generator_module.SigV4QueryAuth + + def _capturing_signer(credentials, *args, **kwargs): + captured["credentials"] = credentials + return real_signer(credentials, *args, **kwargs) + + monkeypatch.setattr(token_generator_module, "SigV4QueryAuth", _capturing_signer) + + _generate_token(_make_rotate_on_read_credentials(), _REGION, 43200) + + signed_with = captured["credentials"] + assert isinstance(signed_with, ReadOnlyCredentials) + # Reading the snapshot repeatedly must be stable, and its access key and + # session token must belong to the same credential version. + assert signed_with.access_key == signed_with.access_key + assert _AKID_TO_VERSION[signed_with.access_key] == _TOKEN_TO_VERSION[signed_with.token] + assert signed_with.secret_key == _AKID_TO_SECRET[signed_with.access_key]