implement trust boundary in compute_engine to support GCE instances.

Author: nbayati

google/auth/compute_engine/credentials.py

@@ -30,11 +30,16 @@
 from google.auth.compute_engine import _metadata
 from google.oauth2 import _client
 
+_TRUST_BOUNDARY_LOOKUP_ENDPOINT = (
+    "https://iamcredentials.{}/v1/projects/-/serviceAccounts/{}/allowedLocations"
+)
+
 
 class Credentials(
     credentials.Scoped,
     credentials.CredentialsWithQuotaProject,
     credentials.CredentialsWithUniverseDomain,
+    credentials.CredentialsWithTrustBoundary,
 ):
     """Compute Engine Credentials.
 
@@ -61,6 +66,7 @@ def __init__(
         scopes=None,
         default_scopes=None,
         universe_domain=None,
+        trust_boundary=None,
     ):
         """
         Args:
@@ -76,6 +82,7 @@ def __init__(
                 provided or None, credential will attempt to fetch the value
                 from metadata server. If metadata server doesn't have universe
                 domain endpoint, then the default googleapis.com will be used.
+            trust_boundary (Mapping[str,str]): A credential trust boundary.
         """
         super(Credentials, self).__init__()
         self._service_account_email = service_account_email
@@ -86,6 +93,7 @@ def __init__(
         if universe_domain:
             self._universe_domain = universe_domain
             self._universe_domain_cached = True
+        self._trust_boundary = trust_boundary
 
     def _metric_header_for_usage(self):
         return metrics.CRED_TYPE_SA_MDS
@@ -111,6 +119,33 @@ def refresh(self, request):
             new_exc = exceptions.RefreshError(caught_exc)
             raise new_exc from caught_exc
 
+        self._refresh_trust_boundary(request)
+
+    def _build_trust_boundary_lookup_url(self):
+        """Builds and returns the URL for the trust boundary lookup API for GCE."""
+        # If the service account email is 'default', we need to get the
+        # actual email address from the metadata server.
+        if self._service_account_email == "default":
+            from google.auth.transport import requests as google_auth_requests
+
+            request = google_auth_requests.Request()
+            try:
+                info = _metadata.get_service_account_info(request, "default")
+                # Cache the fetched email so we don't have to do this again.
+                self._service_account_email = info["email"]
+
+            except exceptions.TransportError as e:
+                # If fetching the service account email fails due to a transport error,
+                # it means we cannot build the trust boundary lookup URL.
+                # Wrap this in a RefreshError so it's caught by _refresh_trust_boundary.
+                raise exceptions.RefreshError(
+                    f"Failed to get service account email for trust boundary lookup: {e}"
+                ) from e
+
+        return _TRUST_BOUNDARY_LOOKUP_ENDPOINT.format(
+            self.universe_domain, self.service_account_email
+        )
+
     @property
     def service_account_email(self):
         """The service account email.
@@ -152,6 +187,7 @@ def with_quota_project(self, quota_project_id):
             quota_project_id=quota_project_id,
             scopes=self._scopes,
             default_scopes=self._default_scopes,
+            trust_boundary=self._trust_boundary,
         )
         creds._universe_domain = self._universe_domain
         creds._universe_domain_cached = self._universe_domain_cached
@@ -167,6 +203,7 @@ def with_scopes(self, scopes, default_scopes=None):
             default_scopes=default_scopes,
             service_account_email=self._service_account_email,
             quota_project_id=self._quota_project_id,
+            trust_boundary=self._trust_boundary,
         )
         creds._universe_domain = self._universe_domain
         creds._universe_domain_cached = self._universe_domain_cached
@@ -179,9 +216,23 @@ def with_universe_domain(self, universe_domain):
             default_scopes=self._default_scopes,
             service_account_email=self._service_account_email,
             quota_project_id=self._quota_project_id,
+            trust_boundary=self._trust_boundary,
             universe_domain=universe_domain,
         )
 
+    @_helpers.copy_docstring(credentials.CredentialsWithTrustBoundary)
+    def with_trust_boundary(self, trust_boundary):
+        creds = self.__class__(
+            service_account_email=self._service_account_email,
+            quota_project_id=self._quota_project_id,
+            scopes=self._scopes,
+            default_scopes=self._default_scopes,
+            trust_boundary=trust_boundary,
+        )
+        creds._universe_domain = self._universe_domain
+        creds._universe_domain_cached = self._universe_domain_cached
+        return creds
+
 
 _DEFAULT_TOKEN_LIFETIME_SECS = 3600  # 1 hour in seconds
 _DEFAULT_TOKEN_URI = "https://www.googleapis.com/oauth2/v4/token"

google/auth/impersonated_credentials.py

@@ -432,10 +432,18 @@ def _make_copy(self):
             lifetime=self._lifetime,
             quota_project_id=self._quota_project_id,
             iam_endpoint_override=self._iam_endpoint_override,
+            trust_boundary=self._trust_boundary,
         )
         cred._cred_file_path = self._cred_file_path
         return cred
 
+    @_helpers.copy_docstring(credentials.CredentialsWithTrustBoundary)
+    def with_trust_boundary(self, trust_boundary):
+        """Returns a copy of these credentials with a modified trust boundary."""
+        cred = self._make_copy()
+        cred._trust_boundary = trust_boundary
+        return cred
+
     @_helpers.copy_docstring(credentials.CredentialsWithQuotaProject)
     def with_quota_project(self, quota_project_id):
         cred = self._make_copy()
@@ -519,9 +527,7 @@ def from_impersonated_service_account_info(cls, info, scopes=None):
 
 
 class IDTokenCredentials(credentials.CredentialsWithQuotaProject):
-    """Open ID Connect ID Token-based service account credentials.
-
-    """
+    """Open ID Connect ID Token-based service account credentials."""
 
     def __init__(
         self,

google/oauth2/_client.py

@@ -622,7 +622,7 @@ def _lookup_trust_boundary_request_no_throw(
           is retryable.
     """
 
-    headers_to_use = {"Authorization", "Bearer {}".format(access_token)}
+    headers_to_use = {"Authorization": "Bearer {}".format(access_token)}
 
     response_data = {}
     retryable_error = False

google/oauth2/service_account.py

@@ -386,6 +386,13 @@ def with_token_uri(self, token_uri):
         cred._token_uri = token_uri
         return cred
 
+    @_helpers.copy_docstring(credentials.CredentialsWithTrustBoundary)
+    def with_trust_boundary(self, trust_boundary):
+        """Returns a copy of these credentials with a modified trust boundary."""
+        cred = self._make_copy()
+        cred._trust_boundary = trust_boundary
+        return cred
+
     def _make_authorization_grant_assertion(self):
         """Create the OAuth 2.0 assertion.