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

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

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:

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."""

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-ownership guard for the SDK-derived default bucket."""
+
+    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)