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

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/_jwt_async.py

@@ -91,7 +91,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 +144,15 @@ 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)
+        from google.auth import _regional_access_boundary_utils
+
+        self._rab_manager.refresh_manager = (
+            _regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
+        )
+
 
 class OnDemandCredentials(
     jwt.OnDemandCredentials, _credentials_async.Signing, _credentials_async.Credentials

packages/google-auth/google/auth/_regional_access_boundary_utils.py

@@ -14,6 +14,7 @@
 
 """Utilities for Regional Access Boundary management."""
 
+import asyncio
 import copy
 import datetime
 import functools
@@ -170,12 +171,11 @@ def apply_headers(self, headers):
         else:
             headers.pop(_REGIONAL_ACCESS_BOUNDARY_HEADER, None)
 
-    def maybe_start_refresh(self, credentials, request):
-        """Starts a background thread to refresh the Regional Access Boundary if needed.
+    def _should_refresh(self):
+        """Checks if the Regional Access Boundary data needs a refresh and is not in cooldown.
 
-        Args:
-            credentials (google.auth.credentials.Credentials): The credentials to refresh.
-            request (google.auth.transport.Request): The object used to make HTTP requests.
+        Returns:
+            bool: True if a refresh is required, False otherwise.
         """
         rab_data = self._data
 
@@ -186,10 +186,22 @@ def maybe_start_refresh(self, credentials, request):
             and _helpers.utcnow()
             < (rab_data.expiry - REGIONAL_ACCESS_BOUNDARY_REFRESH_THRESHOLD)
         ):
-            return
+            return False
 
         # Don't start a new refresh if the cooldown is still in effect.
         if rab_data.cooldown_expiry and _helpers.utcnow() < rab_data.cooldown_expiry:
+            return False
+
+        return True
+
+    def maybe_start_refresh(self, credentials, request):
+        """Starts a background thread to refresh the Regional Access Boundary if needed.
+
+        Args:
+            credentials (google.auth.credentials.Credentials): The credentials to refresh.
+            request (google.auth.transport.Request): The object used to make HTTP requests.
+        """
+        if not self._should_refresh():
             return
 
         # If all checks pass, start the background refresh.
@@ -198,6 +210,22 @@ def maybe_start_refresh(self, credentials, request):
         else:
             self.refresh_manager.start_refresh(credentials, request, self)
 
+    async def maybe_start_refresh_async(self, credentials, request):
+        """Starts a background refresh or performs a blocking refresh asynchronously.
+
+        Args:
+            credentials (google.auth.credentials.Credentials): The credentials to refresh.
+            request (google.auth.aio.transport.Request): The object used to make HTTP requests.
+        """
+        if not self._should_refresh():
+            return
+
+        # If all checks pass, start the refresh.
+        if self._use_blocking_regional_access_boundary_lookup:
+            await self.start_blocking_refresh_async(credentials, request)
+        else:
+            self.refresh_manager.start_refresh(credentials, request, self)
+
     def start_blocking_refresh(self, credentials, request):
         """Initiates a blocking lookup of the Regional Access Boundary.
 
@@ -227,6 +255,37 @@ def start_blocking_refresh(self, credentials, request):
 
         self.process_regional_access_boundary_info(regional_access_boundary_info)
 
+    async def start_blocking_refresh_async(self, credentials, request):
+        """Initiates a blocking lookup of the Regional Access Boundary asynchronously.
+
+        If the lookup raises an exception, it is caught and logged as a warning,
+        and the lookup is treated as a failure (entering cooldown). Exceptions
+        are not propagated to the caller.
+
+        Args:
+            credentials (google.auth.credentials.Credentials): The credentials to refresh.
+            request (google.auth.aio.transport.Request): The object used to make HTTP requests.
+        """
+        try:
+            # The fail_fast parameter is set to True to ensure we don't block the calling
+            # thread for too long. This will do two things: 1) set a timeout to 3s
+            # instead of the default 120s and 2) ensure we do not retry at all
+            regional_access_boundary_info = (
+                await credentials._lookup_regional_access_boundary(
+                    request, fail_fast=True
+                )
+            )
+        except Exception as e:
+            if _helpers.is_logging_enabled(_LOGGER):
+                _LOGGER.warning(
+                    "Regional Access Boundary lookup raised an exception: %s",
+                    e,
+                    exc_info=True,
+                )
+            regional_access_boundary_info = None
+
+        self.process_regional_access_boundary_info(regional_access_boundary_info)
+
     def process_regional_access_boundary_info(self, regional_access_boundary_info):
         """Processes the regional access boundary info and updates the state.
 
@@ -384,3 +443,61 @@ def start_refresh(self, credentials, request, rab_manager):
                 credentials, copied_request, rab_manager
             )
             self._worker.start()
+
+
+class _AsyncRegionalAccessBoundaryRefreshManager(object):
+    """Manages a task for background refreshing of the Regional Access Boundary in async flows."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._worker_task = None
+
+    def __getstate__(self):
+        """Pickle helper that excludes the un-picklable _lock and _worker_task attributes from serialization."""
+        state = self.__dict__.copy()
+        state["_lock"] = None
+        state["_worker_task"] = None
+        return state
+
+    def __setstate__(self, state):
+        """Pickle helper that restores state and re-initializes the _lock and _worker_task attributes."""
+        self.__dict__.update(state)
+        self._lock = threading.Lock()
+        self._worker_task = None
+
+    def start_refresh(self, credentials, request, rab_manager):
+        """
+        Starts a background task to refresh the Regional Access Boundary if one is not already running.
+
+        Args:
+            credentials (CredentialsWithRegionalAccessBoundary): The credentials
+                to refresh.
+            request (google.auth.aio.transport.Request): The object used to make
+                HTTP requests.
+            rab_manager (_RegionalAccessBoundaryManager): The manager container to update.
+        """
+        with self._lock:
+            if self._worker_task and not self._worker_task.done():
+                # A refresh is already in progress.
+                return
+
+            async def _worker():
+                try:
+                    # credentials._lookup_regional_access_boundary should be async in the async creds class
+                    regional_access_boundary_info = (
+                        await credentials._lookup_regional_access_boundary(request)
+                    )
+                except Exception as e:
+                    if _helpers.is_logging_enabled(_LOGGER):
+                        _LOGGER.warning(
+                            "Asynchronous Regional Access Boundary lookup raised an exception: %s",
+                            e,
+                            exc_info=True,
+                        )
+                    regional_access_boundary_info = None
+
+                rab_manager.process_regional_access_boundary_info(
+                    regional_access_boundary_info
+                )
+
+            self._worker_task = asyncio.create_task(_worker())