fix #2640

1 files changed, 20 insertions, 5 deletions

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)