Merge pull request #733 from Kriechi/proxy-refactor-cb

Proxy refactor cb

2 files changed, 41 insertions, 27 deletions

diff --git a/libmproxy/protocol2/http.py b/libmproxy/protocol2/http.py
index db5aabaf..b8b1e8a5 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
@@ -28,6 +28,12 @@ class Http1Layer(Layer):
         self.client_protocol = HTTP1Protocol(self.client_conn)
         self.server_protocol = HTTP1Protocol(self.server_conn)
 
+    def send_to_client(self, message):
+        self.client_conn.send(self.client_protocol.assemble(message))
+
+    def send_to_server(self, message):
+        self.server_conn.send(self.server_protocol.assemble(message))
+
     def connect(self):
         self.ctx.connect()
         self.server_protocol = HTTP1Protocol(self.server_conn)
@@ -51,6 +57,14 @@ class Http2Layer(Layer):
         self.client_protocol = HTTP2Protocol(self.client_conn, is_server=True)
         self.server_protocol = HTTP2Protocol(self.server_conn, is_server=False)
 
+    def send_to_client(self, message):
+        # TODO: implement flow control and WINDOW_UPDATE frames
+        self.client_conn.send(self.client_protocol.assemble(message))
+
+    def send_to_server(self, message):
+        # TODO: implement flow control and WINDOW_UPDATE frames
+        self.server_conn.send(self.server_protocol.assemble(message))
+
     def connect(self):
         self.ctx.connect()
         self.server_protocol = HTTP2Protocol(self.server_conn)
@@ -64,6 +78,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()
 
@@ -165,11 +180,8 @@ class UpstreamConnectLayer(Layer):
         else:
             self.ctx.set_server(address, server_tls, sni, depth-1)
 
-class HttpLayer(Layer):
-    """
-    HTTP 1 Layer
-    """
 
+class HttpLayer(Layer):
     def __init__(self, ctx, mode):
         super(HttpLayer, self).__init__(ctx)
         self.mode = mode
@@ -337,15 +349,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,
@@ -461,11 +476,3 @@ class HttpLayer(Layer):
                     odict.ODictCaseless([[k,v] for k, v in self.config.authenticator.auth_challenge_headers().items()])
                 ))
                 raise InvalidCredentials("Proxy Authentication Required")
-
-    def send_to_server(self, message):
-        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))
diff --git a/libmproxy/protocol2/root_context.py b/libmproxy/protocol2/root_context.py
index f0e5b9a7..9b18f0aa 100644
--- a/libmproxy/protocol2/root_context.py
+++ b/libmproxy/protocol2/root_context.py
@@ -6,6 +6,7 @@ from .rawtcp import RawTcpLayer
 from .tls import TlsLayer
 from .http import Http1Layer, Http2Layer, HttpLayer
 
+from netlib.http.http2 import HTTP2Protocol
 
 class RootContext(object):
     """
@@ -25,11 +26,11 @@ class RootContext(object):
         :return: The next layer.
         """
 
-        d = top_layer.client_conn.rfile.peek(3)
-
         # TODO: Handle ignore and tcp passthrough
 
-        # TLS ClientHello magic, see http://www.moserware.com/2009/06/first-few-milliseconds-of-https.html#client-hello
+        # TLS ClientHello magic, works for SSLv3, TLSv1.0, TLSv1.1, TLSv1.2
+        # http://www.moserware.com/2009/06/first-few-milliseconds-of-https.html#client-hello
+        d = top_layer.client_conn.rfile.peek(3)
         is_tls_client_hello = (
             len(d) == 3 and
             d[0] == '\x16' and
@@ -37,20 +38,26 @@ class RootContext(object):
             d[2] in ('\x00', '\x01', '\x02', '\x03')
         )
 
-        is_ascii = all(x in string.ascii_uppercase for x in d)
+        d = top_layer.client_conn.rfile.peek(3)
+        is_ascii = (
+            len(d) == 3 and
+            all(x in string.ascii_uppercase for x in d)
+        )
 
-        # TODO: build is_http2_magic check here, maybe this is an easy way to detect h2c
+        d = top_layer.client_conn.rfile.peek(len(HTTP2Protocol.CLIENT_CONNECTION_PREFACE))
+        is_http2_magic = (d == HTTP2Protocol.CLIENT_CONNECTION_PREFACE)
 
-        if not d:
-            return iter([])
+        is_alpn_h2_negotiated = (
+            isinstance(top_layer, TlsLayer) and
+            top_layer.client_conn.get_alpn_proto_negotiated() == HTTP2Protocol.ALPN_PROTO_H2
+        )
 
         if is_tls_client_hello:
             return TlsLayer(top_layer, True, True)
-        elif isinstance(top_layer, TlsLayer) and is_ascii:
-            if top_layer.client_conn.get_alpn_proto_negotiated() == 'h2':
-                return Http2Layer(top_layer, 'transparent')
-            else:
-                return Http1Layer(top_layer, "transparent")
+        elif is_alpn_h2_negotiated or is_http2_magic:
+            return Http2Layer(top_layer, 'transparent')
+        elif is_ascii:
+            return Http1Layer(top_layer, 'transparent')
         else:
             return RawTcpLayer(top_layer)