Merge pull request #21 from nibrag/connector-fix

Connector fix

Author: nibrag

.coveragerc

@@ -1,7 +1,7 @@
 [run]
 branch = True
 source = aiosocks, tests
-omit = site-packages
+omit = site-packages,aiosocks/test_utils.py
 
 [html]
 directory = coverage

.travis.yml

@@ -18,7 +18,7 @@ install:
   - pip install --upgrade pip wheel
   - pip install --upgrade setuptools
   - pip install pip
-  - pip install flake8
+  - pip install flake8==3.3.0
   - pip install pyflakes==1.1.0
   - pip install coverage
   - pip install pytest

README.rst

@@ -15,7 +15,7 @@ SOCKS proxy client for asyncio and aiohttp
 Dependencies
 ------------
 python 3.5+
-aiohttp 2.1+
+aiohttp 2.3.2+
 
 Features
 --------
@@ -175,8 +175,10 @@ aiohttp usage
                                proxy_auth=ba) as resp:
           if resp.status == 200:
             print(await resp.text())
-    except aiohttp.ProxyConnectionError:
+    except aiohttp.ClientProxyConnectionError:
       # connection problem
+    except aiohttp.ClientConnectorError:
+      # ssl error, certificate error, etc
     except aiosocks.SocksError:
       # communication problem
 
@@ -185,22 +187,3 @@ aiohttp usage
     loop = asyncio.get_event_loop()
     loop.run_until_complete(load_github_main())
     loop.close()
-
-Proxy from environment
-^^^^^^^^^^^^^^^^^^^^^^
-
-.. code-block:: python
-
-  import os
-  from aiosocks.connector import ProxyConnector, ProxyClientRequest
-
-  os.environ['socks4_proxy'] = 'socks4://127.0.0.1:333'
-  # or
-  os.environ['socks5_proxy'] = 'socks5://127.0.0.1:444'
-
-  conn = ProxyConnector()
-
-  with aiohttp.ClientSession(connector=conn, request_class=ProxyClientRequest) as session:
-        async with session.get('http://github.com/', proxy_from_env=True) as resp:
-          if resp.status == 200:
-            print(await resp.text())

aiosocks/__init__.py

@@ -8,7 +8,7 @@
 )
 from .protocols import Socks4Protocol, Socks5Protocol, DEFAULT_LIMIT
 
-__version__ = '0.2.4'
+__version__ = '0.2.5'
 
 __all__ = ('Socks4Protocol', 'Socks5Protocol', 'Socks4Auth',
            'Socks5Auth', 'Socks4Addr', 'Socks5Addr', 'SocksError',
@@ -79,7 +79,7 @@ def socks_factory():
 
     try:
         await waiter
-    except:
+    except:  # noqa
         transport.close()
         raise
 

aiosocks/connector.py

@@ -1,30 +1,24 @@
 try:
     import aiohttp
     from aiohttp.connector import sentinel
+    from aiohttp.client_exceptions import certificate_errors, ssl_errors
 except ImportError:
     raise ImportError('aiosocks.SocksConnector require aiohttp library')
-
-from yarl import URL
-from urllib.request import getproxies
-
-from .errors import SocksError, SocksConnectionError
+from .errors import SocksConnectionError
 from .helpers import Socks4Auth, Socks5Auth, Socks4Addr, Socks5Addr
 from . import create_connection
 
 __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()
+from distutils.version import StrictVersion
 
-            proxy_url = proxies.get(self.original_url.scheme)
-            if not proxy_url:
-                proxy_url = proxies.get('socks4') or proxies.get('socks5')
+if StrictVersion(aiohttp.__version__) < StrictVersion('2.3.2'):
+    raise RuntimeError('aiosocks.connector depends on aiohttp 2.3.2+')
 
-            proxy = URL(proxy_url) if proxy_url else None
 
+class ProxyClientRequest(aiohttp.ClientRequest):
+    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 +35,9 @@ def update_proxy(self, proxy, proxy_auth, proxy_from_env):
                     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 +63,41 @@ async def _create_proxy_connection(self, req):
         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 +106,37 @@ async def _create_socks_connection(self, req):
                 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

aiosocks/test_utils.py

@@ -113,7 +113,7 @@ async def retranslator(self, reader, writer):
                         data.append(byte[0])
                         writer.write(byte)
                         await writer.drain()
-                    except:
+                    except:  # noqa
                         break
 
         def factory():

tests/test_connector.py

@@ -1,3 +1,5 @@
+import ssl
+
 import aiosocks
 import aiohttp
 import pytest
@@ -9,22 +11,21 @@
 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()
@@ -89,20 +90,40 @@ async def test_connect_locale_resolve(loop):
     conn.close()
 
 
-async def test_proxy_connect_fail(loop):
+@pytest.mark.parametrize('remote_resolve', [True, False])
+async def test_resolve_host_fail(loop, remote_resolve):
+    tr, proto = mock.Mock(name='transport'), mock.Mock(name='protocol')
+
+    with mock.patch('aiosocks.connector.create_connection',
+                    make_mocked_coro((tr, proto))):
+        req = ProxyClientRequest(
+            'GET', URL('http://python.org'), loop=loop,
+            proxy=URL('socks5://proxy.example'))
+        connector = ProxyConnector(loop=loop, remote_resolve=remote_resolve)
+        connector._resolve_host = make_mocked_coro(raise_exception=OSError())
+
+        with pytest.raises(aiohttp.ClientConnectorError):
+            await connector.connect(req)
+
+
+@pytest.mark.parametrize('exc', [
+    (ssl.CertificateError, aiohttp.ClientConnectorCertificateError),
+    (ssl.SSLError, aiohttp.ClientConnectorSSLError),
+    (aiosocks.SocksConnectionError, aiohttp.ClientProxyConnectionError)])
+async def test_proxy_connect_fail(loop, exc):
     loop_mock = mock.Mock()
     loop_mock.getaddrinfo = make_mocked_coro(
         [[0, 0, 0, 0, ['127.0.0.1', 1080]]])
     cc_coro = make_mocked_coro(
-        raise_exception=aiosocks.SocksConnectionError())
+        raise_exception=exc[0]())
 
     with mock.patch('aiosocks.connector.create_connection', cc_coro):
         req = ProxyClientRequest(
             'GET', URL('http://python.org'), loop=loop,
             proxy=URL('socks5://127.0.0.1'))
         connector = ProxyConnector(loop=loop_mock)
 
-        with pytest.raises(aiohttp.ClientConnectionError):
+        with pytest.raises(exc[1]):
             await connector.connect(req)
 
 
@@ -177,38 +198,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')