attach application protocol to connection

5 files changed, 67 insertions, 46 deletions

diff --git a/examples/ignore_websocket.py b/examples/ignore_websocket.py
index 479d0984..bea7e565 100644
--- a/examples/ignore_websocket.py
+++ b/examples/ignore_websocket.py
@@ -30,7 +30,7 @@ def response(context, flow):
     value = flow.response.headers.get_first("Connection", None)
     if value and value.upper() == "UPGRADE":
         # We need to send the response manually now...
-        flow.client_conn.send(flow.client_protocol.assemble(flow.response))
+        flow.client_conn.send(flow.client_conn.protocol.assemble(flow.response))
         # ...and then delegate to tcp passthrough.
         TCPHandler(flow.live.c, log=False).handle_messages()
         flow.reply(KILL)
diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index 5ce2c0b7..1940e390 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -252,7 +252,7 @@ def copy_flow_format_data(part, scope, flow):
                 return None, "Request content is missing"
             with decoded(flow.request):
                 if part == "h":
-                    data += flow.client_protocol.assemble(flow.request)
+                    data += flow.client_conn.protocol.assemble(flow.request)
                 elif part == "c":
                     data += flow.request.content
                 else:
@@ -265,7 +265,7 @@ def copy_flow_format_data(part, scope, flow):
                 return None, "Response content is missing"
             with decoded(flow.response):
                 if part == "h":
-                    data += flow.client_protocol.assemble(flow.response)
+                    data += flow.client_conn.protocol.assemble(flow.response)
                 elif part == "c":
                     data += flow.response.content
                 else:
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 7f1aa78b..35fd7d28 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -9,7 +9,7 @@ from email.utils import parsedate_tz, formatdate, mktime_tz
 
 import netlib
 from netlib import http, tcp, odict, utils
-from netlib.http import cookies, http1
+from netlib.http import cookies, http1, http2
 from netlib.http.semantics import CONTENT_MISSING
 
 from .tcp import TCPHandler
@@ -39,7 +39,7 @@ def send_connect_request(conn, host, port, update_state=True):
         odict.ODictCaseless(),
         ""
     )
-    protocol = http.http1.HTTP1Protocol(conn)
+    protocol = http1.HTTP1Protocol(conn)
     conn.send(protocol.assemble(upstream_request))
     resp = HTTPResponse.from_protocol(protocol, upstream_request.method)
     if resp.status_code != 200:
@@ -177,12 +177,16 @@ class HTTPHandler(ProtocolHandler):
 
         for attempt in (0, 1):
             try:
-                flow.server_protocol = http.http1.HTTP1Protocol(self.c.server_conn)
-                self.c.server_conn.send(flow.server_protocol.assemble(flow.request))
+                if not self.c.server_conn.protocol:
+                    # instantiate new protocol if connection does not have one yet
+                    self.c.server_conn.protocol = http2.HTTP2Protocol(self.c.server_conn)
+                    self.c.server_conn.protocol.perform_connection_preface()
+
+                self.c.server_conn.send(self.c.server_conn.protocol.assemble(flow.request))
 
                 # Only get the headers at first...
                 flow.response = HTTPResponse.from_protocol(
-                    flow.server_protocol,
+                    flow.server_conn.protocol,
                     flow.request.method,
                     body_size_limit=self.c.config.body_size_limit,
                     include_body=False
@@ -220,23 +224,27 @@ class HTTPHandler(ProtocolHandler):
             if flow.response.stream:
                 flow.response.content = CONTENT_MISSING
             else:
-                flow.server_protocol = http1.HTTP1Protocol(self.c.server_conn)
-                flow.response.content = flow.server_protocol.read_http_body(
-                    flow.response.headers,
-                    self.c.config.body_size_limit,
-                    flow.request.method,
-                    flow.response.code,
-                    False
-                )
+                if isinstance(flow.server_conn.protocol, http1.HTTP1Protocol):
+                    flow.response.content = flow.server_conn.protocol.read_http_body(
+                        flow.response.headers,
+                        self.c.config.body_size_limit,
+                        flow.request.method,
+                        flow.response.code,
+                        False
+                    )
         flow.response.timestamp_end = utils.timestamp()
 
     def handle_flow(self):
         flow = HTTPFlow(self.c.client_conn, self.c.server_conn, self.live)
+
         try:
             try:
-                flow.client_protocol = http.http1.HTTP1Protocol(self.c.client_conn)
+                if not flow.client_conn.protocol:
+                    # instantiate new protocol if connection does not have one yet
+                    flow.client_conn.protocol = http1.HTTP1Protocol(self.c.client_conn)
+
                 req = HTTPRequest.from_protocol(
-                    flow.client_protocol,
+                    flow.client_conn.protocol,
                     body_size_limit=self.c.config.body_size_limit
                 )
             except tcp.NetLibError:
@@ -249,9 +257,15 @@ class HTTPHandler(ProtocolHandler):
                 [repr(req)]
             )
             ret = self.process_request(flow, req)
+            if ret:
+                # CONNECT successful - upgrade to HTTP/2
+                # instantiate new protocol if connection does not have one yet
+                flow.client_conn.protocol = http2.HTTP2Protocol(self.c.client_conn, is_server=True)
             if ret is not None:
                 return ret
 
