Add option to limit the memory usage of Lua code

CHANGES.rst

@@ -11,6 +11,9 @@ Lupa change log
   Lua 5.4, LuaJIT 2.0 and LuaJIT 2.1 beta.  Note that this is build specific
   and may depend on the platform.  A normal Python import cascade can be used.
 
+* GH#211: A new option `max_memory` allows to limit the memory usage of Lua code.
+  (patch by Leo Developer)
+
 * GH#171: Python references in Lua are now more safely reference counted
   to prevent garbage collection glitches.
   (patch by Guilherme Dantas)

lupa/_lupa.pyx

@@ -9,6 +9,8 @@ from __future__ import absolute_import
 cimport cython
 
 from libc.string cimport strlen, strchr
+from libc.stdlib cimport malloc, calloc, free, realloc
+from libc.stdio cimport fprintf, stderr, fflush
 from . cimport luaapi as lua
 from .luaapi cimport lua_State
 
@@ -58,7 +60,7 @@ from functools import wraps
 
 
 __all__ = ['LUA_VERSION', 'LUA_MAXINTEGER', 'LUA_MININTEGER',
-            'LuaRuntime', 'LuaError', 'LuaSyntaxError',
+            'LuaRuntime', 'LuaError', 'LuaSyntaxError', 'LuaMemoryError',
            'as_itemgetter', 'as_attrgetter', 'lua_type',
            'unpacks_lua_table', 'unpacks_lua_table_method']
 
@@ -121,6 +123,11 @@ class LuaSyntaxError(LuaError):
     """
 
 
+class LuaMemoryError(LuaError, MemoryError):
+    """Memory error in Lua code.
+    """
+
+
 def lua_type(obj):
     """
     Return the Lua type name of a wrapped object as string, as provided
@@ -217,6 +224,12 @@ cdef class LuaRuntime:
       Normally, it should return the now well-behaved object that can be
       converted/wrapped to a Lua type. If the object cannot be precisely
       represented in Lua, it should raise an ``OverflowError``.
+
+    * ``max_memory``: max memory usage this LuaRuntime can use in bytes.
+      If max_memory is None, the default lua allocator is used and calls to
+      ``set_max_memory(limit)`` will fail with a ``LuaMemoryError``.
+      Note: Not supported on 64bit LuaJIT.
+      (default: None, i.e. no limitation. New in Lupa 2.0)
 
     Example usage::
 
@@ -242,14 +255,24 @@ cdef class LuaRuntime:
     cdef object _attribute_getter
     cdef object _attribute_setter
     cdef bint _unpack_returned_tuples
+    cdef MemoryStatus* _memory_status
 
     def __cinit__(self, encoding='UTF-8', source_encoding=None,
                   attribute_filter=None, attribute_handlers=None,
                   bint register_eval=True, bint unpack_returned_tuples=False,
-                  bint register_builtins=True, overflow_handler=None):
-        cdef lua_State* L = lua.luaL_newstate()
+                  bint register_builtins=True, overflow_handler=None,
+                  max_memory=None):
+        cdef lua_State* L
+
+        if max_memory is None:
+            L = lua.luaL_newstate()
+        else:
+            memory_status = <MemoryStatus*>calloc(1, sizeof(MemoryStatus))
+            self._memory_status = memory_status
+            L = lua.lua_newstate(<lua.lua_Alloc>&_lua_alloc_restricted, <void*>self._memory_status)
         if L is NULL:
             raise LuaError("Failed to initialise Lua runtime")
+
         self._state = L
         self._lock = FastRLock()
         self._pyrefs_in_lua = {}
@@ -276,16 +299,55 @@ cdef class LuaRuntime:
                 raise ValueError("attribute_filter and attribute_handlers are mutually exclusive")
             self._attribute_getter, self._attribute_setter = getter, setter
 
+        lua.lua_atpanic(L, &_lua_panic)
         lua.luaL_openlibs(L)
         self.init_python_lib(register_eval, register_builtins)
