From 04abe6b85b25941eea1b4874e686e1cde7e6ec18 Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Thu, 16 Jul 2020 17:34:42 +0200
Subject: [PATCH] match HTTP-WebSocket handshake flow with ~websocket

fixes #3990
---
 CHANGELOG                         |  1 +
 mitmproxy/flowfilter.py           | 11 ++++++++---
 mitmproxy/master.py               | 13 +++++++++----
 test/mitmproxy/test_flowfilter.py | 11 +++++++++++
 tox.ini                           |  2 +-
 5 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG b/CHANGELOG
index b3a07fa00..f15a84723 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -20,6 +20,7 @@ Unreleased: mitmproxy next
     * Add support for HTTP Trailers to the HTTP/2 protocol (@sanlengjingvv and @Kriechi)
     * Fix certificate runtime error during expire cleanup (@gorogoroumaru)
     * Fixed the DNS Rebind Protection for secure support of IPv6 addresses (@tunnelpr0)
+    * WebSockets: match the HTTP-WebSocket flow for the ~websocket filter (@Kriechi)
 
     * --- TODO: add new PRs above this line ---
 
diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py
index b222d2a89..453fb4071 100644
--- a/mitmproxy/flowfilter.py
+++ b/mitmproxy/flowfilter.py
@@ -43,6 +43,7 @@ from mitmproxy import flow
 from mitmproxy import http
 from mitmproxy import tcp
 from mitmproxy import websocket
+from mitmproxy.net import websockets as net_websockets
 
 
 def only(*types):
@@ -104,11 +105,15 @@ class FHTTP(_Action):
 
 class FWebSocket(_Action):
     code = "websocket"
-    help = "Match WebSocket flows"
+    help = "Match WebSocket flows (and HTTP-WebSocket handshake flows)"
 
-    @only(websocket.WebSocketFlow)
+    @only(http.HTTPFlow, websocket.WebSocketFlow)
     def __call__(self, f):
-        return True
+        m = (
+            (isinstance(f, http.HTTPFlow) and f.request and net_websockets.check_handshake(f.request.headers))
+            or isinstance(f, websocket.WebSocketFlow)
+        )
+        return m
 
 
 class FTCP(_Action):
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index e5683356b..e3867033a 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -152,10 +152,15 @@ class Master:
                 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)
+            hfs = [hf for hf in self.waiting_flows if hf.id == f.metadata['websocket_handshake']]
+            if hfs:
+                hf = hfs[0]
+                f.handshake_flow = hf
+                self.waiting_flows.remove(hf)
+                self._change_reverse_host(f.handshake_flow)
+            else:
+                # this will fail - but at least it will load the remaining flows
+                f.handshake_flow = http.HTTPFlow(None, None)
 
         f.reply = controller.DummyReply()
         for e, o in eventsequence.iterate(f):
diff --git a/test/mitmproxy/test_flowfilter.py b/test/mitmproxy/test_flowfilter.py
index d53cec7d7..e6d39c7f1 100644
--- a/test/mitmproxy/test_flowfilter.py
+++ b/test/mitmproxy/test_flowfilter.py
@@ -439,6 +439,17 @@ class TestMatchingWebSocketFlow:
         assert not self.q("~tcp", f)
         assert not self.q("~http", f)
 
+    def test_handshake(self):
+        f = self.flow().handshake_flow
+        assert self.q("~websocket", f)
+        assert not self.q("~tcp", f)
+        assert self.q("~http", f)
+
+        f = tflow.tflow()
+        assert not self.q("~websocket", f)
+        f = tflow.tflow(resp=True)
+        assert not self.q("~websocket", f)
+
     def test_ferr(self):
         e = self.err()
         assert self.q("~e", e)
diff --git a/tox.ini b/tox.ini
index 353a4d18a..9fba58756 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = py35, py36, py37, flake8, filename_matching, mypy, individual_coverage, docs
+envlist = py35, py36, py37, py38, flake8, filename_matching, mypy, individual_coverage, docs
 skipsdist = True
 toxworkdir={env:TOX_WORK_DIR:.tox}