http2: fix connection preface and wrappers

2 files changed, 17 insertions, 35 deletions

diff --git a/libmproxy/protocol/http_wrappers.py b/libmproxy/protocol/http_wrappers.py
index ed5759ea..e41d65d6 100644
--- a/libmproxy/protocol/http_wrappers.py
+++ b/libmproxy/protocol/http_wrappers.py
@@ -247,24 +247,11 @@ class HTTPRequest(MessageMixin, semantics.Request):
             include_body = include_body,
             body_size_limit = body_size_limit,
         )
-
-        return HTTPRequest(
-            req.form_in,
-            req.method,
-            req.scheme,
-            req.host,
-            req.port,
-            req.path,
-            req.httpversion,
-            req.headers,
-            req.body,
-            req.timestamp_start,
-            req.timestamp_end,
-        )
+        return self.wrap(req)
 
     @classmethod
     def wrap(self, request):
-        return HTTPRequest(
+        req = HTTPRequest(
             form_in=request.form_in,
             method=request.method,
             scheme=request.scheme,
@@ -278,6 +265,9 @@ class HTTPRequest(MessageMixin, semantics.Request):
             timestamp_end=request.timestamp_end,
             form_out=(request.form_out if hasattr(request, 'form_out') else None),
         )
+        if hasattr(request, 'stream_id'):
+            req.stream_id = request.stream_id
+        return req
 
     def __hash__(self):
         return id(self)
@@ -371,20 +361,11 @@ class HTTPResponse(MessageMixin, semantics.Response):
             body_size_limit,
             include_body=include_body
         )
-
-        return HTTPResponse(
-            resp.httpversion,
-            resp.status_code,
-            resp.msg,
-            resp.headers,
-            resp.body,
-            resp.timestamp_start,
-            resp.timestamp_end,
-        )
+        return self.wrap(resp)
 
     @classmethod
     def wrap(self, response):
-        return HTTPResponse(
+        resp = HTTPResponse(
             httpversion=response.httpversion,
             status_code=response.status_code,
             msg=response.msg,
@@ -393,6 +374,9 @@ class HTTPResponse(MessageMixin, semantics.Response):
             timestamp_start=response.timestamp_start,
             timestamp_end=response.timestamp_end,
         )
+        if hasattr(response, 'stream_id'):
+            resp.stream_id = response.stream_id
+        return resp
 
     def _refresh_cookie(self, c, delta):
         """
diff --git a/libmproxy/protocol2/http.py b/libmproxy/protocol2/http.py
index db5aabaf..e73bbb61 100644
--- a/libmproxy/protocol2/http.py
+++ b/libmproxy/protocol2/http.py
@@ -10,7 +10,7 @@ from libmproxy.protocol import KILL
 from libmproxy.protocol.http import HTTPFlow
 from libmproxy.protocol.http_wrappers import HTTPResponse, HTTPRequest
 from netlib import tcp
-from netlib.http import status_codes, http1, HttpErrorConnClosed, HttpError
+from netlib.http import status_codes, http1, http2, HttpErrorConnClosed, HttpError
 from netlib.http.semantics import CONTENT_MISSING
 from netlib import odict
 from netlib.tcp import NetLibError, Address
@@ -64,6 +64,7 @@ class Http2Layer(Layer):
         self.server_protocol = HTTP2Protocol(self.server_conn)
 
     def __call__(self):
+        self.server_protocol.perform_connection_preface()
         layer = HttpLayer(self, self.mode)
         layer()
 
@@ -166,10 +167,6 @@ class UpstreamConnectLayer(Layer):
             self.ctx.set_server(address, server_tls, sni, depth-1)
 
 class HttpLayer(Layer):
-    """
-    HTTP 1 Layer
-    """
-
     def __init__(self, ctx, mode):
         super(HttpLayer, self).__init__(ctx)
         self.mode = mode
@@ -337,15 +334,18 @@ class HttpLayer(Layer):
             self.reconnect()
             get_response()
 
+        if isinstance(self.server_protocol, http2.HTTP2Protocol):
+            flow.response.stream_id = flow.request.stream_id
+
         # call the appropriate script hook - this is an opportunity for an
         # inline script to set flow.stream = True
         flow = self.channel.ask("responseheaders", flow)
         if flow is None or flow == KILL:
             raise Kill()
 
-        if flow.response.stream and isinstance(self.server_protocol, http1.HTTP1Protocol):
+        if flow.response.stream:
             flow.response.content = CONTENT_MISSING
-        else:
+        elif isinstance(self.server_protocol, http1.HTTP1Protocol):
             flow.response.content = self.server_protocol.read_http_body(
                 flow.response.headers,
                 self.config.body_size_limit,
@@ -466,6 +466,4 @@ class HttpLayer(Layer):
         self.server_conn.send(self.server_protocol.assemble(message))
 
     def send_to_client(self, message):
-        # FIXME
-        # - possibly do some http2 stuff here
         self.client_conn.send(self.client_protocol.assemble(message))