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

* feat: IAM signBlob retries * support universe domain and update tests * update test credentials * use ud signing bucket fixture

Author: cojenco

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(

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")

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

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)
 

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)