feat(auth): Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads

packages/google-auth/google/auth/aio/transport/sessions.py

@@ -13,10 +13,14 @@
 # limitations under the License.
 
 import asyncio
+import collections.abc
 from contextlib import asynccontextmanager
 import functools
+import http.client as http_client
+import logging
 import time
 from typing import Mapping, Optional, TYPE_CHECKING, Union
+import urllib.parse
 import warnings
 
 from google.auth import _exponential_backoff, exceptions
@@ -37,6 +41,9 @@
     except (ImportError, AttributeError):
         ClientTimeout = None
 
+_LOGGER = logging.getLogger(__name__)
+MTLS_URL_PREFIXES = ["mtls.googleapis.com", "mtls.sandbox.googleapis.com"]
+
 
 # Tracks the internal aiohttp installation and usage
 try:
@@ -143,11 +150,14 @@ def __init__(
         self._is_mtls = False
         self._mtls_init_task = None
         self._cached_cert = None
+        self._client_cert_callback = None
+        self._old_auth_requests = []  # type: list
         if _auth_request is None:
             raise exceptions.TransportError(
                 "`auth_request` must either be configured or the external package `aiohttp` must be installed to use the default value."
             )
         self._auth_request = _auth_request
+        self._mtls_rotation_lock = None  # type: Optional[asyncio.Lock]
 
     async def configure_mtls_channel(self, client_cert_callback=None):
         """Configure the client certificate and key for SSL connection.
@@ -175,6 +185,7 @@ async def configure_mtls_channel(self, client_cert_callback=None):
                 creation failed for any reason.
         """
         if self._mtls_init_task is None:
+            self._client_cert_callback = client_cert_callback
 
             async def _do_configure():
                 # Run the blocking check in an executor
@@ -204,12 +215,8 @@ async def _do_configure():
 
                             old_auth_request = self._auth_request
                             self._auth_request = AiohttpRequest(session=new_session)
+                            self._old_auth_requests.append(old_auth_request)
 
-                            try:
-                                await old_auth_request.close()
-                            except Exception:
-                                # Suppress so it doesn't abort the mTLS configuration
-                                pass
                         else:
                             is_mtls = False
                             warnings.warn(
@@ -277,7 +284,10 @@ async def request(
                 google.auth.exceptions.TimeoutError: If the method does not complete within
                 the configured `max_allowed_time` or the request exceeds the configured
                 `timeout`.
+                google.auth.exceptions.MutualTLSChannelError: If mutual TLS
+                channel reconfiguration fails for any reason during certificate rotation.
         """
+        _auth_retry_count = kwargs.pop("_auth_retry_count", 0)
         if self._mtls_init_task:
             try:
                 await self._mtls_init_task
@@ -290,6 +300,7 @@ async def request(
         )
         if headers is None:
             headers = {}
+        start_time = time.monotonic()
         async with timeout_guard(max_allowed_time) as with_timeout:
             await with_timeout(
                 # Note: before_request will attempt to refresh credentials if expired.
@@ -310,8 +321,125 @@ async def request(
                         url, method, data, headers, actual_timeout, **kwargs
                     )
                 )
+
                 if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES:
                     break
+
+        if response.status_code == http_client.UNAUTHORIZED:
+            if _auth_retry_count < 2:
+                is_streaming = (
+                    data is not None
+                    and isinstance(
+                        data, (collections.abc.Iterator, collections.abc.AsyncIterable)
+                    )
+                    or hasattr(data, "read")
+                )
+                is_mtls_endpoint = False
+                if getattr(self, "is_mtls", False):
+                    hostname = urllib.parse.urlsplit(url).hostname
+                    if hostname:
+                        is_mtls_endpoint = any(
+                            hostname == prefix or hostname.endswith("." + prefix)
+                            for prefix in MTLS_URL_PREFIXES
+                        )
+                    # Snapshot the stale certificate state BEFORE acquiring the lock.
+                    # This represents the cert that caused the 401 rejection.
+                    if is_mtls_endpoint:
+                        stale_cert = self._cached_cert
+
+                        if self._mtls_rotation_lock is None:
+                            self._mtls_rotation_lock = asyncio.Lock()
+
+                        async with self._mtls_rotation_lock:
+                            # Check Did another coroutine already reconfigure mTLS
+                            if self._cached_cert != stale_cert:
+                                pass
+                            else:
+                                try:
+                                    (
+                                        call_cert_bytes,
+                                        call_key_bytes,
+                                        cached_fingerprint,
+                                        current_cert_fingerprint,
+                                    ) = await mtls._run_in_executor(
+                                        google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response,
+                                        self._cached_cert,
+                                        self._client_cert_callback,
+                                    )
+                                except Exception as e:
+                                    _LOGGER.warning(
+                                        "Failed to check client certificate parameters: %s. Proceeding with original response.",
+                                        e,
+                                    )
+                                    return response
+                                else:
+                                    if cached_fingerprint != current_cert_fingerprint:
+                                        try:
+                                            _LOGGER.info(
+                                                "Client certificate has changed, reconfiguring mTLS "
+                                                "channel."
+                                            )
+                                            if (
+                                                self._mtls_init_task
+                                                and self._mtls_init_task.done()
+                                            ):
+                                                self._mtls_init_task = None
+                                            await self.configure_mtls_channel(
+                                                self._client_cert_callback
+                                            )
+                                        except Exception as e:
+                                            _LOGGER.error(
+                                                "Failed to reconfigure mTLS channel: %s",
+                                                e,
+                                            )
+                                            if hasattr(response, "close"):
+                                                if asyncio.iscoroutinefunction(
+                                                    response.close
+                                                ):
+                                                    await response.close()
+                                                else:
+                                                    response.close()
+                                            raise exceptions.MutualTLSChannelError(
+                                                "Failed to reconfigure mTLS channel"
+                                            ) from e
+                                    else:
+                                        _LOGGER.info(
+                                            "Skipping reconfiguration of mTLS channel because the client"
+                                            " certificate has not changed."
+                                        )
+                if is_streaming:
+                    return response
+                try:
+                    await self._credentials.refresh(self._auth_request)
+                except (
+                    exceptions.RefreshError,
+                    getattr(exceptions, "InvalidOperation", Exception),
+                ) as e:
+                    _LOGGER.debug(
+                        "Credential refresh failed, returning 401 response. Error: %s",
+                        e,
+                    )
+                    return response
+
+                if hasattr(response, "close"):
+                    if asyncio.iscoroutinefunction(response.close):
+                        await response.close()
+                    else:
+                        response.close()
+
+                kwargs["_auth_retry_count"] = _auth_retry_count + 1
+                elapsed_time = time.monotonic() - start_time
+                remaining_time = max(0.0, max_allowed_time - elapsed_time)
+                return await self.request(
+                    method,
+                    url,
+                    data=data,
+                    headers=headers,
+                    max_allowed_time=remaining_time,
+                    timeout=timeout,
+                    total_attempts=total_attempts,
+                    **kwargs,
+                )
         return response
 
     @functools.wraps(request)
@@ -595,3 +723,10 @@ async def close(self) -> None:
             except asyncio.CancelledError:
                 pass
         await self._auth_request.close()
+
+        for old_request in self._old_auth_requests:
+            try:
+                await old_request.close()
+            except Exception:
+                pass
+        self._old_auth_requests.clear()

packages/google-auth/google/auth/transport/_mtls_helper.py

@@ -808,19 +808,24 @@ def check_use_client_cert():
     return False
 
 
-def check_parameters_for_unauthorized_response(cached_cert):
+def check_parameters_for_unauthorized_response(cached_cert, client_cert_callback=None):
     """Returns the cached and current cert fingerprint for reconfiguring mTLS.
 
     Args:
         cached_cert(bytes): The cached client certificate.
+        client_cert_callback(Optional[Callable[[], (bytes, bytes)]]):
+            The optional callback that returns the client certificate and private key bytes.
 
     Returns:
         bytes: The client callback cert bytes.
         bytes: The client callback key bytes.
         str: The base64-encoded SHA256 cached fingerprint.
         str: The base64-encoded SHA256 current cert fingerprint.
     """
-    call_cert_bytes, call_key_bytes = call_client_cert_callback()
+    if client_cert_callback:
+        call_cert_bytes, call_key_bytes = client_cert_callback()
+    else:
+        call_cert_bytes, call_key_bytes = call_client_cert_callback()
     cert_obj = _agent_identity_utils.parse_certificate(call_cert_bytes)
     current_cert_fingerprint = _agent_identity_utils.calculate_certificate_fingerprint(
         cert_obj

packages/google-auth/tests/transport/aio/test_sessions.py

@@ -105,7 +105,9 @@ async def test_timeout_with_simple_async_task_within_bounds(
         self, simple_async_task
     ):
         task = False
-        with patch("time.monotonic", side_effect=[0, 0.25, 0.75]):
+        with patch(
+            "time.monotonic", side_effect=lambda it=iter([0, 0.25]): next(it, 0.75)
+        ):
             with patch("asyncio.wait_for", lambda coro, _: coro):
                 async with self.make_timeout_guard(
                     timeout=self.default_timeout
@@ -255,7 +257,7 @@ async def test_request_raises_transport_error(self):
     async def test_request_max_allowed_time_exceeded_error(self):
         auth_request = MockRequest(side_effect=TransportError)
         authed_session = sessions.AsyncAuthorizedSession(self.credentials, auth_request)
-        with patch("time.monotonic", side_effect=[0, 1, 1]):
+        with patch("time.monotonic", side_effect=[0, 0] + [2] * 10):
             with pytest.raises(TimeoutError):
                 await authed_session.request("GET", self.TEST_URL, max_allowed_time=1)