fix #2640

8 files changed, 63 insertions, 13 deletions

diff --git a/mitmproxy/addons/save.py b/mitmproxy/addons/save.py
index 40cd6f82..7bb55a84 100644
--- a/mitmproxy/addons/save.py
+++ b/mitmproxy/addons/save.py
@@ -74,6 +74,15 @@ class Save:
             self.stream.add(flow)
             self.active_flows.discard(flow)
 
+    def websocket_start(self, flow):
+        if self.stream:
+            self.active_flows.add(flow)
+
+    def websocket_end(self, flow):
+        if self.stream:
+            self.stream.add(flow)
+            self.active_flows.discard(flow)
+
     def response(self, flow):
         if self.stream:
             self.stream.add(flow)
diff --git a/mitmproxy/contrib/wsproto/__init__.py b/mitmproxy/contrib/wsproto/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/mitmproxy/contrib/wsproto/__init__.py
diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py
index 12f2bac8..d1fd8299 100644
--- a/mitmproxy/flowfilter.py
+++ b/mitmproxy/flowfilter.py
@@ -331,6 +331,7 @@ class FDomain(_Rex):
             self.re.search(f.request.pretty_host)
         )
 
+
 class FUrl(_Rex):
     code = "u"
     help = "URL"
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 5997ff6d..9660c137 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -9,6 +9,7 @@ from mitmproxy import eventsequence
 from mitmproxy import exceptions
 from mitmproxy import command
 from mitmproxy import http
+from mitmproxy import websocket
 from mitmproxy import log
 from mitmproxy.net import server_spec
 from mitmproxy.proxy.protocol import http_replay
@@ -41,6 +42,7 @@ class Master:
         self.should_exit = threading.Event()
         self._server = None
         self.first_tick = True
+        self.waiting_flows = []
 
     @property
     def server(self):
@@ -117,15 +119,28 @@ class Master:
         self.should_exit.set()
         self.addons.trigger("done")
 
+    def _change_reverse_host(self, f):
+        if self.options.mode.startswith("reverse:"):
+            _, upstream_spec = server_spec.parse_with_mode(self.options.mode)
+            f.request.host, f.request.port = upstream_spec.address
+            f.request.scheme = upstream_spec.scheme
+
     def load_flow(self, f):
         """
-        Loads a flow
+        Loads a flow and links websocket & handshake flows
         """
+
         if isinstance(f, http.HTTPFlow):
-            if self.options.mode.startswith("reverse:"):
-                _, upstream_spec = server_spec.parse_with_mode(self.options.mode)
-                f.request.host, f.request.port = upstream_spec.address
-                f.request.scheme = upstream_spec.scheme
+            self._change_reverse_host(f)
+            if 'websocket' in f.metadata:
+                self.waiting_flows.append(f)
+
+        if isinstance(f, websocket.WebSocketFlow):
+            hf = [hf for hf in self.waiting_flows if hf.id == f.metadata['websocket_handshake']][0]
+            f.handshake_flow = hf
+            self.waiting_flows.remove(hf)
+            self._change_reverse_host(f.handshake_flow)
+
         f.reply = controller.DummyReply()
         for e, o in eventsequence.iterate(f):
             self.addons.handle_lifecycle(e, o)
diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py
index 57ac0f16..076ffa62 100644
--- a/mitmproxy/proxy/protocol/http.py
+++ b/mitmproxy/proxy/protocol/http.py
@@ -321,6 +321,7 @@ class HttpLayer(base.Layer):
 
         try:
             if websockets.check_handshake(request.headers) and websockets.check_client_version(request.headers):
+                f.metadata['websocket'] = True
                 # We only support RFC6455 with WebSocket version 13
                 # allow inline scripts to manipulate the client handshake
                 self.channel.ask("websocket_handshake", f)
diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py
index 34dcba06..1bd5284d 100644
--- a/mitmproxy/proxy/protocol/websocket.py
+++ b/mitmproxy/proxy/protocol/websocket.py
@@ -1,10 +1,11 @@
 import socket
 from OpenSSL import SSL
 
+
+from mitmproxy.contrib import wsproto
 from mitmproxy.contrib.wsproto import events
 from mitmproxy.contrib.wsproto.connection import ConnectionType, WSConnection
 from mitmproxy.contrib.wsproto.extensions import PerMessageDeflate
