Convert flow.state to an addon

A first step, and already many things are much nicer.

25 files changed, 101 insertions, 158 deletions

diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py
index 5b078129..9e5677da 100644
--- a/mitmproxy/addons.py
+++ b/mitmproxy/addons.py
@@ -25,7 +25,7 @@ class Addons:
     def _options_update(self, options, updated):
         for i in self.chain:
             with self.master.handlecontext():
-                i.configure(options, updated)
+                self.invoke_with_context(i, "configure", options, updated)
 
     def startup(self, s):
         """
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index 1a413990..64761e40 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -222,7 +222,9 @@ class ConsoleMaster(flow.FlowMaster):
     palette = []
 
     def __init__(self, server, options):
-        flow.FlowMaster.__init__(self, options, server, ConsoleState())
+        flow.FlowMaster.__init__(self, options, server)
+        self.state = ConsoleState()
+        self.addons.add(self.state)
         self.stream_path = None
         # This line is just for type hinting
         self.options = self.options  # type: Options
diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index d3f894fa..d3a66876 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -33,7 +33,7 @@ class Options(options.Options):
 class DumpMaster(flow.FlowMaster):
 
     def __init__(self, server, options):
-        flow.FlowMaster.__init__(self, options, server, flow.DummyState())
+        flow.FlowMaster.__init__(self, options, server)
         self.has_errored = False
         self.addons.add(termlog.TermLog())
         self.addons.add(*builtins.default_addons())
diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py
index 9d243f49..07e1de6c 100644
--- a/mitmproxy/flow/__init__.py
+++ b/mitmproxy/flow/__init__.py
@@ -4,10 +4,10 @@ from mitmproxy.flow.master import FlowMaster
 from mitmproxy.flow.modules import (
     AppRegistry
 )
-from mitmproxy.flow.state import State, DummyState, FlowView
+from mitmproxy.flow.state import State, FlowView
 
 __all__ = [
     "export", "modules",
     "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths",
-    "FlowMaster", "AppRegistry", "DummyState", "State", "FlowView",
+    "FlowMaster", "AppRegistry", "State", "FlowView",
 ]
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index 85048b8c..33502dc5 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -47,11 +47,10 @@ class FlowMaster(controller.Master):
         if len(self.servers) > 0:
             return self.servers[0]
 
-    def __init__(self, options, server, state):
+    def __init__(self, options, server):
         super().__init__(options)
         if server:
             self.add_server(server)
-        self.state = state
         self.apps = modules.AppRegistry()
 
     def start_app(self, host, port):
@@ -196,7 +195,7 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def error(self, f):
-        self.state.update_flow(f)
+        pass
 
     @controller.handler
     def requestheaders(self, f):
@@ -216,8 +215,6 @@ class FlowMaster(controller.Master):
                     self.add_log("Error in wsgi app. %s" % err, "error")
                 f.reply.kill()
                 return
-        if f not in self.state.flows:  # don't add again on replay
-            self.state.add_flow(f)
 
     @controller.handler
     def responseheaders(self, f):
@@ -225,7 +222,7 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def response(self, f):
-        self.state.update_flow(f)
+        pass
 
     @controller.handler
     def websocket_handshake(self, f):
@@ -233,8 +230,6 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def tcp_start(self, flow):
-        # TODO: This would break mitmproxy currently.
-        # self.state.add_flow(flow)
         pass
 
     @controller.handler
diff --git a/mitmproxy/flow/state.py b/mitmproxy/flow/state.py
index dddb5ca5..bb7460b6 100644
--- a/mitmproxy/flow/state.py
+++ b/mitmproxy/flow/state.py
@@ -277,13 +277,12 @@ class State:
     def resume(self, f):
         self.update_flow(f)
 
+    def error(self, f):
+        self.update_flow(f)
 
+    def request(self, f):
+        if f not in self.flows:  # don't add again on replay
+            self.add_flow(f)
 
-class DummyState:
-    flows = ()
-
-    def add_flow(self, *args, **kwargs):
-        pass
-
-    def update_flow(self, *args, **kwargs):
-        pass
+    def response(self, f):
+        self.update_flow(f)
diff --git a/mitmproxy/web/master.py b/mitmproxy/web/master.py
index 23c95e57..3f61fd94 100644
--- a/mitmproxy/web/master.py
+++ b/mitmproxy/web/master.py
@@ -133,7 +133,9 @@ class Options(options.Options):
 class WebMaster(flow.FlowMaster):
 
     def __init__(self, server, options):
-        super().__init__(options, server, WebState())
+        super().__init__(options, server)
+        self.state = WebState()
+        self.addons.add(self.state)
         self.addons.add(*builtins.default_addons())
         self.app = app.Application(
             self, self.options.wdebug, self.options.wauthenticator
diff --git a/test/mitmproxy/builtins/test_anticache.py b/test/mitmproxy/builtins/test_anticache.py
index 81041e73..512c90f5 100644
--- a/test/mitmproxy/builtins/test_anticache.py
+++ b/test/mitmproxy/builtins/test_anticache.py
@@ -1,15 +1,13 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import anticache
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
 class TestAntiCache(mastertest.MasterTest):
     def test_simple(self):
-        s = state.State()
         o = options.Options(anticache = True)
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = anticache.AntiCache()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/builtins/test_anticomp.py b/test/mitmproxy/builtins/test_anticomp.py
index f866f33f..1014b9ba 100644
--- a/test/mitmproxy/builtins/test_anticomp.py
+++ b/test/mitmproxy/builtins/test_anticomp.py
@@ -1,15 +1,13 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import anticomp
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
 class TestAntiComp(mastertest.MasterTest):
     def test_simple(self):
-        s = state.State()
         o = options.Options(anticomp = True)
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = anticomp.AntiComp()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/builtins/test_dumper.py b/test/mitmproxy/builtins/test_dumper.py
index 20403eea..3ab75bee 100644
--- a/test/mitmproxy/builtins/test_dumper.py
+++ b/test/mitmproxy/builtins/test_dumper.py
@@ -3,7 +3,6 @@ import io
 from .. import tutils, mastertest
 
 from mitmproxy.builtins import dumper
-from mitmproxy.flow import state
 from mitmproxy import exceptions
 from mitmproxy import dump
 from mitmproxy import models
@@ -71,14 +70,13 @@ class TestContentView(mastertest.MasterTest):
     def test_contentview(self, view_auto):
         view_auto.side_effect = exceptions.ContentViewException("")
 
-        s = state.State()
         sio = io.StringIO()
         o = dump.Options(
             flow_detail=4,
             verbosity=3,
             tfile=sio,
         )
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         d = dumper.Dumper()
         m.addons.add(d)
         m.response(tutils.tflow())
diff --git a/test/mitmproxy/builtins/test_filestreamer.py b/test/mitmproxy/builtins/test_filestreamer.py
index a43ea0b7..6fbeb40b 100644
--- a/test/mitmproxy/builtins/test_filestreamer.py
+++ b/test/mitmproxy/builtins/test_filestreamer.py
@@ -4,7 +4,6 @@ import os.path
 
 from mitmproxy.builtins import filestreamer
 from mitmproxy.flow import master, FlowReader
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
@@ -17,11 +16,10 @@ class TestStream(mastertest.MasterTest):
                 r = FlowReader(open(p, "rb"))
                 return list(r.stream())
 
-            s = state.State()
             o = options.Options(
                 outfile = (p, "wb")
             )
-            m = master.FlowMaster(o, None, s)
+            m = master.FlowMaster(o, None)
             sa = filestreamer.FileStreamer()
 
             m.addons.add(sa)
diff --git a/test/mitmproxy/builtins/test_replace.py b/test/mitmproxy/builtins/test_replace.py
index e9869ab7..03943867 100644
--- a/test/mitmproxy/builtins/test_replace.py
+++ b/test/mitmproxy/builtins/test_replace.py
@@ -1,7 +1,6 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import replace
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
@@ -30,14 +29,13 @@ class TestReplace(mastertest.MasterTest):
         )
 
     def test_simple(self):
-        s = state.State()
         o = options.Options(
             replacements = [
                 ("~q", "foo", "bar"),
                 ("~s", "foo", "bar"),
             ]
         )
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = replace.Replace()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py
index 683f3064..ecf2f3f7 100644
--- a/test/mitmproxy/builtins/test_script.py
+++ b/test/mitmproxy/builtins/test_script.py
@@ -8,7 +8,6 @@ from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy.builtins import script
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 
 from .. import tutils, mastertest
 
@@ -57,9 +56,8 @@ def test_load_script():
 
 class TestScript(mastertest.MasterTest):
     def test_simple(self):
-        s = state.State()
         o = options.Options()
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = script.Script(
             tutils.test_data.path(
                 "data/addonscripts/recorder.py"
@@ -79,9 +77,8 @@ class TestScript(mastertest.MasterTest):
         assert recf[1] == "request"
 
     def test_reload(self):
-        s = state.State()
         o = options.Options()
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         with tutils.tmpdir():
             with open("foo.py", "w"):
                 pass
@@ -98,9 +95,8 @@ class TestScript(mastertest.MasterTest):
             raise AssertionError("Change event not detected.")
 
     def test_exception(self):
-        s = state.State()
         o = options.Options()
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         sc = script.Script(
             tutils.test_data.path("data/addonscripts/error.py")
         )
@@ -113,25 +109,9 @@ class TestScript(mastertest.MasterTest):
         assert re.search('addonscripts/error.py", line \d+, in mkerr', m.event_log[0][1])
         assert m.event_log[0][1].endswith("ValueError: Error!\n")
 
-    def test_duplicate_flow(self):
-        s = state.State()
-        o = options.Options()
-        fm = master.FlowMaster(o, None, s)
-        fm.addons.add(
-            script.Script(
-                tutils.test_data.path("data/addonscripts/duplicate_flow.py")
-            )
-        )
-        f = tutils.tflow()
-        fm.request(f)
-        assert fm.state.flow_count() == 2
-        assert not fm.state.view[0].request.is_replay
-        assert fm.state.view[1].request.is_replay
-
     def test_addon(self):
-        s = state.State()
         o = options.Options()
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = script.Script(
             tutils.test_data.path(
                 "data/addonscripts/addon.py"
@@ -163,9 +143,8 @@ class TestCutTraceback:
 
 class TestScriptLoader(mastertest.MasterTest):
     def test_run_once(self):
-        s = state.State()
         o = options.Options(scripts=[])
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sl = script.ScriptLoader()
         m.addons.add(sl)
 
@@ -188,9 +167,8 @@ class TestScriptLoader(mastertest.MasterTest):
             )
 
     def test_simple(self):
-        s = state.State()
         o = options.Options(scripts=[])
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = script.ScriptLoader()
         m.addons.add(sc)
         assert len(m.addons) == 1
@@ -204,16 +182,14 @@ class TestScriptLoader(mastertest.MasterTest):
         assert len(m.addons) == 1
 
     def test_dupes(self):
-        s = state.State()
         o = options.Options(scripts=["one", "one"])
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = script.ScriptLoader()
         tutils.raises(exceptions.OptionsError, m.addons.add, o, sc)
 
     def test_order(self):
         rec = tutils.test_data.path("data/addonscripts/recorder.py")
 
-        s = state.State()
         o = options.Options(
             scripts = [
                 "%s %s" % (rec, "a"),
@@ -221,7 +197,7 @@ class TestScriptLoader(mastertest.MasterTest):
                 "%s %s" % (rec, "c"),
             ]
         )
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         sc = script.ScriptLoader()
         m.addons.add(sc)
 
diff --git a/test/mitmproxy/builtins/test_serverplayback.py b/test/mitmproxy/builtins/test_serverplayback.py
index 55ce9350..b97c01dd 100644
--- a/test/mitmproxy/builtins/test_serverplayback.py
+++ b/test/mitmproxy/builtins/test_serverplayback.py
@@ -4,7 +4,6 @@ import netlib.tutils
 from mitmproxy.builtins import serverplayback
 from mitmproxy import options
 from mitmproxy import exceptions
-from mitmproxy import flow
 
 
 class TestServerPlayback:
@@ -238,10 +237,9 @@ class TestServerPlayback:
         assert not s._hash(r) == s._hash(r2)
 
     def test_server_playback_full(self):
-        state = flow.State()
         s = serverplayback.ServerPlayback()
         o = options.Options(refresh_server_playback = True, keepserving=False)
-        m = mastertest.RecordingMaster(o, None, state)
+        m = mastertest.RecordingMaster(o, None)
         m.addons.add(s)
 
         f = tutils.tflow()
@@ -268,10 +266,9 @@ class TestServerPlayback:
         assert s.stop
 
     def test_server_playback_kill(self):
-        state = flow.State()
         s = serverplayback.ServerPlayback()
         o = options.Options(refresh_server_playback = True, replay_kill_extra=True)
-        m = mastertest.RecordingMaster(o, None, state)
+        m = mastertest.RecordingMaster(o, None)
         m.addons.add(s)
 
         f = tutils.tflow()
diff --git a/test/mitmproxy/builtins/test_setheaders.py b/test/mitmproxy/builtins/test_setheaders.py
index d148234d..234341f4 100644
--- a/test/mitmproxy/builtins/test_setheaders.py
+++ b/test/mitmproxy/builtins/test_setheaders.py
@@ -1,15 +1,13 @@
 from .. import tutils, mastertest
 
 from mitmproxy.builtins import setheaders
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
 class TestSetHeaders(mastertest.MasterTest):
     def mkmaster(self, **opts):
-        s = state.State()
         o = options.Options(**opts)
-        m = mastertest.RecordingMaster(o, None, s)
+        m = mastertest.RecordingMaster(o, None)
         sh = setheaders.SetHeaders()
         m.addons.add(sh)
         return m, sh
diff --git a/test/mitmproxy/builtins/test_stickyauth.py b/test/mitmproxy/builtins/test_stickyauth.py
index 8b939c24..25141554 100644
--- a/test/mitmproxy/builtins/test_stickyauth.py
+++ b/test/mitmproxy/builtins/test_stickyauth.py
@@ -1,15 +1,13 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import stickyauth
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 
 
 class TestStickyAuth(mastertest.MasterTest):
     def test_simple(self):
-        s = state.State()
         o = options.Options(stickyauth = ".*")
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = stickyauth.StickyAuth()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/builtins/test_stickycookie.py b/test/mitmproxy/builtins/test_stickycookie.py
index 4d3b7cb3..28b05a11 100644
--- a/test/mitmproxy/builtins/test_stickycookie.py
+++ b/test/mitmproxy/builtins/test_stickycookie.py
@@ -1,7 +1,6 @@
 from .. import tutils, mastertest
 from mitmproxy.builtins import stickycookie
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 from mitmproxy import options
 from netlib import tutils as ntutils
 
@@ -13,12 +12,11 @@ def test_domain_match():
 
 class TestStickyCookie(mastertest.MasterTest):
     def mk(self):
-        s = state.State()
         o = options.Options(stickycookie = ".*")
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sc = stickycookie.StickyCookie()
         m.addons.add(sc)
-        return s, m, sc
+        return m, sc
 
     def test_config(self):
         sc = stickycookie.StickyCookie()
@@ -29,7 +27,7 @@ class TestStickyCookie(mastertest.MasterTest):
         )
 
     def test_simple(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
         m.addons.add(sc)
 
         f = tutils.tflow(resp=True)
@@ -47,44 +45,44 @@ class TestStickyCookie(mastertest.MasterTest):
         m.request(f)
         assert f.request.headers["cookie"] == "foo=bar"
 
-    def _response(self, s, m, sc, cookie, host):
+    def _response(self, m, sc, cookie, host):
         f = tutils.tflow(req=ntutils.treq(host=host, port=80), resp=True)
         f.response.headers["Set-Cookie"] = cookie
         m.response(f)
         return f
 
     def test_response(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         c = "SSID=mooo; domain=.google.com, FOO=bar; Domain=.google.com; Path=/; " \
             "Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; "
 
-        self._response(s, m, sc, c, "host")
+        self._response(m, sc, c, "host")
         assert not sc.jar.keys()
 
-        self._response(s, m, sc, c, "www.google.com")
+        self._response(m, sc, c, "www.google.com")
         assert sc.jar.keys()
 
         sc.jar.clear()
         self._response(
-            s, m, sc, "SSID=mooo", "www.google.com"
+            m, sc, "SSID=mooo", "www.google.com"
         )
         assert list(sc.jar.keys())[0] == ('www.google.com', 80, '/')
 
     def test_response_multiple(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         # Test setting of multiple cookies
         c1 = "somecookie=test; Path=/"
         c2 = "othercookie=helloworld; Path=/"
-        f = self._response(s, m, sc, c1, "www.google.com")
+        f = self._response(m, sc, c1, "www.google.com")
         f.response.headers["Set-Cookie"] = c2
         m.response(f)
         googlekey = list(sc.jar.keys())[0]
         assert len(sc.jar[googlekey].keys()) == 2
 
     def test_response_weird(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         # Test setting of weird cookie keys
         f = tutils.tflow(req=ntutils.treq(host="www.google.com", port=80), resp=True)
@@ -100,12 +98,12 @@ class TestStickyCookie(mastertest.MasterTest):
         assert len(sc.jar[googlekey].keys()) == len(cs)
 
     def test_response_overwrite(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         # Test overwriting of a cookie value
         c1 = "somecookie=helloworld; Path=/"
         c2 = "somecookie=newvalue; Path=/"
-        f = self._response(s, m, sc, c1, "www.google.com")
+        f = self._response(m, sc, c1, "www.google.com")
         f.response.headers["Set-Cookie"] = c2
         m.response(f)
         googlekey = list(sc.jar.keys())[0]
@@ -113,19 +111,19 @@ class TestStickyCookie(mastertest.MasterTest):
         assert list(sc.jar[googlekey]["somecookie"].items())[0][1] == "newvalue"
 
     def test_response_delete(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
         # Test that a cookie is be deleted
         # by setting the expire time in the past
-        f = self._response(s, m, sc, "duffer=zafar; Path=/", "www.google.com")
+        f = self._response(m, sc, "duffer=zafar; Path=/", "www.google.com")
         f.response.headers["Set-Cookie"] = "duffer=; Expires=Thu, 01-Jan-1970 00:00:00 GMT"
         m.response(f)
         assert not sc.jar.keys()
 
     def test_request(self):
-        s, m, sc = self.mk()
+        m, sc = self.mk()
 
-        f = self._response(s, m, sc, "SSID=mooo", "www.google.com")
+        f = self._response(m, sc, "SSID=mooo", "www.google.com")
         assert "cookie" not in f.request.headers
         m.request(f)
         assert "cookie" in f.request.headers
diff --git a/test/mitmproxy/builtins/test_streambodies.py b/test/mitmproxy/builtins/test_streambodies.py
index 0e8a82f2..0a306681 100644
--- a/test/mitmproxy/builtins/test_streambodies.py
+++ b/test/mitmproxy/builtins/test_streambodies.py
@@ -1,5 +1,4 @@
 from .. import tutils, mastertest
-from mitmproxy.flow import state
 from mitmproxy.flow import master
 from mitmproxy import options
 
@@ -8,9 +7,8 @@ from mitmproxy.builtins import streambodies
 
 class TestStreamBodies(mastertest.MasterTest):
     def test_simple(self):
-        s = state.DummyState()
         o = options.Options(stream_large_bodies = 10)
-        m = master.FlowMaster(o, None, s)
+        m = master.FlowMaster(o, None)
         sa = streambodies.StreamBodies()
         m.addons.add(sa)
 
diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py
index 5605b1a6..ae11e577 100644
--- a/test/mitmproxy/mastertest.py
+++ b/test/mitmproxy/mastertest.py
@@ -50,8 +50,7 @@ class RecordingMaster(master.FlowMaster):
 
 @contextlib.contextmanager
 def mockctx():
-    state = flow.State()
     o = options.Options(refresh_server_playback = True, keepserving=False)
-    m = RecordingMaster(o, proxy.DummyServer(o), state)
+    m = RecordingMaster(o, proxy.DummyServer(o))
     with m.handlecontext():
         yield
diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py
index 5f943aeb..eb0e83f7 100644
--- a/test/mitmproxy/script/test_concurrent.py
+++ b/test/mitmproxy/script/test_concurrent.py
@@ -3,7 +3,6 @@ from mitmproxy import controller
 from mitmproxy.builtins import script
 from mitmproxy import options
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 import time
 
 
@@ -16,8 +15,7 @@ class Thing:
 class TestConcurrent(mastertest.MasterTest):
     @tutils.skip_appveyor
     def test_concurrent(self):
-        s = state.State()
-        m = master.FlowMaster(options.Options(), None, s)
+        m = master.FlowMaster(options.Options(), None)
         sc = script.Script(
             tutils.test_data.path(
                 "data/addonscripts/concurrent_decorator.py"
@@ -34,8 +32,7 @@ class TestConcurrent(mastertest.MasterTest):
         raise ValueError("Script never acked")
 
     def test_concurrent_err(self):
-        s = state.State()
-        m = mastertest.RecordingMaster(options.Options(), None, s)
+        m = mastertest.RecordingMaster(options.Options(), None)
         sc = script.Script(
             tutils.test_data.path(
                 "data/addonscripts/concurrent_decorator_err.py"
diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py
index 05b7544f..f824584b 100644
--- a/test/mitmproxy/test_examples.py
+++ b/test/mitmproxy/test_examples.py
@@ -6,7 +6,6 @@ from mitmproxy import options
 from mitmproxy import contentviews
 from mitmproxy.builtins import script
 from mitmproxy.flow import master
-from mitmproxy.flow import state
 
 import netlib.utils
 
@@ -32,7 +31,7 @@ class RaiseMaster(master.FlowMaster):
 def tscript(cmd, args=""):
     o = options.Options()
     cmd = example_dir.path(cmd) + " " + args
-    m = RaiseMaster(o, None, state.State())
+    m = RaiseMaster(o, None)
     sc = script.Script(cmd)
     m.addons.add(sc)
     return m, sc
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 02996881..78cfca99 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -130,7 +130,8 @@ class TestHTTPFlow:
     def test_killall(self):
         srv = DummyServer(None)
         s = flow.State()
-        fm = flow.FlowMaster(None, srv, s)
+        fm = flow.FlowMaster(None, srv)
+        fm.addons.add(s)
 
         f = tutils.tflow()
         f.reply.handle()
@@ -382,7 +383,8 @@ class TestSerialize:
     def test_load_flows(self):
         r = self._treader()
         s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
+        fm.addons.add(s)
         fm.load_flows(r)
         assert len(s.flows) == 6
 
@@ -394,7 +396,8 @@ class TestSerialize:
             upstream_server="https://use-this-domain"
         )
         conf = ProxyConfig(opts)
-        fm = flow.FlowMaster(opts, DummyServer(conf), s)
+        fm = flow.FlowMaster(opts, DummyServer(conf))
+        fm.addons.add(s)
         fm.load_flows(r)
         assert s.flows[0].request.host == "use-this-domain"
 
@@ -440,8 +443,7 @@ class TestSerialize:
 class TestFlowMaster:
 
     def test_replay(self):
-        s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
         f = tutils.tflow(resp=True)
         f.request.content = None
         tutils.raises("missing", fm.replay_request, f)
@@ -453,13 +455,13 @@ class TestFlowMaster:
         tutils.raises("live", fm.replay_request, f)
 
     def test_create_flow(self):
-        s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
         assert fm.create_request("GET", "http", "example.com", 80, "/")
 
     def test_all(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
+        fm.addons.add(s)
         f = tutils.tflow(req=None)
         fm.clientconnect(f.client_conn)
         f.request = HTTPRequest.wrap(netlib.tutils.treq())
diff --git a/test/mitmproxy/test_flow_state.py b/test/mitmproxy/test_flow_state.py
index cc05537d..a658e13e 100644
--- a/test/mitmproxy/test_flow_state.py
+++ b/test/mitmproxy/test_flow_state.py
@@ -5,11 +5,13 @@ from . import tutils
 class TestState:
     def test_duplicate_flow(self):
         s = flow.State()
-        fm = flow.FlowMaster(None, None, s)
+        fm = flow.FlowMaster(None, None)
+        fm.addons.add(s)
         f = tutils.tflow(resp=True)
         fm.load_flow(f)
         assert s.flow_count() == 1
-        f2 = fm.state.duplicate_flow(f)
+
+        f2 = s.duplicate_flow(f)
         assert f2.response
         assert s.flow_count() == 2
         assert s.index(f2) == 1
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index e9db210e..0938571b 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -1,13 +1,12 @@
 import os
 import socket
 import time
-import types
 
 import netlib.tutils
 from mitmproxy import controller
 from mitmproxy import options
 from mitmproxy.builtins import script
-from mitmproxy.models import Error, HTTPResponse, HTTPFlow
+from mitmproxy.models import HTTPResponse, HTTPFlow
 from mitmproxy.proxy.config import HostMatcher, parse_server_spec
 from netlib import tcp, http, socks
 from netlib.certutils import SSLCert
@@ -986,6 +985,17 @@ class TestUpstreamProxySSL(
         assert self.chain[1].tmaster.state.flow_count() == 1
 
 
+class RequestKiller:
+    def __init__(self, exclude):
+        self.exclude = exclude
+        self.k = 0
+
+    def request(self, f):
+        self.k += 1
+        if self.k not in self.exclude:
+            f.reply.kill()
+
+
 class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
     ssl = True
 
@@ -995,39 +1005,18 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
         If we have a disconnect on a secure connection that's transparently proxified to
         an upstream http proxy, we need to send the CONNECT request again.
         """
