Support self-signed jwt and refactor refresh to handle refreshing trust boundary in the base class.

Author: nbayati

google/auth/_helpers.py

@@ -295,6 +295,7 @@ def get_bool_from_env(variable_name, default=False):
     (case-insensitive) rules:
     - "true", "1" are considered true.
     - "false", "0" are considered false.
+    Any other values will raise an exception.
 
     Args:
         variable_name (str): The name of the environment variable.

google/auth/compute_engine/credentials.py

@@ -98,7 +98,7 @@ def __init__(
     def _metric_header_for_usage(self):
         return metrics.CRED_TYPE_SA_MDS
 
-    def refresh(self, request):
+    def _refresh_token(self, request):
         """Refresh the access token and scopes.
 
         Args:
@@ -119,8 +119,6 @@ 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

google/auth/credentials.py

@@ -292,6 +292,20 @@ def with_universe_domain(self, universe_domain):
 class CredentialsWithTrustBoundary(Credentials):
     """Abstract base for credentials supporting ``with_trust_boundary`` factory"""
 
+    @abc.abstractmethod
+    def _refresh_token(self, request):
+        """Refreshes the access token.
+
+        Args:
+            request (google.auth.transport.Request): The object used to make
+                HTTP requests.
+
+        Raises:
+            google.auth.exceptions.RefreshError: If the credentials could
+                not be refreshed.
+        """
+        raise NotImplementedError("_refresh_token must be implemented")
+
     def with_trust_boundary(self, trust_boundary):
         """Returns a copy of these credentials with a modified trust boundary.
 
@@ -306,6 +320,29 @@ def with_trust_boundary(self, trust_boundary):
         """
         raise NotImplementedError("This credential does not support trust boundaries.")
 
+    def _is_trust_boundary_lookup_required(self):
+        """Checks if a trust boundary lookup is required.
+
+        A lookup is required if the feature is enabled via an environment
+        variable, the universe domain is supported, and a no-op boundary
+        is not already cached.
+
+        Returns:
+            bool: True if a trust boundary lookup is required, False otherwise.
+        """
+        # 1. Check if the feature is enabled via environment variable.
+        if not _helpers.get_bool_from_env(
+            environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED, default=False
+        ):
+            return False
+
+        # 2. Skip trust boundary flow for non-default universe domains.
+        if self.universe_domain != DEFAULT_UNIVERSE_DOMAIN:
+            return False
+
+        # 3. Do not trigger refresh if credential has a cached no-op trust boundary.
+        return not self._has_no_op_trust_boundary()
+
     def _get_trust_boundary_header(self):
         if self._trust_boundary is not None:
             if self._has_no_op_trust_boundary():
@@ -320,6 +357,15 @@ def apply(self, headers, token=None):
         super().apply(headers, token)
         headers.update(self._get_trust_boundary_header())
 
+    def refresh(self, request):
+        """Refreshes the access token and the trust boundary.
+
+        This method calls the subclass's token refresh logic and then
+        refreshes the trust boundary if applicable.
+        """
+        self._refresh_token(request)
+        self._refresh_trust_boundary(request)
+
     def _refresh_trust_boundary(self, request):
         """Triggers a refresh of the trust boundary and updates the cache if necessary.
 
@@ -331,12 +377,10 @@ def _refresh_trust_boundary(self, request):
             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():
+        if not self._is_trust_boundary_lookup_required():
             return
-        new_trust_boundary = {}
         try:
-            new_trust_boundary = self._lookup_trust_boundary(request)
+            self._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.
@@ -347,8 +391,6 @@ def _refresh_trust_boundary(self, request):
                     "Using cached trust boundary due to refresh error: %s", error
                 )
             return
-        else:
-            self._trust_boundary = new_trust_boundary
 
     def _lookup_trust_boundary(self, request):
         """Calls the trust boundary lookup API to refresh the trust boundary cache.
@@ -366,24 +408,13 @@ def _lookup_trust_boundary(self, request):
         """
         from google.oauth2 import _client
 
