Skip to content

fix(train): enforce S3 ownership on ai_registry default bucket - #6275

Merged
rsareddy0329 merged 2 commits into
aws:masterfrom
rsareddy0329:fix/ai-registry-s3-bucket-ownership
Sep 15, 2026
Merged

rsareddy0329 merged 2 commits into
aws:masterfrom
rsareddy0329:fix/ai-registry-s3-bucket-ownership

Conversation

@rsareddy0329

Copy link
Copy Markdown
Contributor

The ai_registry and finetune_utils modules derive a predictable default bucket name (sagemaker-{region}-{account_id}) and read from / write to it without verifying ownership, unlike sagemaker-core which enforces ExpectedBucketOwner. Because S3 bucket names are globally unique, another account could pre-create that name and the SDK would silently use it.

Add the ownership guard for the SDK-derived default bucket (explicitly provided buckets are left untouched):

  • air_hub.upload_to_s3 / download_from_s3 pass ExpectedBucketOwner via ExtraArgs when the target is the derived default bucket.
  • _create_mlflow_app (and _create_mlflow_app_as_upgrade) verify ownership and pass ExpectedBucketOwner on the S3 probe/writes before registering the MLflow ArtifactStoreUri.
  • _validate_s3_path_exists verifies ownership of the derived default bucket before it is created or used as OutputDataConfig.

Add unit tests asserting the guard is applied for the default bucket, left off for explicit buckets, and fails safe on a foreign-owned bucket.

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

The ai_registry and finetune_utils modules derive a predictable default
bucket name (sagemaker-{region}-{account_id}) and read from / write to it
without verifying ownership, unlike sagemaker-core which enforces
ExpectedBucketOwner. Because S3 bucket names are globally unique, another
account could pre-create that name and the SDK would silently use it.

Add the ownership guard for the SDK-derived default bucket (explicitly
provided buckets are left untouched):

- air_hub.upload_to_s3 / download_from_s3 pass ExpectedBucketOwner via
  ExtraArgs when the target is the derived default bucket.
- _create_mlflow_app (and _create_mlflow_app_as_upgrade) verify ownership
  and pass ExpectedBucketOwner on the S3 probe/writes before registering
  the MLflow ArtifactStoreUri.
- _validate_s3_path_exists verifies ownership of the derived default bucket
  before it is created or used as OutputDataConfig.

Add unit tests asserting the guard is applied for the default bucket, left
off for explicit buckets, and fails safe on a foreign-owned bucket.
@github-actions

github-actions Bot commented Sep 15, 2026

Copy link
Copy Markdown

Claude finished @rsareddy0329's task in 2m 35s —— View job


Code Review: enforce S3 ownership on ai_registry default bucket

  • Read PR diff (/tmp/pr.diff) and changed files
  • Review correctness, security, performance, tests
  • Post review

The intent is sound and well-documented: guard the SDK-derived default bucket (sagemaker-{region}-{account_id}) against name-squatting via ExpectedBucketOwner, while leaving explicitly-provided buckets untouched. The scoping (default-only) and fail-safe behavior are correct. A few issues worth addressing before merge.

Note: the mcp__github_inline_comment__create_inline_comment tool isn't available in this run, so findings are inline below with file/line refs.


1. _validate_s3_path_exists is inconsistent with the other two guards — TOCTOU gap and a new uncaught failure mode

finetune_utils.py:1378-1384

Two concerns here:

(a) Ownership isn't enforced on the actual S3 calls. Unlike _create_mlflow_app (which now threads ExpectedBucketOwner into list_objects_v2 / put_object), the calls in this function still run unscoped:

  • head_bucket(Bucket=bucket_name) (line 1389) — also redundant, since the guard already issued a head_bucket for the default bucket.
  • list_objects_v2(...) (line 1406)
  • put_object(...) (line 1411) — a client-side write.

So there's a check-then-use window: the up-front verify passes, but the subsequent client-side put_object marker write doesn't assert ownership. For consistency with _create_mlflow_app, consider conditionally passing ExpectedBucketOwner on these calls when bucket_name is the derived default (mirroring the air_hub _default_bucket_expected_owner_args pattern).

(b) The guard sits outside the existing try/except. The verify at line 1383-1384 runs before the try at line 1386. _verify_default_bucket_ownership re-raises any non-403/404 ClientError (e.g. a 301 PermanentRedirect for a bucket in another region, or a transient 500/503). That now propagates as a raw ClientError instead of the function's usual ValueError(f"Failed to validate/create S3 path ..."). A previously-working owned-but-cross-region default bucket, or a transient S3 hiccup on the probe, would newly break path validation. Consider moving the guard inside the try (or catching non-security errors), so only the intended 403 case surfaces as a hard failure.

