match HTTP-WebSocket handshake flow with ~websocket

fixes #3990

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 ---
 

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):

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]
+            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):

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)

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}