From a9f169ca86f7c53158c4c6595b44259e622fb24e Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Tue, 1 Aug 2017 23:38:54 +0200
Subject: [PATCH] gracefully handle errors during script load

---
 mitmproxy/addons/script.py                    |  5 ++--
 test/mitmproxy/addons/test_script.py          | 27 ++++++++++++++-----
 .../mitmproxy/data/addonscripts/load_error.py |  2 ++
 3 files changed, 26 insertions(+), 8 deletions(-)
 create mode 100644 test/mitmproxy/data/addonscripts/load_error.py

diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index 58e8cdcd9..2d0303215 100644
--- a/mitmproxy/addons/script.py
+++ b/mitmproxy/addons/script.py
@@ -72,11 +72,12 @@ class Script:
                     ctx.master.addons.remove(self.ns)
                 self.ns = None
                 with addonmanager.safecall():
-                    self.ns = load_script(self.fullpath)
+                    ns = load_script(self.fullpath)
+                    ctx.master.addons.register(ns)
+                    self.ns = ns
                 if self.ns:
                     # We're already running, so we have to explicitly register and
                     # configure the addon
-                    ctx.master.addons.register(self.ns)
                     ctx.master.addons.invoke_addon(self.ns, "running")
                     ctx.master.addons.invoke_addon(
                         self.ns,
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 64fd95058..aa7ca68e0 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -1,15 +1,16 @@
-import traceback
-import sys
 import os
+import sys
+import traceback
+from unittest import mock
+
 import pytest
 
-from unittest import mock
-from mitmproxy.test import tflow
-from mitmproxy.test import tutils
-from mitmproxy.test import taddons
 from mitmproxy import addonmanager
 from mitmproxy import exceptions
 from mitmproxy.addons import script
+from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+from mitmproxy.test import tutils
 
 
 def test_load_script():
@@ -216,6 +217,20 @@ class TestScriptLoader:
             assert not tctx.options.scripts
             assert not sl.addons
 
+    def test_load_err(self):
+        sc = script.ScriptLoader()
+        with taddons.context() as tctx:
+            tctx.configure(sc, scripts=[
+                tutils.test_data.path("mitmproxy/data/addonscripts/load_error.py")
+            ])
+            try:
+                tctx.invoke(sc, "tick")
+            except ValueError:
+                pass  # this is expected and normally guarded.
+            # on the next tick we should not fail however.
+            tctx.invoke(sc, "tick")
+            assert len(tctx.master.addons) == 0
+
     def test_order(self):
         rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder")
         sc = script.ScriptLoader()
diff --git a/test/mitmproxy/data/addonscripts/load_error.py b/test/mitmproxy/data/addonscripts/load_error.py
new file mode 100644
index 000000000..4c05e9ed9
--- /dev/null
+++ b/test/mitmproxy/data/addonscripts/load_error.py
@@ -0,0 +1,2 @@
+def load(_):
+    raise ValueError()