# Copyright 2026 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utilities for Regional Access Boundary management."""
import asyncio
import copy
import datetime
import functools
import inspect
import logging
import threading
from typing import NamedTuple, Optional, TYPE_CHECKING
from google.auth import _helpers
if TYPE_CHECKING: # pragma: NO COVER
import google.auth.credentials
import google.auth.transport
_LOGGER = logging.getLogger(__name__)
# The default lifetime for a cached Regional Access Boundary.
DEFAULT_REGIONAL_ACCESS_BOUNDARY_TTL = datetime.timedelta(hours=6)
# The period of time prior to the boundary's expiration when a background refresh
# is proactively triggered.
REGIONAL_ACCESS_BOUNDARY_REFRESH_THRESHOLD = datetime.timedelta(hours=1)
# The initial cooldown period for a failed Regional Access Boundary lookup.
DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN = datetime.timedelta(minutes=15)
# The maximum cooldown period for a failed Regional Access Boundary lookup.
MAX_REGIONAL_ACCESS_BOUNDARY_COOLDOWN = datetime.timedelta(hours=6)
# The header key used for Regional Access Boundaries.
_REGIONAL_ACCESS_BOUNDARY_HEADER = "x-allowed-locations"
class _RegionalAccessBoundaryData(NamedTuple):
"""Data container for a Regional Access Boundary snapshot.
Attributes:
encoded_locations (Optional[str]): The encoded Regional Access Boundary string.
expiry (Optional[datetime.datetime]): The hard expiration time of the boundary data.
cooldown_expiry (Optional[datetime.datetime]): The time until which further lookups are skipped.
cooldown_duration (datetime.timedelta): The current duration for the exponential cooldown.
"""
encoded_locations: Optional[str]
expiry: Optional[datetime.datetime]
cooldown_expiry: Optional[datetime.datetime]
cooldown_duration: datetime.timedelta
class _RegionalAccessBoundaryManager(object):
"""Manages the Regional Access Boundary state and its background refresh.
The actual data is held in an immutable `_RegionalAccessBoundaryData` object
and is swapped atomically to ensure thread-safe, lock-free reads.
"""
def __init__(self):
self._data = _RegionalAccessBoundaryData(
encoded_locations=None,
expiry=None,
cooldown_expiry=None,
cooldown_duration=DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN,
)
self.refresh_manager = _RegionalAccessBoundaryRefreshManager()
self._update_lock = threading.Lock()
self._use_blocking_regional_access_boundary_lookup = False
def __getstate__(self):
"""Pickle helper that serializes the _update_lock attribute."""
state = self.__dict__.copy()
state["_update_lock"] = None
return state
def __setstate__(self, state):
"""Pickle helper that deserializes the _update_lock attribute."""
self.__dict__.update(state)
self._update_lock = threading.Lock()
def __eq__(self, other):
"""Checks if two managers are equal."""
if not isinstance(other, _RegionalAccessBoundaryManager):
return NotImplemented
return (
self._data == other._data
and self._use_blocking_regional_access_boundary_lookup
== other._use_blocking_regional_access_boundary_lookup
)
def enable_blocking_lookup(self):
"""Enables blocking Regional Access Boundary lookup.
When enabled, the Regional Access Boundary lookup will be performed
synchronously in the calling thread instead of asynchronously in a
background thread.
"""
self._use_blocking_regional_access_boundary_lookup = True
def set_initial_regional_access_boundary(self, encoded_locations=None, expiry=None):
"""Manually sets the regional access boundary to the client provided initial values.
Args:
encoded_locations (Optional[str]): The encoded locations string.
expiry (Optional[datetime.datetime]): The expiry time for the boundary.
If encoded_locations is not provided, expiry is ignored.
"""
if not encoded_locations:
expiry = None
self._data = _RegionalAccessBoundaryData(
encoded_locations=encoded_locations,
expiry=expiry,
cooldown_expiry=None,
cooldown_duration=DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN,
)
def apply_headers(self, headers):
"""Applies the Regional Access Boundary header to the provided dictionary.
If the boundary is valid, the 'x-allowed-locations' header is added
or updated. Otherwise, the header is removed to ensure no stale
data is sent.
Args:
headers (MutableMapping[str, str]): The headers dictionary to update.
"""
rab_data = self._data
if rab_data.encoded_locations and (
rab_data.expiry is not None and _helpers.utcnow() < rab_data.expiry
):
headers[_REGIONAL_ACCESS_BOUNDARY_HEADER] = rab_data.encoded_locations
else:
headers.pop(_REGIONAL_ACCESS_BOUNDARY_HEADER, None)
def _should_refresh(self):
"""Checks if the Regional Access Boundary data needs a refresh and is not in cooldown.
Returns:
bool: True if a refresh is required, False otherwise.
"""
rab_data = self._data
# Don't start a new refresh if the Regional Access Boundary info is still fresh.
if (
rab_data.encoded_locations
and rab_data.expiry
and _helpers.utcnow()
< (rab_data.expiry - REGIONAL_ACCESS_BOUNDARY_REFRESH_THRESHOLD)
):
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.
if self._use_blocking_regional_access_boundary_lookup:
self.start_blocking_refresh(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.
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.transport.Request): The object used to make HTTP requests.
"""
# Async credentials do not support blocking lookups.
if inspect.iscoroutinefunction(credentials._lookup_regional_access_boundary):
_LOGGER.debug(
"Blocking Regional Access Boundary lookup is not supported for async credentials."
)
self.process_regional_access_boundary_info(None)
return
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 = (
credentials._lookup_regional_access_boundary(request, fail_fast=True)
)
except Exception as e:
_LOGGER.debug(
"Blocking 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)
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:
_LOGGER.debug(
"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.
Args:
regional_access_boundary_info (Optional[Mapping[str, str]]): The regional access
boundary info to process.
"""
with self._update_lock:
# Capture the current state before calculating updates.
current_data = self._data
if regional_access_boundary_info:
# On success, update the boundary and its expiry, and clear any cooldown.
encoded_locations = regional_access_boundary_info.get(
"encodedLocations"
)
updated_data = _RegionalAccessBoundaryData(
encoded_locations=encoded_locations,
expiry=_helpers.utcnow() + DEFAULT_REGIONAL_ACCESS_BOUNDARY_TTL,
cooldown_expiry=None,
cooldown_duration=DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN,
)
_LOGGER.debug("Regional Access Boundary lookup successful.")
else:
# On failure, calculate cooldown and update state.
_LOGGER.debug(
"Regional Access Boundary lookup failed. Entering cooldown."
)
next_cooldown_expiry = (
_helpers.utcnow() + current_data.cooldown_duration
)
next_cooldown_duration = min(
current_data.cooldown_duration * 2,
MAX_REGIONAL_ACCESS_BOUNDARY_COOLDOWN,
)
# If the refresh failed, we keep reusing the existing data unless
# it has reached its hard expiration time.
if current_data.expiry and _helpers.utcnow() > current_data.expiry:
next_encoded_locations = None
next_expiry = None
else:
next_encoded_locations = current_data.encoded_locations
next_expiry = current_data.expiry
updated_data = _RegionalAccessBoundaryData(
encoded_locations=next_encoded_locations,
expiry=next_expiry,
cooldown_expiry=next_cooldown_expiry,
cooldown_duration=next_cooldown_duration,
)
# Perform the atomic swap of the state object.
self._data = updated_data
class _RegionalAccessBoundaryRefreshThread(threading.Thread):
"""Thread for background refreshing of the Regional Access Boundary."""
def __init__(
self,
credentials: "google.auth.credentials.CredentialsWithRegionalAccessBoundary", # noqa: F821
request: "google.auth.transport.Request", # noqa: F821
rab_manager: "_RegionalAccessBoundaryManager",
):
super().__init__()
self.daemon = True
self._credentials = credentials
self._request = request
self._rab_manager = rab_manager
def run(self):
"""
Performs the Regional Access Boundary lookup and updates the state.
This method is run in a separate thread. It delegates the actual lookup
to the credentials object's `_lookup_regional_access_boundary` method.
Based on the lookup's outcome (success or complete failure after retries),
it updates the cached Regional Access Boundary information,
its expiry, its cooldown expiry, and its exponential cooldown duration.
"""
# Catch exceptions (e.g., from the underlying transport) to prevent the
# background thread from crashing. This ensures we can gracefully enter
# an exponential cooldown state on failure.
try:
regional_access_boundary_info = (
self._credentials._lookup_regional_access_boundary(self._request)
)
except Exception as e:
_LOGGER.debug(
"Asynchronous Regional Access Boundary lookup raised an exception: %s",
e,
exc_info=True,
)
regional_access_boundary_info = None
self._rab_manager.process_regional_access_boundary_info(
regional_access_boundary_info
)
class _RegionalAccessBoundaryRefreshManager(object):
"""Manages a thread for background refreshing of the Regional Access Boundary."""
def __init__(self):
self._lock = threading.Lock()
self._worker = None
def __getstate__(self):
"""Pickle helper that serializes the _lock and _worker attributes."""
state = self.__dict__.copy()
state["_lock"] = None
state["_worker"] = None
return state
def __setstate__(self, state):
"""Pickle helper that deserializes the _lock and _worker attributes."""
self.__dict__.update(state)
self._lock = threading.Lock()
self._worker = None
def start_refresh(self, credentials, request, rab_manager):
"""
Starts a background thread to refresh the Regional Access Boundary if one is not already running.
Args:
credentials (CredentialsWithRegionalAccessBoundary): The credentials
to refresh.
request (google.auth.transport.Request): The object used to make
HTTP requests.
rab_manager (_RegionalAccessBoundaryManager): The manager container to update.
"""
with self._lock:
if self._worker and self._worker.is_alive():
# A refresh is already in progress.
return
try:
copied_request = copy.deepcopy(request)
except Exception as e:
_LOGGER.debug(
"Could not deepcopy transport for background RAB refresh. "
"Skipping background refresh to avoid thread safety issues. "
"Exception: %s",
e,
)
return
self._worker = _RegionalAccessBoundaryRefreshThread(
credentials, copied_request, rab_manager
)
self._worker.start()
def _prepare_async_lookup_callable(request):
"""Unwraps a request callable, clones the transport, and returns the new callable.
Args:
request: The original request callable (e.g. functools.partial or raw Request).
Returns:
Tuple[Callable, Any, bool]: A tuple containing the new lookup callable, the
underlying request object, and a boolean indicating if it was cloned.
"""
is_partial = isinstance(request, functools.partial)
base_callable = request.func if is_partial else request
if not hasattr(base_callable, "_clone"):
return request, base_callable, False
cloned_callable = base_callable._clone()
is_cloned = cloned_callable is not base_callable
if is_partial:
new_request = functools.partial(
cloned_callable, *request.args, **request.keywords
)
else:
new_request = cloned_callable
return new_request, cloned_callable, is_cloned
async def _close_cloned_request(lookup_request, is_cloned):
"""Safely closes the underlying cloned request transport, if applicable.
Args:
lookup_request (Any): The request object/transport to close.
is_cloned (bool): Whether the request was actually cloned.
"""
if not is_cloned or not hasattr(lookup_request, "close"):
return
is_async = False
try:
maybe_coro = lookup_request.close()
if is_async := inspect.isawaitable(maybe_coro):
await maybe_coro
except Exception as e:
adapter_type = " asynchronous " if is_async else " "
_LOGGER.debug(
"Failed to cleanly close cloned%srequest transport: %s",
adapter_type,
e,
exc_info=True,
)
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
try:
(
lookup_callable,
lookup_request,
is_cloned,
) = _prepare_async_lookup_callable(request)
except Exception as e:
_LOGGER.debug(
"Synchronous cloning of request for Regional Access Boundary lookup failed: %s",
e,
exc_info=True,
)
rab_manager.process_regional_access_boundary_info(None)
return
async def _worker():
try:
regional_access_boundary_info = (
await credentials._lookup_regional_access_boundary(
lookup_callable
)
)
except Exception as e:
_LOGGER.debug(
"Asynchronous Regional Access Boundary lookup raised an exception: %s",
e,
exc_info=True,
)
regional_access_boundary_info = None
finally:
await _close_cloned_request(lookup_request, is_cloned)
rab_manager.process_regional_access_boundary_info(
regional_access_boundary_info
)
coro = _worker()
try:
self._worker_task = asyncio.create_task(coro)
except Exception:
# Clean up cloned request if task creation fails
coro.close()
try:
asyncio.get_running_loop().create_task(
_close_cloned_request(lookup_request, is_cloned)
)
except RuntimeError:
pass
rab_manager.process_regional_access_boundary_info(None)
raise
def _get_domain() -> str:
"""Dynamically determines the domain for IAM credentials based on active mTLS configuration.
Returns:
str: The dynamic domain string.
"""
from google.auth.transport import _mtls_helper
if (
hasattr(_mtls_helper, "check_use_client_cert")
and _mtls_helper.check_use_client_cert()
):
return f"iamcredentials.mtls.{_helpers.DEFAULT_UNIVERSE_DOMAIN}"
else:
return f"iamcredentials.{_helpers.DEFAULT_UNIVERSE_DOMAIN}"
def get_service_account_rab_endpoint(service_account_email: str) -> str:
"""Builds the Regional Access Boundary lookup URL for service accounts.
Args:
service_account_email: The service account email.
Returns:
str: The complete lookup URL.
"""
return f"https://{_get_domain()}/v1/projects/-/serviceAccounts/{service_account_email}/allowedLocations"
def get_workforce_pool_rab_endpoint(pool_id: str) -> str:
"""Builds the Regional Access Boundary lookup URL for workforce pools.
Args:
pool_id: The workforce pool ID.
Returns:
str: The complete lookup URL.
"""
return f"https://{_get_domain()}/v1/locations/global/workforcePools/{pool_id}/allowedLocations"
def get_workload_identity_pool_rab_endpoint(project_number: str, pool_id: str) -> str:
"""Builds the Regional Access Boundary lookup URL for workload identity pools.
Args:
project_number: The Google Cloud project number.
pool_id: The workload identity pool ID.
Returns:
str: The complete lookup URL.
"""
return f"https://{_get_domain()}/v1/projects/{project_number}/locations/global/workloadIdentityPools/{pool_id}/allowedLocations"