[WIP] Extending the capabilities of the HTTP/2.0 server

* Made multithreading more robust, and made it so each stream gets its
own thread

* Requests start getting parsed immediately, and do not wait for the
entire request to get there. This helps with 404 detection early on and
other use cases.

* Created H2Request object

* Improved `write_push` docstring

* Improved flexibility of `write_push`, the user can now choose wether or not to immediately push.

* Using `six.moves` for python compatability with Queue

tools/wptserve/wptserve/request.py

@@ -278,15 +278,14 @@ def __init__(self, request_handler):
 
         self.raw_input = InputFile(request_handler.rfile,
                                    int(self.headers.get("Content-Length", 0)))
+
         self._body = None
 
         self._GET = None
         self._POST = None
         self._cookies = None
         self._auth = None
 
-        self.h2_stream_id = request_handler.h2_stream_id if hasattr(request_handler, 'h2_stream_id') else None
-
         self.server = Server(self)
 
     def __repr__(self):
@@ -349,6 +348,13 @@ def auth(self):
         return self._auth
 
 
+class H2Request(Request):
+    def __init__(self, request_handler):
+        self.h2_stream_id = request_handler.h2_stream_id
+        self.frames = []
+        super(H2Request, self).__init__(request_handler)
+
+
 class RequestHeaders(dict):
     """Dictionary-like API for accessing request headers."""
     def __init__(self, items):

tools/wptserve/wptserve/response.py

@@ -387,7 +387,7 @@ def __init__(self, handler, response):
         self.request = response.request
         self.logger = response.logger
 
-    def write_headers(self, headers, status_code, status_message=None):
+    def write_headers(self, headers, status_code, status_message=None, stream_id=None, last=False):
         formatted_headers = []
         secondary_headers = []  # Non ':' prefixed headers are to be added afterwards
 
@@ -403,13 +403,14 @@ def write_headers(self, headers, status_code, status_message=None):
 
         with self.h2conn as connection:
             connection.send_headers(
-                stream_id=self.request.h2_stream_id,
+                stream_id=self.request.h2_stream_id if stream_id is None else stream_id,
                 headers=formatted_headers,
+                end_stream=last or self.request.method == "HEAD"
             )
 
             self.write(connection)
 
-    def write_content(self, item, last=False):
+    def write_content(self, item, last=False, stream_id=None):
         if isinstance(item, (text_type, binary_type)):
             data = BytesIO(self.encode(item))
         else:
@@ -427,18 +428,48 @@ def write_content(self, item, last=False):
             data_len -= payload_size
             payload_size = self.get_max_payload_size()
 
-        self.write_content_frame(data.read(), last)
+        self.write_content_frame(data.read(), last, stream_id)
 
-    def write_content_frame(self, data, last):
+    def write_content_frame(self, data, last, stream_id=None):
         with self.h2conn as connection:
             connection.send_data(
-                stream_id=self.request.h2_stream_id,
+                stream_id=self.request.h2_stream_id if stream_id is None else stream_id,
                 data=data,
                 end_stream=last,
             )
             self.write(connection)
         self.content_written = last
 
+    def write_push(self, promise_headers, push_stream_id=None, status=None, response_headers=None, response_data=None):
+        """Write a push promise, and optionally write the push content.
+
+        This will write a push promise to the request stream. If you do not provide headers and data for the response,
+        then no response will be pushed, and you should push them yourself using the ID returned from this function
+
+        :param promise_headers: A list of header tuples that matches what the client would use to
+                                request the pushed response
+        :param push_stream_id: The ID of the stream the response should be pushed to. If none given, will
+                               use the next available id.
+        :param status: The status code of the response, REQUIRED if response_headers given
+        :param response_headers: The headers of the response
+        :param response_data: The response data.
+        :return: The ID of the push stream
+        """
+        with self.h2conn as connection:
+            push_stream_id = push_stream_id if push_stream_id is not None else connection.get_next_available_stream_id()
+            connection.push_stream(self.request.h2_stream_id, push_stream_id, promise_headers)
+            self.write(connection)
+
+        has_data = response_data is not None
+        if response_headers is not None:
+            assert status is not None
+            self.write_headers(response_headers, status, stream_id=push_stream_id, last=not has_data)
+
+        if has_data:
+            self.write_content(response_data, last=True, stream_id=push_stream_id)
+
+        return push_stream_id
+
     def get_max_payload_size(self):
         with self.h2conn as connection:
             return min(connection.remote_settings.max_frame_size, connection.local_flow_control_window(self.request.h2_stream_id)) - 9

