From b4e9e55c3420f9bcd27142dfaaf0934e73d90b6f Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Tue, 24 Jul 2012 15:15:41 +1200
Subject: [PATCH] Be more tolerant of corrupted or truncated flows.

We load as far as possible. mitmproxy will only terminate if it was not able to
recover any flows. mitmdump will stop loading as soon as an error is
encountered, but not exit with an error.
---
 libmproxy/console/__init__.py | 8 ++++++--
 libmproxy/dump.py             | 3 +--
 libmproxy/flow.py             | 2 +-
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index bad73e91a..898e858f4 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -548,7 +548,9 @@ class ConsoleMaster(flow.FlowMaster):
 
         if self.options.rfile:
             ret = self.load_flows(self.options.rfile)
-            if ret:
+            if ret and self.state.flow_count():
+                self.add_event("File truncated or corrupted. Loaded as many flows as possible.")
+            else:
                 self.shutdown()
                 print >> sys.stderr, "Could not load file:", ret
                 sys.exit(1)
@@ -653,14 +655,16 @@ class ConsoleMaster(flow.FlowMaster):
             fr = flow.FlowReader(f)
         except IOError, v:
             return v.strerror
+        reterr = None
         try:
             flow.FlowMaster.load_flows(self, fr)
         except flow.FlowReadError, v:
-            return v.strerror
+            reterr = v.strerror
         f.close()
         if self.flow_list_walker:
             self.sync_list_view()
             self.focus_current()
+        return reterr
 
     def path_prompt(self, prompt, text, callback, *args):
         self.statusbar.path_prompt(prompt, text)
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 9ad6c3272..95e5ac0f7 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -128,8 +128,7 @@ class DumpMaster(flow.FlowMaster):
             try:
                 self.load_flows(freader)
             except flow.FlowReadError, v:
-                raise DumpError(v)
-
+                self.add_event("Flow file corrupted. Stopped loading.")
 
     def _readflow(self, path):
         path = os.path.expanduser(path)
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index aa02b5c53..ef2266414 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -1482,7 +1482,7 @@ class FlowReader:
                     raise FlowReadError("Incompatible serialized data version: %s"%v)
                 off = self.fo.tell()
                 yield Flow._from_state(data)
-        except ValueError:
+        except ValueError, v:
             # Error is due to EOF
             if self.fo.tell() == off and self.fo.read() == '':
                 return