fix: check GCS ownership for default pipeline bucket

google/cloud/aiplatform/utils/gcs_utils.py

@@ -299,7 +299,9 @@ def generate_gcs_directory_for_pipeline_artifacts(
     project = project or initializer.global_config.project
     location = location or initializer.global_config.location
 
-    pipelines_bucket_name = project + "-vertex-pipelines-" + location
+    pipelines_bucket_name = (
+        project + "-vertex-pipelines-" + location + "-" + _DEFAULT_STAGING_BUCKET_SALT
+    )
     output_artifacts_gcs_dir = "gs://" + pipelines_bucket_name + "/output_artifacts/"
     return output_artifacts_gcs_dir
 
@@ -375,6 +377,14 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
             f"serviceAccount:{service_account}"
         )
         pipelines_bucket.set_iam_policy(bucket_iam_policy)
+    elif not _verify_bucket_ownership(pipelines_bucket, project, storage_client):
+        raise ValueError(
+            f'Output artifacts bucket "{pipelines_bucket.name}" exists but does '
+            f'not belong to project "{project}". This may indicate a '
+            f"bucket squatting attack. Please provide an explicit "
+            f"output_artifacts_gcs_dir parameter or configure one via "
+            f"aiplatform.init(output_artifacts_gcs_dir='gs://your-bucket')."
+        )
     return output_artifacts_gcs_dir
 
 

tests/unit/aiplatform/test_utils.py

@@ -609,7 +609,8 @@ def test_generate_gcs_directory_for_pipeline_artifacts(self):
         output = gcs_utils.generate_gcs_directory_for_pipeline_artifacts(
             "project", "us-central1"
         )
-        assert output == "gs://project-vertex-pipelines-us-central1/output_artifacts/"
+        assert output.startswith("gs://project-vertex-pipelines-us-central1-")
+        assert output.endswith("/output_artifacts/")
 
     @patch.object(storage.Bucket, "exists", return_value=False)
     @patch.object(storage, "Client")
@@ -627,9 +628,8 @@ def test_create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
         assert mock_storage_client.called
         assert mock_bucket_not_exist.called
         assert mock_get_project_number.called
-        assert (
-            output == "gs://test-project-vertex-pipelines-us-central1/output_artifacts/"
-        )
+        assert output.startswith("gs://test-project-vertex-pipelines-us-central1-")
+        assert output.endswith("/output_artifacts/")
 
     def test_download_from_gcs_dir(
         self, mock_storage_client_list_blobs, mock_storage_blob_download_to_filename

tests/unit/vertexai/tuning/test_tuning.py

@@ -34,6 +34,7 @@
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import utils as aiplatform_utils
 from google.cloud.aiplatform.metadata import experiment_resources
+from google.cloud.aiplatform.utils import gcs_utils
 from google.cloud.aiplatform_v1beta1.services import gen_ai_tuning_service
 from google.cloud.aiplatform_v1beta1.types import job_state
 from google.cloud.aiplatform_v1beta1.types import tuning_job as gca_tuning_job
@@ -283,6 +284,13 @@ def test_genai_tuning_service_service_account(
         attribute="exists",
         new=lambda _: True,
     )
+    @mock.patch.object(
+        target=gcs_utils,
+        attribute="_verify_bucket_ownership",
+        # Override bucket ownership check to allow tests to use a non-existent
+        # bucket.
+        new=lambda *args, **kwargs: True,
+    )
     def test_genai_tuning_service_distillation_distill_model(self):
         distillation_train = _distillation.distill_model