-        lua.lua_atpanic(L, <lua.lua_CFunction>1)
 
         self.set_overflow_handler(overflow_handler)
 
+        # lupa init done, set real limit
+        if max_memory is not None:
+            self._memory_status.base_usage = self._memory_status.used
+            if max_memory > 0:
+                self._memory_status.limit =  self._memory_status.base_usage + <size_t>max_memory
+
     def __dealloc__(self):
         if self._state is not NULL:
             lua.lua_close(self._state)
             self._state = NULL
+        if self._memory_status is not NULL:
+            free(self._memory_status)
+            self._memory_status = NULL
+
+    def get_max_memory(self, count_base=False):
+        """
+        Maximum memory allowed to be used by this LuaRuntime.
+        0 indicates no limit meanwhile None indicates that the default lua
+        allocator is being used and ``set_max_memory()`` cannot be used.
+
+        If ``count_base`` is True, the base memory used by the lua runtime
+        will be included in the limit.
+        """
+        if self._memory_status is NULL:
+            return None
+        elif count_base:
+            return self._memory_status.limit
+        return self._memory_status.limit - self._memory_status.base_usage
+
+    def get_memory_used(self, count_base=False):
+        """
+        Memory currently in use.
+        This is None if the default lua allocator is used and 0 if
+        ``max_memory`` is 0.
+
+        If ``count_base`` is True, the base memory used by the lua runtime
+        will be included.
+        """
+        if self._memory_status is NULL:
+            return None
+        elif count_base:
+            return self._memory_status.used
+        return self._memory_status.used - self._memory_status.base_usage
 
     @property
     def lua_version(self):
@@ -360,7 +422,14 @@ cdef class LuaRuntime:
                 return py_from_lua(self, L, -1)
             else:
                 err = lua.lua_tolstring(L, -1, &size)
-                error = err[:size] if self._encoding is None else err[:size].decode(self._encoding)
+                if self._encoding is None:
+                    error = err[:size]  # bytes
+                    is_memory_error = b"not enough memory" in error
+                else:
+                    error = err[:size].decode(self._encoding)
+                    is_memory_error = u"not enough memory" in error
+                if is_memory_error:
+                    raise LuaMemoryError(error)
                 raise LuaSyntaxError(error)
         finally:
             lua.lua_settop(L, old_top)
@@ -460,6 +529,26 @@ cdef class LuaRuntime:
             lua.lua_settop(L, old_top)
             unlock_runtime(self)
 
+    def set_max_memory(self, size_t max_memory, count_base=False):
+        """Set maximum allowed memory for this LuaRuntime.
+
+        If `max_memory` is 0, there will be no limit.
+        If ``count_base`` is True, the base memory used by the LuaRuntime itself
+        will be included in the memory limit.
+
+        If max_memory was set to None during creation, this will raise a
+        RuntimeError.
+        """
+        cdef size_t used
+        if self._memory_status is NULL:
+            raise RuntimeError("max_memory must be set on LuaRuntime creation")
+        elif max_memory <= 0:
+            self._memory_status.limit = 0
+        elif count_base:
+            self._memory_status.limit = max_memory
+        else:
+            self._memory_status.limit = self._memory_status.base_usage + max_memory
+
     def set_overflow_handler(self, overflow_handler):
         """Set the overflow handler function that is called on failures to pass large numbers to Lua.
         """
