|
20 | 20 | import pytest |
21 | 21 | from unittest.mock import Mock |
22 | 22 |
|
23 | | -from botocore.credentials import Credentials |
| 23 | +import threading |
| 24 | +import urllib.parse |
| 25 | +from datetime import datetime, timezone |
| 26 | + |
| 27 | +import botocore.auth as botocore_auth |
| 28 | +from botocore.auth import SigV4QueryAuth |
| 29 | +from botocore.awsrequest import AWSRequest |
| 30 | +from botocore.credentials import Credentials, ReadOnlyCredentials, RefreshableCredentials |
24 | 31 |
|
25 | 32 | from sagemaker.core.token_generator import generate_token, SageMakerTokenGenerator |
26 | | -from sagemaker.core.token_generator.token_generator import AUTH_PREFIX |
| 33 | +from sagemaker.core.token_generator import token_generator as token_generator_module |
| 34 | +from sagemaker.core.token_generator.token_generator import ( |
| 35 | + AUTH_PREFIX, |
| 36 | + DEFAULT_HOST, |
| 37 | + DEFAULT_URL, |
| 38 | + SERVICE_NAME, |
| 39 | + _generate_token, |
| 40 | +) |
27 | 41 |
|
28 | 42 |
|
29 | 43 | class TestSageMakerTokenGenerator: |
@@ -168,3 +182,205 @@ def test_generate_token_with_custom_expiry_produces_different_token(self): |
168 | 182 | assert token_default != token_custom |
169 | 183 | assert token_default.startswith(AUTH_PREFIX) |
170 | 184 | assert token_custom.startswith(AUTH_PREFIX) |
| 185 | + |
| 186 | + |
| 187 | +# Two internally-consistent credential versions. A generated token is corrupt if |
| 188 | +# it mixes fields (access key / secret / session token) across these versions. |
| 189 | +_KEYPAIRS = [ |
| 190 | + ("AKIDVERSION0000000001", "secretVersion1zzzzzzzzzzzzzzzzzzzzzzzzzz", "sessionTokenV1AAAA"), |
| 191 | + ("AKIDVERSION0000000002", "secretVersion2yyyyyyyyyyyyyyyyyyyyyyyyyy", "sessionTokenV2BBBB"), |
| 192 | +] |
| 193 | +_AKID_TO_VERSION = {kp[0]: i for i, kp in enumerate(_KEYPAIRS)} |
| 194 | +_TOKEN_TO_VERSION = {kp[2]: i for i, kp in enumerate(_KEYPAIRS)} |
| 195 | +_AKID_TO_SECRET = {kp[0]: kp[1] for kp in _KEYPAIRS} |
| 196 | +_REGION = "us-west-2" |
| 197 | + |
| 198 | + |
| 199 | +def _make_rotate_on_read_credentials(): |
| 200 | + """Build RefreshableCredentials that flip to the next version on every read. |
| 201 | +
|
| 202 | + The expiry is always kept just inside botocore's advisory refresh window, so |
| 203 | + ``refresh_needed()`` is perpetually True and every access to ``access_key``, |
| 204 | + ``secret_key``, or ``token`` triggers a refresh. A single unfrozen SigV4 |
| 205 | + signing pass therefore reads its three credential fields from three different |
| 206 | + versions -- deterministically corrupting the signature unless the signer |
| 207 | + snapshots the credentials first. This is the deterministic (100%-reliable) |
| 208 | + analogue of a concurrent credential-refresh race, kept fast for CI. |
| 209 | + """ |
| 210 | + state = {"i": 0} |
| 211 | + lock = threading.Lock() |
| 212 | + |
| 213 | + def _refresh(): |
| 214 | + with lock: |
| 215 | + state["i"] = (state["i"] + 1) % len(_KEYPAIRS) |
| 216 | + access_key, secret_key, session_token = _KEYPAIRS[state["i"]] |
| 217 | + # Expiry inside the advisory window keeps refresh_needed() True without |
| 218 | + # tripping the still-expired RuntimeError in botocore's refresh path. |
| 219 | + expiry = (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat() |
| 220 | + return { |
| 221 | + "access_key": access_key, |
| 222 | + "secret_key": secret_key, |
| 223 | + "token": session_token, |
| 224 | + "expiry_time": expiry, |
| 225 | + } |
| 226 | + |
| 227 | + access_key, secret_key, session_token = _KEYPAIRS[0] |
| 228 | + return RefreshableCredentials( |
| 229 | + access_key=access_key, |
| 230 | + secret_key=secret_key, |
| 231 | + token=session_token, |
| 232 | + expiry_time=datetime.now(timezone.utc) + timedelta(seconds=60), |
| 233 | + refresh_using=_refresh, |
| 234 | + method="test-rotate-on-read", |
| 235 | + ) |
| 236 | + |
| 237 | + |
| 238 | +def _decode_query(token: str) -> dict: |
| 239 | + """Decode a bearer token back into its presigned-URL query parameters.""" |
| 240 | + presigned = base64.b64decode(token[len(AUTH_PREFIX) :]).decode("utf-8") |
| 241 | + query = presigned.split("?", 1)[1] |
| 242 | + return { |
| 243 | + key: value[0] for key, value in urllib.parse.parse_qs(query, keep_blank_values=True).items() |
| 244 | + } |
| 245 | + |
| 246 | + |
| 247 | +def _recompute_signature(akid: str, embedded_token: str, amzdate: str, expires: str, region: str): |
| 248 | + """Recompute the SigV4 signature for the exact request the token embeds. |
| 249 | +
|
| 250 | + Signs an identical canonical request with a coherent static credential whose |
| 251 | + secret is the one that pairs with the embedded access key, pinning the clock |
| 252 | + to the token's own X-Amz-Date so the canonical request is byte-identical. If |
| 253 | + the token was signed with a secret from a different credential version, the |
| 254 | + recomputed signature will not match the embedded one. |
| 255 | + """ |
| 256 | + static = Credentials(akid, _AKID_TO_SECRET[akid], embedded_token) |
| 257 | + fixed_dt = datetime.strptime(amzdate, "%Y%m%dT%H%M%SZ").replace(tzinfo=timezone.utc) |
| 258 | + |
| 259 | + original_now = botocore_auth.get_current_datetime |
| 260 | + botocore_auth.get_current_datetime = lambda: fixed_dt |
| 261 | + try: |
| 262 | + request = AWSRequest( |
| 263 | + method="POST", |
| 264 | + url=DEFAULT_URL, |
| 265 | + headers={"host": DEFAULT_HOST}, |
| 266 | + params={"Action": "CallWithBearerToken"}, |
| 267 | + ) |
| 268 | + SigV4QueryAuth(static, SERVICE_NAME, region, expires=int(expires)).add_auth(request) |
| 269 | + finally: |
| 270 | + botocore_auth.get_current_datetime = original_now |
| 271 | + |
| 272 | + recomputed = { |
| 273 | + key: value[0] |
| 274 | + for key, value in urllib.parse.parse_qs( |
| 275 | + urllib.parse.urlsplit(request.url).query, keep_blank_values=True |
| 276 | + ).items() |
| 277 | + } |
| 278 | + return recomputed["X-Amz-Signature"] |
| 279 | + |
| 280 | + |
| 281 | +def _token_is_corrupt(token: str) -> bool: |
| 282 | + """Return True if a generated token mixes credential versions. |
| 283 | +
|
| 284 | + Two independent checks catch the non-atomic-read race: (1) the embedded |
| 285 | + access key and session token must belong to the same credential version, and |
| 286 | + (2) the embedded signature must match one recomputed with the secret paired |
| 287 | + to the embedded access key. A single credential refresh landing mid-signature |
| 288 | + trips one or both checks. |
| 289 | + """ |
| 290 | + query = _decode_query(token) |
| 291 | + embedded_akid = query["X-Amz-Credential"].split("/", 1)[0] |
| 292 | + embedded_token = query.get("X-Amz-Security-Token") |
| 293 | + |
| 294 | + akid_version = _AKID_TO_VERSION.get(embedded_akid) |
| 295 | + token_version = _TOKEN_TO_VERSION.get(embedded_token) |
| 296 | + if akid_version is None or token_version is None or akid_version != token_version: |
| 297 | + return True |
| 298 | + if embedded_akid not in _AKID_TO_SECRET: |
| 299 | + return True |
| 300 | + |
| 301 | + scope_parts = query["X-Amz-Credential"].split("/") |
| 302 | + region = scope_parts[2] if len(scope_parts) > 2 else _REGION |
| 303 | + expected_signature = _recompute_signature( |
| 304 | + embedded_akid, |
| 305 | + embedded_token, |
| 306 | + query["X-Amz-Date"], |
| 307 | + query["X-Amz-Expires"], |
| 308 | + region, |
| 309 | + ) |
| 310 | + return expected_signature != query["X-Amz-Signature"] |
| 311 | + |
| 312 | + |
| 313 | +class TestGenerateTokenFreezesCredentials: |
| 314 | + """Regression tests for the credential-snapshot fix in ``_generate_token``. |
| 315 | +
|
| 316 | + Without the ``get_frozen_credentials()`` snapshot, refreshable credentials |
| 317 | + that rotate mid-signature yield tokens that fail server-side verification |
| 318 | + with InvalidSignature. These tests reproduce that corruption deterministically |
| 319 | + and assert the fix eliminates it. |
| 320 | + """ |
| 321 | + |
| 322 | + def test_static_credentials_are_never_flagged_corrupt(self): |
| 323 | + """Sanity check: coherent static credentials must produce coherent tokens.""" |
| 324 | + access_key, secret_key, session_token = _KEYPAIRS[0] |
| 325 | + token = _generate_token(Credentials(access_key, secret_key, session_token), _REGION, 43200) |
| 326 | + assert not _token_is_corrupt(token) |
| 327 | + |
| 328 | + def test_rotating_credentials_produce_only_coherent_tokens(self): |
| 329 | + """Rotate-on-read credentials must still yield internally consistent tokens. |
| 330 | +
|
| 331 | + With rotate-on-read credentials, an unfrozen signer deterministically mixes |
| 332 | + credential versions (100% corrupt). The frozen-snapshot fix must bring the |
| 333 | + corrupt count to zero. |
| 334 | + """ |
| 335 | + corrupt = 0 |
| 336 | + for _ in range(200): |
| 337 | + token = _generate_token(_make_rotate_on_read_credentials(), _REGION, 43200) |
| 338 | + if _token_is_corrupt(token): |
| 339 | + corrupt += 1 |
| 340 | + assert corrupt == 0 |
| 341 | + |
| 342 | + def test_rotating_credentials_are_coherent_under_concurrency(self): |
| 343 | + """A shared rotating credential hammered by many threads yields no corruption.""" |
| 344 | + shared = _make_rotate_on_read_credentials() |
| 345 | + tokens = [] |
| 346 | + tokens_lock = threading.Lock() |
| 347 | + |
| 348 | + def _worker(): |
| 349 | + local = [_generate_token(shared, _REGION, 43200) for _ in range(40)] |
| 350 | + with tokens_lock: |
| 351 | + tokens.extend(local) |
| 352 | + |
| 353 | + threads = [threading.Thread(target=_worker) for _ in range(32)] |
| 354 | + for thread in threads: |
| 355 | + thread.start() |
| 356 | + for thread in threads: |
| 357 | + thread.join() |
| 358 | + |
| 359 | + corrupt = sum(1 for token in tokens if _token_is_corrupt(token)) |
| 360 | + assert corrupt == 0 |
| 361 | + |
| 362 | + def test_signer_receives_single_consistent_credential_triple(self, monkeypatch): |
| 363 | + """The SigV4 signer must be constructed with one immutable credential snapshot. |
| 364 | +
|
| 365 | + Captures the credentials object handed to ``SigV4QueryAuth`` and asserts it |
| 366 | + is a frozen ``ReadOnlyCredentials`` whose three fields form a single coherent |
| 367 | + version -- i.e. the signer never sees the live, rotating credential object. |
| 368 | + """ |
| 369 | + captured = {} |
| 370 | + real_signer = token_generator_module.SigV4QueryAuth |
| 371 | + |
| 372 | + def _capturing_signer(credentials, *args, **kwargs): |
| 373 | + captured["credentials"] = credentials |
| 374 | + return real_signer(credentials, *args, **kwargs) |
| 375 | + |
| 376 | + monkeypatch.setattr(token_generator_module, "SigV4QueryAuth", _capturing_signer) |
| 377 | + |
| 378 | + _generate_token(_make_rotate_on_read_credentials(), _REGION, 43200) |
| 379 | + |
| 380 | + signed_with = captured["credentials"] |
| 381 | + assert isinstance(signed_with, ReadOnlyCredentials) |
| 382 | + # Reading the snapshot repeatedly must be stable, and its access key and |
| 383 | + # session token must belong to the same credential version. |
| 384 | + assert signed_with.access_key == signed_with.access_key |
| 385 | + assert _AKID_TO_VERSION[signed_with.access_key] == _TOKEN_TO_VERSION[signed_with.token] |
| 386 | + assert signed_with.secret_key == _AKID_TO_SECRET[signed_with.access_key] |
0 commit comments