From 1b64150cb6f9062e2126dacd8cb19e2e4bed2a2c Mon Sep 17 00:00:00 2001 From: Radhika Agrawal Date: Tue, 13 Oct 2026 21:06:49 +0000 Subject: [PATCH] fix(auth): preserve aiohttp session settings when configuring mTLS in AsyncAuthorizedSession configure_mtls_channel() replaced the transport with a new aiohttp.ClientSession/TCPConnector built with default arguments, so the caller's session settings were dropped on the first mTLS setup and on every certificate rotation triggered by a 401. When the current transport already has a session, the new mTLS session now copies the same settings that aiohttp.Request._clone() copies: connector limit, limit_per_host, force_close and local_addr, and session headers, cookie_jar, auth, timeout, trust_env, auto_decompress, trace_configs and json_serialize. The ssl context is always the new mTLS context. Behavior is unchanged when there is no session yet. Fixes #18314 --- .../google/auth/aio/transport/sessions.py | 55 +++++- .../tests/transport/aio/test_sessions_mtls.py | 167 ++++++++++++++++++ 2 files changed, 220 insertions(+), 2 deletions(-) diff --git a/packages/google-auth/google/auth/aio/transport/sessions.py b/packages/google-auth/google/auth/aio/transport/sessions.py index 980c86f5ea19..03d6e8fd6e58 100644 --- a/packages/google-auth/google/auth/aio/transport/sessions.py +++ b/packages/google-auth/google/auth/aio/transport/sessions.py @@ -26,6 +26,7 @@ import google.auth.credentials import google.auth.transport._mtls_helper from google.auth import _exponential_backoff, _helpers, exceptions +from google.auth.aio import _helpers as _helpers_async from google.auth.aio import transport from google.auth.aio.credentials import Credentials from google.auth.aio.transport import mtls @@ -321,8 +322,58 @@ async def _do_configure(): ssl_context = await mtls._run_in_executor( mtls.make_client_cert_ssl_context, cert, key ) - connector = aiohttp.TCPConnector(ssl=ssl_context) - new_session = aiohttp.ClientSession(connector=connector) + connector_kwargs: dict = {"ssl": ssl_context} + session_kwargs: dict = {} + orig_session = getattr(self._auth_request, "_session", None) + if orig_session is not None: + orig_connector = getattr( + orig_session, "_connector", None + ) + if isinstance(orig_connector, aiohttp.TCPConnector): + connector_kwargs.update( + { + "limit": getattr( + orig_connector, "_limit", 100 + ), + "limit_per_host": getattr( + orig_connector, "_limit_per_host", 0 + ), + "force_close": getattr( + orig_connector, "_force_close", False + ), + "local_addr": _helpers_async._get_local_addr( + orig_connector + ), + } + ) + session_kwargs["auto_decompress"] = getattr( + orig_session, "_auto_decompress", True + ) + session_kwargs["trust_env"] = getattr( + orig_session, "_trust_env", False + ) + trace_configs = getattr( + orig_session, "_trace_configs", None + ) + if trace_configs: + session_kwargs["trace_configs"] = list( + trace_configs + ) + for attr_name, kwarg_name in [ + ("_default_headers", "headers"), + ("_cookie_jar", "cookie_jar"), + ("_default_auth", "auth"), + ("_timeout", "timeout"), + ("_json_serialize", "json_serialize"), + ]: + val = getattr(orig_session, attr_name, None) + if val is not None: + session_kwargs[kwarg_name] = val + + connector = aiohttp.TCPConnector(**connector_kwargs) + new_session = aiohttp.ClientSession( + connector=connector, **session_kwargs + ) old_auth_request = self._auth_request self._auth_request = AiohttpRequest(session=new_session) 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 a96a1e5a404e..77d8b41e1b0a 100644 --- a/packages/google-auth/tests/transport/aio/test_sessions_mtls.py +++ b/packages/google-auth/tests/transport/aio/test_sessions_mtls.py @@ -82,6 +82,173 @@ async def test_configure_mtls_channel(self): mock_session.assert_called_once_with(connector=mock_connector.return_value) await session.close() + @pytest.mark.asyncio + async def test_configure_mtls_channel_preserves_existing_session_config(self): + """ + Tests that configure_mtls_channel preserves connector and session settings + from an existing aiohttp.ClientSession across initial configuration and rotation. + """ + import aiohttp + + orig_connector = aiohttp.TCPConnector( + limit=50, + limit_per_host=10, + force_close=True, + local_addr=("127.0.0.1", 8080), + ) + custom_cookie_jar = aiohttp.CookieJar(unsafe=True) + custom_auth = aiohttp.BasicAuth("user", "pass") + custom_timeout = aiohttp.ClientTimeout(total=42.0) + trace_config = aiohttp.TraceConfig() + + def custom_json_serialize(obj): + return json.dumps(obj) # pragma: NO COVER + + orig_session = aiohttp.ClientSession( + connector=orig_connector, + headers={"X-Custom-Header": "custom-value"}, + cookie_jar=custom_cookie_jar, + auth=custom_auth, + timeout=custom_timeout, + trust_env=True, + auto_decompress=False, + trace_configs=[trace_config], + json_serialize=custom_json_serialize, + ) + auth_request = sessions.AiohttpRequest(session=orig_session) + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=auth_request) + + mock_context_1 = ssl.create_default_context() + mock_context_2 = ssl.create_default_context() + + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch( + "google.auth.transport._mtls_helper.check_use_client_cert", + return_value=True, + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key", + side_effect=[ + (True, b"fake_cert_1", b"fake_key_1"), + (True, b"fake_cert_2", b"fake_key_2"), + ], + ), + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context", + side_effect=[mock_context_1, mock_context_2], + ), + ): + await session.configure_mtls_channel() + + new_session_1 = session._auth_request._session + assert new_session_1 is not orig_session + new_connector_1 = new_session_1.connector + assert isinstance(new_connector_1, aiohttp.TCPConnector) + assert new_connector_1 is not orig_connector + assert new_connector_1._ssl is mock_context_1 + assert new_connector_1.limit == 50 + assert new_connector_1.limit_per_host == 10 + assert new_connector_1.force_close is True + assert sessions._helpers_async._get_local_addr(new_connector_1) == ( + "127.0.0.1", + 8080, + ) + assert new_session_1.headers["X-Custom-Header"] == "custom-value" + assert new_session_1.cookie_jar is custom_cookie_jar + assert new_session_1._default_auth == custom_auth + assert new_session_1.timeout == custom_timeout + assert new_session_1.trust_env is True + assert new_session_1.auto_decompress is False + assert new_session_1._trace_configs == [trace_config] + assert new_session_1._trace_configs is not orig_session._trace_configs + assert new_session_1._json_serialize is custom_json_serialize + + # Simulate certificate rotation by resetting _mtls_init_task and reconfiguring + session._mtls_init_task = None + await session.configure_mtls_channel() + + new_session_2 = session._auth_request._session + assert new_session_2 is not new_session_1 + new_connector_2 = new_session_2.connector + assert isinstance(new_connector_2, aiohttp.TCPConnector) + assert new_connector_2._ssl is mock_context_2 + assert new_connector_2.limit == 50 + assert new_connector_2.limit_per_host == 10 + assert new_connector_2.force_close is True + assert sessions._helpers_async._get_local_addr(new_connector_2) == ( + "127.0.0.1", + 8080, + ) + assert new_session_2.headers["X-Custom-Header"] == "custom-value" + assert new_session_2.cookie_jar is custom_cookie_jar + assert new_session_2._default_auth == custom_auth + assert new_session_2.timeout == custom_timeout + assert new_session_2.trust_env is True + assert new_session_2.auto_decompress is False + assert new_session_2._trace_configs == [trace_config] + assert new_session_2._json_serialize is custom_json_serialize + + await session.close() + + @pytest.mark.asyncio + async def test_configure_mtls_channel_with_minimal_existing_session(self): + """ + Tests that configure_mtls_channel handles an existing session whose + connector is not a TCPConnector and optional attributes are None/empty. + """ + mock_orig_session = mock.MagicMock() + mock_orig_session._connector = None + mock_orig_session._auto_decompress = True + mock_orig_session._trust_env = False + mock_orig_session._trace_configs = [] + mock_orig_session._default_headers = None + mock_orig_session._cookie_jar = None + mock_orig_session._default_auth = None + mock_orig_session._timeout = None + mock_orig_session._json_serialize = None + mock_orig_session.close = mock.AsyncMock() + + auth_request = sessions.AiohttpRequest(session=mock_orig_session) + mock_creds = mock.AsyncMock(spec=credentials.Credentials) + session = sessions.AsyncAuthorizedSession(mock_creds, auth_request=auth_request) + + mock_context = ssl.create_default_context() + with ( + mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}), + mock.patch( + "google.auth.transport._mtls_helper.check_use_client_cert", + return_value=True, + ), + mock.patch( + "google.auth.aio.transport.mtls.get_client_cert_and_key", + return_value=(True, b"fake_cert_data", b"fake_key_data"), + ), + mock.patch( + "google.auth.aio.transport.mtls.make_client_cert_ssl_context", + return_value=mock_context, + ), + mock.patch("aiohttp.ClientSession") as mock_session, + ): + mock_session.return_value.close = mock.AsyncMock() + await session.configure_mtls_channel() + + mock_session.assert_called_once() + _, session_kwargs = mock_session.call_args + assert session_kwargs["connector"]._ssl is mock_context + assert session_kwargs["auto_decompress"] is True + assert session_kwargs["trust_env"] is False + assert "trace_configs" not in session_kwargs + assert "headers" not in session_kwargs + assert "cookie_jar" not in session_kwargs + assert "auth" not in session_kwargs + assert "timeout" not in session_kwargs + assert "json_serialize" not in session_kwargs + await session_kwargs["connector"].close() + + await session.close() + @pytest.mark.asyncio async def test_configure_mtls_channel_disabled(self): """