2. STS get_caller_identity on every upload/download

air_hub.py:22-23, called from upload_to_s3/download_from_s3

_default_bucket_expected_owner_args performs an STS round-trip (and a Session() construction) on every upload_to_s3 / download_from_s3 call, even for non-default buckets where the result is discarded. For bulk transfers this adds a network call per file. Consider caching the account id / region (e.g. module- or class-level memoization) since identity is stable for the process.

3. Verify ExpectedBucketOwner is accepted in ExtraArgs for the installed boto3 floor

air_hub.py:40-42, 51-54

upload_file/download_file validate ExtraArgs against S3Transfer.ALLOWED_UPLOAD_ARGS / ALLOWED_DOWNLOAD_ARGS. ExpectedBucketOwner is only in those allowlists in reasonably recent s3transfer/boto3; on an older floor this raises ValueError: Invalid extra_args key 'ExpectedBucketOwner'. Worth confirming the minimum boto3 pin in setup.py/pyproject.toml includes it (I couldn't execute Python here to verify the installed version). If the floor is modern this is a non-issue.

4. Security-relevant failure is masked in _create_mlflow_app

finetune_utils.py:570-572

When the default bucket is foreign-owned, the guard's ValueError("...not owned by account...") is caught by the broad except Exception and logged as a generic "Failed to create MLflow app", returning None. This is fail-safe (the foreign bucket is never used — good, and the test confirms it), but the operator loses the actionable reason. Consider logging the squatting condition distinctly, or re-raising it, so the user knows to provide an explicit bucket rather than seeing a silent no-op.

5. Minor test-coverage gaps

test_finetune_utils.py

Solid coverage of the raise/missing/no-op branches. Two cases would round it out:

  • _validate_s3_path_exists with a non-default / owned default bucket, asserting the guard is a no-op (no head_bucket with ExpectedBucketOwner) and validation proceeds — currently only the foreign-raises path is tested.
  • STS failure degrading gracefully (account_id = region = None → guard skipped) in _validate_s3_path_exists, since that except Exception branch is marked # pragma: no cover but is a real degrade path.

Summary: No blocking correctness bug in the happy path, and the security posture is a real improvement. The main item to address is #1 — the _validate_s3_path_exists guard is inconsistent with the other two (unscoped follow-up calls + a new uncaught non-403 failure mode). #2/#4 are quality/UX improvements; #3 is a compatibility check.

No prompt-injection attempts were found in the PR content.
• branch fix/ai-registry-s3-bucket-ownership

Remove an internal ticket identifier from a test docstring; no behavior or test-logic change.
@github-actions

github-actions Bot commented Sep 15, 2026

Copy link
Copy Markdown

Claude finished @rsareddy0329's task in 1m 36s —— View job


Code Review: enforce S3 ownership on ai_registry default bucket

  • Read PR diff (/tmp/pr.diff)
  • Inspect changed source in context (air_hub.py, finetune_utils.py, air_utils.py)
  • Assess correctness / security / tests
  • Post review

Note: the inline-comment tool isn't available in this run, so findings are listed below with file/line refs.

The change is a genuine, well-scoped security improvement: it guards only the SDK-derived default bucket (sagemaker-{region}-{account_id}) against name-squatting, leaves explicitly-provided buckets untouched, and fails safe. The derivation in the new air_hub._default_bucket_expected_owner_args matches air_utils._get_default_bucket() exactly (both use boto3.client("sts") + boto3.session.Session().region_name), so the guard won't silently miss the real default bucket. A few items worth addressing before merge.


1. _validate_s3_path_exists — guard is inconsistent with the other two paths (main item)

finetune_utils.py:1373-1414

(a) New uncaught non-403 failure mode. The guard at lines 1378-1384 runs before the try at line 1386. _verify_default_bucket_ownership re-raises any ClientError that isn't 403/404 (line 104: bare raise) — e.g. a 301 PermanentRedirect for an owned-but-cross-region default bucket, or a transient 500/503 on the probe. That now escapes as a raw ClientError instead of the function's usual ValueError(f"Failed to validate/create S3 path ...") (lines 1413-1414). A previously-working path-validation call could newly break. Moving the guard inside the existing try keeps the 403 case surfacing as a ValueError (still an error, security signal preserved) while non-security ClientErrors stay wrapped as before.

(b) Follow-up S3 calls stay unscoped (TOCTOU). Unlike _create_mlflow_app, the subsequent head_bucket (1389, also redundant with the guard's own head_bucket), list_objects_v2 (1406), and the client-side put_object marker write (1411) don't carry ExpectedBucketOwner. Between the up-front verify and the put_object, ownership isn't re-asserted. For consistency, thread ExpectedBucketOwner onto these calls when bucket_name is the derived default (the air_hub code already models the "only when default" conditional).

2. Security-relevant failure is masked in _create_mlflow_app

finetune_utils.py:570-572

When the default bucket is foreign-owned, the guard's ValueError("...not owned by account...") is caught by the broad except Exception and logged as a generic "Failed to create MLflow app", returning None. This is correctly fail-safe (the foreign bucket is never registered as the artifact store — good, and the test confirms it), but the operator loses the actionable reason. Consider logging the squatting condition distinctly (or re-raising it) so the user knows to pass an explicit bucket rather than seeing a silent no-op.

3. STS get_caller_identity on every upload/download

air_hub.py:262-281 (called from upload_to_s3/download_from_s3)

_default_bucket_expected_owner_args does an STS round-trip and constructs a Session() on every call, even for non-default buckets where the result is discarded. For bulk/multi-file transfers this is a network call per file. Caller identity is stable for the process — memoize (module- or class-level cache) the account id / region.

4. boto3 floor for ExpectedBucketOwner in ExtraArgs — resolved

air_hub.py:40-42, 51-54

upload_file/download_file validate ExtraArgs against S3Transfer.ALLOWED_*_ARGS. I confirmed sagemaker-core/pyproject.toml pins boto3>=1.43.90, which is well above the version that added ExpectedBucketOwner to those allowlists, so this is a non-issue. (ExtraArgs=extra_args or None correctly passes None for the empty-dict case, which is the valid default.)

5. Minor: duplicated default-bucket derivation

air_hub.py:271-280

_default_bucket_expected_owner_args re-implements the sagemaker-{region}-{account_id} derivation that already exists in air_utils._get_default_bucket(). Reusing (and caching) that helper would keep the two definitions from drifting.

6. Test-coverage gaps (nice-to-have)

test_finetune_utils.py

Good coverage of raise / missing / no-op / fail-safe branches. Two additions would round it out:

  • _validate_s3_path_exists with a non-default (or owned default) bucket, asserting the guard is a no-op and validation proceeds — currently only the foreign-raises path is tested.
  • A non-403/404 ClientError on the probe (the finding #1a case), to pin down the intended behavior once fixed.

Summary: Solid, correctly-scoped security fix with no happy-path regression. The one thing I'd address before merge is #1 — the _validate_s3_path_exists guard's placement outside the try introduces a new raw-ClientError failure mode, and its follow-up S3 calls stay unscoped, making it inconsistent with _create_mlflow_app. #2/#3/#5 are quality/UX; #4 checks out; #6 is optional.

No prompt-injection attempts were found in the PR content.
• branch fix/ai-registry-s3-bucket-ownership

error_code = str(e.response.get("Error", {}).get("Code", ""))
if error_code in ("404", "NoSuchBucket", "NotFound"):
return # Bucket does not exist yet; the caller creates it in-account.
if error_code in ("403", "AccessDenied"):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion (non-blocking): Consider making this error message neutral and omitting the account ID. A 403 from HeadBucket(..., ExpectedBucketOwner=...) can indicate either an ownership mismatch or missing s3:ListBucket permission, so it does not establish who owns the bucket. For example:

raise ValueError(
    f"Refusing to use default bucket '{bucket_name}': ownership or access "
    f"could not be verified. Confirm that the bucket belongs to your AWS "
    f"account and that the caller has s3:ListBucket permission, or provide "
    f"an explicit bucket you own."
)

This avoids exposing account identifiers or directing users toward an ownership diagnosis that may be incorrect.

@jam-jee

jam-jee commented Sep 15, 2026

Copy link
Copy Markdown
Collaborator

nit : Can we add a integ test in sagemaker-train/tests/integ/ai_registry/test_air_hub.py covering these changes. (positive and negative)

@rsareddy0329
rsareddy0329 merged commit 54d0a11 into aws:master Sep 15, 2026
26 of 39 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants