feat: IAM signBlob retry and universe domain support (#1380)

cojenco · web-flow · commit abc80615ee00 · 2024-11-20T14:50:42.000-08:00
* feat: IAM signBlob retries

* support universe domain and update tests

* update test credentials

* use ud signing bucket fixture
diff --git a/google/cloud/storage/_signing.py b/google/cloud/storage/_signing.py
@@ -28,8 +28,10 @@
 from google.auth import exceptions
 from google.auth.transport import requests
 from google.cloud import _helpers
+from google.cloud.storage._helpers import _DEFAULT_UNIVERSE_DOMAIN
 from google.cloud.storage._helpers import _NOW
 from google.cloud.storage._helpers import _UTC
+from google.cloud.storage.retry import DEFAULT_RETRY
 
 
 # `google.cloud.storage._signing.NOW` is deprecated.
@@ -271,6 +273,7 @@ def generate_signed_url_v2(
     query_parameters=None,
     service_account_email=None,
     access_token=None,
+    universe_domain=None,
 ):
     """Generate a V2 signed URL to provide query-string auth'n to a resource.
 
@@ -384,7 +387,9 @@ def generate_signed_url_v2(
     # See https://github.com/googleapis/google-cloud-python/issues/922
     # Set the right query parameters.
     if access_token and service_account_email:
-        signature = _sign_message(string_to_sign, access_token, service_account_email)
+        signature = _sign_message(
+            string_to_sign, access_token, service_account_email, universe_domain
+        )
         signed_query_params = {
             "GoogleAccessId": service_account_email,
             "Expires": expiration_stamp,
@@ -432,6 +437,7 @@ def generate_signed_url_v4(
     query_parameters=None,
     service_account_email=None,
     access_token=None,
+    universe_domain=None,
     _request_timestamp=None,  # for testing only
 ):
     """Generate a V4 signed URL to provide query-string auth'n to a resource.
@@ -623,7 +629,9 @@ def generate_signed_url_v4(
     string_to_sign = "\n".join(string_elements)
 
     if access_token and service_account_email:
-        signature = _sign_message(string_to_sign, access_token, service_account_email)
+        signature = _sign_message(
+            string_to_sign, access_token, service_account_email, universe_domain
+        )
         signature_bytes = base64.b64decode(signature)
         signature = binascii.hexlify(signature_bytes).decode("ascii")
     else:
@@ -647,7 +655,12 @@ def get_v4_now_dtstamps():
     return timestamp, datestamp
 
 
-def _sign_message(message, access_token, service_account_email):
+def _sign_message(
+    message,
+    access_token,
+    service_account_email,
+    universe_domain=_DEFAULT_UNIVERSE_DOMAIN,
+):
     """Signs a message.
 
     :type message: str
@@ -669,17 +682,22 @@ def _sign_message(message, access_token, service_account_email):
     message = _helpers._to_bytes(message)
 
     method = "POST"
-    url = "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{}:signBlob?alt=json".format(
-        service_account_email
-    )
+    url = f"https://iamcredentials.{universe_domain}/v1/projects/-/serviceAccounts/{service_account_email}:signBlob?alt=json"
     headers = {
         "Authorization": "Bearer " + access_token,
         "Content-type": "application/json",
     }
     body = json.dumps({"payload": base64.b64encode(message).decode("utf-8")})
-
     request = requests.Request()
-    response = request(url=url, method=method, body=body, headers=headers)
+
+    def retriable_request():
+        response = request(url=url, method=method, body=body, headers=headers)
+        return response
+
+    # Apply the default retry object to the signBlob call.
+    retry = DEFAULT_RETRY
+    call = retry(retriable_request)
+    response = call()
 
     if response.status != http.client.OK:
         raise exceptions.TransportError(
diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py
@@ -607,6 +607,9 @@ def generate_signed_url(
             client = self._require_client(client)  # May be redundant, but that's ok.
             credentials = client._credentials
 
+        client = self._require_client(client)
+        universe_domain = client.universe_domain
+
         if version == "v2":
             helper = generate_signed_url_v2
         else:
@@ -638,6 +641,7 @@ def generate_signed_url(
             query_parameters=query_parameters,
             service_account_email=service_account_email,
             access_token=access_token,
+            universe_domain=universe_domain,
         )
 
     @create_trace_span(name="Storage.Blob.exists")
diff --git a/tests/system/conftest.py b/tests/system/conftest.py
@@ -384,3 +384,33 @@ def universe_domain_client(
     )
     with contextlib.closing(ud_storage_client):
         yield ud_storage_client
+
+
+@pytest.fixture(scope="function")
+def universe_domain_bucket(universe_domain_client, test_universe_location):
+    bucket_name = _helpers.unique_name("gcp-systest-ud")
+    bucket = universe_domain_client.create_bucket(
+        bucket_name, location=test_universe_location
+    )
+
+    blob = bucket.blob("README.txt")
+    blob.upload_from_string(_helpers.signing_blob_content)
+
+    yield bucket
+
+    _helpers.delete_bucket(bucket)
+
+
+@pytest.fixture(scope="function")
+def universe_domain_iam_client(
+    test_universe_domain, test_universe_project_id, universe_domain_credential
+):
+    from google.cloud import iam_credentials_v1
+
+    client_options = {"universe_domain": test_universe_domain}
+    iam_client = iam_credentials_v1.IAMCredentialsClient(
+        credentials=universe_domain_credential,
+        client_options=client_options,
+    )
+
+    return iam_client
diff --git a/tests/system/test__signing.py b/tests/system/test__signing.py
@@ -287,6 +287,35 @@ def test_create_signed_read_url_v4_w_access_token(
     )
 
 
+def test_create_signed_read_url_v4_w_access_token_universe_domain(
+    universe_domain_iam_client,
+    universe_domain_client,
+    test_universe_location,
+    universe_domain_credential,
+    universe_domain_bucket,
+    no_mtls,
+):
+    service_account_email = universe_domain_credential.service_account_email
+    name = path_template.expand(
+        "projects/{project}/serviceAccounts/{service_account}",
+        project="-",
+        service_account=service_account_email,
+    )
+    scope = [
+        "https://www.googleapis.com/auth/devstorage.read_write",
+        "https://www.googleapis.com/auth/iam",
+    ]
+    response = universe_domain_iam_client.generate_access_token(name=name, scope=scope)
+
+    _create_signed_read_url_helper(
+        universe_domain_client,
+        universe_domain_bucket,
+        version="v4",
+        service_account_email=service_account_email,
+        access_token=response.access_token,
+    )
+
+
 def _create_signed_delete_url_helper(client, bucket, version="v2", expiration=None):
     expiration = _morph_expiration(version, expiration)
 
diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py
@@ -487,6 +487,8 @@ def _generate_signed_url_helper(
             expected_creds = credentials
             client = self._make_client(_credentials=object())
 
+        expected_universe_domain = client.universe_domain
+
         bucket = _Bucket(client)
         blob = self._make_one(blob_name, bucket=bucket, encryption_key=encryption_key)
 
@@ -564,6 +566,7 @@ def _generate_signed_url_helper(
             "query_parameters": query_parameters,
             "access_token": access_token,
             "service_account_email": service_account_email,
+            "universe_domain": expected_universe_domain,
         }
         signer.assert_called_once_with(expected_creds, **expected_kwargs)
 
