Add connection/read timeout for requests adapter

hyper/contrib.py

@@ -33,7 +33,7 @@ def __init__(self, *args, **kwargs):
         self.connections = {}
 
     def get_connection(self, host, port, scheme, cert=None, verify=True,
-                       proxy=None):
+                       proxy=None, timeout=None):
         """
         Gets an appropriate HTTP/2 connection object based on
         host/port/scheme/cert tuples.
@@ -77,7 +77,8 @@ def get_connection(self, host, port, scheme, cert=None, verify=True,
                 secure=secure,
                 ssl_context=ssl_context,
                 proxy_host=proxy_netloc,
-                proxy_headers=proxy_headers)
+                proxy_headers=proxy_headers,
+                timeout=timeout)
             self.connections[connection_key] = conn
 
         return conn
@@ -92,13 +93,15 @@ def send(self, request, stream=False, cert=None, verify=True, proxies=None,
             proxy = prepend_scheme_if_needed(proxy, 'http')
 
         parsed = urlparse(request.url)
+        timeout = kwargs.get('timeout')
         conn = self.get_connection(
             parsed.hostname,
             parsed.port,
             parsed.scheme,
             cert=cert,
             verify=verify,
-            proxy=proxy)
+            proxy=proxy,
+            timeout=timeout)
 
         # Build the selector.
         selector = parsed.path

hyper/http11/connection.py

@@ -101,7 +101,7 @@ class HTTP11Connection(object):
 
     def __init__(self, host, port=None, secure=None, ssl_context=None,
                  proxy_host=None, proxy_port=None, proxy_headers=None,
-                 **kwargs):
+                 timeout=None, **kwargs):
         if port is None:
             self.host, self.port = to_host_port_tuple(host, default_port=80)
         else:
@@ -150,6 +150,14 @@ def __init__(self, host, port=None, secure=None, ssl_context=None,
         #: the standard hyper parsing interface.
         self.parser = Parser()
 
+        # timeout
+        if isinstance(timeout, tuple):
+            self._connect_timeout = timeout[0]
+            self._read_timeout = timeout[1]
+        else:
+            self._connect_timeout = timeout
+            self._read_timeout = timeout
+
     def connect(self):
         """
         Connect to the server specified when the object was created. This is a
@@ -172,10 +180,11 @@ def connect(self):
                 # Simple http proxy
                 sock = socket.create_connection(
                     (self.proxy_host, self.proxy_port),
-                    5
+                    timeout=self._connect_timeout
                 )
             else:
-                sock = socket.create_connection((self.host, self.port), 5)
+                sock = socket.create_connection((self.host, self.port),
+                                                timeout=self._connect_timeout)
             proto = None
 
             if self.secure:
@@ -184,6 +193,9 @@ def connect(self):
             log.debug("Selected protocol: %s", proto)
             sock = BufferedSocket(sock, self.network_buffer_size)
 
+            # Set read timeout
+            sock.settimeout(self._read_timeout)
+
             if proto not in ('http/1.1', None):
                 raise TLSUpgrade(proto, sock)
 

hyper/http20/connection.py

@@ -102,7 +102,7 @@ class HTTP20Connection(object):
     def __init__(self, host, port=None, secure=None, window_manager=None,
                  enable_push=False, ssl_context=None, proxy_host=None,
                  proxy_port=None, force_proto=None, proxy_headers=None,
-                 **kwargs):
+                 timeout=None, **kwargs):
         """
         Creates an HTTP/2 connection to a specific server.
         """
