Validate Project of staging bucket (#39008)

Author: tarun-google

sdks/python/apache_beam/io/gcp/gcsio.py

@@ -77,6 +77,48 @@ def default_gcs_bucket_name(project, region):
       region, md5(project.encode('utf8')).hexdigest())
 
 
+def _get_project_number(project_id, credentials=None):
+  """Resolves a project ID to its project number using Cloud Resource Manager API."""
+  from google.cloud import resourcemanager_v3
+  client = resourcemanager_v3.ProjectsClient(credentials=credentials)
+  project_info = client.get_project(name=f"projects/{project_id}")
+  # project_info.name is of the form "projects/PROJECT_NUMBER"
+  return int(project_info.name.split('/')[-1])
+
+
+def _validate_bucket_project(bucket, project_id, credentials=None):
+  """Verifies that the GCS bucket is owned by the executing project."""
+  bucket_project_number = bucket.project_number
+
+  # Skip validation if the bucket project number is a mock object
+  if (type(bucket_project_number).__name__.endswith('Mock') or
+      hasattr(bucket_project_number, '_mock_self')):
+    _LOGGER.warning(
+        'Bucket project number is a mock object. Skipping ownership validation.'
+    )
+    return
+
+  if bucket_project_number is None:
+    _LOGGER.warning(
+        'Bucket gs://%s does not have a project number. Skipping ownership validation.',
+        bucket.name)
+    return
+
+  try:
+    project_number = _get_project_number(project_id, credentials=credentials)
+  except Exception as e:
+    _LOGGER.warning(
+        'Failed to resolve project number for project %s: %s. '
+        'Skipping bucket ownership validation.',
+        project_id,
+        e)
+    return
+
+  if bucket_project_number != project_number:
+    raise ValueError(
+        f'Bucket gs://{bucket.name} is not owned by project {project_id}.')
+
+
 def get_or_create_default_gcs_bucket(options):
   """Create a default GCS bucket for this project."""
   if getattr(options, 'dataflow_kms_key', None):
@@ -90,16 +132,18 @@ def get_or_create_default_gcs_bucket(options):
     return None
 
   bucket_name = default_gcs_bucket_name(project, region)
-  bucket = GcsIO(pipeline_options=options).get_bucket(bucket_name)
+  gcs = GcsIO(pipeline_options=options)
+  bucket = gcs.get_bucket(bucket_name)
   if bucket:
+    _validate_bucket_project(
+        bucket, project, credentials=getattr(gcs.client, '_credentials', None))
     return bucket
   else:
     _LOGGER.warning(
         'Creating default GCS bucket for project %s: gs://%s',
         project,
         bucket_name)
-    return GcsIO(pipeline_options=options).create_bucket(
-        bucket_name, project, location=region)
+    return gcs.create_bucket(bucket_name, project, location=region)
 
 
 def create_storage_client(pipeline_options, use_credentials=True):

sdks/python/apache_beam/io/gcp/gcsio_test.py

@@ -395,6 +395,98 @@ def test_default_bucket_name_failure(self):
                 DEFAULT_GCP_PROJECT, "us-central1", kms_key="kmskey!")),
         None)
 
