From 2baea5b1a36e4cd5f9b07d6f3e7ae17ebc61d585 Mon Sep 17 00:00:00 2001
From: Jeff Scudder <jscudder@google.com>
Date: Fri, 14 Aug 2026 19:43:02 -0700
Subject: [PATCH] fix: check GCS ownership for default pipeline bucket

PiperOrigin-RevId: 965016212
---
 google/cloud/aiplatform/utils/gcs_utils.py | 12 +++++++++++-
 tests/unit/aiplatform/test_utils.py        |  8 ++++----
 tests/unit/vertexai/tuning/test_tuning.py  |  8 ++++++++
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/google/cloud/aiplatform/utils/gcs_utils.py b/google/cloud/aiplatform/utils/gcs_utils.py
index 6153504c15..58239f800a 100644
--- a/google/cloud/aiplatform/utils/gcs_utils.py
+++ b/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
 
 
diff --git a/tests/unit/aiplatform/test_utils.py b/tests/unit/aiplatform/test_utils.py
index 0688f78c5a..218b382cb1 100644
--- a/tests/unit/aiplatform/test_utils.py
+++ b/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
diff --git a/tests/unit/vertexai/tuning/test_tuning.py b/tests/unit/vertexai/tuning/test_tuning.py
index 929fa7bcd9..b214e6e6ae 100644
--- a/tests/unit/vertexai/tuning/test_tuning.py
+++ b/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