clean up code, improve DummyReply

21 files changed, 90 insertions, 74 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)
 
diff --git a/test/mitmproxy/builtins/test_anticache.py b/test/mitmproxy/builtins/test_anticache.py
index ac321e26..8897de52 100644
--- a/test/mitmproxy/builtins/test_anticache.py
+++ b/test/mitmproxy/builtins/test_anticache.py
@@ -14,11 +14,11 @@ class TestAntiCache(mastertest.MasterTest):
         m.addons.add(o, sa)
 
         f = tutils.tflow(resp=True)
-        self.invoke(m, "request", f)
+        m.request(f)
 
         f = tutils.tflow(resp=True)
         f.request.headers["if-modified-since"] = "test"
         f.request.headers["if-none-match"] = "test"
-        self.invoke(m, "request", f)
+        m.request(f)
         assert "if-modified-since" not in f.request.headers
         assert "if-none-match" not in f.request.headers
diff --git a/test/mitmproxy/builtins/test_anticomp.py b/test/mitmproxy/builtins/test_anticomp.py
index a5f5a270..af9e4a6a 100644
--- a/test/mitmproxy/builtins/test_anticomp.py
+++ b/test/mitmproxy/builtins/test_anticomp.py
@@ -14,10 +14,10 @@ class TestAntiComp(mastertest.MasterTest):
         m.addons.add(o, sa)
 
         f = tutils.tflow(resp=True)
-        self.invoke(m, "request", f)
+        m.request(f)
 
         f = tutils.tflow(resp=True)
 
         f.request.headers["Accept-Encoding"] = "foobar"
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.headers["Accept-Encoding"] == "identity"
diff --git a/test/mitmproxy/builtins/test_dumper.py b/test/mitmproxy/builtins/test_dumper.py
index 1c7173e0..9b518b53 100644
--- a/test/mitmproxy/builtins/test_dumper.py
+++ b/test/mitmproxy/builtins/test_dumper.py
@@ -80,5 +80,5 @@ class TestContentView(mastertest.MasterTest):
         m = mastertest.RecordingMaster(o, None, s)
         d = dumper.Dumper()
         m.addons.add(o, d)
-        self.invoke(m, "response", tutils.tflow())
+        m.response(tutils.tflow())
         assert "Content viewer failed" in m.event_log[0][1]
diff --git a/test/mitmproxy/builtins/test_filestreamer.py b/test/mitmproxy/builtins/test_filestreamer.py
index 0e69b340..94d68813 100644
--- a/test/mitmproxy/builtins/test_filestreamer.py
+++ b/test/mitmproxy/builtins/test_filestreamer.py
@@ -28,8 +28,8 @@ class TestStream(mastertest.MasterTest):
 
             m.addons.add(o, sa)
             f = tutils.tflow(resp=True)
-            self.invoke(m, "request", f)
-            self.invoke(m, "response", f)
+            m.request(f)
+            m.response(f)
             m.addons.remove(sa)
 
             assert r()[0].response
@@ -38,6 +38,6 @@ class TestStream(mastertest.MasterTest):
 
             m.addons.add(o, sa)
             f = tutils.tflow()
-            self.invoke(m, "request", f)
+            m.request(f)
             m.addons.remove(sa)
             assert not r()[1].response
diff --git a/test/mitmproxy/builtins/test_replace.py b/test/mitmproxy/builtins/test_replace.py
index 5e70ce56..07abcda4 100644
--- a/test/mitmproxy/builtins/test_replace.py
+++ b/test/mitmproxy/builtins/test_replace.py
@@ -43,10 +43,10 @@ class TestReplace(mastertest.MasterTest):
 
         f = tutils.tflow()
         f.request.content = b"foo"
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.content == b"bar"
 
         f = tutils.tflow(resp=True)
         f.response.content = b"foo"
-        self.invoke(m, "response", f)
+        m.response(f)
         assert f.response.content == b"bar"
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index 2870fd17..0bac6ca0 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -69,7 +69,7 @@ class TestScript(mastertest.MasterTest):
 
         sc.ns.call_log = []
         f = tutils.tflow(resp=True)
-        self.invoke(m, "request", f)
+        m.request(f)
 
         recf = sc.ns.call_log[0]
         assert recf[1] == "request"
