Make @controller.handler inheritance-friendly

Use this to adapt mitmweb and mitproxy console

6 files changed, 36 insertions, 18 deletions

diff --git a/mitmproxy/console/__init__.py b/mitmproxy/console/__init__.py
index ce202b39..a87e691e 100644
--- a/mitmproxy/console/__init__.py
+++ b/mitmproxy/console/__init__.py
@@ -713,14 +713,15 @@ class ConsoleMaster(flow.FlowMaster):
             )
 
     def process_flow(self, f):
-        if self.state.intercept and f.match(self.state.intercept) and not f.request.is_replay:
+        should_intercept = any(
+            [
+                self.state.intercept and f.match(self.state.intercept) and not f.request.is_replay,
+                f.intercepted,
+            ]
+        )
+        if should_intercept:
             f.intercept(self)
-        else:
-            # check if flow was intercepted within an inline script by flow.intercept()
-            if f.intercepted:
-                f.intercept(self)
-            else:
-                f.reply()
+            f.reply.take()
         signals.flowlist_change.send(self)
         signals.flow_change.send(self, flow = f)
 
@@ -728,24 +729,29 @@ class ConsoleMaster(flow.FlowMaster):
         self.eventlist[:] = []
 
     # Handlers
+    @controller.handler
     def handle_error(self, f):
         f = flow.FlowMaster.handle_error(self, f)
         if f:
             self.process_flow(f)
         return f
 
+    @controller.handler
     def handle_request(self, f):
         f = flow.FlowMaster.handle_request(self, f)
+        self.add_event(f.reply.acked, "info")
         if f:
             self.process_flow(f)
         return f
 
+    @controller.handler
     def handle_response(self, f):
         f = flow.FlowMaster.handle_response(self, f)
         if f:
             self.process_flow(f)
         return f
 
+    @controller.handler
     def handle_script_change(self, script):
         if super(ConsoleMaster, self).handle_script_change(script):
             signals.status_message.send(message='"{}" reloaded.'.format(script.filename))
diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py
index c125ee4f..57c01f59 100644
--- a/mitmproxy/controller.py
+++ b/mitmproxy/controller.py
@@ -125,7 +125,6 @@ class Channel(object):
             if g == exceptions.Kill:
                 raise exceptions.Kill()
             return g
-
         raise exceptions.Kill()
 
     def tell(self, mtype, m):
@@ -146,6 +145,7 @@ class DummyReply(object):
     def __init__(self):
         self.acked = False
         self.taken = False
+        self.handled = False
 
     def take(self):
         self.taken = True
@@ -164,8 +164,16 @@ def handler(f):
         message = args[-1]
         if not hasattr(message, "reply"):
             raise ControlError("Message %s has no reply attribute"%message)
+
+        handling = False
+        # We're the first handler - ack responsibility is ours
+        if not message.reply.handled:
+            handling = True
+            message.reply.handled = True
+
         ret = f(*args, **kwargs)
-        if not message.reply.acked and not message.reply.taken:
+
+        if handling and not message.reply.acked and not message.reply.taken:
             message.reply()
         return ret
     wrapper.func_dict["handler"] = True
@@ -181,8 +189,12 @@ class Reply(object):
     def __init__(self, obj):
         self.obj = obj
         self.q = queue.Queue()
+        # Has this message been acked?
         self.acked = False
+        # Has the user taken responsibility for ack-ing?
         self.taken = False
+        # Has a handler taken responsibility for ack-ing?
+        self.handled = False
 
     def take(self):
         self.taken = True
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 2a324993..ba8dec5d 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -546,7 +546,8 @@ class FlowStore(FlowList):
 
     def kill_all(self, master):
         for f in self._list:
-            f.kill(master)
+            if not f.reply.acked:
+                f.kill(master)
 
 
 class State(object):
diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py
index 96e7aab6..9b4e2963 100644
--- a/mitmproxy/proxy/root_context.py
+++ b/mitmproxy/proxy/root_context.py
@@ -132,7 +132,6 @@ class RootContext(object):
 
 
 class Log(object):
-
     def __init__(self, msg, level="info"):
         self.msg = msg
         self.level = level
diff --git a/mitmproxy/web/__init__.py b/mitmproxy/web/__init__.py
index 956d221d..d90830d6 100644
--- a/mitmproxy/web/__init__.py
+++ b/mitmproxy/web/__init__.py
@@ -6,7 +6,7 @@ import sys
 
 from netlib.http import authentication
 
-from .. import flow
+from .. import flow, controller
 from ..exceptions import FlowReadException
 from . import app
 
@@ -194,21 +194,24 @@ class WebMaster(flow.FlowMaster):
         if self.state.intercept and self.state.intercept(
                 f) and not f.request.is_replay:
             f.intercept(self)
-        else:
-            f.reply()
+            f.reply.take()
 
+    @controller.handler
     def handle_request(self, f):
         super(WebMaster, self).handle_request(f)
         self._process_flow(f)
 
+    @controller.handler
     def handle_response(self, f):
         super(WebMaster, self).handle_response(f)
         self._process_flow(f)
 
+    @controller.handler
     def handle_error(self, f):
         super(WebMaster, self).handle_error(f)
         self._process_flow(f)
 
+    @controller.handler
     def add_event(self, e, level="info"):
         super(WebMaster, self).add_event(e, level)
         self.state.add_event(e, level)
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 62f23ac8..da8b8ddd 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -470,10 +470,7 @@ class TestFlow(object):
         fm = flow.FlowMaster(None, s)
 
         f = tutils.tflow()
-        fm.handle_request(f)
-
-        f = tutils.tflow()
-        fm.handle_request(f)
+        f.intercept(fm)
 
         s.killall(fm)
         for i in s.view: