Improve timeouts handling

romis2012 · Jul 16, 2020 · 32792a5 · 32792a5
1 parent a2c61b2
commit 32792a5
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 13 deletions.
diff --git a/httpx_socks/__init__.py b/httpx_socks/__init__.py
@@ -1,5 +1,5 @@
 __title__ = 'httpx-socks'
-__version__ = '0.2.4'
+__version__ = '0.2.5'
 
 from .core_socks import (
     ProxyError,

diff --git a/httpx_socks/core_socks/_proxy_async_aio.py b/httpx_socks/core_socks/_proxy_async_aio.py
@@ -11,6 +11,8 @@
 from ._proto_http_async import HttpProto
 from ._proto_socks4_async import Socks4Proto
 
+DEFAULT_TIMEOUT = 60
+
 
 class Proxy:
     @classmethod
@@ -82,10 +84,23 @@ def __init__(self, loop: asyncio.AbstractEventLoop,
         self._stream = SocketStream(loop=loop)
 
     async def connect(self, dest_host, dest_port, timeout=None, _socket=None):
+        if timeout is None:
+            timeout = DEFAULT_TIMEOUT
+
         self._dest_host = dest_host
         self._dest_port = dest_port
         self._timeout = timeout
 
+        try:
+            await self._connect(_socket=_socket)
+        except asyncio.TimeoutError as e:
+            raise ProxyTimeoutError(
+                'Proxy connection timed out: %s'
+                % self._timeout) from e
+
+        return self._stream.socket
+
+    async def _connect(self, _socket=None):
         async with async_timeout.timeout(self._timeout):
             try:
                 await self._stream.open_connection(
@@ -99,26 +114,22 @@ async def connect(self, dest_host, dest_port, timeout=None, _socket=None):
                 msg = ('Can not connect to proxy %s:%s [%s]' %
                        (self._proxy_host, self._proxy_port, e.strerror))
                 raise ProxyConnectionError(e.errno, msg) from e
-            except asyncio.CancelledError as e:  # pragma: no cover
+            except Exception:  # pragma: no cover
                 await self._stream.close()
-                raise ProxyTimeoutError('Proxy connection timed out: %s'
-                                        % self._timeout) from e
+                raise
 
             try:
                 await self._negotiate()
-            except asyncio.CancelledError as e:  # pragma: no cover
+            except asyncio.CancelledError:  # pragma: no cover
                 # https://bugs.python.org/issue30064
                 # https://bugs.python.org/issue34795
                 if self._can_be_closed_safely():
                     await self._stream.close()
-                raise ProxyTimeoutError('Proxy connection timed out: %s'
-                                        % self._timeout) from e
+                raise
             except Exception:
                 await self._stream.close()
                 raise
 
-        return self._stream.socket
-
     def _can_be_closed_safely(self):  # pragma: no cover
         def is_proactor_event_loop():
             try:

diff --git a/httpx_socks/core_socks/_proxy_async_trio.py b/httpx_socks/core_socks/_proxy_async_trio.py
@@ -10,12 +10,14 @@
 from ._proto_http_async import HttpProto
 from ._proto_socks4_async import Socks4Proto
 
+DEFAULT_TIMEOUT = 60
+
 
 class Proxy:
     @classmethod
     def create(cls, proxy_type: ProxyType, host: str, port: int,
                username: str = None, password: str = None,
-               rdns: bool = None) -> 'BaseProxy':
+               rdns: bool = None) -> AsyncProxy:
 
         if proxy_type == ProxyType.SOCKS4:
             return Socks4Proxy(
@@ -46,7 +48,7 @@ def create(cls, proxy_type: ProxyType, host: str, port: int,
                          % proxy_type)
 
     @classmethod
-    def from_url(cls, url: str, **kwargs) -> 'BaseProxy':
+    def from_url(cls, url: str, **kwargs) -> AsyncProxy:
         proxy_type, host, port, username, password = parse_proxy_url(url)
         return cls.create(
             proxy_type=proxy_type,
@@ -71,7 +73,7 @@ def __init__(self, proxy_host, proxy_port):
 
     async def connect(self, dest_host, dest_port, timeout=None, _socket=None):
         if timeout is None:
-            timeout = 5
+            timeout = DEFAULT_TIMEOUT
 
         self._dest_host = dest_host
         self._dest_port = dest_port

diff --git a/httpx_socks/core_socks/_proxy_sync.py b/httpx_socks/core_socks/_proxy_sync.py
@@ -9,6 +9,8 @@
 from ._proto_http_sync import HttpProto
 from ._proto_socks4_sync import Socks4Proto
 
+DEFAULT_TIMEOUT = 60
+
 
 class SyncProxy:
     def connect(self, dest_host, dest_port, timeout=None, _socket=None):
@@ -82,6 +84,9 @@ def __init__(self, proxy_host, proxy_port):
         self._stream = SyncSocketStream()
 
     def connect(self, dest_host, dest_port, timeout=None, _socket=None):
+        if timeout is None:
+            timeout = DEFAULT_TIMEOUT
+
         self._dest_host = dest_host
         self._dest_port = dest_port
         self._timeout = timeout

diff --git a/httpx_socks/core_socks/_version.py b/httpx_socks/core_socks/_version.py
@@ -1,2 +1,2 @@
 __title__ = 'core-socks'
-__version__ = '0.1.3'
+__version__ = '0.1.4'