@@ -102,7 +102,7 @@ class TestScript(mastertest.MasterTest):
         )
         m.addons.add(o, sc)
         f = tutils.tflow(resp=True)
-        self.invoke(m, "request", f)
+        m.request(f)
         assert m.event_log[0][0] == "error"
 
     def test_duplicate_flow(self):
diff --git a/test/mitmproxy/builtins/test_setheaders.py b/test/mitmproxy/builtins/test_setheaders.py
index 41c18360..63685177 100644
--- a/test/mitmproxy/builtins/test_setheaders.py
+++ b/test/mitmproxy/builtins/test_setheaders.py
@@ -33,12 +33,12 @@ class TestSetHeaders(mastertest.MasterTest):
         )
         f = tutils.tflow()
         f.request.headers["one"] = "xxx"
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.headers["one"] == "two"
 
         f = tutils.tflow(resp=True)
         f.response.headers["one"] = "xxx"
-        self.invoke(m, "response", f)
+        m.response(f)
         assert f.response.headers["one"] == "three"
 
         m, sh = self.mkmaster(
@@ -50,7 +50,7 @@ class TestSetHeaders(mastertest.MasterTest):
         f = tutils.tflow(resp=True)
         f.request.headers["one"] = "xxx"
         f.response.headers["one"] = "xxx"
-        self.invoke(m, "response", f)
+        m.response(f)
         assert f.response.headers.get_all("one") == ["two", "three"]
 
         m, sh = self.mkmaster(
@@ -61,5 +61,5 @@ class TestSetHeaders(mastertest.MasterTest):
         )
         f = tutils.tflow()
         f.request.headers["one"] = "xxx"
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.headers.get_all("one") == ["two", "three"]
diff --git a/test/mitmproxy/builtins/test_stickyauth.py b/test/mitmproxy/builtins/test_stickyauth.py
index 5757fb2d..00b12072 100644
--- a/test/mitmproxy/builtins/test_stickyauth.py
+++ b/test/mitmproxy/builtins/test_stickyauth.py
@@ -15,10 +15,10 @@ class TestStickyAuth(mastertest.MasterTest):
 
         f = tutils.tflow(resp=True)
         f.request.headers["authorization"] = "foo"
-        self.invoke(m, "request", f)
+        m.request(f)
 
         assert "address" in sa.hosts
 
         f = tutils.tflow(resp=True)
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.headers["authorization"] == "foo"
diff --git a/test/mitmproxy/builtins/test_stickycookie.py b/test/mitmproxy/builtins/test_stickycookie.py
index e9d92c83..13cea6a2 100644
--- a/test/mitmproxy/builtins/test_stickycookie.py
+++ b/test/mitmproxy/builtins/test_stickycookie.py
@@ -34,23 +34,23 @@ class TestStickyCookie(mastertest.MasterTest):
 
         f = tutils.tflow(resp=True)
         f.response.headers["set-cookie"] = "foo=bar"
-        self.invoke(m, "request", f)
+        m.request(f)
 
         f.reply.acked = False
-        self.invoke(m, "response", f)
+        m.response(f)
 
         assert sc.jar
         assert "cookie" not in f.request.headers
 
         f = f.copy()
         f.reply.acked = False
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.headers["cookie"] == "foo=bar"
 
     def _response(self, s, m, sc, cookie, host):
         f = tutils.tflow(req=ntutils.treq(host=host, port=80), resp=True)
         f.response.headers["Set-Cookie"] = cookie
-        self.invoke(m, "response", f)
+        m.response(f)
         return f
 
     def test_response(self):
@@ -79,7 +79,7 @@ class TestStickyCookie(mastertest.MasterTest):
         c2 = "othercookie=helloworld; Path=/"
         f = self._response(s, m, sc, c1, "www.google.com")
         f.response.headers["Set-Cookie"] = c2
-        self.invoke(m, "response", f)
+        m.response(f)
         googlekey = list(sc.jar.keys())[0]
         assert len(sc.jar[googlekey].keys()) == 2
 
@@ -96,7 +96,7 @@ class TestStickyCookie(mastertest.MasterTest):
         ]
         for c in cs:
             f.response.headers["Set-Cookie"] = c
-            self.invoke(m, "response", f)
+            m.response(f)
         googlekey = list(sc.jar.keys())[0]
         assert len(sc.jar[googlekey].keys()) == len(cs)
 
