aiohttp connector new api supporting

aiosocks/connector.py

@@ -1,12 +1,11 @@
+from aiohttp.client_exceptions import certificate_errors, ssl_errors
+
 try:
     import aiohttp
     from aiohttp.connector import sentinel
 except ImportError:
     raise ImportError('aiosocks.SocksConnector require aiohttp library')
 
-from yarl import URL
-from urllib.request import getproxies
-
 from .errors import SocksError, SocksConnectionError
 from .helpers import Socks4Auth, Socks5Auth, Socks4Addr, Socks5Addr
 from . import create_connection
@@ -15,16 +14,7 @@ __all__ = ('ProxyConnector', 'ProxyClientRequest')
 
 
 class ProxyClientRequest(aiohttp.ClientRequest):
-    def update_proxy(self, proxy, proxy_auth, proxy_from_env):
-        if proxy_from_env and not proxy:
-            proxies = getproxies()
-
-            proxy_url = proxies.get(self.original_url.scheme)
-            if not proxy_url:
-                proxy_url = proxies.get('socks4') or proxies.get('socks5')
-
-            proxy = URL(proxy_url) if proxy_url else None
-
+    def update_proxy(self, proxy, proxy_auth, proxy_headers):
         if proxy and proxy.scheme not in ['http', 'socks4', 'socks5']:
             raise ValueError(
                 "Only http, socks4 and socks5 proxies are supported")
@@ -41,9 +31,9 @@ class ProxyClientRequest(aiohttp.ClientRequest):
                     not isinstance(proxy_auth, Socks5Auth):
                 raise ValueError("proxy_auth must be None or Socks5Auth() "
                                  "tuple for socks5 proxy")
-
         self.proxy = proxy
         self.proxy_auth = proxy_auth
+        self.proxy_headers = proxy_headers
 
 
 class ProxyConnector(aiohttp.TCPConnector):
@@ -69,20 +59,36 @@ class ProxyConnector(aiohttp.TCPConnector):
         else:
             return await self._create_socks_connection(req)
 
+    async def _wrap_create_socks_connection(self, *args, req, **kwargs):
+        try:
+            return await create_connection(*args, **kwargs)
+        except certificate_errors as exc:
+            raise aiohttp.ClientConnectorCertificateError(
+                req.connection_key, exc) from exc
+        except ssl_errors as exc:
+            raise aiohttp.ClientConnectorSSLError(req.connection_key, exc) from exc
+        except (OSError, SocksConnectionError) as exc:
+            raise aiohttp.ClientProxyConnectionError(req.connection_key, exc) from exc
+
     async def _create_socks_connection(self, req):
-        if req.ssl:
-            sslcontext = self.ssl_context
-        else:
-            sslcontext = None
+        sslcontext = self._get_ssl_context(req)
+        fingerprint, hashfunc = self._get_fingerprint_and_hashfunc(req)
 
         if not self._remote_resolve:
-            dst_hosts = list(await self._resolve_host(req.host, req.port))
-            dst = dst_hosts[0]['host'], dst_hosts[0]['port']
+            try:
+                dst_hosts = list(await self._resolve_host(req.host, req.port))
+                dst = dst_hosts[0]['host'], dst_hosts[0]['port']
+            except OSError as exc:
+                raise aiohttp.ClientConnectorError(req.connection_key, exc) from exc
         else:
             dst = req.host, req.port
 
-        proxy_hosts = await self._resolve_host(req.proxy.host, req.proxy.port)
-        exc = None
+        try:
+            proxy_hosts = await self._resolve_host(req.proxy.host, req.proxy.port)
+        except OSError as exc:
+            raise aiohttp.ClientConnectorError(req.connection_key, exc) from exc
+
+        last_exc = None
 
         for hinfo in proxy_hosts:
             if req.proxy.scheme == 'socks4':
@@ -91,45 +97,37 @@ class ProxyConnector(aiohttp.TCPConnector):
                 proxy = Socks5Addr(hinfo['host'], hinfo['port'])
 
             try:
-                transp, proto = await create_connection(
+                transp, proto = await self._wrap_create_socks_connection(
                     self._factory, proxy, req.proxy_auth, dst,
                     loop=self._loop, remote_resolve=self._remote_resolve,
                     ssl=sslcontext, family=hinfo['family'],
                     proto=hinfo['proto'], flags=hinfo['flags'],
-                    local_addr=self._local_addr,
+                    local_addr=self._local_addr, req=req,
                     server_hostname=req.host if sslcontext else None)
-
-                self._validate_ssl_fingerprint(transp, req.host, req.port)
-                return transp, proto
-            except (OSError, SocksError, SocksConnectionError) as e:
-                exc = e
+            except aiohttp.ClientConnectorError as exc:
+                last_exc = exc
+                continue
+
+            has_cert = transp.get_extra_info('sslcontext')
+            if has_cert and fingerprint:
+                sock = transp.get_extra_info('socket')
+                if not hasattr(sock, 'getpeercert'):
+                    # Workaround for asyncio 3.5.0
+                    # Starting from 3.5.1 version
+                    # there is 'ssl_object' extra info in transport
+                    sock = transp._ssl_protocol._sslpipe.ssl_object
+                # gives DER-encoded cert as a sequence of bytes (or None)
+                cert = sock.getpeercert(binary_form=True)
+                assert cert
+                got = hashfunc(cert).digest()
+                expected = fingerprint
+                if got != expected:
+                    transp.close()
+                    if not self._cleanup_closed_disabled:
+                        self._cleanup_closed_transports.append(transp)
+                    last_exc = aiohttp.ServerFingerprintMismatch(
+                        expected, got, req.host, req.port)
+                    continue
+            return transp, proto
         else:
-            if isinstance(exc, SocksConnectionError):
-                raise aiohttp.ClientProxyConnectionError(*exc.args)
-            if isinstance(exc, SocksError):
-                raise exc
-            else:
-                raise aiohttp.ClientOSError(
-                    exc.errno, 'Can not connect to %s:%s [%s]' %
-                               (req.host, req.port, exc.strerror)) from exc
-
-    def _validate_ssl_fingerprint(self, transp, host, port):
-        has_cert = transp.get_extra_info('sslcontext')
-        if has_cert and self._fingerprint:
-            sock = transp.get_extra_info('socket')
-            if not hasattr(sock, 'getpeercert'):
-                # Workaround for asyncio 3.5.0
-                # Starting from 3.5.1 version
-                # there is 'ssl_object' extra info in transport
-                sock = transp._ssl_protocol._sslpipe.ssl_object
-            # gives DER-encoded cert as a sequence of bytes (or None)
-            cert = sock.getpeercert(binary_form=True)
-            assert cert
-            got = self._hashfunc(cert).digest()
-            expected = self._fingerprint
-            if got != expected:
-                transp.close()
-                if not self._cleanup_closed_disabled:
-                    self._cleanup_closed_transports.append(transp)
-                raise aiohttp.ServerFingerprintMismatch(
-                    expected, got, host, port)
+            raise last_exc

tests/test_connector.py

@@ -9,22 +9,21 @@ from aiosocks.connector import ProxyConnector, ProxyClientRequest
 from aiosocks.helpers import Socks4Auth, Socks5Auth
 
 
-async def test_connect_proxy_ip():
+async def test_connect_proxy_ip(loop):
     tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol')
 
     with mock.patch('aiosocks.connector.create_connection',
                     make_mocked_coro((tr, proto))):
-        loop_mock = mock.Mock()
-        loop_mock.getaddrinfo = make_mocked_coro(
-            [[0, 0, 0, 0, ['127.0.0.1', 1080]]])
+        loop.getaddrinfo = make_mocked_coro(
+             [[0, 0, 0, 0, ['127.0.0.1', 1080]]])
 
         req = ProxyClientRequest(
-            'GET', URL('http://python.org'), loop=loop_mock,
+            'GET', URL('http://python.org'), loop=loop,
             proxy=URL('socks5://proxy.org'))
-        connector = ProxyConnector(loop=loop_mock)
+        connector = ProxyConnector(loop=loop)
         conn = await connector.connect(req)
 
-    assert loop_mock.getaddrinfo.called
+    assert loop.getaddrinfo.called
     assert conn.protocol is proto
 
     conn.close()
@@ -177,38 +176,3 @@ def test_proxy_client_request_invalid(loop):
             proxy=URL('socks5://proxy.org'), proxy_auth=Socks4Auth('l'))
     assert 'proxy_auth must be None or Socks5Auth() ' \
            'tuple for socks5 proxy' in str(cm)
-
-
-def test_proxy_from_env_http(loop):
-    proxies = {'http': 'http://proxy.org'}
-
-    with mock.patch('aiosocks.connector.getproxies',  return_value=proxies):
-        req = ProxyClientRequest('GET', URL('http://python.org'), loop=loop)
-        req.update_proxy(None, None, True)
-        assert req.proxy == URL('http://proxy.org')
-
-        req.original_url = URL('https://python.org')
-        req.update_proxy(None, None, True)
-        assert req.proxy is None
-
-        proxies.update({'https': 'http://proxy.org',
-                        'socks4': 'socks4://127.0.0.1:33',
-                        'socks5': 'socks5://localhost:44'})
-        req.update_proxy(None, None, True)
-        assert req.proxy == URL('http://proxy.org')
-
-
-def test_proxy_from_env_socks(loop):
-    proxies = {'socks4': 'socks4://127.0.0.1:33',
-               'socks5': 'socks5://localhost:44'}
-
-    with mock.patch('aiosocks.connector.getproxies', return_value=proxies):
-        req = ProxyClientRequest('GET', URL('http://python.org'), loop=loop)
-
-        req.update_proxy(None, None, True)
-        assert req.proxy == URL('socks4://127.0.0.1:33')
-
-        del proxies['socks4']
-
-        req.update_proxy(None, None, True)
-        assert req.proxy == URL('socks5://localhost:44')