From 04a1ff4a42d4801014c77e77e4ecef3aeedac9bd Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Tue, 3 Aug 2021 14:34:12 +0200
Subject: [PATCH] catch TypeError when reading flows, fix #4705

---
 mitmproxy/io/io.py           |  2 +-
 test/mitmproxy/io/test_io.py | 51 ++++++++++++++++++++----------------
 2 files changed, 30 insertions(+), 23 deletions(-)

diff --git a/mitmproxy/io/io.py b/mitmproxy/io/io.py
index 5fdf03993..b57d3d6ee 100644
--- a/mitmproxy/io/io.py
+++ b/mitmproxy/io/io.py
@@ -46,7 +46,7 @@ class FlowReader:
                 if mdata["type"] not in FLOW_TYPES:
                     raise exceptions.FlowReadException("Unknown flow type: {}".format(mdata["type"]))
                 yield FLOW_TYPES[mdata["type"]].from_state(mdata)
-        except ValueError as e:
+        except (ValueError, TypeError) as e:
             if str(e) == "not a tnetstring: empty file":
                 return  # Error is due to EOF
             raise exceptions.FlowReadException("Invalid data format.")
diff --git a/test/mitmproxy/io/test_io.py b/test/mitmproxy/io/test_io.py
index a868055a1..ea67cbbe0 100644
--- a/test/mitmproxy/io/test_io.py
+++ b/test/mitmproxy/io/test_io.py
@@ -1,28 +1,35 @@
-from unittest import mock
+import io
 
 import pytest
+from hypothesis import example, given
+from hypothesis.strategies import binary
 
-from mitmproxy import exceptions
-from mitmproxy.io import FlowReader
+from mitmproxy import exceptions, version
+from mitmproxy.io import FlowReader, tnetstring
 
 
 class TestFlowReader:
-    @mock.patch("mitmproxy.io.tnetstring.load")
-    @mock.patch("mitmproxy.io.compat.migrate_flow")
-    def test_stream_with_exception(self, mock1, mock2):
-        with open("./abc", "rb") as fp:
-            reader = FlowReader(fp)
-            mock2.side_effect = ValueError()
-            with pytest.raises(exceptions.FlowReadException, match="Invalid data format."):
-                for i in reader.stream():
-                    pass
-            mock2.side_effect = None
-            mock1.side_effect = ValueError("TestException")
-            with pytest.raises(exceptions.FlowReadException, match="TestException"):
-                for i in reader.stream():
-                    pass
-            mock1.side_effect = None
-            mock1.return_value = {"type": "test_type"}
-            with pytest.raises(exceptions.FlowReadException, match="Unknown flow type: test_type"):
-                for i in reader.stream():
-                    pass
+    @given(binary())
+    @example(b'51:11:12345678901#4:this,8:true!0:~,4:true!0:]4:\\x00,~}')
+    def test_fuzz(self, data):
+        f = io.BytesIO(data)
+        reader = FlowReader(f)
+        try:
+            for _ in reader.stream():
+                pass
+        except exceptions.FlowReadException:
+            pass  # should never raise anything else.
+
+    def test_empty(self):
+        assert list(FlowReader(io.BytesIO(b"")).stream()) == []
+
+    def test_unknown_type(self):
+        with pytest.raises(exceptions.FlowReadException, match="Unknown flow type"):
+            weird_flow = tnetstring.dumps({"type": "unknown", "version": version.FLOW_FORMAT_VERSION})
+            for _ in FlowReader(io.BytesIO(weird_flow)).stream():
+                pass
+
+    def test_cannot_migrate(self):
+        with pytest.raises(exceptions.FlowReadException, match="cannot read files with flow format version 0"):
+            for _ in FlowReader(io.BytesIO(b"14:7:version;1:0#}")).stream():
+                pass