@@ -108,7 +108,7 @@ class TestStickyCookie(mastertest.MasterTest):
         c2 = "somecookie=newvalue; Path=/"
         f = self._response(s, m, sc, c1, "www.google.com")
         f.response.headers["Set-Cookie"] = c2
-        self.invoke(m, "response", f)
+        m.response(f)
         googlekey = list(sc.jar.keys())[0]
         assert len(sc.jar[googlekey].keys()) == 1
         assert list(sc.jar[googlekey]["somecookie"].items())[0][1] == "newvalue"
@@ -120,7 +120,7 @@ class TestStickyCookie(mastertest.MasterTest):
         # by setting the expire time in the past
         f = self._response(s, m, sc, "duffer=zafar; Path=/", "www.google.com")
         f.response.headers["Set-Cookie"] = "duffer=; Expires=Thu, 01-Jan-1970 00:00:00 GMT"
-        self.invoke(m, "response", f)
+        m.response(f)
         assert not sc.jar.keys()
 
     def test_request(self):
@@ -128,5 +128,5 @@ class TestStickyCookie(mastertest.MasterTest):
 
         f = self._response(s, m, sc, "SSID=mooo", "www.google.com")
         assert "cookie" not in f.request.headers
-        self.invoke(m, "request", f)
+        m.request(f)
         assert "cookie" in f.request.headers
diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py
index 95597d2c..08659d19 100644
--- a/test/mitmproxy/mastertest.py
+++ b/test/mitmproxy/mastertest.py
@@ -1,5 +1,3 @@
-import mock
-
 from . import tutils
 import netlib.tutils
 
@@ -8,24 +6,19 @@ from mitmproxy import flow, proxy, models, controller
 
 
 class MasterTest:
-    def invoke(self, master, handler, *message):
-        with master.handlecontext():
-            func = getattr(master, handler)
-            func(*message)
 
     def cycle(self, master, content):
         f = tutils.tflow(req=netlib.tutils.treq(content=content))
         l = proxy.Log("connect")
         l.reply = controller.DummyReply()
         master.log(l)
-        self.invoke(master, "clientconnect", f.client_conn)
-        self.invoke(master, "clientconnect", f.client_conn)
-        self.invoke(master, "serverconnect", f.server_conn)
-        self.invoke(master, "request", f)
+        master.clientconnect(f.client_conn)
+        master.serverconnect(f.server_conn)
+        master.request(f)
         if not f.error:
             f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content))
-            self.invoke(master, "response", f)
-        self.invoke(master, "clientdisconnect", f)
+            master.response(f)
+        master.clientdisconnect(f)
         return f
 
     def dummy_cycle(self, master, n, content):
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
index 07ba1c57..c4f1e9ae 100644
--- a/test/mitmproxy/script/test_concurrent.py
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -25,8 +25,8 @@ class TestConcurrent(mastertest.MasterTest):
         )
         m.addons.add(m.options, sc)
         f1, f2 = tutils.tflow(), tutils.tflow()
-        self.invoke(m, "request", f1)
-        self.invoke(m, "request", f2)
+        m.request(f1)
+        m.request(f2)
         start = time.time()
         while time.time() - start < 5:
             if f1.reply.state == f2.reply.state == "committed":
diff --git a/test/mitmproxy/test_controller.py b/test/mitmproxy/test_controller.py
index 8fe2453d..abf66b6c 100644
--- a/test/mitmproxy/test_controller.py
+++ b/test/mitmproxy/test_controller.py
@@ -187,6 +187,7 @@ class TestDummyReply(object):
             reply.ack()
             reply.take()
             reply.commit()
+            reply.mark_reset()
             reply.reset()
         assert reply.state == "unhandled"
 
@@ -196,9 +197,11 @@ class TestDummyReply(object):
         reply.ack()
         reply.take()
         reply.commit()
+        reply.mark_reset()
+        assert reply.state == "committed"
         reply.reset()
         assert reply.state == "unhandled"
 
     def test_del(self):
         reply = controller.DummyReply()
-        reply.__del__()
\ No newline at end of file
+        reply.__del__()
diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py
index 34fcc261..6c24ace5 100644
--- a/test/mitmproxy/test_examples.py
+++ b/test/mitmproxy/test_examples.py
@@ -39,7 +39,7 @@ class TestScripts(mastertest.MasterTest):
     def test_add_header(self):
         m, _ = tscript("add_header.py")
         f = tutils.tflow(resp=netutils.tresp())
