From da8dec982373b5137c5700818505887fe0249e98 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Wed, 19 Oct 2016 12:01:08 +1300
Subject: [PATCH] addons: add AddonHalt exception

This can be raised from an addon handler to stop further processing of a flow.

Use this to prevent further handling of web app requests.
---
 mitmproxy/addons.py           | 5 ++++-
 mitmproxy/builtins/wsgiapp.py | 2 ++
 mitmproxy/console/master.py   | 2 +-
 mitmproxy/exceptions.py       | 4 ++++
 mitmproxy/web/master.py       | 2 +-
 5 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py
index ea6411121..34ca8e554 100644
--- a/mitmproxy/addons.py
+++ b/mitmproxy/addons.py
@@ -86,4 +86,7 @@ class Addons:
 
     def __call__(self, name, *args, **kwargs):
         for i in self.chain:
-            self.invoke(i, name, *args, **kwargs)
+            try:
+                self.invoke(i, name, *args, **kwargs)
+            except exceptions.AddonHalt:
+                return
diff --git a/mitmproxy/builtins/wsgiapp.py b/mitmproxy/builtins/wsgiapp.py
index 07ec8a6f9..9444fcec5 100644
--- a/mitmproxy/builtins/wsgiapp.py
+++ b/mitmproxy/builtins/wsgiapp.py
@@ -1,4 +1,5 @@
 from mitmproxy import ctx
+from mitmproxy import exceptions
 
 from netlib import wsgi
 from netlib import version
@@ -30,6 +31,7 @@ class WSGIApp:
         if err:
             ctx.log.warn("Error in wsgi app. %s" % err, "error")
         flow.reply.kill()
+        raise exceptions.AddonHalt()
 
     def request(self, f):
         if (f.request.pretty_host, f.request.port) == (self.host, self.port):
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index 05602e1a3..bee5954b0 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -224,7 +224,6 @@ class ConsoleMaster(flow.FlowMaster):
     def __init__(self, server, options):
         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
@@ -252,6 +251,7 @@ class ConsoleMaster(flow.FlowMaster):
         signals.push_view_state.connect(self.sig_push_view_state)
         signals.sig_add_log.connect(self.sig_add_log)
         self.addons.add(*builtins.default_addons())
+        self.addons.add(self.state)
 
     def __setattr__(self, name, value):
         self.__dict__[name] = value
diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py
index b696c91f6..64cc457a4 100644
--- a/mitmproxy/exceptions.py
+++ b/mitmproxy/exceptions.py
@@ -96,3 +96,7 @@ class OptionsError(Exception):
 
 class AddonError(Exception):
     pass
+
+
+class AddonHalt(Exception):
+    pass
diff --git a/mitmproxy/web/master.py b/mitmproxy/web/master.py
index 3d2182a8d..a05d97806 100644
--- a/mitmproxy/web/master.py
+++ b/mitmproxy/web/master.py
@@ -135,8 +135,8 @@ class WebMaster(flow.FlowMaster):
     def __init__(self, server, options):
         super().__init__(options, server)
         self.state = WebState()
-        self.addons.add(self.state)
         self.addons.add(*builtins.default_addons())
+        self.addons.add(self.state)
         self.app = app.Application(
             self, self.options.wdebug, self.options.wauthenticator
         )