feat(auth): implement regional access boundary support for standalone JWT and async service accounts (#17025)

This PR implements the following changes:

- Add RAB support to async service account and jwt credential types, by providing
async manager and fetching methods.
- Update unit tests to accept both mtls and standard lookup endpoint
urls.
- Refactor before_request to use a _after_refresh hook so we don't have
to override the method.
- Add RAb support for self signed jwt flow through jwt.py
- some small enhancements for test coverage and backward compatibility

Author: nbayati

packages/google-auth/google/auth/_credentials_async.py

@@ -18,6 +18,7 @@
 import abc
 import inspect
 
+from google.auth import _regional_access_boundary_utils
 from google.auth import credentials
 
 
@@ -64,8 +65,28 @@ async def before_request(self, request, method, url, headers):
                 await self.refresh(request)
             else:
                 self.refresh(request)
+
+        if inspect.iscoroutinefunction(self._after_refresh):
+            await self._after_refresh(request, method, url, headers)
+        else:
+            self._after_refresh(request, method, url, headers)
+
         self.apply(headers)
 
+    def _after_refresh(self, request, method, url, headers):
+        """Hook for subclasses to perform actions after refresh but before
+        applying credentials to headers.
+
+        Args:
+            request (google.auth.transport.Request): The object used to make
+                HTTP requests.
+            method (str): The request's HTTP method or the RPC method being
+                invoked.
+            url (str): The request's URI or the RPC service's URI.
+            headers (Mapping[str, str]): The request's headers.
+        """
+        pass
+
 
 class CredentialsWithQuotaProject(credentials.CredentialsWithQuotaProject):
     """Abstract base for credentials supporting ``with_quota_project`` factory"""
@@ -169,3 +190,74 @@ def with_scopes_if_required(credentials, scopes):
 
 class Signing(credentials.Signing, metaclass=abc.ABCMeta):
     """Interface for credentials that can cryptographically sign messages."""
+
+
+class CredentialsWithRegionalAccessBoundary(
+    Credentials, credentials.CredentialsWithRegionalAccessBoundary
+):
+    """Async base for credentials supporting regional access boundary configuration."""
+
+    def __init__(self):
+        super().__init__()
+        self._rab_manager.refresh_manager = (
+            _regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
+        )
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        self._rab_manager.refresh_manager = (
+            _regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
+        )
+
+    async def _after_refresh(self, request, method, url, headers):
+        """Triggers the Regional Access Boundary lookup asynchronously if necessary."""
+        await self._maybe_start_regional_access_boundary_refresh_async(request, url)
+
+    async def _maybe_start_regional_access_boundary_refresh_async(self, request, url):
+        """Starts a background refresh or performs a blocking refresh asynchronously.
+
+        Args:
+            request (google.auth.aio.transport.Request): The object used to make
+                HTTP requests.
+            url (str): The URL of the request.
+        """
+        # Do not perform a lookup if the request is for a regional endpoint.
+        if self._is_regional_endpoint(url):
+            return
+
+        # A refresh is only needed if the feature is enabled.
+        if not self._is_regional_access_boundary_lookup_required():
+            return
+
+        # Trigger background or blocking refresh if needed.
+        await self._rab_manager.maybe_start_refresh_async(self, request)
+
+    async def _lookup_regional_access_boundary(self, request, fail_fast=False):
+        """Calls the Regional Access Boundary lookup API asynchronously.
+
+        Args:
+            request (google.auth.aio.transport.Request): The object used to make
+                HTTP requests.
+            fail_fast (bool): Whether the lookup should fail fast (short timeout, no retries).
+
+        Returns:
+            Optional[Dict[str, str]]: The Regional Access Boundary information
+                returned by the lookup API, or None if the lookup failed.
+        """
+        url_builder = self._build_regional_access_boundary_lookup_url
+        if inspect.iscoroutinefunction(url_builder):
+            url = await url_builder(request=request)
+        else:
+            url = url_builder(request=request)
+
+        if not url:
+            return None
+
+        headers = {}
+        self._apply(headers)
+
+        from google.oauth2 import _client_async
+
+        return await _client_async._lookup_regional_access_boundary(
+            request, url, headers=headers, fail_fast=fail_fast
+        )

packages/google-auth/google/auth/_helpers.py

@@ -28,6 +28,8 @@
 from google.auth import exceptions
 
 
+DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
+
 # _BASE_LOGGER_NAME is the base logger for all google-based loggers.
 _BASE_LOGGER_NAME = "google"
 

packages/google-auth/google/auth/_jwt_async.py

@@ -44,6 +44,8 @@
 """
 
 from google.auth import _credentials_async
+from google.auth import _helpers
+from google.auth import _regional_access_boundary_utils
 from google.auth import jwt
 
 
@@ -91,7 +93,9 @@ def decode(token, certs=None, verify=True, audience=None):
 
 
 class Credentials(
-    jwt.Credentials, _credentials_async.Signing, _credentials_async.Credentials
+    jwt.Credentials,
+    _credentials_async.Signing,
+    _credentials_async.CredentialsWithRegionalAccessBoundary,
 ):
     """Credentials that use a JWT as the bearer token.
 
@@ -142,6 +146,14 @@ class Credentials(
         new_credentials = credentials.with_claims(audience=new_audience)
     """
 
+    def __setstate__(self, state):
+        """Restores the credential state and ensures the async refresh manager is attached."""
+        super().__setstate__(state)
+
+        self._rab_manager.refresh_manager = (
+            _regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
+        )
+
 
 class OnDemandCredentials(
     jwt.OnDemandCredentials, _credentials_async.Signing, _credentials_async.Credentials
@@ -162,3 +174,7 @@ class OnDemandCredentials(
 
     .. _grpc: http://www.grpc.io/
     """
+
+    @_helpers.copy_docstring(jwt.OnDemandCredentials)
+    async def before_request(self, request, method, url, headers):
+        super(OnDemandCredentials, self).before_request(request, method, url, headers)