Fix use of HTTPS proxies #1070 (#1077)

* fix use of proxies #1070
* use `contextlib.AsyncExitStack` to simplify proxy session handling in `AIOHTTPSession`
* extract session creation into `AIOHTTPSession._get_session()`
* extract `chunked` detection into `AIOHTTPSession._chunked()`

Author: jakob-keller

CHANGES.rst

@@ -1,6 +1,10 @@
 Changes
 -------
 
+2.12.1 (2024-03-04)
+^^^^^^^^^^^^^^^^^^^
+* fix use of proxies #1070
+
 2.12.0 (2024-02-28)
 ^^^^^^^^^^^^^^^^^^^
 * bump botocore dependency specification

aiobotocore/__init__.py

@@ -1 +1 @@
-__version__ = '2.12.0'
+__version__ = '2.12.1'

aiobotocore/httpsession.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import io
 import os
 import socket
@@ -54,8 +55,11 @@ def __init__(
         proxies_config=None,
         connector_args=None,
     ):
+        self._exit_stack = contextlib.AsyncExitStack()
+
         # TODO: handle socket_options
-        self._session: Optional[aiohttp.ClientSession] = None
+        # keep track of sessions by proxy url (if any)
+        self._sessions: Dict[Optional[str], aiohttp.ClientSession] = {}
         self._verify = verify
         self._proxy_config = ProxyConfiguration(
             proxies=proxies, proxies_settings=proxies_config
@@ -93,53 +97,17 @@ def __init__(
         # it also pools by host so we don't need a manager, and can pass proxy via
         # request so don't need proxy manager
 
-        ssl_context = None
-        if bool(verify):
-            if proxies:
-                proxies_settings = self._proxy_config.settings
-                ssl_context = self._setup_proxy_ssl_context(proxies_settings)
-                # TODO: add support for
-                #    proxies_settings.get('proxy_use_forwarding_for_https')
-            else:
-                ssl_context = self._get_ssl_context()
-
-                # inline self._setup_ssl_cert
-                ca_certs = get_cert_path(verify)
-                if ca_certs:
-                    ssl_context.load_verify_locations(ca_certs, None, None)
-
-        self._create_connector = lambda: aiohttp.TCPConnector(
-            limit=max_pool_connections,
-            verify_ssl=bool(verify),
-            ssl=ssl_context,
-            **self._connector_args
-        )
-        self._connector = None
-
     async def __aenter__(self):
-        assert not self._session and not self._connector
+        assert not self._sessions
 
-        self._connector = self._create_connector()
-
-        self._session = aiohttp.ClientSession(
-            connector=self._connector,
-            timeout=self._timeout,
-            skip_auto_headers={'CONTENT-TYPE'},
-            auto_decompress=False,
-        )
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        if self._session:
-            await self._session.__aexit__(exc_type, exc_val, exc_tb)
-            self._session = None
-            self._connector = None
+        self._sessions.clear()
+        await self._exit_stack.aclose()
 
     def _get_ssl_context(self):
-        ssl_context = create_urllib3_context()
-        if self._cert_file:
-            ssl_context.load_cert_chain(self._cert_file, self._key_file)
-        return ssl_context
+        return create_urllib3_context()
 
     def _setup_proxy_ssl_context(self, proxy_url):
         proxies_settings = self._proxy_config.settings
@@ -167,6 +135,58 @@ def _setup_proxy_ssl_context(self, proxy_url):
         except (OSError, LocationParseError) as e:
             raise InvalidProxiesConfigError(error=e)
 
+    def _chunked(self, headers):
+        transfer_encoding = headers.get('Transfer-Encoding', '')
+        if chunked := transfer_encoding.lower() == 'chunked':
+            # aiohttp wants chunking as a param, and not a header
+            del headers['Transfer-Encoding']
+        return chunked or None
+
+    def _create_connector(self, proxy_url):
+        ssl_context = None
+        if bool(self._verify):
+            if proxy_url:
+                ssl_context = self._setup_proxy_ssl_context(proxy_url)
+                # TODO: add support for
+                #    proxies_settings.get('proxy_use_forwarding_for_https')
+            else:
+                ssl_context = self._get_ssl_context()
+
+            if ssl_context:
+                if self._cert_file:
+                    ssl_context.load_cert_chain(
+                        self._cert_file,
+                        self._key_file,
+                    )
+
+                # inline self._setup_ssl_cert
+                ca_certs = get_cert_path(self._verify)
+                if ca_certs:
+                    ssl_context.load_verify_locations(ca_certs, None, None)
+
+        return aiohttp.TCPConnector(
+            limit=self._max_pool_connections,
+            verify_ssl=bool(self._verify),
+            ssl=ssl_context,
+            **self._connector_args,
+        )
+
+    async def _get_session(self, proxy_url):
+        if not (session := self._sessions.get(proxy_url)):
+            connector = self._create_connector(proxy_url)
+            self._sessions[
+                proxy_url
+            ] = session = await self._exit_stack.enter_async_context(
+                aiohttp.ClientSession(
+                    connector=connector,
+                    timeout=self._timeout,
+                    skip_auto_headers={'CONTENT-TYPE'},
+                    auto_decompress=False,
+                ),
+            )
+
+        return session
+
     async def close(self):
         await self.__aexit__(None, None, None)
 
@@ -195,20 +215,15 @@ async def send(self, request):
             # https://github.com/boto/botocore/issues/1255
             headers_['Accept-Encoding'] = 'identity'
 
-            chunked = None
-            if headers_.get('Transfer-Encoding', '').lower() == 'chunked':
-                # aiohttp wants chunking as a param, and not a header
-                headers_.pop('Transfer-Encoding', '')
-                chunked = True
-
             if isinstance(data, io.IOBase):
                 data = _IOBaseWrapper(data)
 
             url = URL(url, encoded=True)
-            response = await self._session.request(
+            session = await self._get_session(proxy_url)
+            response = await session.request(
                 request.method,
                 url=url,
-                chunked=chunked,
+                chunked=self._chunked(headers_),
                 headers=headers_,
                 data=data,
                 proxy=proxy_url,