clean up code, improve DummyReply

1 files changed, 19 insertions, 4 deletions

diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index 5ef81f8c..72374f31 100644
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -203,6 +203,12 @@ def handler(f):
         if not hasattr(message, "reply"):
             raise exceptions.ControlException("Message %s has no reply attribute" % message)
 
+        # DummyReplys may be reused multiple times.
+        # We only clear them up on the next handler so that we can access value and
+        # state in the meantime.
+        if isinstance(message.reply, DummyReply):
+            message.reply.reset()
+
         # The following ensures that inheritance with wrapped handlers in the
         # base class works. If we're the first handler, then responsibility for
         # acking is ours. If not, it's someone else's and we ignore it.
@@ -228,7 +234,7 @@ def handler(f):
 
             # DummyReplys may be reused multiple times.
             if isinstance(message.reply, DummyReply):
-                message.reply.reset()
+                message.reply.mark_reset()
         return ret
     # Mark this function as a handler wrapper
     wrapper.__dict__["__handler"] = True
@@ -269,6 +275,10 @@ class Reply(object):
     def has_message(self):
         return self.value != NO_REPLY
 
+    @property
+    def done(self):
+        return self.state == "committed"
+
     def handle(self):
         """
         Reply are handled by controller.handlers, which may be nested. The first handler takes
@@ -329,12 +339,17 @@ class DummyReply(Reply):
     """
     def __init__(self):
         super(DummyReply, self).__init__(None)
+        self._should_reset = False
 
-    def reset(self):
+    def mark_reset(self):
         if self.state != "committed":
             raise exceptions.ControlException("Uncommitted reply: %s" % self.obj)
-        self._state = "unhandled"
-        self.value = NO_REPLY
+        self._should_reset = True
+
+    def reset(self):
+        if self._should_reset:
+            self._state = "unhandled"
+            self.value = NO_REPLY
 
     def __del__(self):
         pass