Merge pull request #68 from theseion/performance-improvement-http

Rewrote http.py to be faster

ftw/http.py

@@ -9,8 +9,8 @@
 import socket
 import ssl
 import sys
-import time
 import zlib
+import select
 
 import brotli
 from IPy import IP
@@ -19,6 +19,9 @@
 from . import util
 
 
+SOCKET_TIMEOUT = .3
+
+
 class HttpResponse(object):
     def __init__(self, http_response, user_agent):
         self.response = util.ensure_binary(http_response)
@@ -263,7 +266,6 @@ def __init__(self):
             'ADH-AES256-SHA:ECDHE-ECDSA-AES128-GCM-SHA256:' \
             'ECDHE-RSA-AES128-GCM-SHA256:AES128-GCM-SHA256:AES128-SHA256:HIGH:'
         self.CRLF = '\r\n'
-        self.HTTP_TIMEOUT = .3
         self.RECEIVE_BYTES = 8192
         self.SOCKET_TIMEOUT = 5
 
@@ -463,38 +465,77 @@ def get_response(self):
         """
         Get the response from the socket
         """
-        self.sock.setblocking(0)
         our_data = []
-        # Beginning time
-        begin = time.time()
+        self.sock.setblocking(False)
+        try:
+            our_data = self.read_response_from_socket()
+        finally:
+            try:
+                self.sock.shutdown(socket.SHUT_WR)
+                self.sock.close()
+            except OSError as err:
+                raise errors.TestError(
+                    'We were unable to close the socket as expected.',
+                    {
+                        'msg': err,
+                        'function': 'http.HttpUA.get_response'
+                    })
+            else:
+                self.response_object = HttpResponse(b''.join(our_data), self)
+            finally:
+                if not b''.join(our_data):
+                    raise errors.TestError(
+                        'No response from server.'
+                        + ' Request likely timed out.',
+                        {
+                            'host': self.request_object.dest_addr,
+                            'port': self.request_object.port,
+                            'proto': self.request_object.protocol,
+                            'msg': 'Please send the request and check'
+                            + ' Wireshark',
+                            'function': 'http.HttpUA.get_response'
+                        })
+
+    def read_response_from_socket(self):
+        # wait for socket to become ready
+        ready_sock, _, _ = select.select(
+            [self.sock], [], [self.sock], self.SOCKET_TIMEOUT)
+        if not ready_sock:
+            raise errors.TestError(
+                f'No response from server within {self.SOCKET_TIMEOUT}s',
+                {
+                    'host': self.request_object.dest_addr,
+                    'port': self.request_object.port,
+                    'proto': self.request_object.protocol,
+                    'msg': 'Please send the request and check Wireshark',
+                    'function': 'http.HttpUA.get_response'
+                })
+
+        our_data = []
         while True:
-            # If we have data then if we're passed the timeout break
-            if our_data and time.time() - begin > self.HTTP_TIMEOUT:
-                break
-            # If we're dataless wait just a bit
-            elif time.time() - begin > self.HTTP_TIMEOUT * 2:
-                break
-            # Recv data
             try:
                 data = self.sock.recv(self.RECEIVE_BYTES)
-                if data:
-                    our_data.append(util.ensure_binary(data))
-                    begin = time.time()
-                else:
-                    # Sleep for sometime to indicate a gap
-                    time.sleep(self.HTTP_TIMEOUT)
-            except socket.error as err:
-                # Check if we got a timeout
-                if err.errno == errno.EAGAIN:
-                    pass
+                if len(data) == 0:
+                    # we're done
+                    break
+                our_data.append(util.ensure_binary(data))
+            except BlockingIOError as e:
+                # If we can't handle the error here, pass it on
+                if e.errno == socket.EAGAIN or e.errno == socket.EWOULDBLOCK:
+                    # we're done
+                    break
+            except OSError as err:
                 # SSL will return SSLWantRead instead of EAGAIN
-                elif sys.platform == 'win32' and \
-                        err.errno == errno.WSAEWOULDBLOCK:
+                if (sys.platform == 'win32' and
+                        err.errno == errno.WSAEWOULDBLOCK):
                     pass
                 elif (self.request_object.protocol == 'https' and
-                      err.args[0] == ssl.SSL_ERROR_WANT_READ):
-                    continue
-                # If we didn't it's an error
+                        err.args[0] == ssl.SSL_ERROR_WANT_READ):
+                    ready_sock, _, _ = select.select(
+                        [self.sock], [], [self.sock], SOCKET_TIMEOUT)
+                    if not ready_sock:
+                        break
+                # It's an error
                 else:
                     raise errors.TestError(
                         'Failed to connect to server',
@@ -505,26 +546,4 @@ def get_response(self):
                             'message': err,
                             'function': 'http.HttpUA.get_response'
                         })
-        try:
-            self.sock.shutdown(socket.SHUT_WR)
-            self.sock.close()
-        except socket.error as err:
-            raise errors.TestError(
-                'We were unable to close the socket as expected.',
-                {
-                    'msg': err,
-                    'function': 'http.HttpUA.get_response'
-                })
-        else:
-            self.response_object = HttpResponse(b''.join(our_data), self)
-        finally:
-            if not b''.join(our_data):
-                raise errors.TestError(
-                    'No response from server. Request likely timed out.',
-                    {
-                        'host': self.request_object.dest_addr,
-                        'port': self.request_object.port,
-                        'proto': self.request_object.protocol,
-                        'msg': 'Please send the request and check Wireshark',
-                        'function': 'http.HttpUA.get_response'
-                    })
+        return our_data