Skip to content

Commit c0bc32f

Browse files
HollowTubetttritin
andauthored
fix(core): freeze credentials before signing bearer token (#6360)
_generate_token passed botocore RefreshableCredentials straight into SigV4QueryAuth.add_auth(), which reads access_key, token, and secret_key as three separate, non-atomic property accesses. A credential refresh landing mid-signature mixes key material from two versions, producing a presigned URL whose signature fails verification with InvalidSignature (intermittent under concurrency; ~58% of tokens in a 96-thread repro). Snapshot the credentials with get_frozen_credentials() before signing so all three reads come from a single credential version. The read is taken under botocore's refresh lock. Static credentials also expose the method and return a coherent snapshot, so the change is behavior-preserving for non-refreshable credentials; a hasattr guard leaves any custom credential object without the method unchanged. Adds regression tests: a deterministic rotate-on-read case (100% corrupt without the fix), a concurrency case, and a signer-snapshot assertion. Co-authored-by: tttritin <tttritin@amazon.com>
1 parent db90567 commit c0bc32f

2 files changed

Lines changed: 229 additions & 2 deletions

File tree

‎sagemaker-core/src/sagemaker/core/token_generator/token_generator.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,17 @@ def _generate_token(credentials: Credentials, region: str, expires: int) -> str:
3838
params={"Action": "CallWithBearerToken"},
3939
)
4040

41+
# Take an atomic snapshot of the credentials before signing. botocore's
42+
# SigV4 signer reads access_key, token, and secret_key as three separate,
43+
# non-atomic property accesses. With refreshable credentials, a refresh that
44+
# lands mid-signature would mix key material from two credential versions and
45+
# produce a token that fails verification with InvalidSignature. Freezing
46+
# once (under botocore's refresh lock) guarantees all three reads come from a
47+
# single credential version. Static credentials also expose this method and
48+
# return a coherent snapshot, so this is safe for every credential type.
49+
if hasattr(credentials, "get_frozen_credentials"):
50+
credentials = credentials.get_frozen_credentials()
51+
4152
auth = SigV4QueryAuth(credentials, SERVICE_NAME, region, expires=expires)
4253
auth.add_auth(request)
4354

‎sagemaker-core/tests/unit/test_token_generator.py‎

Lines changed: 218 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,24 @@
2020
import pytest
2121
from unittest.mock import Mock
2222

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
2431

2532
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+
)
2741

2842

2943
class TestSageMakerTokenGenerator:
@@ -168,3 +182,205 @@ def test_generate_token_with_custom_expiry_produces_different_token(self):
168182
assert token_default != token_custom
169183
assert token_default.startswith(AUTH_PREFIX)
170184
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

Comments
 (0)