2525import os
2626import pathlib
2727import tempfile
28- from typing import Optional , TYPE_CHECKING
28+ from typing import Optional , Tuple , TYPE_CHECKING
2929import warnings
3030
3131from google .auth import credentials as auth_credentials
6363_DEFAULT_STAGING_BUCKET_SALT = str (uuid .uuid4 ())
6464
6565
66+ class BucketOwnershipError (ValueError ):
67+ """Raised when a GCS bucket does not belong to the expected project."""
68+
69+
70+ def _get_bucket_project_numbers (
71+ bucket : storage .Bucket ,
72+ expected_project : str ,
73+ client : storage .Client ,
74+ ) -> Tuple [Optional [str ], Optional [str ]]:
75+ """Gets the project number of a GCS bucket and of the expected project.
76+
77+ Args:
78+ bucket: The GCS bucket to check.
79+ expected_project: The project ID that should own the bucket.
80+ client: Storage client instance.
81+
82+ Returns:
83+ A tuple of the bucket's project number and the expected project's
84+ number. Each value is None if it could not be retrieved.
85+ """
86+ bucket_project_number = None
87+ expected_project_number = None
88+ try :
89+ bucket .reload (client = client )
90+ if bucket .project_number is not None :
91+ bucket_project_number = str (bucket .project_number )
92+ except Exception :
93+ pass
94+ try :
95+ expected_project_number = str (
96+ resource_manager_utils .get_project_number (
97+ expected_project , credentials = getattr (client , "_credentials" , None )
98+ )
99+ )
100+ except Exception :
101+ pass
102+ return bucket_project_number , expected_project_number
103+
104+
66105def _verify_bucket_ownership (
67106 bucket : storage .Bucket ,
68107 expected_project : str ,
@@ -81,15 +120,70 @@ def _verify_bucket_ownership(
81120 Returns:
82121 True if the bucket belongs to the expected project, False otherwise.
83122 """
84- try :
85- bucket .reload (client = client )
86- bucket_project_number = str (bucket .project_number )
87- expected_project_number = str (
88- resource_manager_utils .get_project_number (expected_project )
123+ bucket_project_number , expected_project_number = _get_bucket_project_numbers (
124+ bucket , expected_project , client
125+ )
126+ return (
127+ bucket_project_number is not None
128+ and bucket_project_number == expected_project_number
129+ )
130+
131+
132+ def _check_pipeline_bucket_ownership (
133+ bucket : storage .Bucket ,
134+ expected_project : str ,
135+ client : storage .Client ,
136+ is_sdk_generated : bool = False ,
137+ ) -> None :
138+ """Warns, or raises, if a bucket is not owned by the project.
139+
140+ Args:
141+ bucket: The existing GCS bucket for pipeline artifacts.
142+ expected_project: The project ID that should own the bucket.
143+ client: Storage client instance.
144+ is_sdk_generated: Whether the SDK generated the bucket name. A
145+ generated name can be predictable, so a mismatch always raises.
146+
147+ Raises:
148+ BucketOwnershipError: If the bucket is not owned by `expected_project`,
149+ or its owner cannot be verified, and either `is_sdk_generated` is
150+ True or `aiplatform.init(enforce_bucket_ownership=True)` is set.
151+ """
152+ bucket_project_number , expected_project_number = _get_bucket_project_numbers (
153+ bucket , expected_project , client
154+ )
155+ if bucket_project_number is not None and (
156+ bucket_project_number == expected_project_number
157+ ):
158+ return
159+
160+ if bucket_project_number is None or expected_project_number is None :
161+ message = (
162+ f'Could not verify that GCS bucket "{ bucket .name } " is owned by '
163+ f'project "{ expected_project } ".'
89164 )
90- return bucket_project_number == expected_project_number
91- except Exception :
92- return False
165+ else :
166+ message = (
167+ f'GCS bucket "{ bucket .name } " is owned by project number '
168+ f'"{ bucket_project_number } ", not by project "{ expected_project } " '
169+ f'(project number "{ expected_project_number } ").'
170+ )
171+ message += (
172+ " Pipeline artifacts will be written to this bucket. If you do not "
173+ "own this bucket, other users may be able to read or modify the "
174+ "artifacts. Use a bucket owned by your project, for example by "
175+ "passing pipeline_root or by setting "
176+ "aiplatform.init(staging_bucket='gs://your-bucket')."
177+ )
178+
179+ if is_sdk_generated or initializer .global_config .enforce_bucket_ownership :
180+ raise BucketOwnershipError (message )
181+ warnings .warn (
182+ f"{ message } To raise an error instead, set "
183+ "aiplatform.init(enforce_bucket_ownership=True)." ,
184+ UserWarning ,
185+ stacklevel = 2 ,
186+ )
93187
94188
95189def blob_from_uri (uri : str , client : storage .Client ) -> storage .Blob :
@@ -335,12 +429,12 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
335429 service_account = service_account or initializer .global_config .service_account
336430 credentials = credentials or initializer .global_config .credentials
337431
432+ default_output_artifacts_gcs_dir = generate_gcs_directory_for_pipeline_artifacts (
433+ project = project ,
434+ location = location ,
435+ )
338436 output_artifacts_gcs_dir = (
339- output_artifacts_gcs_dir
340- or generate_gcs_directory_for_pipeline_artifacts (
341- project = project ,
342- location = location ,
343- )
437+ output_artifacts_gcs_dir or default_output_artifacts_gcs_dir
344438 )
345439
346440 # Creating the bucket if needed
@@ -377,13 +471,21 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
377471 f"serviceAccount:{ service_account } "
378472 )
379473 pipelines_bucket .set_iam_policy (bucket_iam_policy )
380- elif not _verify_bucket_ownership (pipelines_bucket , project , storage_client ):
381- raise ValueError (
382- f'Output artifacts bucket "{ pipelines_bucket .name } " exists but does '
383- f'not belong to project "{ project } ". This may indicate a '
384- f"bucket squatting attack. Please provide an explicit "
385- f"output_artifacts_gcs_dir parameter or configure one via "
386- f"aiplatform.init(output_artifacts_gcs_dir='gs://your-bucket')."
474+ else :
475+ # PipelineJob generates its default pipeline_root before calling this
476+ # function, so compare bucket names instead of checking for None.
477+ is_sdk_generated = (
478+ pipelines_bucket .name
479+ == bucket_from_uri (
480+ uri = default_output_artifacts_gcs_dir ,
481+ client = storage_client ,
482+ ).name
483+ )
484+ _check_pipeline_bucket_ownership (
485+ pipelines_bucket ,
486+ project ,
487+ storage_client ,
488+ is_sdk_generated = is_sdk_generated ,
387489 )
388490 return output_artifacts_gcs_dir
389491
0 commit comments