diff --git a/google/cloud/aiplatform/utils/gcs_utils.py b/google/cloud/aiplatform/utils/gcs_utils.py index 5bebd9ee01..c5faa7011b 100644 --- a/google/cloud/aiplatform/utils/gcs_utils.py +++ b/google/cloud/aiplatform/utils/gcs_utils.py @@ -377,6 +377,16 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist( f"serviceAccount:{service_account}" ) pipelines_bucket.set_iam_policy(bucket_iam_policy) + else: + # Verify bucket ownership to prevent bucket squatting attacks. + if not _verify_bucket_ownership(pipelines_bucket, project, storage_client): + raise ValueError( + f'Pipeline 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 a pipeline " + f"root via PipelineJob(pipeline_root='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 e40aae37b8..670bde96b0 100644 --- a/tests/unit/aiplatform/test_utils.py +++ b/tests/unit/aiplatform/test_utils.py @@ -603,6 +603,19 @@ def test_create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist( output == "gs://test-project-vertex-pipelines-us-central1/output_artifacts/" ) + @patch.object( + gcs_utils, "_verify_bucket_ownership", return_value=False + ) + @patch.object(storage.Bucket, "exists", return_value=True) + @patch.object(storage, "Client") + def test_create_gcs_bucket_for_pipeline_artifacts_rejects_foreign_bucket( + self, mock_storage_client, mock_bucket_exists, mock_verify + ): + with pytest.raises(ValueError, match="bucket squatting"): + gcs_utils.create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist( + project="test-project", location="us-central1" + ) + def test_download_from_gcs_dir( self, mock_storage_client_list_blobs, mock_storage_blob_download_to_filename ):