catch TypeError when reading flows, fix #4705

This commit is contained in:
Maximilian Hils 2021-08-03 14:34:12 +02:00
parent 331061c6f2
commit 04a1ff4a42
2 changed files with 30 additions and 23 deletions

View File

@ -46,7 +46,7 @@ class FlowReader:
if mdata["type"] not in FLOW_TYPES: if mdata["type"] not in FLOW_TYPES:
raise exceptions.FlowReadException("Unknown flow type: {}".format(mdata["type"])) raise exceptions.FlowReadException("Unknown flow type: {}".format(mdata["type"]))
yield FLOW_TYPES[mdata["type"]].from_state(mdata) 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": if str(e) == "not a tnetstring: empty file":
return # Error is due to EOF return # Error is due to EOF
raise exceptions.FlowReadException("Invalid data format.") raise exceptions.FlowReadException("Invalid data format.")

View File

@ -1,28 +1,35 @@
from unittest import mock import io
import pytest import pytest
from hypothesis import example, given
from hypothesis.strategies import binary
from mitmproxy import exceptions from mitmproxy import exceptions, version
from mitmproxy.io import FlowReader from mitmproxy.io import FlowReader, tnetstring
class TestFlowReader: class TestFlowReader:
@mock.patch("mitmproxy.io.tnetstring.load") @given(binary())
@mock.patch("mitmproxy.io.compat.migrate_flow") @example(b'51:11:12345678901#4:this,8:true!0:~,4:true!0:]4:\\x00,~}')
def test_stream_with_exception(self, mock1, mock2): def test_fuzz(self, data):
with open("./abc", "rb") as fp: f = io.BytesIO(data)
reader = FlowReader(fp) reader = FlowReader(f)
mock2.side_effect = ValueError() try:
with pytest.raises(exceptions.FlowReadException, match="Invalid data format."): for _ in reader.stream():
for i in reader.stream(): pass
pass except exceptions.FlowReadException:
mock2.side_effect = None pass # should never raise anything else.
mock1.side_effect = ValueError("TestException")
with pytest.raises(exceptions.FlowReadException, match="TestException"): def test_empty(self):
for i in reader.stream(): assert list(FlowReader(io.BytesIO(b"")).stream()) == []
pass
mock1.side_effect = None def test_unknown_type(self):
mock1.return_value = {"type": "test_type"} with pytest.raises(exceptions.FlowReadException, match="Unknown flow type"):
with pytest.raises(exceptions.FlowReadException, match="Unknown flow type: test_type"): weird_flow = tnetstring.dumps({"type": "unknown", "version": version.FLOW_FORMAT_VERSION})
for i in reader.stream(): for _ in FlowReader(io.BytesIO(weird_flow)).stream():
pass 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