tools/wptserve/wptserve/server.py

@@ -12,16 +12,19 @@
 import uuid
 from collections import OrderedDict
 
+from six.moves.queue import Queue
+from copy import copy
+
 from h2.config import H2Configuration
 from h2.connection import H2Connection
-from h2.events import RequestReceived, ConnectionTerminated
+from h2.events import RequestReceived, ConnectionTerminated, DataReceived, StreamReset
 
 from six.moves.urllib.parse import urlsplit, urlunsplit
 
 from . import routes as default_routes
 from .config import ConfigBuilder
 from .logger import get_logger
-from .request import Server, Request
+from .request import Server, Request, H2Request
 from .response import Response, H2Response
 from .router import Router
 from .utils import HTTPException
@@ -220,11 +223,12 @@ def __init__(self, *args, **kwargs):
         self.logger = get_logger()
         BaseHTTPServer.BaseHTTPRequestHandler.__init__(self, *args, **kwargs)
 
-    def finish_handling(self, request_line_is_valid, response_cls):
-            self.server.rewriter.rewrite(self)
+    def finish_handling(self, request_line_is_valid, response_cls, request=None, response=None):
+            if request is None or response is None:
+                self.server.rewriter.rewrite(self)
 
-            request = Request(self)
-            response = response_cls(self, request)
+                request = Request(self)
+                response = response_cls(self, request)
 
             if request.method == "CONNECT":
                 self.handle_connect(response)
@@ -254,6 +258,7 @@ def finish_handling(self, request_line_is_valid, response_cls):
                 time.sleep(latency / 1000.)
 
             if handler is None:
+                self.logger.debug("No Handler found!")
                 response.set_error(404)
             else:
                 try:
@@ -322,9 +327,9 @@ def handle_one_request(self):
         self.close_connection = False
 
         # Generate a UUID to make it easier to distinguish different H2 connection debug messages
-        uid = uuid.uuid4()
+        self.uid = str(uuid.uuid4())[:8]
 
-        self.logger.debug('(%s) Initiating h2 Connection' % uid)
+        self.logger.debug('(%s) Initiating h2 Connection' % self.uid)
 
         with self.conn as connection:
             connection.initiate_connection()
@@ -337,48 +342,113 @@ def handle_one_request(self):
         # TODO Need to do some major work on multithreading. Current idea is to have a thread per stream
         # so that processing of the request can start from the first frame.
 
+        self.stream_queues = {}
+
         while not self.close_connection:
             try:
                 # This size may need to be made variable based on remote settings?
                 data = self.request.recv(65535)
 
                 with self.conn as connection:
-                    events = connection.receive_data(data)
-
-                self.logger.debug('(%s) Events: ' % (uid) + str(events))
-
-                for event in events:
-                    if isinstance(event, RequestReceived):
-                        self.logger.debug('(%s) Parsing RequestReceived' % (uid))
-                        self._h2_parse_request(event)
-                        t = threading.Thread(target=BaseWebTestRequestHandler.finish_handling, args=(self, True, H2Response))
-                        self.request_threads.append(t)
-                        t.start()
-                    if isinstance(event, ConnectionTerminated):
-                        self.logger.debug('(%s) Connection terminated by remote peer ' % (uid))
+                    frames = connection.receive_data(data)
+
+                self.logger.debug('(%s) Frames Received: ' % (self.uid) + str(frames))
+
+                for frame in frames:
+                    if isinstance(frame, ConnectionTerminated):
+                        self.logger.debug('(%s) Connection terminated by remote peer ' % (self.uid))
                         self.close_connection = True
 