-        # Verify the trust boundary feature flag is enabled.
-        if not _helpers.get_bool_from_env(
-            environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED, default=False
-        ):
-            # Skip the lookup and return early if it's not explicitly enabled.
-            return None
-
-        # Skip trust boundary flow for non-gdu universe domain.
-        if self.universe_domain != DEFAULT_UNIVERSE_DOMAIN:
-            return None
-
         url = self._build_trust_boundary_lookup_url()
         if not url:
             raise exceptions.InvalidValue("Failed to build trust boundary lookup URL.")
 
         headers = {}
         self.apply(headers)
-        return _client.lookup_trust_boundary(request, url, headers=headers)
+        return _client._lookup_trust_boundary(request, url, headers=headers)
 
     @abc.abstractmethod
     def _build_trust_boundary_lookup_url(self):

google/auth/impersonated_credentials.py

@@ -264,12 +264,7 @@ def __init__(
     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):
+    def _refresh_token(self, request):
         """Updates credentials with a new access_token representing
         the impersonated account.
 

google/oauth2/_client.py

@@ -515,7 +515,7 @@ def refresh_grant(
     return _handle_refresh_grant_response(response_data, refresh_token)
 
 
-def lookup_trust_boundary(request, url, headers=None):
+def _lookup_trust_boundary(request, url, headers=None):
     """Implements the global lookup of a credential trust boundary.
     For the lookup, we send a request to the global lookup endpoint and then
     parse the response. Service account credentials, workload identity
@@ -556,9 +556,7 @@ def lookup_trust_boundary(request, url, headers=None):
     return response_data
 
 
-def _lookup_trust_boundary_request(
-    request, url, can_retry=True, headers=None, **kwargs
-):
+def _lookup_trust_boundary_request(request, url, can_retry=True, headers=None):
     """Makes a request to the trust boundary lookup endpoint.
 
     Args:
@@ -567,13 +565,6 @@ def _lookup_trust_boundary_request(
         url (str): The trust boundary lookup url.
         can_retry (bool): Enable or disable request retry behavior. Defaults to true.
         headers (Optional[Mapping[str, str]]): The headers for the request.
-        kwargs: Additional arguments passed on to the request method. The
-            kwargs will be passed to `requests.request` method, see:
-            https://docs.python-requests.org/en/latest/api/#requests.request.
-            For example, you can use `cert=("cert_pem_path", "key_pem_path")`
-            to set up client side SSL certificate, and use
-            `verify="ca_bundle_path"` to set up the CA certificates for sever
-            side SSL certificate verification.
 
     Returns:
         Mapping[str, str]: The JSON-decoded response data.
@@ -583,18 +574,14 @@ def _lookup_trust_boundary_request(
             an error.
     """
     response_status_ok, response_data, retryable_error = (
-        _lookup_trust_boundary_request_no_throw(
-            request, url, can_retry, headers, **kwargs
-        )
+        _lookup_trust_boundary_request_no_throw(request, url, can_retry, headers)
     )
     if not response_status_ok:
         _handle_error_response(response_data, retryable_error)
     return response_data
 
 
