Merge pull request #2492 from mhils/addon-contains

Gracefully handle errors during script load

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,

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()

test/mitmproxy/data/addonscripts/load_error.py

@@ -0,0 +1,2 @@
+def load(_):
+    raise ValueError()