fix: Add bucket ownership verification to prevent bucket squatting in Model.upload()

PiperOrigin-RevId: 899642927

Author: vertex-sdk-bot

google/cloud/aiplatform/utils/gcs_utils.py

@@ -63,6 +63,35 @@
 _DEFAULT_STAGING_BUCKET_SALT = str(uuid.uuid4())
 
 
+def _verify_bucket_ownership(
+    bucket: storage.Bucket,
+    expected_project: str,
+    client: storage.Client,
+) -> bool:
+    """Verifies that a GCS bucket belongs to the expected project.
+
+    This check mitigates bucket squatting attacks where an attacker creates a
+    bucket with a predictable name in their own project before the victim does.
+
+    Args:
+        bucket: The GCS bucket to verify.
+        expected_project: The project ID that should own the bucket.
+        client: Storage client instance.
+
+    Returns:
+        True if the bucket belongs to the expected project, False otherwise.
+    """
+    try:
+        bucket.reload(client=client)
+        bucket_project_number = str(bucket.project_number)
+        expected_project_number = str(
+            resource_manager_utils.get_project_number(expected_project)
+        )
+        return bucket_project_number == expected_project_number
+    except Exception:
+        return False
+
+
 def blob_from_uri(uri: str, client: storage.Client) -> storage.Blob:
     """Create a Blob from a GCS URI, compatible with v2 and v3.
 
@@ -221,6 +250,17 @@ def stage_local_data_in_gcs(
                 project=project,
                 location=location,
             )
+        else:
+            # Verify bucket ownership to prevent bucket squatting attacks.
+            # See b/469987320 for details.
+            if not _verify_bucket_ownership(staging_bucket, project, client):
+                raise ValueError(
+                    f'Staging bucket "{staging_bucket_name}" exists but does '
+                    f'not belong to project "{project}". This may indicate a '
+                    f"bucket squatting attack. Please provide an explicit "
+                    f"staging_bucket parameter or configure one via "
+                    f"aiplatform.init(staging_bucket='gs://your-bucket')."
+                )
         staging_gcs_dir = "gs://" + staging_bucket_name
 
     timestamp = datetime.datetime.now().isoformat(sep="-", timespec="milliseconds")

tests/unit/aiplatform/test_models.py

@@ -899,6 +899,9 @@ def mock_storage_blob_upload_from_filename():
         "google.cloud.storage.Blob.upload_from_filename"
     ) as mock_blob_upload_from_filename, patch(
         "google.cloud.storage.Bucket.exists", return_value=True
+    ), patch(
+        "google.cloud.aiplatform.utils.gcs_utils._verify_bucket_ownership",
+        return_value=True,
     ):
         yield mock_blob_upload_from_filename
 

tests/unit/aiplatform/test_utils.py

@@ -668,6 +668,79 @@ def test_validate_gcs_path(self):
         with pytest.raises(ValueError, match=err_msg):
             gcs_utils.validate_gcs_path(test_invalid_path)
 
+    @patch.object(
+        gcs_utils.resource_manager_utils,
+        "get_project_number",
+        return_value=12345,
+    )
+    @patch.object(storage.Bucket, "reload")
+    def test_verify_bucket_ownership_matching_project(
+        self, mock_reload, mock_get_project_number
+    ):
+        mock_client = mock.MagicMock(spec=storage.Client)
+        mock_bucket = mock.MagicMock(spec=storage.Bucket)
+        mock_bucket.project_number = 12345
+        assert gcs_utils._verify_bucket_ownership(
+            mock_bucket, "test-project", mock_client
+        )
+
+    @patch.object(
+        gcs_utils.resource_manager_utils,
+        "get_project_number",
+        return_value=12345,
+    )
+    @patch.object(storage.Bucket, "reload")
+    def test_verify_bucket_ownership_different_project(
+        self, mock_reload, mock_get_project_number
+    ):
+        mock_client = mock.MagicMock(spec=storage.Client)
+        mock_bucket = mock.MagicMock(spec=storage.Bucket)
+        mock_bucket.project_number = 99999
+        assert not gcs_utils._verify_bucket_ownership(
+            mock_bucket, "test-project", mock_client
+        )
+
+    @patch.object(storage.Bucket, "exists", return_value=True)
+    @patch.object(storage, "Client")
+    @patch.object(gcs_utils, "_verify_bucket_ownership", return_value=False)
+    def test_stage_local_data_in_gcs_rejects_squatted_bucket(
+        self, mock_verify, mock_storage_client, mock_bucket_exists, json_file
+    ):
+        mock_config = mock.MagicMock()
+        mock_config.project = "victim-project"
+        mock_config.location = "us-central1"
+        mock_config.staging_bucket = None
+        mock_config.credentials = None
+        with patch.object(gcs_utils.initializer, "global_config", mock_config):
+            with pytest.raises(
+                ValueError,
+                match="bucket squatting",
+            ):
+                gcs_utils.stage_local_data_in_gcs(json_file)
+
+    @patch.object(storage.Bucket, "exists", return_value=True)
+    @patch.object(storage, "Client")
+    @patch.object(gcs_utils, "_verify_bucket_ownership", return_value=True)
+    @patch("google.cloud.storage.Blob.upload_from_filename")
+    def test_stage_local_data_in_gcs_accepts_owned_bucket(
+        self,
+        mock_upload,
+        mock_verify,
+        mock_storage_client,
+        mock_bucket_exists,
+        json_file,
+        mock_datetime,
+    ):
+        mock_config = mock.MagicMock()
+        mock_config.project = "my-project"
+        mock_config.location = "us-central1"
+        mock_config.staging_bucket = None
+        mock_config.credentials = None
+        with patch.object(gcs_utils.initializer, "global_config", mock_config):
+            result = gcs_utils.stage_local_data_in_gcs(json_file)
+            assert result.startswith("gs://")
+            mock_verify.assert_called_once()
+
 
 class TestPipelineUtils:
     SAMPLE_JOB_SPEC = {