@@ -584,7 +673,7 @@ cdef int check_lua_stack(lua_State* L, int extra) except -1:
     """
     assert extra >= 0
     if not lua.lua_checkstack(L, extra):
-        raise MemoryError
+        raise LuaMemoryError
     return 0
 
 
@@ -1558,9 +1647,12 @@ cdef int raise_lua_error(LuaRuntime runtime, lua_State* L, int result) except -1
     if result == 0:
         return 0
     elif result == lua.LUA_ERRMEM:
-        raise MemoryError()
+        raise LuaMemoryError()
     else:
-        raise LuaError(build_lua_error_message(runtime, L))
+        error_message = build_lua_error_message(runtime, L)
+        if u"not enough memory" in error_message:
+            raise LuaMemoryError(error_message)
+        raise LuaError(error_message)
 
 
 cdef bint _looks_like_traceback_line(unicode line):
@@ -1597,9 +1689,8 @@ cdef unicode _reorder_lua_stack_trace(unicode error_message):
     return error_message
 
 
-cdef build_lua_error_message(LuaRuntime runtime, lua_State* L, unicode err_message=None, int stack_index=-1):
+cdef build_lua_error_message(LuaRuntime runtime, lua_State* L, int stack_index=-1):
     """Removes the string at the given stack index ``n`` to build an error message.
