fix(google-auth): keep PoolManager settings in AuthorizedHttp.configure_mtls_channel

packages/google-auth/google/auth/transport/urllib3.py

@@ -19,6 +19,7 @@
 import http.client as http_client
 import logging
 import warnings
+from collections.abc import Mapping
 
 # Certifi is Mozilla's certificate bundle. Urllib3 needs a certificate bundle
 # to verify HTTPS requests, and certifi is the recommended and most reliable
@@ -153,20 +154,61 @@ def __call__(
             raise new_exc from caught_exc
 
 
-def _make_default_http():
+# Connection pool settings unrelated to TLS that are carried over when
+# ``AuthorizedHttp.configure_mtls_channel`` replaces the ``urllib3.PoolManager``.
+_CARRIED_OVER_POOL_SETTINGS = ("retries", "timeout", "maxsize", "block")
+_DEFAULT_NUM_POOLS = 10
+
+
+def _pool_manager_settings(http):
+    """Get the non-TLS settings of a ``urllib3.PoolManager``.
+
+    Args:
+        http (Any): The HTTP object to read the settings from.
+
+    Returns:
+        Mapping[str, Any]: The ``urllib3.PoolManager`` keyword arguments for the
+            settings that were set on ``http``. Empty if ``http`` is not a
+            ``urllib3.PoolManager``.
+    """
+    if not isinstance(http, urllib3.PoolManager):
+        return {}
+
+    settings = {}
+    pool_kw = getattr(http, "connection_pool_kw", None)
+    if isinstance(pool_kw, dict):
+        settings.update(
+            (name, pool_kw[name])
+            for name in _CARRIED_OVER_POOL_SETTINGS
+            if name in pool_kw
+        )
+    headers = getattr(http, "headers", None)
+    if isinstance(headers, Mapping) and headers:
+        settings["headers"] = headers
+    num_pools = getattr(getattr(http, "pools", None), "_maxsize", None)
+    if isinstance(num_pools, int) and num_pools != _DEFAULT_NUM_POOLS:
+        settings["num_pools"] = num_pools
+    return settings
+
+
+def _make_default_http(**pool_kwargs):
     if certifi is not None:
-        return urllib3.PoolManager(cert_reqs="CERT_REQUIRED", ca_certs=certifi.where())
+        return urllib3.PoolManager(
+            cert_reqs="CERT_REQUIRED", ca_certs=certifi.where(), **pool_kwargs
+        )
     else:
-        return urllib3.PoolManager()
+        return urllib3.PoolManager(**pool_kwargs)
 
 
-def _make_mutual_tls_http(cert, key):
+def _make_mutual_tls_http(cert, key, **pool_kwargs):
     """Create a mutual TLS HTTP connection with the given client cert and key.
     See https://github.com/urllib3/urllib3/issues/474#issuecomment-253168415
 
     Args:
         cert (bytes): client certificate in PEM format
         key (bytes): client private key in PEM format
+        pool_kwargs: Additional keyword arguments, such as ``retries`` or
+            ``maxsize``, passed to the ``urllib3.PoolManager``.
 
     Returns:
         urllib3.PoolManager: Mutual TLS HTTP connection.
@@ -198,7 +240,7 @@ def _make_mutual_tls_http(cert, key):
             "Failed to configure client certificate and key for mTLS."
         ) from exc
 
-    http = urllib3.PoolManager(ssl_context=ctx)
+    http = urllib3.PoolManager(ssl_context=ctx, **pool_kwargs)
     return http
 
 
@@ -334,6 +376,10 @@ def configure_mtls_channel(self, client_cert_callback=None):
                 If the callback is None, application default SSL credentials
                 will be used.
 
+        The new `urllib3.PoolManager` keeps the ``retries``, ``timeout``,
+        ``maxsize``, ``block``, ``headers`` and ``num_pools`` settings of the
+        current one.
+
         .. warning::
             Calling this method mutates the underlying `urllib3.PoolManager`.
             It is not thread-safe to call this explicitly while other
@@ -351,16 +397,17 @@ def configure_mtls_channel(self, client_cert_callback=None):
         if not use_client_cert:
             return False
 
+        pool_settings = _pool_manager_settings(self.http)
         try:
             found_cert_key, cert, key = transport._mtls_helper.get_client_cert_and_key(
                 client_cert_callback
             )
 
             if found_cert_key:
-                new_http = _make_mutual_tls_http(cert, key)
+                new_http = _make_mutual_tls_http(cert, key, **pool_settings)
                 new_is_mtls = True
             else:
-                new_http = _make_default_http()
+                new_http = _make_default_http(**pool_settings)
                 new_is_mtls = False
         except (
             exceptions.ClientCertError,

packages/google-auth/tests/transport/test_urllib3.py

@@ -113,6 +113,36 @@ def test_setup_error_raises_mutual_tls_channel_error(self, mock_secure_paths):
         assert "Failed to configure client certificate" in str(exc_info.value)
         assert isinstance(exc_info.value.__cause__, OSError)
 
+    def test_success_with_pool_kwargs(self):
+        retries = urllib3.util.Retry(total=3)
+        http = google.auth.transport.urllib3._make_mutual_tls_http(
+            pytest.public_cert_bytes,
+            pytest.private_key_bytes,
+            retries=retries,
+            maxsize=7,
+        )
+        assert isinstance(http, urllib3.PoolManager)
+        assert http.connection_pool_kw["ssl_context"] is not None
+        assert http.connection_pool_kw["retries"] is retries
+        assert http.connection_pool_kw["maxsize"] == 7
+
+
+def _custom_pool_manager():
+    return urllib3.PoolManager(
+        num_pools=3,
+        retries=urllib3.util.Retry(total=3),
+        timeout=urllib3.util.Timeout(connect=5.0),
+        maxsize=7,
+        block=True,
+    )
+
+
+def _assert_same_pool_settings(new_http, old_http):
+    for name in ("retries", "timeout", "maxsize", "block"):
+        assert new_http.connection_pool_kw[name] is old_http.connection_pool_kw[name]
+    assert new_http.headers == {"x-custom": "value"}
+    assert new_http.pools._maxsize == 3
+
 
 class TestAuthorizedHttp(object):
     TEST_URL = "http://example.com"
@@ -215,6 +245,50 @@ def test_configure_mtls_channel_with_callback(self, mock_make_mutual_tls_http):
             cert=pytest.public_cert_bytes, key=pytest.private_key_bytes
         )
 
+    def test_configure_mtls_channel_preserves_pool_settings(self):
+        callback = mock.Mock()
+        callback.return_value = (pytest.public_cert_bytes, pytest.private_key_bytes)
+        old_http = _custom_pool_manager()
+        authed_http = google.auth.transport.urllib3.AuthorizedHttp(
+            credentials=mock.Mock(), http=old_http
+        )
+        authed_http.headers = {"x-custom": "value"}
+
+        with pytest.warns(UserWarning):
+            with mock.patch.dict(
+                os.environ, {environment_vars.GOOGLE_API_USE_CLIENT_CERTIFICATE: "true"}
+            ):
+                is_mtls = authed_http.configure_mtls_channel(callback)
+
+        assert is_mtls
+        assert authed_http.http is not old_http
+        assert authed_http.http.connection_pool_kw["ssl_context"] is not None
+        _assert_same_pool_settings(authed_http.http, old_http)
+        assert authed_http.headers == {"x-custom": "value"}
+
+    @mock.patch(
+        "google.auth.transport._mtls_helper.get_client_cert_and_key", autospec=True
+    )
+    def test_configure_mtls_channel_non_mtls_preserves_pool_settings(
+        self, mock_get_client_cert_and_key
+    ):
+        mock_get_client_cert_and_key.return_value = (False, None, None)
+        old_http = _custom_pool_manager()
+        authed_http = google.auth.transport.urllib3.AuthorizedHttp(
+            credentials=mock.Mock(), http=old_http
+        )
+        authed_http.headers = {"x-custom": "value"}
+
+        with pytest.warns(UserWarning):
+            with mock.patch.dict(
+                os.environ, {environment_vars.GOOGLE_API_USE_CLIENT_CERTIFICATE: "true"}
+            ):
+                is_mtls = authed_http.configure_mtls_channel()
+
+        assert not is_mtls
+        assert authed_http.http is not old_http
+        _assert_same_pool_settings(authed_http.http, old_http)
+
     @mock.patch("google.auth.transport.urllib3._make_mutual_tls_http", autospec=True)
     @mock.patch(
         "google.auth.transport._mtls_helper.get_client_cert_and_key", autospec=True