feat: [aiohttp] Add mTLS reconfiguration logic when certificate mismatch for existing credentials & Agent Identity workloads

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

@@ -15,6 +15,8 @@
 import asyncio
 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 warnings
@@ -26,6 +28,8 @@
 from google.auth.exceptions import TimeoutError
 import google.auth.transport._mtls_helper
 
+_LOGGER = logging.getLogger(__name__)
+
 if TYPE_CHECKING:  # pragma: NO COVER
     import aiohttp
     from aiohttp import ClientTimeout  # type: ignore
@@ -310,6 +314,45 @@ async def request(
                         url, method, data, headers, actual_timeout, **kwargs
                     )
                 )
+                if response.status_code == http_client.UNAUTHORIZED:
+                    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
+                        )
+                        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(
+                                    lambda: (call_cert_bytes, call_key_bytes)
+                                )
+                                continue
+                            except Exception as e:
+                                _LOGGER.warning(
+                                    "Failed to reconfigure mTLS channel: %s. Proceeding with original response.",
+                                    e
+                                )
+                        else:
+                             _LOGGER.info(
+                                "Skipping reconfiguration of mTLS channel because the client"
+                                " certificate has not changed."
+                            )
+                    except Exception as e:
+                        _LOGGER.warning(
+                            "Failed to check client certificate parameters: %s. Proceeding with original response.",
+                            e,
+                        )
+
                 if response.status_code not in transport.DEFAULT_RETRYABLE_STATUS_CODES:
                     break
         return response

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

@@ -344,3 +344,78 @@ async def test_configure_mtls_channel_close_exception_does_not_abort(self):
             assert session._is_mtls is True
             assert session._cached_cert == b"fake_cert_data"
             await session.close()
+
+    @pytest.mark.asyncio
+    async def test_cert_rotation_failure_logs(self):
+        mock_creds = mock.AsyncMock(spec=credentials.Credentials)
+        mock_creds.before_request = mock.AsyncMock(return_value=None)
+
+        mock_auth_req = mock.AsyncMock()
+        mock_resp = mock.Mock()
+        import http.client as http_client
+        mock_resp.status_code = http_client.UNAUTHORIZED
+        mock_auth_req.return_value = mock_resp
+
+        session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req)
+        session._is_mtls = True
+        session._cached_cert = b"old_cert"
+
+        new_cert = b"new_cert"
+        new_key = b"new_key"
+
+        with mock.patch("google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response") as mock_check, \
+             mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf:
+
+            mock_check.return_value = (new_cert, new_key, b"old_fp", b"new_fp")
+            mock_conf.side_effect = Exception("Failed to reconfigure")
+
+            resp = await session.request("GET", "http://example.com")
+            assert resp == mock_resp
+
+            mock_check.assert_called_once()
+            mock_conf.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_cert_rotation_check_params_fails(self):
+        mock_creds = mock.AsyncMock(spec=credentials.Credentials)
+        mock_auth_req = mock.AsyncMock()
+        mock_resp = mock.Mock()
+        import http.client as http_client
+        mock_resp.status_code = http_client.UNAUTHORIZED
+        mock_auth_req.return_value = mock_resp
+
+        session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req)
+        session._is_mtls = True
+        session._cached_cert = b"cached_cert"
+
+        with mock.patch(
+            "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response",
+            side_effect=Exception("check_params failed"),
+        ) as mock_check_params:
+            resp = await session.request("GET", "http://example.com")
+            assert resp == mock_resp
+            mock_check_params.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_no_cert_rotation_when_cert_match_and_mTLS_enabled(self):
+        mock_creds = mock.AsyncMock(spec=credentials.Credentials)
+        mock_auth_req = mock.AsyncMock()
+        mock_resp = mock.Mock()
+        import http.client as http_client
+        mock_resp.status_code = http_client.UNAUTHORIZED
+        mock_auth_req.return_value = mock_resp
+
+        session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=mock_auth_req)
+        session._is_mtls = True
+        session._cached_cert = b"old_cert"
+
+        with mock.patch(
+            "google.auth.transport._mtls_helper.check_parameters_for_unauthorized_response",
+        ) as mock_check, mock.patch.object(session, "configure_mtls_channel", new_callable=mock.AsyncMock) as mock_conf:
+            # same fingerprint, so no call to configure_mtls_channel
+            mock_check.return_value = (b"new_cert", b"new_key", b"same_fp", b"same_fp")
+
+            await session.request("GET", "http://example.com")
+
+            mock_check.assert_called_once()
+            mock_conf.assert_not_called()