-    If ``err_message`` is provided, it is used as a %-format string to build the error message.
     """
     cdef size_t size = 0
     cdef const char *s = lua.lua_tolstring(L, stack_index, &size)
@@ -1615,9 +1706,6 @@ cdef build_lua_error_message(LuaRuntime runtime, lua_State* L, unicode err_messa
     if u"stack traceback:" in py_ustring:
         py_ustring = _reorder_lua_stack_trace(py_ustring)
 
-    if err_message is not None:
-        py_ustring = err_message % py_ustring
-
     return py_ustring
 
 
@@ -1631,8 +1719,10 @@ cdef run_lua(LuaRuntime runtime, bytes lua_code, tuple args):
     try:
         check_lua_stack(L, 1)
         if lua.luaL_loadbuffer(L, lua_code, len(lua_code), '<python>'):
-            raise LuaSyntaxError(build_lua_error_message(
-                runtime, L, err_message=u"error loading code: %s"))
+            error = build_lua_error_message(runtime, L)
+            if error.startswith("not enough memory"):
+                raise LuaMemoryError(error)
+            raise LuaSyntaxError(u"error loading code: " + error)
         return call_lua(runtime, L, args)
     finally:
         lua.lua_settop(L, old_top)
@@ -1725,6 +1815,53 @@ cdef tuple unpack_multiple_lua_results(LuaRuntime runtime, lua_State *L, int nar
     return args
 
 
+# bounded memory allocation
+
+cdef struct MemoryStatus:
+    size_t used
+    size_t base_usage
+    size_t limit
+
+cdef void* _lua_alloc_restricted(void* ud, void* ptr, size_t old_size, size_t new_size) nogil:
+    # adapted from https://stackoverflow.com/a/9672205
+    # print(<size_t>ud, <size_t>ptr, old_size, new_size)
+    cdef MemoryStatus* memory_status = <MemoryStatus*>ud
+    # print("  ", memory_status.used, memory_status.base_usage, memory_status.limit)
+
+    if ptr is NULL:
+        # <http://www.lua.org/manual/5.2/manual.html#lua_Alloc>:
+        # When ptr is NULL, old_size encodes the kind of object that Lua is allocating.
+        # Since we don’t care about that, just mark it as 0.
+        old_size = 0
+
+    cdef void* new_ptr
+    if new_size == 0:
+        free(ptr)
+        memory_status.used -= old_size  # add deallocated old size to available memory
+        return NULL
+    elif new_size == old_size:
+        return ptr
+
+    if memory_status.limit > 0 and new_size > old_size and memory_status.limit <= memory_status.used + new_size - old_size:  # reached the limit
+        # print("REACHED LIMIT")
+        return NULL
+    # print("  realloc()...")
+    new_ptr = realloc(ptr, new_size)
+    # print("  ", memory_status.used, new_size - old_size, memory_status.used + new_size - old_size)
+    if new_ptr is not NULL:
+        memory_status.used += new_size - old_size
+    return new_ptr
+
+cdef int _lua_panic(lua_State *L) nogil:
+    cdef const char* msg = lua.lua_tostring(L, -1)
+    if msg == NULL:
+        msg = "error object is not a string"
+    cdef char* message = "PANIC: unprotected error in call to Lua API (%s)\n"
+    fprintf(stderr, message, msg)
+    fflush(stderr)
+    return 0  # return to Lua to abort
+
+
 ################################################################################
 # Python support in Lua
 

lupa/tests/test.py

@@ -2995,6 +2995,62 @@ def test_tostring_err(self):
         self.assertRaises(self.lupa.LuaError, str, self.lua.eval('setmetatable({}, {__tostring = function() error() end})'))
 
 
+################################################################################
+# test LuaRuntime max_memory
+
+class TestMaxMemory(SetupLuaRuntimeMixin, LupaTestCase):
+    lua_runtime_kwargs = {"max_memory": 10000}
+
+    def setUp(self):
+        # need to test in here because the creation of the LuaRuntime fails
+        if "luajit" in self.lupa.LuaRuntime().lua_implementation.lower():
+            return self.skipTest("not supported in LuaJIT")
+        return super(TestMaxMemory, self).setUp()
+
+    def test_getters(self):
+        self.assertEqual(self.lua.get_memory_used(), 0)
+        self.assertGreater(self.lua.get_memory_used(count_base=True), 0)
+        self.assertEqual(self.lua.get_max_memory(), 10000)
+        self.assertGreater(self.lua.get_max_memory(count_base=True), 10000)
+        self.lua.set_max_memory(1000000)
+        self.assertEqual(self.lua.get_memory_used(), 0)
+        self.assertGreater(self.lua.get_memory_used(count_base=True), 0)
+        self.assertEqual(self.lua.get_max_memory(), 1000000)
+        self.assertGreater(self.lua.get_max_memory(count_base=True), 1000000)
+        self.lua.set_max_memory(1000000, count_base=True)
+        self.assertEqual(self.lua.get_max_memory(count_base=True), 1000000)
+        self.assertLess(self.lua.get_max_memory(), 1000000)
+
+    def test_not_enough_memory(self):
+        self.lua.eval("('a'):rep(50)")
+        self.assertRaises(self.lupa.LuaMemoryError, self.lua.eval, "('a'):rep(50000)")
+
+    def test_decrease_memory(self):
+        self.lua.set_max_memory(1000000)
+        self.lua.execute("a = ('a'):rep(50000)")
+        self.lua.set_max_memory(10000)
+        self.assertEqual(self.lua.get_max_memory(), 10000)
+        self.assertGreaterEqual(self.lua.get_memory_used(), 50000)
+        self.assertRaises(self.lupa.LuaMemoryError, self.lua.eval, "('b'):rep(10)")
+        del self.lua.globals()["a"]
+        self.lua.eval("('b'):rep(10)")
+
+    def test_compile_not_enough_memory(self):
+        self.lua.set_max_memory(10)
+        self.assertRaises(self.lupa.LuaMemoryError, self.lua.compile, "_G.a = function() return 'test abcdef' end")
+
+    def test_unlimited_memory(self):
+        self.lua.set_max_memory(0)
+        self.lua.execute("a = ('a'):rep(50000)")
+
+
+class TestMaxMemoryWithoutSettingIt(SetupLuaRuntimeMixin, LupaTestCase):
+    def test_property(self):
+        self.assertEqual(self.lua.get_max_memory(), None)
+
+    def test_set_max(self):
+        self.assertRaises(RuntimeError, self.lua.set_max_memory, 10000)
+
 
 ################################################################################
 # Load tests for different Lua version modules
@@ -3003,6 +3059,7 @@ def load_tests(loader, standard_tests, pattern):
     return lupa.tests.build_suite_for_modules(loader, globals())
 
 
+
 if __name__ == '__main__':
     def print_version():
         version = lupa.LuaRuntime().lua_implementation