Merge pull request #1797 from Kriechi/websocket++

Pass Autobahn WebSocket Test Suite

3 files changed, 67 insertions, 22 deletions

diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py
index 3a44efd1..304670db 100644
--- a/mitmproxy/flowfilter.py
+++ b/mitmproxy/flowfilter.py
@@ -37,6 +37,7 @@ import sys
 import functools
 
 from mitmproxy import http
+from mitmproxy import websocket
 from mitmproxy import tcp
 from mitmproxy import flow
 
@@ -99,6 +100,14 @@ class FHTTP(_Action):
         return True
 
 
+class FWebSocket(_Action):
+    code = "websocket"
+    help = "Match WebSocket flows"
+
+    @only(websocket.WebSocketFlow)
+    def __call__(self, f):
+        return True
+
 class FTCP(_Action):
     code = "tcp"
     help = "Match TCP flows"
@@ -245,7 +254,7 @@ class FBod(_Rex):
     help = "Body"
     flags = re.DOTALL
 
-    @only(http.HTTPFlow, tcp.TCPFlow)
+    @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow)
     def __call__(self, f):
         if isinstance(f, http.HTTPFlow):
             if f.request and f.request.raw_content:
@@ -254,7 +263,7 @@ class FBod(_Rex):
             if f.response and f.response.raw_content:
                 if self.re.search(f.response.get_content(strict=False)):
                     return True
-        elif isinstance(f, tcp.TCPFlow):
+        elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow):
             for msg in f.messages:
                 if self.re.search(msg.content):
                     return True
@@ -266,13 +275,13 @@ class FBodRequest(_Rex):
     help = "Request body"
     flags = re.DOTALL
 
-    @only(http.HTTPFlow, tcp.TCPFlow)
+    @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow)
     def __call__(self, f):
         if isinstance(f, http.HTTPFlow):
             if f.request and f.request.raw_content:
                 if self.re.search(f.request.get_content(strict=False)):
                     return True
-        elif isinstance(f, tcp.TCPFlow):
+        elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow):
             for msg in f.messages:
                 if msg.from_client and self.re.search(msg.content):
                     return True
@@ -283,13 +292,13 @@ class FBodResponse(_Rex):
     help = "Response body"
     flags = re.DOTALL
 
-    @only(http.HTTPFlow, tcp.TCPFlow)
+    @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow)
     def __call__(self, f):
         if isinstance(f, http.HTTPFlow):
             if f.response and f.response.raw_content:
                 if self.re.search(f.response.get_content(strict=False)):
                     return True
-        elif isinstance(f, tcp.TCPFlow):
+        elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow):
             for msg in f.messages:
                 if not msg.from_client and self.re.search(msg.content):
                     return True
diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py
index 15d9a288..d0b12540 100644
--- a/mitmproxy/proxy/protocol/websocket.py
+++ b/mitmproxy/proxy/protocol/websocket.py
@@ -59,30 +59,56 @@ class WebSocketLayer(base.Layer):
         fb.append(frame)
 
         if frame.header.fin:
-            if frame.header.opcode == websockets.OPCODE.TEXT:
+            payload = b''.join(f.payload for f in fb)
+            original_chunk_sizes = [len(f.payload) for f in fb]
+            message_type = fb[0].header.opcode
+            compressed_message = fb[0].header.rsv1
+            fb.clear()
+
+            if message_type == websockets.OPCODE.TEXT:
                 t = WebSocketTextMessage
             else:
                 t = WebSocketBinaryMessage
 
-            payload = b''.join(f.payload for f in fb)
-            fb.clear()
-
             websocket_message = t(self.flow, not is_server, payload)
+            length = len(websocket_message.content)
             self.flow.messages.append(websocket_message)
             self.channel.ask("websocket_message", self.flow)
 
