clean up code, improve DummyReply

7 files changed, 30 insertions, 12 deletions

diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index 571119ed..7e69e098 100644
--- a/mitmproxy/console/flowlist.py
+++ b/mitmproxy/console/flowlist.py
@@ -182,7 +182,7 @@ class ConnectionItem(urwid.WidgetWrap):
             self.flow.accept_intercept(self.master)
             signals.flowlist_change.send(self)
         elif key == "d":
-            if self.flow.reply and self.flow.reply.state != "committed":
+            if self.flow.killable:
                 self.flow.kill(self.master)
             self.state.delete_flow(self.flow)
             signals.flowlist_change.send(self)
@@ -246,7 +246,7 @@ class ConnectionItem(urwid.WidgetWrap):
                 callback = self.save_flows_prompt,
             )
         elif key == "X":
-            if self.flow.reply and self.flow.reply.state != "committed":
+            if self.flow.killable:
                 self.flow.kill(self.master)
         elif key == "enter":
             if self.flow.request:
diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py
index 22314396..6d74be65 100644
--- a/mitmproxy/console/flowview.py
+++ b/mitmproxy/console/flowview.py
@@ -8,7 +8,6 @@ import urwid
 from typing import Optional, Union  # noqa
 
 from mitmproxy import contentviews
-from mitmproxy import controller
 from mitmproxy import models
 from mitmproxy import utils
 from mitmproxy.console import common
@@ -537,7 +536,7 @@ class FlowView(tabs.Tabs):
             else:
                 self.view_next_flow(self.flow)
             f = self.flow
-            if not f.reply.acked:
+            if f.killable:
                 f.kill(self.master)
             self.state.delete_flow(f)
         elif key == "D":
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
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index e02cdc15..0475ef4e 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -234,7 +234,7 @@ class FlowMaster(controller.Master):
             pb = self.do_server_playback(f)
             if not pb and self.kill_nonreplay:
                 self.add_log("Killed {}".format(f.request.url), "info")
-                f.kill(self)
+                f.reply.kill()
 
     def replay_request(self, f, block=False):
         """
diff --git a/mitmproxy/flow/state.py b/mitmproxy/flow/state.py
index 6f7e6c03..8576fadc 100644
--- a/mitmproxy/flow/state.py
+++ b/mitmproxy/flow/state.py
@@ -178,7 +178,7 @@ class FlowStore(FlowList):
 
     def kill_all(self, master):
         for f in self._list:
-            if f.reply.state != "committed":
+            if f.killable:
                 f.kill(master)
 
 
diff --git a/mitmproxy/models/flow.py b/mitmproxy/models/flow.py
index 5d61e47f..fc673274 100644
--- a/mitmproxy/models/flow.py
+++ b/mitmproxy/models/flow.py
@@ -149,13 +149,17 @@ class Flow(stateobject.StateObject):
             self.set_state(self._backup)
             self._backup = None
 
+    @property
+    def killable(self):
+        return self.reply and self.reply.state in {"handled", "taken"}
+
     def kill(self, master):
         """
             Kill this request.
         """
         self.error = Error("Connection killed")
         self.intercepted = False
-        # reply.state should only be handled or taken here.
+        # reply.state should only be "handled" or "taken" here.
         # if none of this is the case, .take() will raise an exception.
         if self.reply.state != "taken":
             self.reply.take()
diff --git a/mitmproxy/web/app.py b/mitmproxy/web/app.py
index 097de634..5bd6f274 100644
--- a/mitmproxy/web/app.py
+++ b/mitmproxy/web/app.py
@@ -234,7 +234,7 @@ class AcceptFlow(RequestHandler):
 class FlowHandler(RequestHandler):
 
     def delete(self, flow_id):
-        if self.flow.reply.state != "committed":
+        if self.flow.killable:
             self.flow.kill(self.master)
         self.state.delete_flow(self.flow)