Make mitmdump handle invalid serialized data gracefully.

3 files changed, 21 insertions, 5 deletions

diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 2eb0035e..97232ac7 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -93,7 +93,7 @@ class DumpMaster(flow.FlowMaster):
         try:
             f = file(path, "r")
             flows = list(flow.FlowReader(f).stream())
-        except IOError, v:
+        except (IOError, flow.FlowReadError), v:
             raise DumpError(v.strerror)
         return flows
 
@@ -182,7 +182,6 @@ class DumpMaster(flow.FlowMaster):
         return f
 
 
-
 # begin nocover
     def run(self):
         try:
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 60a6ebd2..d555c926 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -560,6 +560,10 @@ class FlowWriter:
         s = json.dumps(d)
         self.ns.write(s)
 
+class FlowReadError(Exception):
+    @property
+    def strerror(self):
+        return self.args[0]
 
 class FlowReader:
     def __init__(self, fo):
@@ -570,7 +574,10 @@ class FlowReader:
         """
             Yields Flow objects from the dump.
         """
-        for i in self.ns:
-            data = json.loads(i)
-            yield Flow.from_state(data)
+        try:
+            for i in self.ns:
+                data = json.loads(i)
+                yield Flow.from_state(data)
+        except netstring.DecoderError:
+            raise FlowReadError("Invalid data format.")
 
diff --git a/test/test_flow.py b/test/test_flow.py
index 78892b5d..d2cb85dc 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -356,6 +356,16 @@ class uSerialize(libpry.AutoTree):
         assert len(l) == 1
         assert l[0] == f
 
+    def test_error(self):
+        sio = StringIO()
+        sio.write("bogus")
+        sio.seek(0)
+        r = flow.FlowReader(sio)
+        libpry.raises(flow.FlowReadError, list, r.stream())
+
+        f = flow.FlowReadError("foo")
+        assert f.strerror == "foo"
+
 
 class uFlowMaster(libpry.AutoTree):
     def test_all(self):