Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 53 additions & 2 deletions packages/google-auth/google/auth/aio/transport/sessions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Comment on lines +335 to +371

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Relying on private attributes (prefixed with _) of aiohttp.ClientSession and aiohttp.TCPConnector is risky because these are internal implementation details that can change across versions without notice. Instead, use the public properties where available. Accessing private fields is bad practice and creates maintenance risk.

"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 [
    ("headers", "headers"),
    ("cookie_jar", "cookie_jar"),
    ("timeout", "timeout"),
]:
    val = getattr(orig_session, attr_name, None)
    if val is not None:
        session_kwargs[kwarg_name] = val
References
  1. Avoid accessing private attributes to bypass public API methods, as accessing private fields is bad practice and the minor performance gains do not justify the maintenance risk.


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)
Comment on lines 378 to 379

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The old session and its underlying connector/connections are not closed when a new session is created during initial configuration or certificate rotation. This will cause resource leaks (unclosed client sessions and connectors) and trigger warnings in aiohttp. Awaiting close() on the old AiohttpRequest resolves this leak.

old_auth_request = self._auth_request
self._auth_request = AiohttpRequest(session=new_session)
if old_auth_request is not None:
    await old_auth_request.close()

Expand Down
167 changes: 167 additions & 0 deletions packages/google-auth/tests/transport/aio/test_sessions_mtls.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
Loading