+  @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+  @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+  def test_get_or_create_default_gcs_bucket_ownership_match(
+      self, mock_gcsio_class, mock_crm_class):
+    mock_gcsio = mock_gcsio_class.return_value
+    mock_bucket = mock.Mock()
+    mock_bucket.project_number = 123456789
+    mock_gcsio.get_bucket.return_value = mock_bucket
+
+    mock_crm_client = mock_crm_class.return_value
+    mock_project_info = mock.Mock()
+    mock_project_info.name = 'projects/123456789'
+    mock_crm_client.get_project.return_value = mock_project_info
+
+    options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+    bucket = gcsio.get_or_create_default_gcs_bucket(options)
+
+    self.assertEqual(bucket, mock_bucket)
+    mock_gcsio.get_bucket.assert_called_once_with(
+        'dataflow-staging-us-central1-77b801c0838aee13391c0d1885860494')
+
+  @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+  @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+  def test_get_or_create_default_gcs_bucket_ownership_mismatch(
+      self, mock_gcsio_class, mock_crm_class):
+    mock_gcsio = mock_gcsio_class.return_value
+    mock_bucket = mock.Mock()
+    mock_bucket.project_number = 999999999
+    mock_gcsio.get_bucket.return_value = mock_bucket
+
+    mock_crm_client = mock_crm_class.return_value
+    mock_project_info = mock.Mock()
+    mock_project_info.name = 'projects/123456789'
+    mock_crm_client.get_project.return_value = mock_project_info
+
+    options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+    with self.assertRaises(ValueError) as ctx:
+      gcsio.get_or_create_default_gcs_bucket(options)
+
+    self.assertIn(
+        f'is not owned by project {DEFAULT_GCP_PROJECT}', str(ctx.exception))
+
+  @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+  @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+  def test_get_or_create_default_gcs_bucket_ownership_no_bucket_number(
+      self, mock_gcsio_class, mock_crm_class):
+    mock_gcsio = mock_gcsio_class.return_value
+    mock_bucket = mock.Mock()
+    mock_bucket.project_number = None
+    mock_bucket.name = 'dataflow-staging-us-central1-77b801c0838aee13391c0d1885860494'
+    mock_gcsio.get_bucket.return_value = mock_bucket
+
+    options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+    bucket = gcsio.get_or_create_default_gcs_bucket(options)
+
+    self.assertEqual(bucket, mock_bucket)
+    mock_crm_class.assert_not_called()
+
+  @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+  @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+  def test_get_or_create_default_gcs_bucket_ownership_crm_error(
+      self, mock_gcsio_class, mock_crm_class):
+    mock_gcsio = mock_gcsio_class.return_value
+    mock_bucket = mock.Mock()
+    mock_bucket.project_number = 123456789
+    mock_gcsio.get_bucket.return_value = mock_bucket
+
+    mock_crm_client = mock_crm_class.return_value
+    mock_crm_client.get_project.side_effect = Exception("API Disabled")
+
+    options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+    bucket = gcsio.get_or_create_default_gcs_bucket(options)
+
+    # Should fall back to success (warn only)
+    self.assertEqual(bucket, mock_bucket)
+
+  @mock.patch('google.cloud.resourcemanager_v3.ProjectsClient')
+  @mock.patch('apache_beam.io.gcp.gcsio.GcsIO')
+  def test_get_or_create_default_gcs_bucket_ownership_mock_project_number(
+      self, mock_gcsio_class, mock_crm_class):
+    mock_gcsio = mock_gcsio_class.return_value
+    # Creating a mock bucket without setting project_number (it returns another mock object)
+    mock_bucket = mock.Mock()
+    mock_gcsio.get_bucket.return_value = mock_bucket
+
+    options = SampleOptions(DEFAULT_GCP_PROJECT, 'us-central1')
+    bucket = gcsio.get_or_create_default_gcs_bucket(options)
+
+    # Verification should be skipped, returning the mock bucket and never calling CRM
+    self.assertEqual(bucket, mock_bucket)
+    mock_crm_class.assert_not_called()
+
   def test_exists(self):
     file_name = 'gs://gcsio-test/dummy_file'
     file_size = 1234

sdks/python/setup.py

@@ -531,6 +531,7 @@ def get_portability_package_data():
               'google-cloud-datastore>=2.0.0,<3',
               'google-cloud-pubsub>=2.1.0,<3',
               'google-cloud-storage>=2.18.2,<4',
+              'google-cloud-resource-manager>=1.12.0,<2',
               'google-cloud-dataflow-client>=0.13.0,<0.14.0',
               # GCP packages required by tests
               'google-cloud-bigquery>=2.0.0,<4',