feat: adds trust boundary lookup support for SA and impersonated credentials

google/auth/credentials.py

@@ -17,7 +17,9 @@
 
 import abc
 from enum import Enum
+import json
 import os
+import typing
 
 from google.auth import _helpers, environment_vars
 from google.auth import exceptions
@@ -26,6 +28,8 @@
 from google.auth._refresh_worker import RefreshThreadManager
 
 DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
+NO_OP_TRUST_BOUNDARY_LOCATIONS: "typing.Tuple[str]" = ()
+NO_OP_TRUST_BOUNDARY_ENCODED_LOCATIONS = "0x0"
 
 
 class Credentials(_BaseCredentials):
@@ -178,22 +182,7 @@ def apply(self, headers, token=None):
             token (Optional[str]): If specified, overrides the current access
                 token.
         """
-        self._apply(headers, token=token)
-        """Trust boundary value will be a cached value from global lookup.
-
-        The response of trust boundary will be a list of regions and a hex
-        encoded representation.
-
-        An example of global lookup response:
-        {
-          "locations": [
-            "us-central1", "us-east1", "europe-west1", "asia-east1"
-          ]
-          "encoded_locations": "0xA30"
-        }
-        """
-        if self._trust_boundary is not None:
-            headers["x-allowed-locations"] = self._trust_boundary["encoded_locations"]
+        self._apply(headers, token)
         if self.quota_project_id:
             headers["x-goog-user-project"] = self.quota_project_id
 
@@ -299,6 +288,100 @@ def with_universe_domain(self, universe_domain):
         )
 
 
+class CredentialsWithTrustBoundary(Credentials):
+    """Abstract base for credentials supporting ``with_trust_boundary`` factory"""
+
+    def with_trust_boundary(self, trust_boundary):
+        """Returns a copy of these credentials with a modified trust boundary.
+
+        Args:
+            trust_boundary Mapping[str, str]: The trust boundary to use for the
+            credential. This should be a map with a "locations" key that maps to
+            a list of GCP regions, and a "encodedLocations" key that maps to a
+            hex string.
+
+        Returns:
+            google.auth.credentials.Credentials: A new credentials instance.
+        """
+        raise NotImplementedError("This credential does not support trust boundaries.")
+
+    def apply(self, headers, token=None):
+        """Apply the token to the authentication header."""
+        super().apply(headers, token)
+        if self._trust_boundary is not None:
+            headers["x-allowed-locations"] = self._trust_boundary["encodedLocations"]
+
+    def _refresh_trust_boundary(self, request):
+        """Triggers a refresh of the trust boundary and updates the cache if necessary.
+
+        Args:
+            request (google.auth.transport.Request): The object used to make
+                HTTP requests.
+
+        Raises:
+            google.auth.exceptions.RefreshError: If the trust boundary could
+                not be refreshed and no cached value is available.
+        """
+        # Do not trigger refresh if credential has a cached no-op trust boundary.
+        if self._has_no_op_trust_boundary():
+            return
+        new_trust_boundary = {}
+        try:
+            new_trust_boundary = self._lookup_trust_boundary(request)
+        except exceptions.RefreshError as error:
+            # If the call to the lookup API failed, check if there is a trust boundary
+            # already cached. If there is, do nothing. If not, then throw the error.
+            if self._trust_boundary is None:
+                raise (error)
+            return
+        else:
+            self._trust_boundary = new_trust_boundary
+
+    @abc.abstractmethod
+    def _lookup_trust_boundary(self, request):
+        """Calls the trust boundary lookup API to refresh the trust boundary cache.
+
+        Args:
+            request (google.auth.transport.Request): The object used to make
+                HTTP requests.
+
+        Returns:
+            trust_boundary (dict): The trust boundary object returned by the lookup API.
+
+        Raises:
+            google.auth.exceptions.RefreshError: If the trust boundary could not be
+                retrieved.
+        """
+        # pylint: disable=missing-raises-doc
+        # (pylint doesn't recognize that this is abstract)
+        raise NotImplementedError("_lookup_trust_boundary must be implemented")
+
+    @staticmethod
+    def _parse_trust_boundary(trust_boundary_string: str):
+        try:
+            trust_boundary = json.loads(trust_boundary_string)
+            if (
+                "locations" not in trust_boundary
+                or "encodedLocations" not in trust_boundary
+            ):
+                raise exceptions.MalformedError
+            return trust_boundary
+        except Exception:
+            raise exceptions.MalformedError(
+                "Cannot parse trust boundary {}".format(trust_boundary_string)
+            )
+
+    def _has_no_op_trust_boundary(self):
+        if (
+            self._trust_boundary != None
+            and self._trust_boundary["locations"] == NO_OP_TRUST_BOUNDARY_LOCATIONS
+            and self._trust_boundary["encodedLocations"]
+            == NO_OP_TRUST_BOUNDARY_ENCODED_LOCATIONS
+        ):
+            return True
+        return False
+
+
 class AnonymousCredentials(Credentials):
     """Credentials that do not provide any authentication information.
 

google/auth/external_account.py