+            print("still here: %s" % flow.client_conn.protocol.__class__)
+
             # Be careful NOT to assign the request to the flow before
             # process_request completes. This is because the call can raise an
             # exception. If the request object is already attached, this results
@@ -375,30 +389,31 @@ class HTTPHandler(ProtocolHandler):
             pass
 
     def send_error(self, code, message, headers):
-        response = http.status_codes.RESPONSES.get(code, "Unknown")
-        html_content = """
-            <html>
-                <head>
-                    <title>%d %s</title>
-                </head>
-                <body>%s</body>
-            </html>
-        """ % (code, response, message)
-        self.c.client_conn.wfile.write("HTTP/1.1 %s %s\r\n" % (code, response))
-        self.c.client_conn.wfile.write(
-            "Server: %s\r\n" % self.c.config.server_version
-        )
-        self.c.client_conn.wfile.write("Content-type: text/html\r\n")
-        self.c.client_conn.wfile.write(
-            "Content-Length: %d\r\n" % len(html_content)
-        )
-        if headers:
-            for key, value in headers.items():
-                self.c.client_conn.wfile.write("%s: %s\r\n" % (key, value))
-        self.c.client_conn.wfile.write("Connection: close\r\n")
-        self.c.client_conn.wfile.write("\r\n")
-        self.c.client_conn.wfile.write(html_content)
-        self.c.client_conn.wfile.flush()
+        raise NotImplementedError("todo - adapt for HTTP/2 - make use of make_error_reponse from pathod")
+        # response = http.status_codes.RESPONSES.get(code, "Unknown")
+        # html_content = """
+        #     <html>
+        #         <head>
+        #             <title>%d %s</title>
+        #         </head>
+        #         <body>%s</body>
+        #     </html>
+        # """ % (code, response, message)
+        # self.c.client_conn.wfile.write("HTTP/1.1 %s %s\r\n" % (code, response))
+        # self.c.client_conn.wfile.write(
+        #     "Server: %s\r\n" % self.c.config.server_version
+        # )
+        # self.c.client_conn.wfile.write("Content-type: text/html\r\n")
+        # self.c.client_conn.wfile.write(
+        #     "Content-Length: %d\r\n" % len(html_content)
+        # )
+        # if headers:
+        #     for key, value in headers.items():
+        #         self.c.client_conn.wfile.write("%s: %s\r\n" % (key, value))
+        # self.c.client_conn.wfile.write("Connection: close\r\n")
+        # self.c.client_conn.wfile.write("\r\n")
+        # self.c.client_conn.wfile.write(html_content)
+        # self.c.client_conn.wfile.flush()
 
     def process_request(self, flow, request):
         """
@@ -554,7 +569,7 @@ class HTTPHandler(ProtocolHandler):
             # no streaming:
             # we already received the full response from the server and can
             # send it to the client straight away.
-            self.c.client_conn.send(flow.client_protocol.assemble(flow.response))
+            self.c.client_conn.send(self.c.client_conn.protocol.assemble(flow.response))
         else:
             raise NotImplementedError("HTTP streaming is currently not supported.")
             # TODO: implement it according to new protocols and messages
@@ -731,12 +746,11 @@ class RequestReplayThread(threading.Thread):
                         )
                     r.form_out = "relative"
 
-                server.send(self.flow.server_protocol.assemble(r))
+                server.send(self.flow.server_conn.protocol.assemble(r))
                 self.flow.server_conn = server
 
-                self.flow.server_protocol = http.http1.HTTP1Protocol(self.flow.server_conn)
                 self.flow.response = HTTPResponse.from_protocol(
-                    self.flow.server_protocol,
+                    self.flow.server_conn.protocol,
                     r.method,
                     body_size_limit=self.config.body_size_limit,
                 )
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index a9193c5f..92fc95e5 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -167,6 +167,7 @@ class Flow(stateobject.StateObject):
         master.handle_accept_intercept(self)
 
 
+
 class ProtocolHandler(object):
     """
     A ProtocolHandler implements an application-layer protocol, e.g. HTTP.
diff --git a/libmproxy/proxy/connection.py b/libmproxy/proxy/connection.py
index 54b3688e..a0bf2af9 100644
--- a/libmproxy/proxy/connection.py
+++ b/libmproxy/proxy/connection.py
@@ -23,6 +23,7 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
         self.timestamp_start = utils.timestamp()
         self.timestamp_end = None
         self.timestamp_ssl_setup = None
+        self.protocol = None
 
     def __repr__(self):
         return "<ClientConnection: {ssl}{host}:{port}>".format(
@@ -58,6 +59,8 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
         return copy.copy(self)
 
     def send(self, message):
+        if isinstance(message, list):
+            message = b''.join(message)
         self.wfile.write(message)
         self.wfile.flush()
 
@@ -93,6 +96,7 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
         self.timestamp_end = None
         self.timestamp_tcp_setup = None
         self.timestamp_ssl_setup = None
+        self.protocol = None
 
     def __repr__(self):
         if self.ssl_established and self.sni:
@@ -157,6 +161,8 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
         self.timestamp_tcp_setup = utils.timestamp()
 
     def send(self, message):
+        if isinstance(message, list):
+            message = b''.join(message)
         self.wfile.write(message)
         self.wfile.flush()