Merge remote-tracking branch 'Kriechi/proxy-refactor' into proxy-refactor

Conflicts:
	libmproxy/protocol2/http.py

5 files changed, 57 insertions, 73 deletions

diff --git a/libmproxy/protocol2/http.py b/libmproxy/protocol2/http.py
index 90784666..eadde3b3 100644
--- a/libmproxy/protocol2/http.py
+++ b/libmproxy/protocol2/http.py
@@ -4,18 +4,47 @@ from .. import version
 from ..exceptions import InvalidCredentials, HttpException, ProtocolException
 from .layer import Layer
 from libmproxy import utils
-from libmproxy.proxy.connection import ServerConnection
 from .messages import SetServer, Connect, Reconnect, Kill
 from libmproxy.protocol import KILL
 
 from libmproxy.protocol.http import HTTPFlow
 from libmproxy.protocol.http_wrappers import HTTPResponse, HTTPRequest
-from libmproxy.protocol2.http_protocol_mock import HTTP1
 from netlib import tcp
 from netlib.http import status_codes, http1, HttpErrorConnClosed
 from netlib.http.semantics import CONTENT_MISSING
 from netlib import odict
 from netlib.tcp import NetLibError, Address
+from netlib.http.http1 import HTTP1Protocol
+from netlib.http.http2 import HTTP2Protocol
+
+# TODO: The HTTP2 layer is missing multiplexing, which requires a major rewrite.
+
+
+class Http1Layer(Layer):
+    def __init__(self, ctx, mode):
+        super(Http1Layer, self).__init__(ctx)
+        self.mode = mode
+        self.client_protocol = HTTP1Protocol(self.client_conn)
+        self.server_protocol = HTTP1Protocol(self.server_conn)
+
+    def __call__(self):
+        layer = HttpLayer(self, self.mode)
+        for message in layer():
+            yield message
+
+
+class Http2Layer(Layer):
+    def __init__(self, ctx, mode):
+        super(Http2Layer, self).__init__(ctx)
+        self.mode = mode
+        self.client_protocol = HTTP2Protocol(self.client_conn, is_server=True)
+        self.server_protocol = HTTP2Protocol(self.server_conn, is_server=False)
+
+    def __call__(self):
+        # FIXME: Handle Reconnect etc.
+        layer = HttpLayer(self, self.mode)
+        for message in layer():
+            yield message
 
 
 def make_error_response(status_code, message, headers=None):
@@ -96,8 +125,8 @@ class HttpLayer(Layer):
         while True:
             try:
                 try:
-                    request = HTTP1.read_request(
-                        self.client_conn,
+                    request = HTTPRequest.from_protocol(
+                        self.client_protocol,
                         body_size_limit=self.config.body_size_limit
                     )
                 except tcp.NetLibError:
@@ -219,12 +248,12 @@ class HttpLayer(Layer):
             # streaming:
             # First send the headers and then transfer the response
             # incrementally:
-            h = HTTP1._assemble_response_first_line(flow.response)
+            h = self.client_protocol._assemble_response_first_line(flow.response)
             self.send_to_client(h + "\r\n")
-            h = HTTP1._assemble_response_headers(flow.response, preserve_transfer_encoding=True)
+            h = self.client_protocol._assemble_response_headers(flow.response, preserve_transfer_encoding=True)
             self.send_to_client(h + "\r\n")
 
-            chunks = HTTP1.read_http_body_chunked(
+            chunks = self.client_protocol.read_http_body_chunked(
                 flow.response.headers,
                 self.config.body_size_limit,
                 flow.request.method,
@@ -247,8 +276,8 @@ class HttpLayer(Layer):
         # TODO: Add second attempt.
         self.send_to_server(flow.request)
 
-        flow.response = HTTP1.read_response(
-            self.server_conn,
+        flow.response = HTTPResponse.from_protocol(
+            self.server_protocol,
             flow.request.method,
             body_size_limit=self.config.body_size_limit,
             include_body=False,
@@ -260,10 +289,10 @@ class HttpLayer(Layer):
         if flow is None or flow == KILL:
             yield Kill()
 
-        if flow.response.stream:
+        if flow.response.stream and isinstance(self.server_protocol, http1.HTTP1Protocol):
             flow.response.content = CONTENT_MISSING
         else:
-            flow.response.content = HTTP1.read_http_body(
+            flow.response.content = self.server_protocol.read_http_body(
                 self.server_conn,
                 flow.response.headers,
                 self.config.body_size_limit,
@@ -381,9 +410,9 @@ class HttpLayer(Layer):
                 raise InvalidCredentials("Proxy Authentication Required")
 
     def send_to_server(self, message):
-        self.server_conn.send(HTTP1.assemble(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(HTTP1.assemble(message))
+        self.client_conn.send(self.client_protocol.assemble(message))
diff --git a/libmproxy/protocol2/http_protocol_mock.py b/libmproxy/protocol2/http_protocol_mock.py
deleted file mode 100644
index dd3643f6..00000000
--- a/libmproxy/protocol2/http_protocol_mock.py
+++ /dev/null
@@ -1,50 +0,0 @@
-"""
-Temporary mock to sort out API discrepancies
-"""
-from libmproxy.protocol.http_wrappers import HTTPResponse, HTTPRequest
-from netlib.http.http1 import HTTP1Protocol
-
-
-class HTTP1(object):
-    @staticmethod
-    def read_request(connection, *args, **kwargs):
-        """
-        :type connection: object
-        """
-        return HTTPRequest.from_protocol(HTTP1Protocol(connection), *args, **kwargs)
-
-    @staticmethod
-    def read_response(connection, *args, **kwargs):
-        """
-        :type connection: object
-        """
-        return HTTPResponse.from_protocol(HTTP1Protocol(connection), *args, **kwargs)
-
-    @staticmethod
-    def read_http_body(connection, *args, **kwargs):
-        """
-        :type connection: object
-        """
-        return HTTP1Protocol(connection).read_http_body(*args, **kwargs)
-
-
-    @staticmethod
-    def _assemble_response_first_line(*args, **kwargs):
-        return HTTP1Protocol()._assemble_response_first_line(*args, **kwargs)
-
-
-    @staticmethod
-    def _assemble_response_headers(*args, **kwargs):
-        return HTTP1Protocol()._assemble_response_headers(*args, **kwargs)
-
-
-    @staticmethod
-    def read_http_body_chunked(connection, *args, **kwargs):
-        """
-        :type connection: object
-        """
-        return HTTP1Protocol(connection).read_http_body_chunked(*args, **kwargs)
-
-    @staticmethod
-    def assemble(*args, **kwargs):
-        return HTTP1Protocol().assemble(*args, **kwargs)
\ No newline at end of file
diff --git a/libmproxy/protocol2/http_proxy.py b/libmproxy/protocol2/http_proxy.py
index 3cc7fee2..19b5f7ef 100644
--- a/libmproxy/protocol2/http_proxy.py
+++ b/libmproxy/protocol2/http_proxy.py
@@ -1,12 +1,12 @@
 from __future__ import (absolute_import, print_function, division)
 
 from .layer import Layer, ServerConnectionMixin
-from .http import HttpLayer
+from .http import Http1Layer
 
 
 class HttpProxy(Layer, ServerConnectionMixin):
     def __call__(self):
-        layer = HttpLayer(self, "regular")
+        layer = Http1Layer(self, "regular")
         for message in layer():
             if not self._handle_server_message(message):
                 yield message
@@ -17,7 +17,7 @@ class HttpUpstreamProxy(Layer, ServerConnectionMixin):
         super(HttpUpstreamProxy, self).__init__(ctx, server_address=server_address)
 
     def __call__(self):
-        layer = HttpLayer(self, "upstream")
+        layer = Http1Layer(self, "upstream")
         for message in layer():
             if not self._handle_server_message(message):
-                yield message
\ No newline at end of file
+                yield message
diff --git a/libmproxy/protocol2/root_context.py b/libmproxy/protocol2/root_context.py
index f369fd48..6ba6ca9a 100644
--- a/libmproxy/protocol2/root_context.py
+++ b/libmproxy/protocol2/root_context.py
@@ -3,7 +3,7 @@ from __future__ import (absolute_import, print_function, division)
 from .messages import Kill
 from .rawtcp import RawTcpLayer
 from .tls import TlsLayer
-from .http import HttpLayer
+from .http import Http1Layer, Http2Layer, HttpLayer
 
 
 class RootContext(object):
@@ -36,13 +36,17 @@ class RootContext(object):
             d[2] in ('\x00', '\x01', '\x02', '\x03')
         )
 
+        # TODO: build is_http2_magic check here, maybe this is an easy way to detect h2c
+
         if not d:
             return iter([])
 
         if is_tls_client_hello:
             return TlsLayer(top_layer, True, True)
-        elif isinstance(top_layer, TlsLayer) and isinstance(top_layer.ctx, HttpLayer):
-            return HttpLayer(top_layer, "transparent")
+        elif isinstance(top_layer, TlsLayer) and top_layer.client_conn.get_alpn_proto_negotiated() == 'h2':
+                return Http2Layer(top_layer, 'transparent')
+        elif isinstance(top_layer, TlsLayer) and isinstance(top_layer.ctx, Http1Layer):
+            return Http1Layer(top_layer, "transparent")
         else:
             return RawTcpLayer(top_layer)
 
diff --git a/libmproxy/protocol2/tls.py b/libmproxy/protocol2/tls.py
index 12c67f4e..970abe62 100644
--- a/libmproxy/protocol2/tls.py
+++ b/libmproxy/protocol2/tls.py
@@ -124,7 +124,7 @@ class TlsLayer(Layer):
 
             if old_upstream_sni != self.sni_for_upstream_connection:
                 # Perform reconnect
-                if self._server_tls:
+                if self.server_conn and self._server_tls:
                     self.yield_from_callback(Reconnect())
 
             if self.client_sni:
@@ -151,13 +151,14 @@ class TlsLayer(Layer):
         alpn_preference = netlib.http.http2.HTTP2Protocol.ALPN_PROTO_H2
         ###
 
+        # TODO: Not
         if self.client_alpn_protos != options:
             # Perform reconnect
-            if self._server_tls:
+            # TODO: Avoid double reconnect.
+            if self.server_conn and self._server_tls:
                 self.yield_from_callback(Reconnect())
 
         self.client_alpn_protos = options
-        print("foo: %s" % options)
 
         if alpn_preference in options:
             return bytes(alpn_preference)