-            # chunk payload into multiple 10kB frames, and send them
-            payload = websocket_message.content
-            chunk_size = 10240  # 10kB
-            chunks = range(0, len(payload), chunk_size)
+            def get_chunk(payload):
+                if len(payload) == length:
+                    # message has the same length, we can reuse the same sizes
+                    pos = 0
+                    for s in original_chunk_sizes:
+                        yield payload[pos:pos + s]
+                        pos += s
+                else:
+                    # just re-chunk everything into 10kB frames
+                    chunk_size = 10240
+                    chunks = range(0, len(payload), chunk_size)
+                    for i in chunks:
+                        yield payload[i:i + chunk_size]
+
             frms = [
                 websockets.Frame(
-                    payload=payload[i:i + chunk_size],
+                    payload=chunk,
                     opcode=frame.header.opcode,
                     mask=(False if is_server else 1),
-                    masking_key=(b'' if is_server else os.urandom(4))) for i in chunks
+                    masking_key=(b'' if is_server else os.urandom(4)))
+                for chunk in get_chunk(websocket_message.content)
             ]
-            frms[-1].header.fin = 1
+
+            if len(frms) > 0:
+                frms[-1].header.fin = True
+            else:
+                frms.append(websockets.Frame(
+                    fin=True,
+                    opcode=websockets.OPCODE.CONTINUE,
+                    mask=(False if is_server else 1),
+                    masking_key=(b'' if is_server else os.urandom(4))))
+
+            frms[0].header.opcode = message_type
+            frms[0].header.rsv1 = compressed_message
 
             for frm in frms:
                 other_conn.send(bytes(frm))
@@ -105,7 +131,7 @@ class WebSocketLayer(base.Layer):
 
         other_conn.send(bytes(frame))
 
-        # close the connection
+        # initiate close handshake
         return False
 
     def _handle_unknown_frame(self, frame, source_conn, other_conn, is_server):
@@ -126,10 +152,11 @@ class WebSocketLayer(base.Layer):
         client = self.client_conn.connection
         server = self.server_conn.connection
         conns = [client, server]
+        close_received = False
 
         try:
             while not self.channel.should_exit.is_set():
-                r = tcp.ssl_read_select(conns, 1)
+                r = tcp.ssl_read_select(conns, 0.1)
                 for conn in r:
                     source_conn = self.client_conn if conn == client else self.server_conn
                     other_conn = self.server_conn if conn == client else self.client_conn
@@ -137,10 +164,15 @@ class WebSocketLayer(base.Layer):
 
                     frame = websockets.Frame.from_file(source_conn.rfile)
 
-                    if not self._handle_frame(frame, source_conn, other_conn, is_server):
-                        return
+                    cont = self._handle_frame(frame, source_conn, other_conn, is_server)
+                    if not cont:
+                        if close_received:
+                            return
+                        else:
+                            close_received = True
         except (socket.error, exceptions.TcpException, SSL.Error) as e:
-            self.flow.error = flow.Error("WebSocket connection closed unexpectedly: {}".format(repr(e)))
+            s = 'server' if is_server else 'client'
+            self.flow.error = flow.Error("WebSocket connection closed unexpectedly by {}: {}".format(s, repr(e)))
             self.channel.tell("websocket_error", self.flow)
         finally:
             self.channel.tell("websocket_end", self.flow)
diff --git a/test/mitmproxy/protocol/test_websocket.py b/test/mitmproxy/protocol/test_websocket.py
index e1c3e49a..e42250e0 100644
--- a/test/mitmproxy/protocol/test_websocket.py
+++ b/test/mitmproxy/protocol/test_websocket.py
@@ -276,6 +276,7 @@ class TestClose(_WebSocketTest):
     def handle_websockets(cls, rfile, wfile):
         frame = websockets.Frame.from_file(rfile)
         wfile.write(bytes(frame))
+        wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
         wfile.flush()
 
         with pytest.raises(exceptions.TcpDisconnect):
@@ -287,6 +288,7 @@ class TestClose(_WebSocketTest):
         client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
         client.wfile.flush()
 
+        websockets.Frame.from_file(client.rfile)
         with pytest.raises(exceptions.TcpDisconnect):
             websockets.Frame.from_file(client.rfile)
 
@@ -296,6 +298,7 @@ class TestClose(_WebSocketTest):
         client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42')))
         client.wfile.flush()
 
+        websockets.Frame.from_file(client.rfile)
         with pytest.raises(exceptions.TcpDisconnect):
             websockets.Frame.from_file(client.rfile)
 
@@ -305,6 +308,7 @@ class TestClose(_WebSocketTest):
         client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42foobar')))
         client.wfile.flush()
 
+        websockets.Frame.from_file(client.rfile)
         with pytest.raises(exceptions.TcpDisconnect):
             websockets.Frame.from_file(client.rfile)