close response on connection errors during read

test/with_dummyserver/test_socketlevel.py

@@ -297,6 +297,64 @@ def socket_handler(listener):
         self.assertEqual(response.status, 200)
         self.assertEqual(response.data, b'foo')
 
+    def test_connection_cleanup_on_read_timeout(self):
+        timed_out = Event()
+
+        def socket_handler(listener):
+            sock = listener.accept()[0]
+            buf = b''
+            body = 'Hi'
+            while not buf.endswith(b'\r\n\r\n'):
+                buf = sock.recv(65536)
+            sock.send(('HTTP/1.1 200 OK\r\n'
+                       'Content-Type: text/plain\r\n'
+                       'Content-Length: %d\r\n'
+                       '\r\n' % len(body)).encode('utf-8'))
+
+            timed_out.wait()
+            sock.send(body.encode('utf-8'))
+            sock.close()
+
+        self._start_server(socket_handler)
+        with HTTPConnectionPool(self.host, self.port) as pool:
+            poolsize = pool.pool.qsize()
+            response = pool.urlopen('GET', '/', retries=0, preload_content=False,
+                                    timeout=Timeout(connect=1, read=0.001))
+            try:
+                self.assertRaises(ReadTimeoutError, response.read)
+                self.assertEqual(poolsize, pool.pool.qsize())
+            finally:
+                timed_out.set()
+
+    def test_connection_cleanup_on_protocol_error_during_read(self):
+        body = 'Response'
+        partial_body = body[:2]
+
+        def socket_handler(listener):
+            sock = listener.accept()[0]
+
+            # Consume request
+            buf = b''
+            while not buf.endswith(b'\r\n\r\n'):
+                buf = sock.recv(65536)
+
+            # Send partial response and close socket.
+            sock.send((
+                'HTTP/1.1 200 OK\r\n'
+                'Content-Type: text/plain\r\n'
+                'Content-Length: %d\r\n'
+                '\r\n'
+                '%s' % (len(body), partial_body)).encode('utf-8')
+            )
+            sock.close()
+
+        self._start_server(socket_handler)
+        with HTTPConnectionPool(self.host, self.port) as pool:
+            poolsize = pool.pool.qsize()
+            response = pool.request('GET', '/', retries=0, preload_content=False)
+
+            self.assertRaises(ProtocolError, response.read)
+            self.assertEqual(poolsize, pool.pool.qsize())
 
 
 class TestProxyManager(SocketDummyServerTestCase):

urllib3/response.py

@@ -231,57 +231,74 @@ def read(self, amt=None, decode_content=None, cache_content=False):
             return
 
         flush_decoder = False
+        data = None
 
         try:
-            try:
-                if amt is None:
-                    # cStringIO doesn't like amt=None
-                    data = self._fp.read()
+            if amt is None:
+                # cStringIO doesn't like amt=None
+                data = self._fp.read()
+                flush_decoder = True
+            else:
+                cache_content = False
+                data = self._fp.read(amt)
+                if amt != 0 and not data:  # Platform-specific: Buggy versions of Python.
+                    # Close the connection when no data is returned
+                    #
+                    # This is redundant to what httplib/http.client _should_
+                    # already do.  However, versions of python released before
+                    # December 15, 2012 (http://bugs.python.org/issue16298) do
+                    # not properly close the connection in all cases. There is
+                    # no harm in redundantly calling close.
+                    self._fp.close()
                     flush_decoder = True
-                else:
-                    cache_content = False
-                    data = self._fp.read(amt)
-                    if amt != 0 and not data:  # Platform-specific: Buggy versions of Python.
-                        # Close the connection when no data is returned
-                        #
-                        # This is redundant to what httplib/http.client _should_
-                        # already do.  However, versions of python released before
-                        # December 15, 2012 (http://bugs.python.org/issue16298) do
-                        # not properly close the connection in all cases. There is
-                        # no harm in redundantly calling close.
-                        self._fp.close()
-                        flush_decoder = True
-
-            except SocketTimeout:
-                # FIXME: Ideally we'd like to include the url in the ReadTimeoutError but
-                # there is yet no clean way to get at it from this context.
-                raise ReadTimeoutError(self._pool, None, 'Read timed out.')
-
-            except BaseSSLError as e:
-                # FIXME: Is there a better way to differentiate between SSLErrors?
-                if 'read operation timed out' not in str(e):  # Defensive:
-                    # This shouldn't happen but just in case we're missing an edge
-                    # case, let's avoid swallowing SSL errors.
-                    raise
-
-                raise ReadTimeoutError(self._pool, None, 'Read timed out.')
-
-            except HTTPException as e:
-                # This includes IncompleteRead.
-                raise ProtocolError('Connection broken: %r' % e, e)
 
+        except SocketTimeout:
+            # the response may not be closed but we're not going to use it anymore
+            # so close it now to ensure that the connection is released back to the pool
+            if self._original_response and not self._original_response.isclosed():
+                self._original_response.close()
+
+            # FIXME: Ideally we'd like to include the url in the ReadTimeoutError but
+            # there is yet no clean way to get at it from this context.
+            raise ReadTimeoutError(self._pool, None, 'Read timed out.')
+
+        except BaseSSLError as e:
+            # the response may not be closed but we're not going to use it anymore
+            # so close it now to ensure that the connection is released back to the pool
+            if self._original_response and not self._original_response.isclosed():
+                self._original_response.close()
+
+            # FIXME: Is there a better way to differentiate between SSLErrors?
+            if 'read operation timed out' not in str(e):  # Defensive:
+                # This shouldn't happen but just in case we're missing an edge
+                # case, let's avoid swallowing SSL errors.
+                raise
+
+            raise ReadTimeoutError(self._pool, None, 'Read timed out.')
+
+        except HTTPException as e:
+            # the response may not be closed but we're not going to use it anymore
+            # so close it now to ensure that the connection is released back to the pool
+            if self._original_response and not self._original_response.isclosed():
+                self._original_response.close()
+
+            # This includes IncompleteRead.
+            raise ProtocolError('Connection broken: %r' % e, e)
+
+        finally:
+            if self._original_response and self._original_response.isclosed():
+                self.release_conn()
+
+        if data:
             self._fp_bytes_read += len(data)
 
             data = self._decode(data, decode_content, flush_decoder)
 
             if cache_content:
                 self._body = data
 
-            return data
+        return data
 
-        finally:
-            if self._original_response and self._original_response.isclosed():
-                self.release_conn()
 
     def stream(self, amt=2**16, decode_content=None):
         """