+                        # Flood all the streams with connection terminated, this will cause them to stop
+                        for stream_id, queue in self.stream_queues.items():
+                            queue.put(frame)
+                    elif hasattr(frame, 'stream_id'):
+                        if frame.stream_id not in self.stream_queues:
+                            self.stream_queues[frame.stream_id] = Queue()
+                            self.start_stream_thread(frame)
+                        self.stream_queues[frame.stream_id].put(frame)
+
             except (socket.timeout, socket.error) as e:
-                self.logger.debug('(%s) ERROR - Closing Connection - \n%s' % (uid, str(e)))
+                self.logger.debug('(%s) ERROR - Closing Connection - \n%s' % (self.uid, str(e)))
                 self.close_connection = True
                 for t in self.request_threads:
                     t.join()
 
-    def _h2_parse_request(self, event):
-        self.headers = H2Headers(event.headers)
+    def start_stream_thread(self, frame):
+        t = threading.Thread(
+            target=Http2WebTestRequestHandler._stream_thread,
+            args=(self, frame.stream_id)
+        )
+        self.request_threads.append(t)
+        t.start()
+
+    def _stream_thread(self, stream_id):
+        """
+        This thread processes frames for a specific stream. It waits for frames to be placed
+        in the queue, and processes them. When it receives a request frame, it will start processing
+        immediately, even if there are data frames to follow. One of the reasons for this is that it
+        can detect invalid requests before needing to read the rest of the frames.
+        """
+
+        # The file-like pipe object that will be used to share data to request object if data is received
+        wfile = None
+        request = None
+        while not self.close_connection:
+            # Wait for next frame, blocking
+            frame = self.stream_queues[stream_id].get(True, None)
+
+            self.logger.debug('(%s - %s) %s' % (self.uid, stream_id, str(frame)))
+
+            if isinstance(frame, RequestReceived):
+                # Create a shallow copy of this instance, to avoid issues with multiple threads editing the same object
+                self_copy = copy(self)
+                self_copy._h2_parse_request(frame)
+
+                rfile, wfile = os.pipe()
+                rfile, wfile = os.fdopen(rfile, 'rb'), os.fdopen(wfile, 'wb')
+                self_copy.rfile = rfile
+
+                self_copy.server.rewriter.rewrite(self_copy)
+                request = H2Request(self_copy)
+                request.frames.append(frame)
+                response = H2Response(self_copy, request)
+
+                # Begin processing the response in another thread,
+                # and continue listening for more data here in this one
+                t = threading.Thread(
+                    target=BaseWebTestRequestHandler.finish_handling,
+                    args=(self, True, H2Response,),
+                    kwargs={'request':request, 'response':response}
+                )
+                t.start()
+            elif isinstance(frame, DataReceived):
+                request.frames.append(frame)
+                wfile.write(frame.data)
+                if frame.stream_ended:
+                    wfile.close()
+            elif isinstance(frame, (StreamReset, ConnectionTerminated)):
+                assert self.stream_queues[stream_id].empty()
+                del self.stream_queues[stream_id]
+                self.logger.debug('(%s - %s) Stream Reset, Thread Closing' % (self.uid, stream_id))
+                break
+
+    def _h2_parse_request(self, frame):
+        self.headers = H2Headers(frame.headers)
         self.command = self.headers['method']
         self.path = self.headers['path']
-        self.h2_stream_id = event.stream_id
+        self.h2_stream_id = frame.stream_id
 
         # TODO Need to figure out what to do with this thing as it is no longer used
         # For now I can just leave it be as it does not affect anything
         self.raw_requestline = ''
 
 
 class H2ConnectionGuard(object):
+    """H2Connection objects are not threadsafe, so this keeps thread safety"""
     lock = threading.Lock()
 
     def __init__(self, obj):
+        assert isinstance(obj, H2Connection)
         self.obj = obj
 
     def __enter__(self):