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..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 @@ -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)