Commit 403218ec by Ned Batchelder

If sandboxed code raises an exception, the exception will be cached.

parent 001ef7b0
"""Capa's specialized use of codejail.safe_exec.""" """Capa's specialized use of codejail.safe_exec."""
from codejail.safe_exec import safe_exec as 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 . import lazymod
from statsd import statsd from statsd import statsd
...@@ -63,20 +63,34 @@ def safe_exec(code, globals_dict, random_seed=None, python_path=None, cache=None ...@@ -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()) key = "safe_exec.%r.%s" % (random_seed, md5er.hexdigest())
cached = cache.get(key) cached = cache.get(key)
if cached is not None: 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 return
# Create the complete code we'll run. # Create the complete code we'll run.
code_prolog = CODE_PROLOG % random_seed code_prolog = CODE_PROLOG % random_seed
# Run the code! Results are side effects in globals_dict. # Run the code! Results are side effects in globals_dict.
try:
codejail_safe_exec( codejail_safe_exec(
code_prolog + LAZY_IMPORTS + code, globals_dict, code_prolog + LAZY_IMPORTS + code, globals_dict,
python_path=python_path, 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 # Put the result back in the cache. This is complicated by the fact that
# the globals dict might not be entirely serializable. # the globals dict might not be entirely serializable.
if cache: if cache:
cleaned_results = json_safe(globals_dict) 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
...@@ -5,6 +5,8 @@ import random ...@@ -5,6 +5,8 @@ import random
import unittest import unittest
from capa.safe_exec import safe_exec from capa.safe_exec import safe_exec
from codejail.safe_exec import SafeExecException
class TestSafeExec(unittest.TestCase): class TestSafeExec(unittest.TestCase):
def test_set_values(self): def test_set_values(self):
...@@ -57,6 +59,12 @@ class TestSafeExec(unittest.TestCase): ...@@ -57,6 +59,12 @@ class TestSafeExec(unittest.TestCase):
g, python_path=[pylib] 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): class DictCache(object):
"""A cache implementation over a simple dict, for testing.""" """A cache implementation over a simple dict, for testing."""
...@@ -86,10 +94,10 @@ class TestSafeExecCaching(unittest.TestCase): ...@@ -86,10 +94,10 @@ class TestSafeExecCaching(unittest.TestCase):
safe_exec("a = int(math.pi)", g, cache=DictCache(cache)) safe_exec("a = int(math.pi)", g, cache=DictCache(cache))
self.assertEqual(g['a'], 3) self.assertEqual(g['a'], 3)
# A result has been cached # 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. # Fiddle with the cache, then try it again.
cache[cache.keys()[0]] = {'a': 17} cache[cache.keys()[0]] = (None, {'a': 17})
g = {} g = {}
safe_exec("a = int(math.pi)", g, cache=DictCache(cache)) safe_exec("a = int(math.pi)", g, cache=DictCache(cache))
...@@ -104,3 +112,32 @@ class TestSafeExecCaching(unittest.TestCase): ...@@ -104,3 +112,32 @@ class TestSafeExecCaching(unittest.TestCase):
cache = {} cache = {}
safe_exec(code, g, cache=DictCache(cache)) safe_exec(code, g, cache=DictCache(cache))
self.assertEqual(g['a'], 12345) 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)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment