diff --git a/packages/google-auth/google/auth/aio/transport/mtls.py b/packages/google-auth/google/auth/aio/transport/mtls.py index a7d1baf7355d..267a1e401ff0 100644 --- a/packages/google-auth/google/auth/aio/transport/mtls.py +++ b/packages/google-auth/google/auth/aio/transport/mtls.py @@ -22,7 +22,7 @@ import ssl from typing import Optional -from google.auth import exceptions +from google.auth import _agent_identity_utils, exceptions from google.auth.transport._mtls_helper import secure_cert_key_paths import google.auth.transport.mtls @@ -177,3 +177,44 @@ async def get_client_cert_and_key(client_cert_callback=None): has_cert, cert, key, _ = await get_client_ssl_credentials() return has_cert, cert, key + + +async def check_parameters_for_unauthorized_response( + cached_cert, client_cert_callback=None +): + """Async helper to retrieve certs and compute fingerprints for mTLS rotation. + + Args: + cached_cert (bytes): The cached client certificate. + client_cert_callback (Optional[Callable[[], (bytes, bytes)]]): An + optional callback which returns client certificate bytes and private + key bytes both in PEM format. + + Returns: + Tuple[Optional[bytes], Optional[bytes], Optional[str], Optional[str]]: + call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint. + Returns (None, None, None, None) if mTLS is disabled or no client certificate is present. + """ + is_mtls, call_cert_bytes, call_key_bytes = await get_client_cert_and_key( + client_cert_callback + ) + if not is_mtls or not call_cert_bytes: + return None, None, None, None + + def _fetch_fingerprints(): + cert_obj = _agent_identity_utils.parse_certificate(call_cert_bytes) + current_fingerprint = _agent_identity_utils.calculate_certificate_fingerprint( + cert_obj + ) + if cached_cert: + cached_fingerprint = _agent_identity_utils.get_cached_cert_fingerprint( + cached_cert + ) + else: + cached_fingerprint = None + return cached_fingerprint, current_fingerprint + + cached_fingerprint, current_cert_fingerprint = await _run_in_executor( + _fetch_fingerprints + ) + return call_cert_bytes, call_key_bytes, cached_fingerprint, current_cert_fingerprint diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index d88162667bda..4224808c52a4 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -13,10 +13,15 @@ # limitations under the License. import asyncio +import collections.abc from contextlib import asynccontextmanager import functools +import http.client as http_client +import inspect +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 +42,12 @@ except (ImportError, AttributeError): ClientTimeout = None +_LOGGER = logging.getLogger(__name__) +_MTLS_URL_PREFIXES = [ + "mtls.googleapis.com", + "mtls.sandbox.googleapis.com", + "p.googleapis.com", +] # Tracks the internal aiohttp installation and usage try: @@ -66,6 +77,8 @@ async def timeout_guard(timeout): total_timeout = timeout def _remaining_time(): + if total_timeout is None: + return None elapsed = time.monotonic() - start remaining = total_timeout - elapsed if remaining <= 0: @@ -143,11 +156,17 @@ def __init__( self._is_mtls = False self._mtls_init_task = None self._cached_cert = None + self._client_cert_callback = None + self._old_auth_requests: list[transport.Request] = [] 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: Optional[asyncio.Lock] = None + self._mtls_check_counter = 0 + self._refresh_lock: Optional[asyncio.Lock] = None + self._refresh_counter = 0 async def configure_mtls_channel(self, client_cert_callback=None): """Configure the client certificate and key for SSL connection. @@ -175,6 +194,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 +224,19 @@ 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) + + while len(self._old_auth_requests) > 2: + oldest_auth_request = self._old_auth_requests[0] + try: + if hasattr(oldest_auth_request, "close"): + res = oldest_auth_request.close() + if inspect.isawaitable(res): + await res + except Exception: + pass + self._old_auth_requests.pop(0) - 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 +304,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 @@ -288,13 +318,15 @@ async def request( retries = _exponential_backoff.AsyncExponentialBackoff( total_attempts=total_attempts, ) - if headers is None: - headers = {} + request_headers = dict(headers) if headers is not None else {} + start_time = time.monotonic() + refresh_counter_at_error = self._refresh_counter + check_counter_at_error = self._mtls_check_counter async with timeout_guard(max_allowed_time) as with_timeout: await with_timeout( # Note: before_request will attempt to refresh credentials if expired. self._credentials.before_request( - self._auth_request, method, url, headers + self._auth_request, method, url, request_headers ) ) actual_timeout: float = 0.0 @@ -307,11 +339,195 @@ async def request( async for _ in retries: # pragma: no branch response = await with_timeout( self._auth_request( - url, method, data, headers, actual_timeout, **kwargs + url, method, data, request_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: + try: + if max_allowed_time is not None: + elapsed = time.monotonic() - start_time + remaining_time = max(0.0, max_allowed_time - elapsed) + if remaining_time == 0.0: + raise google.auth.exceptions.TimeoutError( + "Timeout exceeded before credential refresh could begin" + ) + else: + remaining_time = None + is_streaming = data is not None and ( + isinstance( + data, + (collections.abc.Iterator, collections.abc.AsyncIterable), + ) + or hasattr(data, "read") + ) + + async def _recover_auth_state(): + is_mtls_endpoint = False + if self._is_mtls: + 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: + if self._mtls_rotation_lock is None: + self._mtls_rotation_lock = asyncio.Lock() + async with self._mtls_rotation_lock: + # Check if another coroutine already reconfigured mTLS or + # ran the validation check. + if ( + self._mtls_check_counter + > check_counter_at_error + ): + pass + else: + try: + ( + call_cert_bytes, + call_key_bytes, + cached_fingerprint, + current_cert_fingerprint, + ) = await mtls.check_parameters_for_unauthorized_response( + self._cached_cert, + self._client_cert_callback, + ) + except ( + exceptions.ClientCertError, + exceptions.MutualTLSChannelError, + OSError, + ValueError, + ImportError, + ) as e: + _LOGGER.warning( + "Failed to check client certificate parameters: %s. Proceeding with original response.", + e, + ) + else: + if ( + current_cert_fingerprint is not None + and cached_fingerprint + != current_cert_fingerprint + ): + saved_callback = ( + self._client_cert_callback + ) + try: + _LOGGER.info( + "Client certificate has changed, reconfiguring mTLS " + "channel." + ) + if self._mtls_init_task is not None: + if not self._mtls_init_task.done(): + try: + await self._mtls_init_task + except Exception: + pass + self._mtls_init_task = None + await self.configure_mtls_channel( + lambda: (call_cert_bytes, call_key_bytes) + ) + except Exception as e: + _LOGGER.error( + "Failed to reconfigure mTLS channel: %s", + e, + ) + raise exceptions.MutualTLSChannelError( + "Failed to reconfigure mTLS channel" + ) from e + finally: + self._client_cert_callback = ( + saved_callback + ) + else: + _LOGGER.info( + "Skipping reconfiguration of mTLS channel because the client" + " certificate has not changed." + ) + # Always increment so waiting tasks skip the check block + self._mtls_check_counter += 1 + if self._refresh_lock is None: + self._refresh_lock = asyncio.Lock() + + async with self._refresh_lock: + # Check if another task already refreshed credentials while we were waiting + if self._refresh_counter > refresh_counter_at_error: + _LOGGER.debug( + "Credentials were already refreshed by a concurrent task. Skipping duplicate refresh." + ) + else: + try: + await self._credentials.refresh(self._auth_request) + except NotImplementedError: + _LOGGER.debug("Credentials do not implement refresh().") + return response + except ( + exceptions.RefreshError, + getattr(exceptions, "InvalidOperation", Exception), + ) as e: + _LOGGER.debug( + "Credential refresh failed, returning 401 response. Error: %s", + e, + ) + return response + else: + self._refresh_counter += 1 + + if is_streaming: + return response + # Return None to explicitly signal successful recovery & trigger retry if needed + return None + + async with timeout_guard(remaining_time) as auth_with_timeout: + early_return_response = await auth_with_timeout( + _recover_auth_state() + ) + except (Exception, asyncio.CancelledError): + if hasattr(response, "close"): + try: + res = response.close() + if inspect.isawaitable(res): + await res + except Exception: + pass + raise + # If it returned a response (meaning streaming or error), bail out + if early_return_response is not None: + return early_return_response + if hasattr(response, "close"): + try: + res = response.close() + if inspect.isawaitable(res): + await res + except Exception: + pass + if max_allowed_time is not None: + remaining_time = max( + 0.0, max_allowed_time - (time.monotonic() - start_time) + ) + if remaining_time == 0.0: + raise google.auth.exceptions.TimeoutError( + "Timeout exceeded before retrying the request" + ) + kwargs["_auth_retry_count"] = _auth_retry_count + 1 + 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) @@ -594,4 +810,18 @@ async def close(self) -> None: await self._mtls_init_task except asyncio.CancelledError: pass - await self._auth_request.close() + try: + if hasattr(self._auth_request, "close"): + res = self._auth_request.close() + if inspect.isawaitable(res): + await res + finally: + for old_request in self._old_auth_requests: + try: + if hasattr(old_request, "close"): + res = old_request.close() + if inspect.isawaitable(res): + await res + except Exception: + pass + self._old_auth_requests.clear() diff --git a/packages/google-auth/tests/transport/aio/test_mtls.py b/packages/google-auth/tests/transport/aio/test_mtls.py new file mode 100644 index 000000000000..7d583dcfb108 --- /dev/null +++ b/packages/google-auth/tests/transport/aio/test_mtls.py @@ -0,0 +1,226 @@ +# Copyright 2026 Google LLC +# +# 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. + +from unittest import mock + +import pytest + +from google.auth import exceptions +from google.auth.aio.transport import mtls + +CERT_BYTES = ( + b"-----BEGIN CERTIFICATE-----\nMIID...CERT1...=\n-----END CERTIFICATE-----\n" +) +KEY_BYTES = ( + b"-----BEGIN PRIVATE KEY-----\nMIIE...KEY1...==\n-----END PRIVATE KEY-----\n" +) +NEW_CERT_BYTES = ( + b"-----BEGIN CERTIFICATE-----\nMIID...CERT2...=\n-----END CERTIFICATE-----\n" +) +NEW_KEY_BYTES = ( + b"-----BEGIN PRIVATE KEY-----\nMIIE...KEY2...==\n-----END PRIVATE KEY-----\n" +) + + +@pytest.mark.asyncio +async def test_check_parameters_no_client_cert(): + """Test when no certificate is discovered (has_cert is False).""" + with mock.patch.object( + mtls, "get_client_cert_and_key", return_value=(False, None, None) + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( + cached_cert=b"stale_cert", client_cert_callback=None + ) + + assert cert is None + assert key is None + assert cached_fp is None + assert current_fp is None + + +@pytest.mark.asyncio +async def test_check_parameters_cert_matched(): + """Test when newly retrieved certificate matches the cached certificate.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate") as mock_parse, + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_A", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FINGERPRINT_A", + ) as mock_get_cached, + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=callback + ) + + assert cert == CERT_BYTES + assert key == KEY_BYTES + assert cached_fp == "FINGERPRINT_A" + assert current_fp == "FINGERPRINT_A" + assert cached_fp == current_fp + + # These assertions will now work correctly using the bound mocks + mock_parse.assert_called_once_with(CERT_BYTES) + mock_get_cached.assert_called_once_with(CERT_BYTES) + + +@pytest.mark.asyncio +async def test_check_parameters_cert_mismatch_rotation(): + """Test when newly retrieved certificate differs from the cached certificate (rotation occurred).""" + + def callback(): + return NEW_CERT_BYTES, NEW_KEY_BYTES + + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate") as mock_parse, + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_NEW", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FINGERPRINT_OLD", + ) as mock_get_cached, + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=callback + ) + + assert cert == NEW_CERT_BYTES + assert key == NEW_KEY_BYTES + assert cached_fp == "FINGERPRINT_OLD" + assert current_fp == "FINGERPRINT_NEW" + assert cached_fp != current_fp + + # The fix: mock_parse is called with the NEW cert from the callback + mock_parse.assert_called_once_with(NEW_CERT_BYTES) + # mock_get_cached is called with the old cached cert + mock_get_cached.assert_called_once_with(CERT_BYTES) + + +@pytest.mark.asyncio +async def test_check_parameters_without_cached_cert(): + """Test when cached_cert is None.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with ( + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FINGERPRINT_CURRENT", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint" + ) as mock_get_cached, + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( + cached_cert=None, client_cert_callback=callback + ) + + assert cert == CERT_BYTES + assert key == KEY_BYTES + assert cached_fp == "FINGERPRINT_CURRENT" + assert current_fp == "FINGERPRINT_CURRENT" + mock_get_cached.assert_not_called() + + +@pytest.mark.asyncio +async def test_check_parameters_executor_fingerprint_computation(): + """Test that fingerprint computation is properly offloaded to the executor.""" + + def callback(): + return CERT_BYTES, KEY_BYTES + + with ( + mock.patch.object( + mtls, "_run_in_executor", wraps=mtls._run_in_executor + ) as mock_run_in_executor, + mock.patch("google.auth._agent_identity_utils.parse_certificate"), + mock.patch( + "google.auth._agent_identity_utils.calculate_certificate_fingerprint", + return_value="FP_CURRENT", + ), + mock.patch( + "google.auth._agent_identity_utils.get_cached_cert_fingerprint", + return_value="FP_CACHED", + ), + ): + ( + cert, + key, + cached_fp, + current_fp, + ) = await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=callback + ) + + assert mock_run_in_executor.called + assert cert == CERT_BYTES + assert cached_fp == "FP_CACHED" + assert current_fp == "FP_CURRENT" + + +@pytest.mark.asyncio +async def test_check_parameters_callback_exception_propagation(): + """Test that exceptions raised by client_cert_callback propagate cleanly.""" + + def failing_callback(): + raise exceptions.ClientCertError("Client cert provider failed") + + with pytest.raises(exceptions.ClientCertError, match="Client cert provider failed"): + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=failing_callback + ) + + +@pytest.mark.asyncio +async def test_check_parameters_async_callback_exception_propagation(): + """Test that exceptions raised in an async client_cert_callback propagate cleanly.""" + + async def failing_async_callback(): + raise OSError("Disk read error while loading certificates") + + with pytest.raises(OSError, match="Disk read error while loading certificates"): + await mtls.check_parameters_for_unauthorized_response( + cached_cert=CERT_BYTES, client_cert_callback=failing_async_callback + ) diff --git a/packages/google-auth/tests/transport/aio/test_sessions.py b/packages/google-auth/tests/transport/aio/test_sessions.py index de283b7b2e7f..ecdb60e4af5a 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions.py +++ b/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) 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..0972a6180299 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import http.client as http_client import json import os import ssl @@ -23,12 +25,16 @@ from google.auth.aio import credentials from google.auth.aio import transport from google.auth.aio.transport import sessions +from google.auth.exceptions import TimeoutError # This is the valid "workload" format the library expects VALID_WORKLOAD_CONFIG = { "version": 1, "cert_configs": { - "workload": {"cert_path": "/tmp/mock_cert.pem", "key_path": "/tmp/mock_key.pem"} + "workload": { + "cert_path": "/tmp/mock_cert.pem", + "key_path": "/tmp/mock_key.pem", + } }, } @@ -36,23 +42,23 @@ class TestSessionsMtls: @pytest.mark.asyncio async def test_configure_mtls_channel(self): - """ - Tests that the mTLS channel configures correctly when a - valid workload config is mocked. - """ - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, mock.patch( - "aiohttp.TCPConnector" - ) as mock_connector, mock.patch( - "aiohttp.ClientSession" - ) as mock_session: + """Tests that the mTLS channel configures correctly when a valid workload config is mocked.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector") as mock_connector, + mock.patch("aiohttp.ClientSession") as mock_session, + ): mock_session.return_value.close = mock.AsyncMock() mock_exists.return_value = True mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") @@ -75,12 +81,11 @@ async def test_configure_mtls_channel(self): @pytest.mark.asyncio async def test_configure_mtls_channel_disabled(self): - """ - Tests behavior when the config file does not exist. - """ - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists: + """Tests behavior when the config file does not exist.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + ): mock_exists.return_value = False mock_creds = mock.AsyncMock(spec=credentials.Credentials) session = sessions.AsyncAuthorizedSession(mock_creds) @@ -90,13 +95,13 @@ async def test_configure_mtls_channel_disabled(self): @pytest.mark.asyncio async def test_configure_mtls_channel_invalid_format(self): - """ - Verifies that the MutualTLSChannelError is raised for bad formats. - """ - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data='{"invalid": "format"}') + """Verifies that the MutualTLSChannelError is raised for bad formats.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", mock.mock_open(read_data='{"invalid": "format"}') + ), ): mock_exists.return_value = True mock_creds = mock.AsyncMock(spec=credentials.Credentials) @@ -107,14 +112,14 @@ async def test_configure_mtls_channel_invalid_format(self): await session.close() @pytest.mark.asyncio - async def test_configure_mtls_channel_invalud_fields(self): - """ - If cert is missing expected keys, it should fail gracefully - """ - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data='{"cert_configs": {}}') + async def test_configure_mtls_channel_invalid_fields(self): + """If cert is missing expected keys, it should fail gracefully.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", mock.mock_open(read_data='{"cert_configs": {}}') + ), ): mock_exists.return_value = True mock_creds = mock.AsyncMock(spec=credentials.Credentials) @@ -125,25 +130,23 @@ async def test_configure_mtls_channel_invalud_fields(self): @pytest.mark.asyncio async def test_configure_mtls_channel_mock_callback(self): - """ - Tests mTLS configuration using bytes-returning callback. - """ + """Tests mTLS configuration using bytes-returning callback.""" def mock_callback(): return (b"fake_cert_bytes", b"fake_key_bytes") - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch( - "google.auth.transport.mtls.has_default_client_cert_source", - return_value=True, - ), mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, mock.patch( - "aiohttp.TCPConnector" - ) as mock_connector, mock.patch( - "aiohttp.ClientSession" - ) as mock_session: + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch( + "google.auth.transport.mtls.has_default_client_cert_source", + return_value=True, + ), + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector") as mock_connector, + mock.patch("aiohttp.ClientSession") as mock_session, + ): mock_session.return_value.close = mock.AsyncMock() mock_context = mock.Mock(spec=ssl.SSLContext) mock_make_context.return_value = mock_context @@ -163,18 +166,21 @@ def mock_callback(): @pytest.mark.asyncio async def test_configure_mtls_channel_custom_request(self): - """Tests that if _auth_request is not an AiohttpRequest, _is_mtls is set to False - because we can't configure the custom request with mTLS. - """ - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context: + """Tests that if _auth_request is not an AiohttpRequest, _is_mtls is set to False.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + ): mock_exists.return_value = True mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") @@ -190,27 +196,27 @@ async def test_configure_mtls_channel_custom_request(self): with pytest.warns(UserWarning, match="Attempted to establish mTLS"): await session.configure_mtls_channel() - # If the request handler is not an AiohttpRequest, the library cannot configure - # the connection to use mTLS, so _is_mtls must be False to reflect this unconfigured state. assert session._is_mtls is False mock_make_context.assert_not_called() await session.close() @pytest.mark.asyncio async def test_configure_mtls_channel_exception_resets_flag(self): - """ - Tests that self._is_mtls is reset to False if an exception is raised - during configuration. - """ - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context: + """Tests that self._is_mtls is reset to False if an exception is raised during configuration.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + ): mock_exists.return_value = True mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") mock_make_context.side_effect = exceptions.ClientCertError("Mock error") @@ -226,19 +232,21 @@ async def test_configure_mtls_channel_exception_resets_flag(self): @pytest.mark.asyncio async def test_configure_mtls_channel_transport_error_resets_flag(self): - """ - Tests that self._is_mtls is reset to False if a TransportError is raised - during configuration. - """ - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context: + """Tests that self._is_mtls is reset to False if a TransportError is raised.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + ): mock_exists.return_value = True mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") mock_make_context.side_effect = exceptions.TransportError("Mock error") @@ -254,27 +262,31 @@ async def test_configure_mtls_channel_transport_error_resets_flag(self): @pytest.mark.asyncio async def test_configure_mtls_channel_atomic_on_exception(self): - """ - Tests that if configure_mtls_channel has already successfully configured mTLS, - a subsequent attempt that raises an exception will preserve the original mTLS state. - """ + """Tests that if configure_mtls_channel already succeeded, a subsequent failure preserves state.""" # Step 1: Successful configuration - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, mock.patch( - "aiohttp.TCPConnector" - ), mock.patch( - "aiohttp.ClientSession" - ) as mock_session: + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector"), + mock.patch("aiohttp.ClientSession") as mock_session, + ): mock_session.return_value.close = mock.AsyncMock() mock_exists.return_value = True - mock_helper.return_value = (True, b"fake_cert_data_1", b"fake_key_data_1") + mock_helper.return_value = ( + True, + b"fake_cert_data_1", + b"fake_key_data_1", + ) mock_context = mock.Mock(spec=ssl.SSLContext) mock_make_context.return_value = mock_context @@ -288,16 +300,12 @@ async def test_configure_mtls_channel_atomic_on_exception(self): first_auth_request = session._auth_request # Step 2: Failed subsequent configuration attempt - # Reset task so we trigger a new configuration run session._mtls_init_task = None - - # Patch context generator to fail this time mock_make_context.side_effect = exceptions.ClientCertError("Mock error") with pytest.raises(exceptions.MutualTLSChannelError): await session.configure_mtls_channel() - # Verify that the state remains unchanged from the first successful configuration assert session._is_mtls is True assert session._cached_cert == b"fake_cert_data_1" assert session._auth_request is first_auth_request @@ -305,24 +313,23 @@ async def test_configure_mtls_channel_atomic_on_exception(self): @pytest.mark.asyncio async def test_configure_mtls_channel_close_exception_does_not_abort(self): - """ - Tests that if old_auth_request.close() raises an exception, the mTLS - configuration is still considered successful, and is_mtls remains True - without raising MutualTLSChannelError. - """ - with mock.patch.dict( - os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"} - ), mock.patch("os.path.exists") as mock_exists, mock.patch( - "builtins.open", mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)) - ), mock.patch( - "google.auth.aio.transport.mtls.get_client_cert_and_key" - ) as mock_helper, mock.patch( - "google.auth.aio.transport.mtls.make_client_cert_ssl_context" - ) as mock_make_context, mock.patch( - "aiohttp.TCPConnector" - ), mock.patch( - "aiohttp.ClientSession" - ) as mock_session: + """Tests that an exception in old_auth_request.close() does not abort configuration.""" + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch("os.path.exists") as mock_exists, + mock.patch( + "builtins.open", + mock.mock_open(read_data=json.dumps(VALID_WORKLOAD_CONFIG)), + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key" + ) as mock_helper, + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context" + ) as mock_make_context, + mock.patch("aiohttp.TCPConnector"), + mock.patch("aiohttp.ClientSession") as mock_session, + ): mock_session.return_value.close = mock.AsyncMock() mock_exists.return_value = True mock_helper.return_value = (True, b"fake_cert_data", b"fake_key_data") @@ -333,14 +340,652 @@ async def test_configure_mtls_channel_close_exception_does_not_abort(self): mock_creds = mock.AsyncMock(spec=credentials.Credentials) session = sessions.AsyncAuthorizedSession(mock_creds) - # Mock close() of the initial self._auth_request to raise an exception session._auth_request.close = mock.AsyncMock( side_effect=Exception("Mock close error") ) - # Should complete successfully without raising MutualTLSChannelError await session.configure_mtls_channel() 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_raises_error(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(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.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) 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") + + with pytest.raises(exceptions.MutualTLSChannelError): + await session.request("GET", "https://pubsub.mtls.googleapis.com/test") + + mock_check.assert_called_once() + mock_conf.assert_called_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_check_params_fails(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock() + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(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.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.side_effect = exceptions.MutualTLSChannelError( + "Failed to check params" + ) + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + mock_check.assert_called_once() + mock_creds.refresh.assert_not_called() + mock_conf.assert_not_called() + + await session.close() + + @pytest.mark.asyncio + async def test_no_cert_rotation_when_cert_matches_and_mtls_enabled(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + # 401 on initial request, 200 on retry after refresh + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + 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.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + # Matching fingerprints mean no mTLS rotation is needed + mock_check.return_value = (new_cert, new_key, b"old_fp", b"old_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_not_called() + mock_creds.refresh.assert_called_once() + assert mock_auth_req.call_count == 2 + mock_resp_401.close.assert_awaited_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_success_and_retry(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + 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.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) 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") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_called_once_with(mock.ANY) + mock_creds.refresh.assert_called_once() + assert mock_creds.before_request.call_count == 2 + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_lock_contention(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock( + side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 + ) + + 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" + + async def mock_configure_mtls_channel(*args, **kwargs): + await asyncio.sleep(0.01) + session._cached_cert = new_cert + + async def mock_check_side_effect(cached_cert, callback=None): + if cached_cert == b"old_cert": + return (new_cert, new_key, b"old_fp", b"new_fp") + return (new_cert, new_key, b"new_fp", b"new_fp") + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.side_effect = mock_check_side_effect + mock_conf.side_effect = mock_configure_mtls_channel + + tasks = [ + session.request("GET", "https://pubsub.mtls.googleapis.com/test") + for _ in range(3) + ] + responses = await asyncio.gather(*tasks) + + for resp in responses: + assert resp == mock_resp_200 + + mock_check.assert_called_once() + mock_conf.assert_called_once() + assert mock_creds.refresh.call_count == 1 + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_lock_contention_no_cert_change(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock( + side_effect=[mock_resp_401] * 3 + [mock_resp_200] * 3 + ) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + async def mock_check_side_effect(cached_cert, callback=None): + await asyncio.sleep(0.01) + return (b"old_cert", b"old_key", b"old_fp", b"old_fp") + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.side_effect = mock_check_side_effect + + tasks = [ + session.request("GET", "https://pubsub.mtls.googleapis.com/test") + for _ in range(3) + ] + responses = await asyncio.gather(*tasks) + + for resp in responses: + assert resp == mock_resp_200 + + mock_check.assert_called_once() + mock_conf.assert_not_called() + # Concurrent 401s properly deduplicate to 1 refresh + assert mock_creds.refresh.call_count == 1 + + await session.close() + + @pytest.mark.asyncio + async def test_psc_endpoint_triggers_cert_rotation(self): + """Verifies that PSC endpoints (*.p.googleapis.com) are recognized as mTLS endpoints.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + + mock_auth_req = mock.AsyncMock(side_effect=[mock_resp_401, mock_resp_200]) + + 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.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) 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") + + resp = await session.request("GET", "https://pubsub.p.googleapis.com/test") + + assert resp == mock_resp_200 + mock_check.assert_called_once() + mock_conf.assert_called_once_with(mock.ANY) + + await session.close() + + @pytest.mark.asyncio + async def test_non_mtls_url_bypasses_rotation(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + + 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.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + resp = await session.request("GET", "https://example.com/test") + + assert resp == mock_resp_401 + mock_check.assert_not_called() + mock_conf.assert_not_called() + assert mock_creds.refresh.call_count == 2 + assert mock_auth_req.call_count == 3 + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_skips_retry_for_streaming(self): + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock() + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(return_value=mock_resp) + + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + session._is_mtls = True + session._cached_cert = b"old_cert" + + class MockStream: + def read(self): + pass + + with ( + mock.patch( + "google.auth.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ) as mock_conf, + ): + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test", data=MockStream() + ) + + assert resp == mock_resp + mock_conf.assert_called_once() + + mock_creds.refresh.assert_called_once() + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_credential_refresh_fails(self): + """Covers the except block for RefreshError when credentials fail to refresh.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock( + side_effect=exceptions.RefreshError("Refresh failed") + ) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_auth_req = mock.AsyncMock(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.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ), + ): + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + mock_creds.refresh.assert_called_once() + mock_auth_req.assert_called_once() + + await session.close() + + @pytest.mark.asyncio + async def test_cert_rotation_max_retries_exceeded(self): + """Covers the `if _auth_retry_count < 2:` max retry limit.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp = mock.Mock() + mock_resp.status_code = http_client.UNAUTHORIZED + mock_resp.close = mock.AsyncMock() + mock_auth_req = mock.AsyncMock(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.aio.transport.mtls.check_parameters_for_unauthorized_response", + new_callable=mock.AsyncMock, + ) as mock_check, + mock.patch.object( + session, "configure_mtls_channel", new_callable=mock.AsyncMock + ), + ): + mock_check.return_value = (b"new", b"new", b"old_fp", b"new_fp") + + resp = await session.request( + "GET", "https://pubsub.mtls.googleapis.com/test" + ) + + assert resp == mock_resp + assert mock_auth_req.call_count == 3 + assert mock_check.call_count == 2 + assert mock_resp.close.call_count == 2 + + await session.close() + + @pytest.mark.asyncio + async def test_session_close_cleans_old_auth_requests(self): + """Covers the loop in the `close()` method that drains `_old_auth_requests`.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock.AsyncMock() + ) + + mock_old_req_1 = mock.AsyncMock() + mock_old_req_2 = mock.AsyncMock() + mock_old_req_3_fails = mock.AsyncMock() + mock_old_req_3_fails.close.side_effect = Exception("Close error") + + session._old_auth_requests.extend( + [mock_old_req_1, mock_old_req_2, mock_old_req_3_fails] + ) + + await session.close() + + mock_old_req_1.close.assert_called_once() + mock_old_req_2.close.assert_called_once() + mock_old_req_3_fails.close.assert_called_once() + assert len(session._old_auth_requests) == 0 + + @pytest.mark.asyncio + async def test_request_401_streaming_refreshes_creds_and_returns_open_response( + self, + ): + """Verifies that streaming requests refresh credentials but return the unclosed response.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + mock_creds.refresh = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + streaming_data = (chunk for chunk in [b"chunk1", b"chunk2"]) + response = await session.request( + "POST", "https://example.com", data=streaming_data + ) + + assert response == mock_resp_401 + mock_creds.refresh.assert_awaited_once() + mock_resp_401.close.assert_not_called() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_timeout_during_recovery(self): + """Verifies that response is closed when auth_with_timeout times out during _recover_auth_state.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + async def slow_refresh(*args, **kwargs): + await asyncio.sleep(10) + + mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + with pytest.raises(TimeoutError): + await session.request("GET", "https://example.com", max_allowed_time=0.01) + + mock_resp_401.close.assert_awaited_once() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_closes_response_on_cancellation(self): + """Verifies that response is closed and CancelledError propagated if task is cancelled.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + refresh_started = asyncio.Event() + + async def cancel_on_refresh(*args, **kwargs): + refresh_started.set() + await asyncio.sleep(10) + + mock_creds.refresh = mock.AsyncMock(side_effect=cancel_on_refresh) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock(return_value=mock_resp_401) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + task = asyncio.create_task(session.request("GET", "https://example.com")) + await refresh_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + mock_resp_401.close.assert_awaited_once() + await session.close() + + @pytest.mark.asyncio + async def test_request_401_concurrent_refreshes_are_deduplicated(self): + """Verifies that concurrent 401s execute only one credentials.refresh call.""" + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + mock_creds.before_request = mock.AsyncMock(return_value=None) + + mock_resp_401 = mock.Mock() + mock_resp_401.status_code = http_client.UNAUTHORIZED + mock_resp_401.close = mock.AsyncMock() + + mock_resp_200 = mock.Mock() + mock_resp_200.status_code = http_client.OK + mock_resp_200.close = mock.AsyncMock() + + mock_auth_req = mock.AsyncMock( + side_effect=[ + mock_resp_401, + mock_resp_401, + mock_resp_200, + mock_resp_200, + ] + ) + session = sessions.AsyncAuthorizedSession( + mock_creds, auth_request=mock_auth_req + ) + + refresh_count = 0 + + async def slow_refresh(*args, **kwargs): + nonlocal refresh_count + refresh_count += 1 + await asyncio.sleep(0.05) + + mock_creds.refresh = mock.AsyncMock(side_effect=slow_refresh) + + results = await asyncio.gather( + session.request("GET", "https://example.com/1"), + session.request("GET", "https://example.com/2"), + ) + + assert all(r.status_code == 200 for r in results) + assert refresh_count == 1 + await session.close()