From 403218ec6b5e5eb46b0a134f7641b6385df3bdb5 Mon Sep 17 00:00:00 2001
From: Ned Batchelder <ned@nedbatchelder.com>
Date: Thu, 2 May 2013 12:46:35 -0400
Subject: [PATCH] If sandboxed code raises an exception, the exception will be
 cached.

---
 common/lib/capa/capa/safe_exec/safe_exec.py   | 28 +++++++++----
 .../capa/safe_exec/tests/test_safe_exec.py    | 41 ++++++++++++++++++-
 2 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/common/lib/capa/capa/safe_exec/safe_exec.py b/common/lib/capa/capa/safe_exec/safe_exec.py
index 745704373b..34cca0593e 100644
--- a/common/lib/capa/capa/safe_exec/safe_exec.py
+++ b/common/lib/capa/capa/safe_exec/safe_exec.py
@@ -1,7 +1,7 @@
 """Capa's specialized use of codejail.safe_exec."""
 
 from codejail.safe_exec import safe_exec as codejail_safe_exec
-from codejail.safe_exec import json_safe
+from codejail.safe_exec import json_safe, SafeExecException
 from . import lazymod
 from statsd import statsd
 
@@ -63,20 +63,34 @@ def safe_exec(code, globals_dict, random_seed=None, python_path=None, cache=None
         key = "safe_exec.%r.%s" % (random_seed, md5er.hexdigest())
         cached = cache.get(key)
         if cached is not None:
-            globals_dict.update(cached)
+            # We have a cached result.  The result is a pair: the exception
+            # message, if any, else None; and the resulting globals dictionary.
+            emsg, cleaned_results = cached
+            globals_dict.update(cleaned_results)
+            if emsg:
+                raise SafeExecException(emsg)
             return
 
     # Create the complete code we'll run.
     code_prolog = CODE_PROLOG % random_seed
 
     # Run the code!  Results are side effects in globals_dict.
-    codejail_safe_exec(
-        code_prolog + LAZY_IMPORTS + code, globals_dict,
-        python_path=python_path,
-    )
+    try:
+        codejail_safe_exec(
+            code_prolog + LAZY_IMPORTS + code, globals_dict,
+            python_path=python_path,
+        )
+    except SafeExecException as e:
+        emsg = e.message
+    else:
+        emsg = None
 
     # Put the result back in the cache.  This is complicated by the fact that
     # the globals dict might not be entirely serializable.
     if cache:
         cleaned_results = json_safe(globals_dict)
-        cache.set(key, cleaned_results)
+        cache.set(key, (emsg, cleaned_results))
+
+    # If an exception happened, raise it now.
+    if emsg:
+        raise e
diff --git a/common/lib/capa/capa/safe_exec/tests/test_safe_exec.py b/common/lib/capa/capa/safe_exec/tests/test_safe_exec.py
index bd960f331c..b648672daf 100644
--- a/common/lib/capa/capa/safe_exec/tests/test_safe_exec.py
+++ b/common/lib/capa/capa/safe_exec/tests/test_safe_exec.py
@@ -5,6 +5,8 @@ import random
 import unittest
 
 from capa.safe_exec import safe_exec
+from codejail.safe_exec import SafeExecException
+
 
 class TestSafeExec(unittest.TestCase):
     def test_set_values(self):
@@ -57,6 +59,12 @@ class TestSafeExec(unittest.TestCase):
             g, python_path=[pylib]
         )
 
+    def test_raising_exceptions(self):
+        g = {}
+        with self.assertRaises(SafeExecException) as cm:
+            safe_exec("1/0", g)
+        self.assertIn("ZeroDivisionError", cm.exception.message)
+
 
 class DictCache(object):
     """A cache implementation over a simple dict, for testing."""
@@ -86,10 +94,10 @@ class TestSafeExecCaching(unittest.TestCase):
         safe_exec("a = int(math.pi)", g, cache=DictCache(cache))
         self.assertEqual(g['a'], 3)
         # A result has been cached
-        self.assertEqual(cache.values(), [{'a': 3}])
+        self.assertEqual(cache.values()[0], (None, {'a': 3}))
 
         # Fiddle with the cache, then try it again.
-        cache[cache.keys()[0]] = {'a': 17}
+        cache[cache.keys()[0]] = (None, {'a': 17})
 
         g = {}
         safe_exec("a = int(math.pi)", g, cache=DictCache(cache))
@@ -104,3 +112,32 @@ class TestSafeExecCaching(unittest.TestCase):
         cache = {}
         safe_exec(code, g, cache=DictCache(cache))
         self.assertEqual(g['a'], 12345)
+
+    def test_cache_exceptions(self):
+        # Used to be that running code that raised an exception didn't cache
+        # the result.  Check that now it does.
+        code = "1/0"
+        g = {}
+        cache = {}
+        with self.assertRaises(SafeExecException):
+            safe_exec(code, g, cache=DictCache(cache))
+
+        # The exception should be in the cache now.
+        self.assertEqual(len(cache), 1)
+        cache_exc_msg, cache_globals = cache.values()[0]
+        self.assertIn("ZeroDivisionError", cache_exc_msg)
+
+        # Change the value stored in the cache, the result should change.
+        cache[cache.keys()[0]] = ("Hey there!", {})
+
+        with self.assertRaises(SafeExecException):
+            safe_exec(code, g, cache=DictCache(cache))
+
+        self.assertEqual(len(cache), 1)
+        cache_exc_msg, cache_globals = cache.values()[0]
+        self.assertEqual("Hey there!", cache_exc_msg)
+
+        # Change it again, now no exception!
+        cache[cache.keys()[0]] = (None, {'a': 17})
+        safe_exec(code, g, cache=DictCache(cache))
+        self.assertEqual(g['a'], 17)