Add "max_memory" option to limit the memory usage of Lua code (GH-212)

Closes #211

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)

README.rst

@@ -913,6 +913,52 @@ setter function implementations for a ``LuaRuntime``:
         AttributeError: not allowed to write attribute "noway"
 
 
+Restricting Lua Memory Usage
+----------------------------
+
+Lupa provides a simple mechanism to control the maximum memory
+usage of the Lua Runtime since version 2.0.
+By default Lupa does not interfere with Lua's memory allocation, to opt-in
+you must set the ``max_memory`` when creating the LuaRuntime.
+
+The ``LuaRuntime`` provides three methods for controlling and reading the
+memory usage:
+
+1. ``get_memory_used(total=False)`` to get the current memory
+   usage of the LuaRuntime.
+
+2. ``get_max_memory(total=False)`` to get the current memory limit.
+   ``0`` means there is no memory limitation.
+
+3. ``set_max_memory(max_memory, total=False)`` to change the memory limit.
+   Values below or equal to 0 mean no limit.
+
+There is always some memory used by the LuaRuntime itself (around ~20KiB,
+depending on your lua version and other factors) which is excluded from all
+calculations unless you specify ``total=True``.
+
+.. code:: python
+
+        >>> lua = LuaRuntime(max_memory=0)  # 0 for unlimited, default is None
+        >>> lua.get_memory_used()  # memory used by your code
+        0
+        >>> total_lua_memory = lua.get_memory_used(total=True)  # includes memory used by the runtime itself
+        >>> assert total_lua_memory > 0  # exact amount depends on your lua version and other factors
+
+
+Lua code hitting the memory limit will receive memory errors:
+
+.. code:: python
+      
+        >>> lua.set_max_memory(100)
+        >>> lua.eval("string.rep('a', 1000)")   # doctest: +IGNORE_EXCEPTION_DETAIL
+        Traceback (most recent call last):
+         ...
+        lupa.LuaMemoryError: not enough memory
+
+``LuaMemoryError`` inherits from ``LuaError`` and ``MemoryError``.
+
+
 Importing Lua binary modules
 ----------------------------
 

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 . 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']
 
@@ -111,6 +113,12 @@ else:  # probably not smaller
     LUA_MININTEGER, LUA_MAXINTEGER = (CHAR_MIN, CHAR_MAX)
 
 
+cdef struct MemoryStatus:
+    size_t used
+    size_t base_usage
+    size_t limit
+
+
 class LuaError(Exception):
     """Base class for errors in the Lua runtime.
     """
@@ -121,6 +129,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 +230,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 +261,23 @@ 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()
+            self._memory_status.limit = <size_t> -1
+        else:
+            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,17 +304,56 @@ 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
+                # Prevent accidental (or deliberate) usage of our special value.
+                if self._memory_status.limit == <size_t> -1:
+                    self._memory_status.limit -= 1
+
     def __dealloc__(self):
         if self._state is not NULL:
             lua.lua_close(self._state)
             self._state = NULL
 
+    def get_max_memory(self, total=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 ``total`` is True, the base memory used by the lua runtime
+        will be included in the limit.
+        """
+        if self._memory_status.limit == <size_t> -1:
+            return None
+        elif total:
+            return self._memory_status.limit
+        return self._memory_status.limit - self._memory_status.base_usage
+
+    def get_memory_used(self, total=False):
+        """
+        Memory currently in use.
+        This is None if the default lua allocator is used and 0 if
+        ``max_memory`` is 0.
+
+        If ``total`` is True, the base memory used by the lua runtime
+        will be included.
+        """
+        if self._memory_status.limit == <size_t> -1:
+            return None
+        elif total:
+            return self._memory_status.used
+        return self._memory_status.used - self._memory_status.base_usage
+
     @property
     def lua_version(self):
         """
@@ -360,7 +427,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 +534,29 @@ cdef class LuaRuntime:
             lua.lua_settop(L, old_top)
             unlock_runtime(self)
 
+    def set_max_memory(self, size_t max_memory, total=False):
+        """Set maximum allowed memory for this LuaRuntime.
+
+        If `max_memory` is 0, there will be no limit.
+        If ``total`` 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.limit == <size_t> -1:
+            raise RuntimeError("max_memory must be set on LuaRuntime creation")
+        elif max_memory == 0:
+            self._memory_status.limit = 0
+        elif total:
+            self._memory_status.limit = max_memory
+        else:
+            self._memory_status.limit = self._memory_status.base_usage + max_memory
+            # Prevent accidental (or deliberate) usage of our special value.
+            if self._memory_status.limit == <size_t> -1:
+                self._memory_status.limit -= 1
+
     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 +681,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 +1655,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 +1697,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 +1714,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 +1727,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 +1823,48 @@ cdef tuple unpack_multiple_lua_results(LuaRuntime runtime, lua_State *L, int nar
     return args
 
 
+# bounded memory allocation
+
+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