diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index d88162667bda..c4ef00528d0a 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/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 diff --git a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py index b68766ca5b5d..0b0a0ec87229 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/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()