-def _lookup_trust_boundary_request_no_throw(
-    request, url, can_retry=True, headers=None, **kwargs
-):
+def _lookup_trust_boundary_request_no_throw(request, url, can_retry=True, headers=None):
     """Makes a request to the trust boundary lookup endpoint. This
         function doesn't throw on response errors.
 
@@ -604,13 +591,6 @@ def _lookup_trust_boundary_request_no_throw(
         url (str): The trust boundary lookup url.
         can_retry (bool): Enable or disable request retry behavior. Defaults to true.
         headers (Optional[Mapping[str, str]]): The headers for the request.
-        kwargs: Additional arguments passed on to the request method. The
-            kwargs will be passed to `requests.request` method, see:
-            https://docs.python-requests.org/en/latest/api/#requests.request.
-            For example, you can use `cert=("cert_pem_path", "key_pem_path")`
-            to set up client side SSL certificate, and use
-            `verify="ca_bundle_path"` to set up the CA certificates for sever
-            side SSL certificate verification.
 
     Returns:
         Tuple(bool, Mapping[str, str], Optional[bool]): A boolean indicating
@@ -624,7 +604,7 @@ def _lookup_trust_boundary_request_no_throw(
 
     retries = _exponential_backoff.ExponentialBackoff()
     for _ in retries:
-        response = request(method="GET", url=url, headers=headers, **kwargs)
+        response = request(method="GET", url=url, headers=headers)
         response_body = (
             response.data.decode("utf-8")
             if hasattr(response.data, "decode")

google/oauth2/service_account.py

@@ -436,7 +436,7 @@ def _metric_header_for_usage(self):
         return metrics.CRED_TYPE_SA_ASSERTION
 
     @_helpers.copy_docstring(credentials.Credentials)
-    def refresh(self, request):
+    def _refresh_token(self, request):
         if self._always_use_jwt_access and not self._jwt_credentials:
             # If self signed jwt should be used but jwt credential is not
             # created, try to create one with scopes
@@ -461,7 +461,50 @@ def refresh(self, request):
             )
             self.token = access_token
             self.expiry = expiry
-        self._refresh_trust_boundary(request)
+
+    def refresh(self, request):
+        """Refreshes the credential's access token.
+
+        This method is overridden to provide special handling for credentials that
+        use a self-signed JWT and have a trust boundary configured. In this
+        scenario, it first generates a temporary, IAM-specific self-signed JWT
+        to perform the trust boundary lookup, and then generates the final
+        self-signed JWT for the target API.
+
+        For all other cases, it falls back to the standard refresh behavior
+        from the parent class.
+
+        Args:
+            request (google.auth.transport.Request): The object used to make
+                HTTP requests.
+
+        Raises:
+            google.auth.exceptions.RefreshError: If the credentials could
+                not be refreshed.
+        """
+        # Determine if we're going to use a self-signed JWT.
+        use_ssjwt = self._use_self_signed_jwt()
+
+        if use_ssjwt and self._is_trust_boundary_lookup_required():
+            # Special case: self-signed JWT with trust boundary.
+            # 1. Create a temporary self-signed JWT for the IAM API.
+            iam_audience = "https://iamcredentials.{}/".format(self._universe_domain)
+            iam_jwt_creds = jwt.Credentials.from_signing_credentials(self, iam_audience)
+            iam_jwt_creds.refresh(request)
+
+            # 2. Use this JWT to perform the trust boundary lookup.
+            # We temporarily set self.token for the base lookup method.
+            # The base lookup method will call self.apply() which adds the
+            # authorization header.
+            with _helpers.update_property(self, "token", iam_jwt_creds.token.decode()):
+                # This will call _lookup_trust_boundary and set self._trust_boundary
+                self._refresh_trust_boundary(request)
+
+            # 3. Now, refresh the original self-signed JWT for the target API.
+            self._refresh_token(request)
+        else:
+            # For all other cases, use the standard refresh mechanism.
+            super(Credentials, self).refresh(request)
 
     def _create_self_signed_jwt(self, audience):
         """Create a self-signed JWT from the credentials if requirements are met.
@@ -798,6 +841,12 @@ def with_trust_boundary(self, trust_boundary):
         cred._trust_boundary = trust_boundary
         return cred
 
+    def _refresh_token(self, request):
+        """Not used by this class, which overrides refresh() directly."""
+        # This is required to satisfy the abstract base class, but this
+        # class's refresh() method is called directly and does not use this.
+        pass
+
     def _build_trust_boundary_lookup_url(self):
         """Builds and returns the URL for the trust boundary lookup API.