Repository: mitmproxy/mitmproxy

Issue:
~websocket filter when reading from a file does not work
#### Problem Description
It appears that the `~websocket` filter, which is supposed to match all WebSocket flows per <https://docs.mitmproxy.org/stable/concepts-filters/>, does not match anything when reading from an existing dump.

#### Steps to reproduce the behavior:
1. Capture WebSocket traffic: `mitmdump -w infile ~websocket`
2. Apply a filter of `~websocket`: `mitmdump -nr infile -w outfile ~websocket`

The result is a non-empty `infile`, which contains the expected WebSocket messages, and an **empty** `outfile` even though they should contain the same data.

In addition, I also noticed that `mitmdump -nr infile` produces no output whatsoever, even with `-v` to increase the verbosity. I'm aware that the UI does not properly support WebSocket (#1020), but I'd expect it to at least report that messages exist as it does during dumping (i.e. `127.0.0.1:57466 -> WebSocket 1 message -> URL` etc.). This does not happen when the initial capturing is done without a filter, i.e. those "message" log lines appear in that case. This might point at a more fundamental issue with reading back the WebSocket data than just the filtering.

#### System Information
```
$ mitmproxy --version
Mitmproxy: 5.1.1 (+2, commit a8ca63f)
Python:    3.8.2
OpenSSL:   OpenSSL 1.1.1g  21 Apr 2020
Platform:  Linux-5.6.0-1-amd64-x86_64-with-glibc2.29
```

(Regarding the version number, see #3987; this is a clean 5.1.1 installation.)


Discussion:
The HTTP-WebSocket handshake flow is missing from the dump file, which seems to cause a few missing internal links...
I'll send a PR to fix this by also matching all HTTP-WebSocket handshake flows with the `~websocket` filter.



Resolving patch:
diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py
--- a/mitmproxy/flowfilter.py
+++ b/mitmproxy/flowfilter.py
@@ -43,6 +43,7 @@
 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 @@ def __call__(self, f):
 
 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
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -152,10 +152,15 @@ async def load_flow(self, f):
                 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):


Test patch:
diff --git a/test/mitmproxy/test_flowfilter.py b/test/mitmproxy/test_flowfilter.py
--- a/test/mitmproxy/test_flowfilter.py
+++ b/test/mitmproxy/test_flowfilter.py
@@ -439,6 +439,17 @@ def test_websocket(self):
         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)