Add option to limit the memory usage of Lua code

CHANGES.rst

@@ -4,6 +4,9 @@ Lupa change log
 2.0a1 (2022-??-??)
 ------------------
 
+* GH#211: Option to limit memory usage
+  (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, free, realloc
+from libc.stdio cimport fprintf, stderr, fflush
 from lupa cimport lua
 from .lua cimport lua_State
 
@@ -217,6 +219,10 @@ 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.
+      Values below 40kb may cause lua panics.
+      (default: None, new in Lupa 2.0)
 
     Example usage::
 
@@ -242,12 +248,24 @@ cdef class LuaRuntime:
     cdef object _attribute_getter
     cdef object _attribute_setter
     cdef bint _unpack_returned_tuples
+    cdef size_t _max_memory
+    cdef size_t* _memory_left
 
     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
+        cdef size_t* memory_left
+        if max_memory is None:
+            L = lua.luaL_newstate()
+        else:
+            memory_left = <size_t*>malloc(sizeof(size_t))
+            if memory_left is NULL:
+                raise LuaError("Failed to allocate Lua runtime memory limiter")
+            memory_left[0] = <size_t>max_memory
+            L = lua.lua_newstate(<lua.lua_Alloc>&_lua_alloc_restricted, memory_left)
         if L is NULL:
             raise LuaError("Failed to initialise Lua runtime")
         self._state = L
@@ -259,6 +277,8 @@ cdef class LuaRuntime:
             raise ValueError("attribute_filter must be callable")
         self._attribute_filter = attribute_filter
         self._unpack_returned_tuples = unpack_returned_tuples
+        self._max_memory = max_memory
+        self._memory_left = memory_left
 
         if attribute_handlers:
             raise_error = False
@@ -276,9 +296,11 @@ cdef class LuaRuntime:
                 raise ValueError("attribute_filter and attribute_handlers are mutually exclusive")
             self._attribute_getter, self._attribute_setter = getter, setter
 
+        if max_memory is not None:
+            # luaL_newstate already registers a panic handler
+            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)
 
@@ -287,6 +309,10 @@ cdef class LuaRuntime:
             lua.lua_close(self._state)
             self._state = NULL
 
+    @property
+    def max_memory(self):
+        return self._max_memory
+
     @property
     def lua_version(self):
         """
@@ -460,6 +486,13 @@ cdef class LuaRuntime:
             lua.lua_settop(L, old_top)
             unlock_runtime(self)
 
+    def set_max_memory(self, max_memory):
+        if self._max_memory is None:
+            raise RuntimeError("max_memory must be set on LuaRuntime creation")
+        used = self._max_memory - self._memory_left[0]
+        self._memory_left[0] = max_memory - used
+        self._max_memory = 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.
         """
@@ -1609,6 +1642,40 @@ cdef call_lua(LuaRuntime runtime, lua_State *L, tuple args):
     push_lua_arguments(runtime, L, args)
     return execute_lua_call(runtime, L, len(args))
 
+# adapted from https://stackoverflow.com/a/9672205
+cdef void* _lua_alloc_restricted(void* ud, void* ptr, size_t osize, size_t nsize) nogil:
+    cdef size_t* left = <size_t*>ud
+
+    if ptr is NULL:
+        # <http://www.lua.org/manual/5.2/manual.html#lua_Alloc>:
+        # When ptr is NULL, osize encodes the kind of object that Lua is
+        # allocating.
+        # Since we don’t care about that, just mark it as 0.
+        osize = 0
+
+    if nsize == 0:
+        free(ptr)
+        left[0] += osize # add old size to available memory
+        return NULL
+    elif nsize == osize:
+        return ptr
+    else:
+        if nsize > osize and left[0] < nsize - osize: # reached the limit
+            return NULL
+        ptr = realloc(ptr, nsize)
+        if ptr: # reallocation successful?
+            left[0] -= nsize + osize
+        return ptr
+
+cdef int _lua_panic(lua_State *L) nogil:
+    cdef 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
+
 cdef object execute_lua_call(LuaRuntime runtime, lua_State *L, Py_ssize_t nargs):
     cdef int result_status
     cdef object result