From fbd4418b5def9ff29a2b370c095d344c4cb412e0 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Mon, 14 Sep 2026 17:22:25 -0700 Subject: [PATCH 1/2] fix(train): enforce S3 ownership on ai_registry default bucket 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. --- .../src/sagemaker/ai_registry/air_hub.py | 32 ++++- .../train/common_utils/finetune_utils.py | 84 +++++++++++- .../tests/unit/ai_registry/test_air_hub.py | 62 +++++++-- .../train/common_utils/test_finetune_utils.py | 124 ++++++++++++++++++ 4 files changed, 288 insertions(+), 14 deletions(-) diff --git a/sagemaker-train/src/sagemaker/ai_registry/air_hub.py b/sagemaker-train/src/sagemaker/ai_registry/air_hub.py index f8fad16a77..6673e62d30 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/air_hub.py +++ b/sagemaker-train/src/sagemaker/ai_registry/air_hub.py @@ -259,6 +259,28 @@ def delete_hub_content(cls, hub_content_type: str, hub_content_name: str, hub_co } return client.delete_hub_content(**request) + @staticmethod + def _default_bucket_expected_owner_args(bucket: str) -> dict: + """Return ExtraArgs enforcing bucket ownership for the SDK-derived default bucket. + + The AI Registry derives a predictable default bucket name + ``sagemaker-{region}-{account_id}``. Because S3 bucket names are globally + unique, another account could pre-create that name and grant the caller + access; without a check the SDK would silently read from / write to that + foreign-owned bucket. Passing ``ExpectedBucketOwner`` makes S3 reject the + call with 403 instead. The guard is applied ONLY when ``bucket`` matches the + derived default, so explicitly-provided (possibly cross-account) buckets are + left untouched. + """ + try: + account_id = boto3.client("sts").get_caller_identity()["Account"] + region = boto3.session.Session().region_name + except Exception: # pragma: no cover - identity resolution is best-effort + return {} + if bucket == f"sagemaker-{region}-{account_id}": + return {"ExpectedBucketOwner": account_id} + return {} + @staticmethod @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="AIRHub.upload_to_s3") def upload_to_s3(bucket: str, prefix: str, local_file_path: str) -> str: @@ -272,7 +294,10 @@ def upload_to_s3(bucket: str, prefix: str, local_file_path: str) -> str: Returns: S3 URI of uploaded file """ - AIRHub._s3_client.upload_file(local_file_path, bucket, prefix) + extra_args = AIRHub._default_bucket_expected_owner_args(bucket) + AIRHub._s3_client.upload_file( + local_file_path, bucket, prefix, ExtraArgs=extra_args or None + ) return f"s3://{bucket}/{prefix}" @staticmethod @@ -287,4 +312,7 @@ def download_from_s3(s3_uri: str, local_path: str) -> None: parsed = urlparse(s3_uri) bucket = parsed.netloc key = parsed.path.lstrip("/") - AIRHub._s3_client.download_file(bucket, key, local_path) + extra_args = AIRHub._default_bucket_expected_owner_args(bucket) + AIRHub._s3_client.download_file( + bucket, key, local_path, ExtraArgs=extra_args or None + ) diff --git a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py b/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py index f3a8ddf950..a3580dfbbc 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py @@ -7,6 +7,7 @@ import json from typing import Any, Dict, Optional, Union import boto3 +from botocore.exceptions import ClientError from sagemaker.core.resources import ModelPackage, ModelPackageGroup from sagemaker.core.helper.session_helper import Session from sagemaker.core.s3.utils import resolve_s3_uri_placeholders @@ -422,6 +423,42 @@ def _wait_for_mlflow_app_ready_boto(sm_client, arn: str, timeout: int = 600) -> return None +def _default_bucket_name(region: str, account_id: str) -> str: + """Return the SDK-derived default bucket name ``sagemaker-{region}-{account_id}``.""" + return f"sagemaker-{region}-{account_id}" + + +def _verify_default_bucket_ownership(s3_client, bucket_name: str, account_id: str, region: str) -> None: + """Refuse to use the SDK-derived default bucket if another account owns it. + + The default bucket name ``sagemaker-{region}-{account_id}`` is predictable, so a + third party could pre-create it in a region the account has not used yet and grant + the account access; the SDK would then silently read/write a foreign-owned bucket + (and, via the MLflow artifact store, expose a pickle-deserialization vector). + + This check is applied ONLY when ``bucket_name`` matches the derived default, so + explicitly-provided (possibly cross-account) buckets are left untouched. A missing + bucket is allowed because the caller creates it in-account. A 403/AccessDenied from + the ownership-scoped ``head_bucket`` means the bucket exists under a different owner + and MUST NOT be used, so a clear error is raised instead of proceeding. + """ + if bucket_name != _default_bucket_name(region, account_id): + return + try: + s3_client.head_bucket(Bucket=bucket_name, ExpectedBucketOwner=account_id) + except ClientError as e: + 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"): + raise ValueError( + f"Refusing to use default bucket '{bucket_name}': it exists but is not " + f"owned by account {account_id}. Another account may have pre-created " + f"this predictable bucket name. Provide an explicit bucket you own." + ) from e + raise + + def _create_mlflow_app_as_upgrade( sagemaker_session, old_app: dict, domain_id: Optional[str] ) -> Optional[str]: @@ -442,6 +479,16 @@ def _create_mlflow_app_as_upgrade( artifact_store_uri = old_app.get("ArtifactStoreUri") or \ f"s3://sagemaker-{region}-{account_id}/mlflow-artifacts" + # If we fell back to the predictable default bucket, refuse it when another + # account owns it before registering it as the MLflow ArtifactStoreUri. + default_bucket = _default_bucket_name(region, account_id) + if artifact_store_uri == f"s3://{default_bucket}/mlflow-artifacts": + _verify_default_bucket_ownership( + sagemaker_session.boto_session.client("s3"), + default_bucket, + account_id, + region, + ) role_arn = old_app.get("RoleArn") or \ TrainDefaults.get_role(role=None, sagemaker_session=sagemaker_session) old_name = old_app.get("Name", "mlflow-app") @@ -481,10 +528,22 @@ def _create_mlflow_app(sagemaker_session) -> Optional[str]: s3_client = sagemaker_session.boto_session.client('s3') bucket_name = f"sagemaker-{region}-{account_id}" + # Refuse the predictable default bucket if it exists under another owner. + _verify_default_bucket_ownership(s3_client, bucket_name, account_id, region) + try: - response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix="mlflow-artifacts/", MaxKeys=1) + response = s3_client.list_objects_v2( + Bucket=bucket_name, + Prefix="mlflow-artifacts/", + MaxKeys=1, + ExpectedBucketOwner=account_id, + ) if 'Contents' not in response: - s3_client.put_object(Bucket=bucket_name, Key="mlflow-artifacts/") + s3_client.put_object( + Bucket=bucket_name, + Key="mlflow-artifacts/", + ExpectedBucketOwner=account_id, + ) except s3_client.exceptions.NoSuchBucket: if region == 'us-east-1': s3_client.create_bucket(Bucket=bucket_name) @@ -493,7 +552,11 @@ def _create_mlflow_app(sagemaker_session) -> Optional[str]: Bucket=bucket_name, CreateBucketConfiguration={'LocationConstraint': region} ) - s3_client.put_object(Bucket=bucket_name, Key="mlflow-artifacts/") + s3_client.put_object( + Bucket=bucket_name, + Key="mlflow-artifacts/", + ExpectedBucketOwner=account_id, + ) resp = sm_client.create_mlflow_app( Name=app_name, @@ -1306,7 +1369,20 @@ def _validate_s3_path_exists(s3_path: str, sagemaker_session): prefix = s3_parts[1] if len(s3_parts) > 1 else "" s3_client = sagemaker_session.boto_session.client('s3') - + + # Refuse the predictable default bucket if another account owns it, before we + # create it or pass it to the training job as OutputDataConfig. The training + # write itself is performed service-side under the execution role, so + # ExpectedBucketOwner cannot be attached to it; verifying ownership of the + # derived bucket up front is the applicable guard. + try: + account_id = sagemaker_session.boto_session.client('sts').get_caller_identity()['Account'] + region = sagemaker_session.boto_session.region_name + except Exception: # pragma: no cover - identity resolution is best-effort + account_id = region = None + if account_id and region: + _verify_default_bucket_ownership(s3_client, bucket_name, account_id, region) + try: # Check if bucket exists, create if it doesn't try: diff --git a/sagemaker-train/tests/unit/ai_registry/test_air_hub.py b/sagemaker-train/tests/unit/ai_registry/test_air_hub.py index 6c405ab65b..b90706c8c1 100644 --- a/sagemaker-train/tests/unit/ai_registry/test_air_hub.py +++ b/sagemaker-train/tests/unit/ai_registry/test_air_hub.py @@ -143,24 +143,70 @@ def test_delete_hub_content(self, mock_boto3): def test_upload_to_s3(self, mock_boto3): mock_s3_client = MagicMock() mock_boto3.client.return_value = mock_s3_client - + AIRHub._s3_client = mock_s3_client - + result = AIRHub.upload_to_s3("test-bucket", "test/key", "/local/path") - + assert result == "s3://test-bucket/test/key" - mock_s3_client.upload_file.assert_called_once_with("/local/path", "test-bucket", "test/key") + # Non-default bucket: no ExpectedBucketOwner enforced (explicit buckets untouched). + mock_s3_client.upload_file.assert_called_once_with( + "/local/path", "test-bucket", "test/key", ExtraArgs=None + ) @patch('sagemaker.ai_registry.air_hub.boto3') def test_download_from_s3(self, mock_boto3): mock_s3_client = MagicMock() mock_boto3.client.return_value = mock_s3_client - + AIRHub._s3_client = mock_s3_client - + AIRHub.download_from_s3("s3://test-bucket/test/key", "/local/path") - - mock_s3_client.download_file.assert_called_once_with("test-bucket", "test/key", "/local/path") + + # Non-default bucket: no ExpectedBucketOwner enforced. + mock_s3_client.download_file.assert_called_once_with( + "test-bucket", "test/key", "/local/path", ExtraArgs=None + ) + + @patch('sagemaker.ai_registry.air_hub.boto3') + def test_upload_to_s3_default_bucket_enforces_owner(self, mock_boto3): + """Uploading to the SDK-derived default bucket passes ExpectedBucketOwner.""" + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "111122223333"} + mock_boto3.client.return_value = mock_sts + mock_boto3.session.Session.return_value.region_name = "us-west-2" + mock_s3_client = MagicMock() + AIRHub._s3_client = mock_s3_client + + default_bucket = "sagemaker-us-west-2-111122223333" + AIRHub.upload_to_s3(default_bucket, "test/key", "/local/path") + + mock_s3_client.upload_file.assert_called_once_with( + "/local/path", + default_bucket, + "test/key", + ExtraArgs={"ExpectedBucketOwner": "111122223333"}, + ) + + @patch('sagemaker.ai_registry.air_hub.boto3') + def test_download_from_s3_default_bucket_enforces_owner(self, mock_boto3): + """Downloading from the SDK-derived default bucket passes ExpectedBucketOwner.""" + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "111122223333"} + mock_boto3.client.return_value = mock_sts + mock_boto3.session.Session.return_value.region_name = "us-west-2" + mock_s3_client = MagicMock() + AIRHub._s3_client = mock_s3_client + + default_bucket = "sagemaker-us-west-2-111122223333" + AIRHub.download_from_s3(f"s3://{default_bucket}/test/key", "/local/path") + + mock_s3_client.download_file.assert_called_once_with( + default_bucket, + "test/key", + "/local/path", + ExtraArgs={"ExpectedBucketOwner": "111122223333"}, + ) def test_generate_hub_names_no_padding(self): """Test that generated hub names don't contain = padding characters.""" diff --git a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py index 32463e3f58..5eb2b38d9c 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py @@ -1784,3 +1784,127 @@ def test_list_hyperparameters_accepts_enum_values(self, mock_boto_client, mock_g ) assert result.learning_rate == 0.0001 + + +class TestDefaultBucketOwnershipGuard: + """Bucket-squatting guard for the SDK-derived default bucket (ticket D524303300).""" + + def test_verify_ownership_foreign_bucket_raises(self): + from botocore.exceptions import ClientError + from sagemaker.train.common_utils.finetune_utils import _verify_default_bucket_ownership + + s3 = Mock() + s3.head_bucket.side_effect = ClientError( + {"Error": {"Code": "403", "Message": "Forbidden"}}, "HeadBucket" + ) + with pytest.raises(ValueError, match="not owned by account"): + _verify_default_bucket_ownership( + s3, "sagemaker-us-west-2-111122223333", "111122223333", "us-west-2" + ) + s3.head_bucket.assert_called_once_with( + Bucket="sagemaker-us-west-2-111122223333", ExpectedBucketOwner="111122223333" + ) + + def test_verify_ownership_missing_bucket_ok(self): + from botocore.exceptions import ClientError + from sagemaker.train.common_utils.finetune_utils import _verify_default_bucket_ownership + + s3 = Mock() + s3.head_bucket.side_effect = ClientError( + {"Error": {"Code": "404", "Message": "Not Found"}}, "HeadBucket" + ) + # Missing bucket is allowed (caller creates it in-account); must not raise. + _verify_default_bucket_ownership( + s3, "sagemaker-us-west-2-111122223333", "111122223333", "us-west-2" + ) + + def test_verify_ownership_non_default_bucket_noop(self): + from sagemaker.train.common_utils.finetune_utils import _verify_default_bucket_ownership + + s3 = Mock() + # Explicit / non-default bucket: guard is a no-op and never probes S3. + _verify_default_bucket_ownership(s3, "my-explicit-bucket", "111122223333", "us-west-2") + s3.head_bucket.assert_not_called() + + @patch('sagemaker.train.common_utils.finetune_utils._wait_for_mlflow_app_ready_boto') + @patch('sagemaker.train.common_utils.finetune_utils.TrainDefaults.get_role') + @patch('sagemaker.train.common_utils.finetune_utils._get_prod_sm_client') + def test_create_mlflow_app_passes_expected_owner(self, mock_get_client, mock_get_role, mock_wait): + from sagemaker.train.common_utils.finetune_utils import _create_mlflow_app + + mock_session = Mock() + mock_session.boto_session.region_name = "us-east-1" + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_s3 = Mock() + mock_s3.list_objects_v2.return_value = {"Contents": [{"Key": "mlflow-artifacts/"}]} + + def _client(service_name): + return mock_sts if service_name == "sts" else mock_s3 + + mock_session.boto_session.client.side_effect = _client + mock_get_role.return_value = "arn:aws:iam::123456789012:role/test-role" + mock_sm = Mock() + mock_sm.create_mlflow_app.return_value = {"Arn": "arn:app"} + mock_get_client.return_value = mock_sm + mock_wait.return_value = "arn:app" + + _create_mlflow_app(mock_session) + + mock_s3.head_bucket.assert_called_once_with( + Bucket="sagemaker-us-east-1-123456789012", ExpectedBucketOwner="123456789012" + ) + mock_s3.list_objects_v2.assert_called_once_with( + Bucket="sagemaker-us-east-1-123456789012", + Prefix="mlflow-artifacts/", + MaxKeys=1, + ExpectedBucketOwner="123456789012", + ) + + @patch('sagemaker.train.common_utils.finetune_utils._wait_for_mlflow_app_ready_boto') + @patch('sagemaker.train.common_utils.finetune_utils.TrainDefaults.get_role') + @patch('sagemaker.train.common_utils.finetune_utils._get_prod_sm_client') + def test_create_mlflow_app_foreign_bucket_returns_none(self, mock_get_client, mock_get_role, mock_wait): + from botocore.exceptions import ClientError + from sagemaker.train.common_utils.finetune_utils import _create_mlflow_app + + mock_session = Mock() + mock_session.boto_session.region_name = "us-east-1" + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_s3 = Mock() + mock_s3.head_bucket.side_effect = ClientError({"Error": {"Code": "403"}}, "HeadBucket") + + def _client(service_name): + return mock_sts if service_name == "sts" else mock_s3 + + mock_session.boto_session.client.side_effect = _client + mock_get_role.return_value = "arn:aws:iam::123456789012:role/test-role" + mock_sm = Mock() + mock_get_client.return_value = mock_sm + + result = _create_mlflow_app(mock_session) + + # Foreign-owned default bucket: app is NOT created; call fails safe. + assert result is None + mock_sm.create_mlflow_app.assert_not_called() + + @patch('boto3.client') + def test_validate_s3_path_foreign_default_bucket_raises(self, _mock_boto_client): + from botocore.exceptions import ClientError + from sagemaker.train.common_utils.finetune_utils import _validate_s3_path_exists + + mock_session = Mock() + mock_session.boto_session.region_name = "us-west-2" + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "111122223333"} + mock_s3 = Mock() + mock_s3.head_bucket.side_effect = ClientError({"Error": {"Code": "403"}}, "HeadBucket") + + def _client(service_name): + return mock_sts if service_name == "sts" else mock_s3 + + mock_session.boto_session.client.side_effect = _client + + with pytest.raises(ValueError, match="not owned by account"): + _validate_s3_path_exists("s3://sagemaker-us-west-2-111122223333/output", mock_session) From e44c0c310526f546b503d0ce4e0afca7c30309a8 Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Tue, 15 Sep 2026 09:55:41 -0700 Subject: [PATCH 2/2] chore(train): drop internal tracking ref from test docstring Remove an internal ticket identifier from a test docstring; no behavior or test-logic change. --- .../tests/unit/train/common_utils/test_finetune_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py index 5eb2b38d9c..1789550b97 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py @@ -1787,7 +1787,7 @@ def test_list_hyperparameters_accepts_enum_values(self, mock_boto_client, mock_g class TestDefaultBucketOwnershipGuard: - """Bucket-squatting guard for the SDK-derived default bucket (ticket D524303300).""" + """Bucket-ownership guard for the SDK-derived default bucket.""" def test_verify_ownership_foreign_bucket_raises(self): from botocore.exceptions import ClientError