If a message has been acked, all other processors are skipped

This applies the constraint, but does to clumsily. When we've unified modules
and processors it will be much nicer. We also make some exceptions for the
master processors that we may want to re-evaluate down the track.

3 files changed, 40 insertions, 22 deletions

diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index 31475f5b..289102a1 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -103,9 +103,10 @@ class FlowMaster(controller.Master):
             except script.ScriptException as e:
                 self.add_event("Script error:\n{}".format(e), "error")
 
-    def run_script_hook(self, name, *args, **kwargs):
+    def run_scripts(self, name, msg):
         for script_obj in self.scripts:
-            self._run_single_script_hook(script_obj, name, *args, **kwargs)
+            if not msg.reply.acked:
+                self._run_single_script_hook(script_obj, name, msg)
 
     def get_ignore_filter(self):
         return self.server.config.check_ignore.patterns
@@ -373,28 +374,28 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def clientconnect(self, root_layer):
-        self.run_script_hook("clientconnect", root_layer)
+        self.run_scripts("clientconnect", root_layer)
 
     @controller.handler
     def clientdisconnect(self, root_layer):
-        self.run_script_hook("clientdisconnect", root_layer)
+        self.run_scripts("clientdisconnect", root_layer)
 
     @controller.handler
     def serverconnect(self, server_conn):
-        self.run_script_hook("serverconnect", server_conn)
+        self.run_scripts("serverconnect", server_conn)
 
     @controller.handler
     def serverdisconnect(self, server_conn):
-        self.run_script_hook("serverdisconnect", server_conn)
+        self.run_scripts("serverdisconnect", server_conn)
 
     @controller.handler
     def next_layer(self, top_layer):
-        self.run_script_hook("next_layer", top_layer)
+        self.run_scripts("next_layer", top_layer)
 
     @controller.handler
     def error(self, f):
         self.state.update_flow(f)
-        self.run_script_hook("error", f)
+        self.run_scripts("error", f)
         if self.client_playback:
             self.client_playback.clear(f)
         return f
@@ -416,10 +417,14 @@ class FlowMaster(controller.Master):
         if f not in self.state.flows:  # don't add again on replay
             self.state.add_flow(f)
         self.active_flows.add(f)
-        self.replacehooks.run(f)
-        self.setheaders.run(f)
-        self.process_new_request(f)
-        self.run_script_hook("request", f)
+        if not f.reply.acked:
+            self.replacehooks.run(f)
+        if not f.reply.acked:
+            self.setheaders.run(f)
+        if not f.reply.acked:
+            self.process_new_request(f)
+        if not f.reply.acked:
+            self.run_scripts("request", f)
         return f
 
     @controller.handler
@@ -430,18 +435,21 @@ class FlowMaster(controller.Master):
         except netlib.exceptions.HttpException:
             f.reply.kill()
             return
-        self.run_script_hook("responseheaders", f)
+        self.run_scripts("responseheaders", f)
         return f
 
     @controller.handler
     def response(self, f):
         self.active_flows.discard(f)
         self.state.update_flow(f)
-        self.replacehooks.run(f)
-        self.setheaders.run(f)
-        self.run_script_hook("response", f)
-        if self.client_playback:
-            self.client_playback.clear(f)
+        if not f.reply.acked:
+            self.replacehooks.run(f)
+        if not f.reply.acked:
+            self.setheaders.run(f)
+        self.run_scripts("response", f)
+        if not f.reply.acked:
+            if self.client_playback:
+                self.client_playback.clear(f)
         self.process_new_response(f)
         if self.stream:
             self.stream.add(f)
@@ -487,11 +495,11 @@ class FlowMaster(controller.Master):
         # TODO: This would break mitmproxy currently.
         # self.state.add_flow(flow)
         self.active_flows.add(flow)
-        self.run_script_hook("tcp_open", flow)
+        self.run_scripts("tcp_open", flow)
 
     @controller.handler
     def tcp_message(self, flow):
-        self.run_script_hook("tcp_message", flow)
+        self.run_scripts("tcp_message", flow)
         message = flow.messages[-1]
         direction = "->" if message.from_client else "<-"
         self.add_event("{client} {direction} tcp {direction} {server}".format(
@@ -507,14 +515,14 @@ class FlowMaster(controller.Master):
             repr(flow.server_conn.address),
             flow.error
         ), "info")
-        self.run_script_hook("tcp_error", flow)
+        self.run_scripts("tcp_error", flow)
 
     @controller.handler
     def tcp_close(self, flow):
         self.active_flows.discard(flow)
         if self.stream:
             self.stream.add(flow)
-        self.run_script_hook("tcp_close", flow)
+        self.run_scripts("tcp_close", flow)
 
     def shutdown(self):
         super(FlowMaster, self).shutdown()
diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py
index 9bb8826d..4d04f337 100644
--- a/test/mitmproxy/mastertest.py
+++ b/test/mitmproxy/mastertest.py
@@ -16,7 +16,9 @@ class MasterTest:
         master.request(f)
         if not f.error:
             f.response = models.HTTPResponse.wrap(netlib.tutils.tresp(content=content))
+            f.reply.acked = False
             f = master.response(f)
+        f.client_conn.reply.acked = False
         master.clientdisconnect(f.client_conn)
         return f
 
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 1b1f03f9..af8256c4 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -807,17 +807,22 @@ class TestFlowMaster:
         fm.load_script(tutils.test_data.path("data/scripts/all.py"))
         f = tutils.tflow(resp=True)
 
+        f.client_conn.acked = False
         fm.clientconnect(f.client_conn)
         assert fm.scripts[0].ns["log"][-1] == "clientconnect"
+        f.server_conn.acked = False
         fm.serverconnect(f.server_conn)
         assert fm.scripts[0].ns["log"][-1] == "serverconnect"
+        f.reply.acked = False
         fm.request(f)
         assert fm.scripts[0].ns["log"][-1] == "request"
+        f.reply.acked = False
         fm.response(f)
         assert fm.scripts[0].ns["log"][-1] == "response"
         # load second script
         fm.load_script(tutils.test_data.path("data/scripts/all.py"))
         assert len(fm.scripts) == 2
+        f.server_conn.reply.acked = False
         fm.clientdisconnect(f.server_conn)
         assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
         assert fm.scripts[1].ns["log"][-1] == "clientdisconnect"
@@ -828,6 +833,7 @@ class TestFlowMaster:
         fm.load_script(tutils.test_data.path("data/scripts/all.py"))
 
         f.error = tutils.terr()
+        f.reply.acked = False
         fm.error(f)
         assert fm.scripts[0].ns["log"][-1] == "error"
 
@@ -977,10 +983,12 @@ class TestFlowMaster:
         f = tutils.tflow(resp=True)
         f.response.headers["set-cookie"] = "foo=bar"
         fm.request(f)
+        f.reply.acked = False
         fm.response(f)
         assert fm.stickycookie_state.jar
         assert "cookie" not in f.request.headers
         f = f.copy()
+        f.reply.acked = False
         fm.request(f)
         assert f.request.headers["cookie"] == "foo=bar"