Convert flow.state to an addon

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

7 files changed, 20 insertions, 22 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