-
-        def kill_requests(master, attr, exclude):
-            k = [0]  # variable scope workaround: put into array
-            _func = getattr(master, attr)
-
-            @controller.handler
-            def handler(*args):
-                f = args[-1]
-                k[0] += 1
-                if not (k[0] in exclude):
-                    f.client_conn.finish()
-                    f.error = Error("terminated")
-                    f.reply.kill()
-                return _func(f)
-
-            setattr(master, attr, types.MethodType(handler, master))
-
-        kill_requests(
-            self.chain[1].tmaster,
-            "request",
-            exclude = [
-                # fail first request
-                2,  # allow second request
-            ]
+        self.chain[1].tmaster.addons.add(
+            RequestKiller([2])
+        )
+        self.chain[0].tmaster.addons.add(
+            RequestKiller(
+                [
+                    1,  # CONNECT
+                    3,  # reCONNECT
+                    4   # request
+                ]
+            )
         )
-
-        kill_requests(self.chain[0].tmaster, "request",
-                      exclude=[
-                          1,  # CONNECT
-                          # fail first request
-                          3,  # reCONNECT
-                          4,  # request
-        ])
 
         p = self.pathoc()
         with p.connect():
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 431e0f90..c4230d6f 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -7,6 +7,7 @@ import sys
 
 from mitmproxy.proxy.config import ProxyConfig
 from mitmproxy.proxy.server import ProxyServer
+from mitmproxy.flow import state
 import pathod.test
 import pathod.pathoc
 from mitmproxy import flow, controller, options
@@ -34,8 +35,9 @@ class TestMaster(flow.FlowMaster):
 
     def __init__(self, opts, config):
         s = ProxyServer(config)
-        state = flow.State()
-        flow.FlowMaster.__init__(self, opts, s, state)
+        flow.FlowMaster.__init__(self, opts, s)
+        self.state = state.State()
+        self.addons.add(self.state)
         self.addons.add(*builtins.default_addons())
         self.apps.add(testapp, "testapp", 80)
         self.apps.add(errapp, "errapp", 80)
@@ -116,8 +118,8 @@ class ProxyTestBase:
                 raise
 
     def setup(self):
-        self.master.clear_log()
         self.master.state.clear()
+        self.master.clear_log()
         self.server.clear_log()
         self.server2.clear_log()