@@ -151,6 +151,14 @@ def __init__(self, host, port=None, secure=None, window_manager=None,
         self.__wm_class = window_manager or FlowControlManager
         self.__init_state()
 
+        # timeout
+        if isinstance(timeout, tuple):
+            self._connect_timeout = timeout[0]
+            self._read_timeout = timeout[1]
+        else:
+            self._connect_timeout = timeout
+            self._read_timeout = timeout
+
         return
 
     def __init_state(self):
@@ -355,10 +363,12 @@ def connect(self):
             elif self.proxy_host:
                 # Simple http proxy
                 sock = socket.create_connection(
-                    (self.proxy_host, self.proxy_port)
+                    (self.proxy_host, self.proxy_port),
+                    timeout=self._connect_timeout
                 )
             else:
-                sock = socket.create_connection((self.host, self.port))
+                sock = socket.create_connection((self.host, self.port),
+                                                timeout=self._connect_timeout)
 
             if self.secure:
                 sock, proto = wrap_socket(sock, self.host, self.ssl_context,
@@ -374,6 +384,9 @@ def connect(self):
 
             self._sock = BufferedSocket(sock, self.network_buffer_size)
 
+            # Set read timeout
+            self._sock.settimeout(self._read_timeout)
+
             self._send_preamble()
 
     def _connect_upgrade(self, sock):

test/server.py

@@ -108,12 +108,13 @@ class SocketLevelTest(object):
     A test-class that defines a few helper methods for running socket-level
     tests.
     """
-    def set_up(self, secure=True, proxy=False):
+    def set_up(self, secure=True, proxy=False, timeout=None):
         self.host = None
         self.port = None
         self.socket_security = SocketSecuritySetting(secure)
         self.proxy = proxy
         self.server_thread = None
+        self.timeout = timeout
 
     def _start_server(self, socket_handler):
         """
@@ -146,18 +147,22 @@ def secure(self, value):
     def get_connection(self):
         if self.h2:
             if not self.proxy:
-                return HTTP20Connection(self.host, self.port, self.secure)
+                return HTTP20Connection(self.host, self.port, self.secure,
+                                        timeout=self.timeout)
             else:
                 return HTTP20Connection('http2bin.org', secure=self.secure,
                                         proxy_host=self.host,
-                                        proxy_port=self.port)
+                                        proxy_port=self.port,
+                                        timeout=self.timeout)
         else:
             if not self.proxy:
-                return HTTP11Connection(self.host, self.port, self.secure)
+                return HTTP11Connection(self.host, self.port, self.secure,
+                                        timeout=self.timeout)
             else:
                 return HTTP11Connection('httpbin.org', secure=self.secure,
                                         proxy_host=self.host,
-                                        proxy_port=self.port)
+                                        proxy_port=self.port,
+                                        timeout=self.timeout)
 
     def get_encoder(self):
         """

test/test_http11.py

@@ -110,6 +110,18 @@ def test_initialization_with_ipv6_addresses_proxy_inline_port(self):
         assert c.proxy_host == 'ffff:aaaa::1'
         assert c.proxy_port == 8443
 
+    def test_initialization_timeout(self):
+        c = HTTP11Connection('httpbin.org', timeout=30)
+
+        assert c._connect_timeout == 30
+        assert c._read_timeout == 30
+
+    def test_initialization_tuple_timeout(self):
+        c = HTTP11Connection('httpbin.org', timeout=(5, 60))
+
+        assert c._connect_timeout == 5
+        assert c._read_timeout == 60
+
     def test_basic_request(self):
         c = HTTP11Connection('httpbin.org')
         c._sock = sock = DummySocket()

test/test_integration.py

@@ -13,6 +13,7 @@
 import hyper
 import hyper.http11.connection
 import pytest
+from socket import timeout as SocketTimeout
 from contextlib import contextmanager
 from mock import patch
 from concurrent.futures import ThreadPoolExecutor, TimeoutError
@@ -1230,6 +1231,114 @@ def do_connect(conn):
 
         self.tear_down()
 
+    def test_connection_timeout(self):
+        self.set_up(timeout=0.5)
+
+        def socket_handler(listener):
+            time.sleep(1)
+            sock = listener.accept()[0]
+            sock.close()
+
+        self._start_server(socket_handler)
+        conn = self.get_connection()
+        try:
+            conn.connect()
+            assert False
+        except (SocketTimeout, ssl.SSLError):
+            # Py2 raises this as a BaseSSLError,
+            # Py3 raises it as socket timeout.
+            # assert 'timed out' in e.message
+            pass
+
+        self.tear_down()
+
+    def test_read_timeout(self):
+        self.set_up(timeout=0.5)
+
+        req_event = threading.Event()
+        recv_event = threading.Event()
+
+        def socket_handler(listener):
+            sock = listener.accept()[0]
+
+            # We get two messages for the connection open and then a HEADERS
+            # frame.
+            receive_preamble(sock)
+            sock.recv(65535)
+
+            # Wait for request
+            req_event.wait(5)
+            # Now, send the headers for the response. This response has no body
+            time.sleep(1)
+
+            # Now, send the headers for the response. This response has no body
+            f = build_headers_frame(
+                [(':status', '204'), ('content-length', '0')]
+            )
+            f.flags.add('END_STREAM')
+            f.stream_id = 1
+            sock.send(f.serialize())
+
+            # Wait for the message from the main thread.
+            recv_event.wait(5)
+            sock.close()
+
+        self._start_server(socket_handler)
+        conn = self.get_connection()
+        conn.request('GET', '/')
+        req_event.set()
+
+        try:
+            conn.get_response()
+            assert False
+        except (SocketTimeout, ssl.SSLError):
+            # Py2 raises this as a BaseSSLError,
+            # Py3 raises it as socket timeout.
+            # assert 'timed out' in e.message
+            pass
+
+        # Awesome, we're done now.
+        recv_event.set()
+        self.tear_down()
+
+    def test_default_connection_timeout(self):
+        self.set_up(timeout=None)
+
+        # Confirm that we send the connection upgrade string and the initial
+        # SettingsFrame.
+        data = []
+        send_event = threading.Event()
+
+        def socket_handler(listener):
+            time.sleep(1)
+            sock = listener.accept()[0]
+
+            # We should get one big chunk.
+            first = sock.recv(65535)
+            data.append(first)
+
+            # We need to send back a SettingsFrame.
+            f = SettingsFrame(0)
+            sock.send(f.serialize())
+
+            send_event.set()
+            sock.close()
+
+        self._start_server(socket_handler)
+        conn = self.get_connection()
+        try:
+            conn.connect()
+        except (SocketTimeout, ssl.SSLError):
+            # Py2 raises this as a BaseSSLError,
+            # Py3 raises it as socket timeout.
+            assert False
+
+        send_event.wait(5)
+
+        assert data[0].startswith(b'PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n')
+
+        self.tear_down()
+
 
 @patch('hyper.http20.connection.H2_NPN_PROTOCOLS', PROTOCOLS)
 class TestRequestsAdapter(SocketLevelTest):
@@ -1290,7 +1399,8 @@ def socket_handler(listener):
 
         s = requests.Session()
         s.mount('https://%s' % self.host, HTTP20Adapter())
-        r = s.get('https://%s:%s/some/path' % (self.host, self.port))
+        r = s.get('https://%s:%s/some/path' % (self.host, self.port),
+                  timeout=(10, 60))
 
         # Assert about the received values.
         assert r.status_code == 200