Merge pull request #145 from python-discord/fix-eval-exploits

1 files changed, 15 insertions, 8 deletions

diff --git a/resources/unittest_template.py b/resources/unittest_template.py
index 2410278..05730ce 100644
--- a/resources/unittest_template.py
+++ b/resources/unittest_template.py
@@ -19,6 +19,9 @@ class RunnerTestCase(unittest.TestCase):
 ### UNIT CODE
 
 
+normal_exit = False
+
+
 def _exit_sandbox(code: int) -> NoReturn:
     """
     Exit the sandbox by printing the result to the actual stdout and exit with the provided code.
@@ -32,6 +35,8 @@ def _exit_sandbox(code: int) -> NoReturn:
     137 can also be generated by NsJail when killing the process.
     """
     print(RESULT.getvalue(), file=ORIGINAL_STDOUT, end="")
+    global normal_exit
+    normal_exit = True
     sys.exit(code)
 
 
@@ -72,19 +77,21 @@ try:
     DEVNULL = SimpleNamespace(write=lambda *_: None, flush=lambda *_: None)
 
     RESULT = io.StringIO()
-    ORIGINAL_STDOUT = sys.stdout
+    ORIGINAL_STDOUT = sys.__stdout__
 
     # stdout/err is patched in order to control what is outputted by the runner
-    sys.stdout = DEVNULL
-    sys.stderr = DEVNULL
-    
+    sys.__stdout__ = sys.stdout = DEVNULL
+    sys.__stderr__ = sys.stderr = DEVNULL
+
     # Load the user code as a global module variable
     try:
         module = _load_user_module()
-    except Exception:
-        RESULT.write("Uncaught exception while loading user code.")
+    except BaseException as e:
+        RESULT.write(f"Uncaught exception while loading user code: {e}")
         _exit_sandbox(6)
     _main()
-except Exception:
-    RESULT.write("Uncaught exception inside runner.")
+except BaseException as e:
+    if isinstance(e, SystemExit) and normal_exit:
+        raise e from None
+    RESULT.write(f"Uncaught exception inside runner: {e}")
     _exit_sandbox(99)