@@ -81,6 +81,7 @@ class Credentials(
     credentials.Scoped,
     credentials.CredentialsWithQuotaProject,
     credentials.CredentialsWithTokenUri,
+    credentials.CredentialsWithTrustBoundary,
     metaclass=abc.ABCMeta,
 ):
     """Base class for all external account credentials.
@@ -133,14 +134,14 @@ def __init__(
                 authorization grant.
             default_scopes (Optional[Sequence[str]]): Default scopes passed by a
                 Google client library. Use 'scopes' for user-defined scopes.
-            workforce_pool_user_project (Optona[str]): The optional workforce pool user
+            workforce_pool_user_project (Optonal[str]): The optional workforce pool user
                 project number when the credential corresponds to a workforce pool and not
                 a workload identity pool. The underlying principal must still have
                 serviceusage.services.use IAM permission to use the project for
                 billing/quota.
             universe_domain (str): The universe domain. The default universe
                 domain is googleapis.com.
-            trust_boundary (str): String representation of trust boundary meta.
+            trust_boundary (str): String representation of trust boundary metadata.
         Raises:
             google.auth.exceptions.RefreshError: If the generateAccessToken
                 endpoint returned an error.
@@ -167,9 +168,9 @@ def __init__(
         self._default_scopes = default_scopes
         self._workforce_pool_user_project = workforce_pool_user_project
         self._trust_boundary = {
-            "locations": [],
-            "encoded_locations": "0x0",
-        }  # expose a placeholder trust boundary value.
+            "locations": credentials.NO_OP_TRUST_BOUNDARY_LOCATIONS,
+            "encodedLocations": credentials.NO_OP_TRUST_BOUNDARY_ENCODED_LOCATIONS,
+        }  # Sets a no-op trust boundary value.
 
         if self._client_id:
             self._client_auth = utils.ClientAuthentication(
@@ -456,6 +457,12 @@ def refresh(self, request):
 
             self.expiry = now + lifetime
 
+    def _lookup_trust_boundary(self, request):
+        """Trust boundary lookup for external account. Currently a no-op because the lookup
+            endpoint does not support external account lookup.
+        """
+        return
+
     def _make_copy(self):
         kwargs = self._constructor_args()
         new_cred = self.__class__(**kwargs)

google/auth/impersonated_credentials.py

@@ -46,6 +46,9 @@
 _DEFAULT_TOKEN_LIFETIME_SECS = 3600  # 1 hour in seconds
 
 _GOOGLE_OAUTH2_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token"
+_TRUST_BOUNDARY_LOOKUP_ENDPOINT = (
+    "https://iamcredentials.{}/v1/projects/-/serviceAccounts/{}/allowedLocations"
+)
 
 
 def _make_iam_token_request(
@@ -111,7 +114,10 @@ def _make_iam_token_request(
 
 
 class Credentials(
-    credentials.Scoped, credentials.CredentialsWithQuotaProject, credentials.Signing
+    credentials.Scoped,
+    credentials.CredentialsWithQuotaProject,
+    credentials.Signing,
+    credentials.CredentialsWithTrustBoundary,
 ):
     """This module defines impersonated credentials which are essentially
     impersonated identities.
@@ -184,6 +190,7 @@ def __init__(
         lifetime=_DEFAULT_TOKEN_LIFETIME_SECS,
         quota_project_id=None,
         iam_endpoint_override=None,
+        trust_boundary=None,
     ):
         """
         Args:
@@ -214,6 +221,7 @@ def __init__(
             subject (Optional[str]): sub field of a JWT. This field should only be set
                 if you wish to impersonate as a user. This feature is useful when
                 using domain wide delegation.
+            trust_boundary (Mapping[str,str]): A credential trust boundary.
         """
 
         super(Credentials, self).__init__()
@@ -245,13 +253,15 @@ def __init__(
         self._quota_project_id = quota_project_id
         self._iam_endpoint_override = iam_endpoint_override
         self._cred_file_path = None
+        self._trust_boundary = trust_boundary
 
     def _metric_header_for_usage(self):
         return metrics.CRED_TYPE_SA_IMPERSONATE
 
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
         self._update_token(request)
+        self._refresh_trust_boundary(request)
 
     def _update_token(self, request):
         """Updates credentials with a new access_token representing
@@ -325,6 +335,19 @@ def _update_token(self, request):
             iam_endpoint_override=self._iam_endpoint_override,
         )
 
+    def _lookup_trust_boundary(self, request):
+        """Trust boundary lookup for service account using endpoint:
+            iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{service_account_email}/allowedLocations
+            And we are using a fresh access token as basic auth.
+        """
+        # Skip trust boundary flow for non-gdu universe domain.
+        if self.universe_domain == credentials.DEFAULT_UNIVERSE_DOMAIN:
+            return
+        url = _TRUST_BOUNDARY_LOOKUP_ENDPOINT.format(
+            self.universe_domain, self.service_account_email
+        )
+        return _client.lookup_trust_boundary(request, url, self.token)
+
     def sign_bytes(self, message):
         from google.auth.transport.requests import AuthorizedSession