From f8ab713e1ed48c466ad88d7c35f420ad2a1ae9a7 Mon Sep 17 00:00:00 2001 From: Leo Developer Date: Mon, 3 Apr 2023 15:16:15 +0200 Subject: [PATCH] Add "max_memory" option to limit the memory usage of Lua code (GH-212) Closes https://github.com/scoder/lupa/issues/211 --- CHANGES.rst | 3 + README.rst | 46 ++++++++++++ lupa/_lupa.pyx | 170 +++++++++++++++++++++++++++++++++++++++++---- lupa/tests/test.py | 57 +++++++++++++++ 4 files changed, 261 insertions(+), 15 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index 8df4572c..8d185de2 100644 --- a/CHANGES.rst +++ b/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) diff --git a/README.rst b/README.rst index 2c53f0fc..40ec9a0c 100644 --- a/README.rst +++ b/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 ---------------------------- diff --git a/lupa/_lupa.pyx b/lupa/_lupa.pyx index ba7303cf..9310e354 100644 --- a/lupa/_lupa.pyx +++ b/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 = -1 + else: + L = lua.lua_newstate(&_lua_alloc_restricted, &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, 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 + max_memory + # Prevent accidental (or deliberate) usage of our special value. + if self._memory_status.limit == -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 == -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 == -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 == -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 == -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), ''): - 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(ud, ptr, old_size, new_size) + cdef MemoryStatus* memory_status = ud + # print(" ", memory_status.used, memory_status.base_usage, memory_status.limit) + + if ptr is NULL: + # : + # 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 diff --git a/lupa/tests/test.py b/lupa/tests/test.py index 67873011..1ce505ae 100644 --- a/lupa/tests/test.py +++ b/lupa/tests/test.py @@ -2996,6 +2996,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(total=True), 0) + self.assertEqual(self.lua.get_max_memory(), 10000) + self.assertGreater(self.lua.get_max_memory(total=True), 10000) + self.lua.set_max_memory(1000000) + self.assertEqual(self.lua.get_memory_used(), 0) + self.assertGreater(self.lua.get_memory_used(total=True), 0) + self.assertEqual(self.lua.get_max_memory(), 1000000) + self.assertGreater(self.lua.get_max_memory(total=True), 1000000) + self.lua.set_max_memory(1000000, total=True) + self.assertEqual(self.lua.get_max_memory(total=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 @@ -3004,6 +3060,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