-from mitmproxy.contrib.wsproto.frame_protocol import Opcode
 
 from mitmproxy import exceptions
 from mitmproxy import flow
@@ -93,11 +94,14 @@ class WebSocketLayer(base.Layer):
 
         if event.message_finished:
             original_chunk_sizes = [len(f) for f in fb]
-            message_type = Opcode.TEXT if isinstance(event, events.TextReceived) else Opcode.BINARY
-            if message_type == Opcode.TEXT:
+
+            if isinstance(event, events.TextReceived):
+                message_type = wsproto.frame_protocol.Opcode.TEXT
                 payload = ''.join(fb)
             else:
+                message_type = wsproto.frame_protocol.Opcode.BINARY
                 payload = b''.join(fb)
+
             fb.clear()
 
             websocket_message = WebSocketMessage(message_type, not is_server, payload)
diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py
index 6c1e7000..ade23732 100644
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -1,6 +1,8 @@
 import time
 from typing import List, Optional
 
+from mitmproxy.contrib import wsproto
+
 from mitmproxy import flow
 from mitmproxy.net import websockets
 from mitmproxy.coretypes import serializable
@@ -11,7 +13,7 @@ class WebSocketMessage(serializable.Serializable):
     def __init__(
         self, type: int, from_client: bool, content: bytes, timestamp: Optional[int]=None
     ) -> None:
-        self.type = type
+        self.type = wsproto.frame_protocol.Opcode(type)
         self.from_client = from_client
         self.content = content
         self.timestamp = timestamp or int(time.time())  # type: int
@@ -21,13 +23,14 @@ class WebSocketMessage(serializable.Serializable):
         return cls(*state)
 
     def get_state(self):
-        return self.type, self.from_client, self.content, self.timestamp
+        return int(self.type), self.from_client, self.content, self.timestamp
 
     def set_state(self, state):
         self.type, self.from_client, self.content, self.timestamp = state
+        self.type = wsproto.frame_protocol.Opcode(self.type)  # replace enum with bare int
 
     def __repr__(self):
-        if self.type == websockets.OPCODE.TEXT:
+        if self.type == wsproto.frame_protocol.Opcode.TEXT:
             return "text message: {}".format(repr(self.content))
         else:
             return "binary message: {}".format(strutils.bytes_to_escaped_str(self.content))
@@ -42,7 +45,7 @@ class WebSocketFlow(flow.Flow):
         super().__init__("websocket", client_conn, server_conn, live)
         self.messages = []  # type: List[WebSocketMessage]
         self.close_sender = 'client'
-        self.close_code = '(status code missing)'
+        self.close_code = wsproto.frame_protocol.CloseReason.NORMAL_CLOSURE
         self.close_message = '(message missing)'
         self.close_reason = 'unknown status code'
         self.stream = False
@@ -69,7 +72,7 @@ class WebSocketFlow(flow.Flow):
     _stateobject_attributes.update(dict(
         messages=List[WebSocketMessage],
         close_sender=str,
-        close_code=str,
+        close_code=int,
         close_message=str,
         close_reason=str,
         client_key=str,
@@ -83,6 +86,11 @@ class WebSocketFlow(flow.Flow):
         # dumping the handshake_flow will include the WebSocketFlow too.
     ))
 
+    def get_state(self):
+        d = super().get_state()
+        d['close_code'] = int(d['close_code'])  # replace enum with bare int
+        return d
+
     @classmethod
     def from_state(cls, state):
         f = cls(None, None, None)
diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py
index a4e425cd..84564157 100644
--- a/test/mitmproxy/addons/test_save.py
+++ b/test/mitmproxy/addons/test_save.py
@@ -44,6 +44,18 @@ def test_tcp(tmpdir):
         assert rd(p)
 
 
+def test_websocket(tmpdir):
+    sa = save.Save()
+    with taddons.context() as tctx:
+        p = str(tmpdir.join("foo"))
+        tctx.configure(sa, save_stream_file=p)
+
+        f = tflow.twebsocketflow()
+        sa.websocket_start(f)
+        tctx.configure(sa, save_stream_file=None)
+        assert rd(p)
+
+
 def test_save_command(tmpdir):
     sa = save.Save()
     with taddons.context() as tctx: