Merge pull request #3050 from kira0204/error-script

Handling user script's exceptions, fix #2839

3 files changed, 29 insertions, 10 deletions

diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py
index f5c4d599..51341f2f 100644
--- a/mitmproxy/addons/script.py
+++ b/mitmproxy/addons/script.py
@@ -25,6 +25,7 @@ def load_script(path: str) -> types.ModuleType:
     sys.modules.pop(fullname, None)
     oldpath = sys.path
     sys.path.insert(0, os.path.dirname(path))
+    m = None
     try:
         loader = importlib.machinery.SourceFileLoader(fullname, path)
         spec = importlib.util.spec_from_loader(fullname, loader=loader)
@@ -32,9 +33,11 @@ def load_script(path: str) -> types.ModuleType:
         loader.exec_module(m)
         if not getattr(m, "name", None):
             m.name = path  # type: ignore
-        return m
+    except Exception as e:
+        script_error_handler(path, e, msg=str(e))
     finally:
         sys.path[:] = oldpath
+        return m
 
 
 def script_error_handler(path, exc, msg="", tb=False):
@@ -48,11 +51,11 @@ def script_error_handler(path, exc, msg="", tb=False):
     lineno = ""
     if hasattr(exc, "lineno"):
         lineno = str(exc.lineno)
-    log_msg = "in Script {}:{} {}".format(path, lineno, exception)
+    log_msg = "in script {}:{} {}".format(path, lineno, exception)
     if tb:
         etype, value, tback = sys.exc_info()
         tback = addonmanager.cut_traceback(tback, "invoke_addon")
-        log_msg = log_msg.join(["\n"] + traceback.format_exception(etype, value, tback))
+        log_msg = log_msg + "\n" + "".join(traceback.format_exception(etype, value, tback))
     ctx.log.error(log_msg)
 
 
diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py
index 3f0ce68c..96e19841 100644
--- a/test/mitmproxy/addons/test_script.py
+++ b/test/mitmproxy/addons/test_script.py
@@ -12,18 +12,27 @@ from mitmproxy.test import tflow
 from mitmproxy.test import tutils
 
 
-def test_load_script():
-    ns = script.load_script(
-        tutils.test_data.path(
-            "mitmproxy/data/addonscripts/recorder/recorder.py"
+@pytest.mark.asyncio
+async def test_load_script():
+    with taddons.context() as tctx:
+        ns = script.load_script(
+            tutils.test_data.path(
+                "mitmproxy/data/addonscripts/recorder/recorder.py"
+            )
         )
-    )
-    assert ns.addons
+        assert ns.addons
 
-    with pytest.raises(FileNotFoundError):
         script.load_script(
             "nonexistent"
         )
+        assert await tctx.master.await_log("No such file or directory")
+
+        script.load_script(
+            tutils.test_data.path(
+                "mitmproxy/data/addonscripts/recorder/error.py"
+            )
+        )
+        assert await tctx.master.await_log("invalid syntax")
 
 
 def test_load_fullname():
diff --git a/test/mitmproxy/data/addonscripts/recorder/error.py b/test/mitmproxy/data/addonscripts/recorder/error.py
new file mode 100644
index 00000000..2e7e648a
--- /dev/null
+++ b/test/mitmproxy/data/addonscripts/recorder/error.py
@@ -0,0 +1,7 @@
+"""
+This file is intended to have syntax errors for test purposes
+"""
+
+impotr recorder # Intended Syntax Error
+
+addons = [recorder.Recorder("e")]