-        self.invoke(m, "response", f)
+        m.response(f)
         assert f.response.headers["newheader"] == "foo"
 
     def test_custom_contentviews(self):
@@ -54,9 +54,9 @@ class TestScripts(mastertest.MasterTest):
             tscript("iframe_injector.py")
 
         m, sc = tscript("iframe_injector.py", "http://example.org/evil_iframe")
-        flow = tutils.tflow(resp=netutils.tresp(content=b"<html>mitmproxy</html>"))
-        self.invoke(m, "response", flow)
-        content = flow.response.content
+        f = tutils.tflow(resp=netutils.tresp(content=b"<html>mitmproxy</html>"))
+        m.response(f)
+        content = f.response.content
         assert b'iframe' in content and b'evil_iframe' in content
 
     def test_modify_form(self):
@@ -64,23 +64,23 @@ class TestScripts(mastertest.MasterTest):
 
         form_header = Headers(content_type="application/x-www-form-urlencoded")
         f = tutils.tflow(req=netutils.treq(headers=form_header))
-        self.invoke(m, "request", f)
+        m.request(f)
 
         assert f.request.urlencoded_form[b"mitmproxy"] == b"rocks"
 
         f.request.headers["content-type"] = ""
-        self.invoke(m, "request", f)
+        m.request(f)
         assert list(f.request.urlencoded_form.items()) == [(b"foo", b"bar")]
 
     def test_modify_querystring(self):
         m, sc = tscript("modify_querystring.py")
         f = tutils.tflow(req=netutils.treq(path="/search?q=term"))
 
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.query["mitmproxy"] == "rocks"
 
         f.request.path = "/"
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.query["mitmproxy"] == "rocks"
 
     def test_modify_response_body(self):
@@ -89,13 +89,13 @@ class TestScripts(mastertest.MasterTest):
 
         m, sc = tscript("modify_response_body.py", "mitmproxy rocks")
         f = tutils.tflow(resp=netutils.tresp(content=b"I <3 mitmproxy"))
-        self.invoke(m, "response", f)
+        m.response(f)
         assert f.response.content == b"I <3 rocks"
 
     def test_redirect_requests(self):
         m, sc = tscript("redirect_requests.py")
         f = tutils.tflow(req=netutils.treq(host="example.org"))
-        self.invoke(m, "request", f)
+        m.request(f)
         assert f.request.host == "mitmproxy.org"
 
     def test_har_extractor(self):
@@ -119,7 +119,7 @@ class TestScripts(mastertest.MasterTest):
                 req=netutils.treq(**times),
                 resp=netutils.tresp(**times)
             )
-            self.invoke(m, "response", f)
+            m.response(f)
             m.addons.remove(sc)
 
             with open(path, "rb") as f:
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 256ee124..1caeb100 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -3,9 +3,9 @@ import io
 
 import netlib.utils
 from netlib.http import Headers
-from mitmproxy import filt, controller, flow, options
+from mitmproxy import filt, flow, options
 from mitmproxy.contrib import tnetstring
-from mitmproxy.exceptions import FlowReadException
+from mitmproxy.exceptions import FlowReadException, Kill
 from mitmproxy.models import Error
 from mitmproxy.models import Flow
 from mitmproxy.models import HTTPFlow
@@ -372,14 +372,16 @@ class TestHTTPFlow(object):
         assert f.get_state() == f2.get_state()
 
     def test_kill(self):
-        s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = mock.Mock()
         f = tutils.tflow()
         f.reply.handle()
-        f.intercept(mock.Mock())
+        f.intercept(fm)
+        assert fm.handle_intercept.called
+        assert f.killable
         f.kill(fm)
-        for i in s.view:
-            assert "killed" in str(i.error)
+        assert not f.killable
+        assert fm.error.called
+        assert f.reply.value == Kill
 
     def test_killall(self):
         s = flow.State()
@@ -835,8 +837,8 @@ class TestFlowMaster:
 
         f = tutils.tflow()
         f.request.host = "nonexistent"
-        fm.process_new_request(f)
-        assert "killed" in f.error.msg
+        fm.request(f)
+        assert f.reply.value == Kill
 
 
 class TestRequest: