From d9aac08b81d0829211684c758d6e0ca85a6e70d6 Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Sun, 26 Mar 2023 00:02:35 +0000 Subject: [PATCH] Support setting memory limit for Lua 5.1/JIT/Luau Other versions already support this feature. Closes #119 --- src/function.rs | 3 +- src/lib.rs | 1 + src/lua.rs | 116 +++++++++--------------------------- src/memory.rs | 153 ++++++++++++++++++++++++++++++++++++++++++++++++ src/util.rs | 20 +++++-- tests/memory.rs | 15 +++-- 6 files changed, 209 insertions(+), 99 deletions(-) create mode 100644 src/memory.rs diff --git a/src/function.rs b/src/function.rs index 9302cbb6..92a78eb1 100644 --- a/src/function.rs +++ b/src/function.rs @@ -5,6 +5,7 @@ use std::slice; use crate::error::{Error, Result}; use crate::ffi; +use crate::memory::MemoryState; use crate::types::LuaRef; use crate::util::{ assert_stack, check_stack, error_traceback, pop_error, ptr_to_cstr_bytes, StackGuard, @@ -118,7 +119,7 @@ impl<'lua> Function<'lua> { let _sg = StackGuard::new(state); check_stack(state, nargs + 3)?; - ffi::lua_pushcfunction(state, error_traceback); + MemoryState::relax_limit_with(state, || ffi::lua_pushcfunction(state, error_traceback)); let stack_start = ffi::lua_gettop(state); lua.push_ref(&self.0); for arg in args.drain_all() { diff --git a/src/lib.rs b/src/lib.rs index c574f4a1..ca5d635e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -90,6 +90,7 @@ mod hook; mod lua; #[cfg(feature = "luau")] mod luau; +mod memory; mod multi; mod scope; mod stdlib; diff --git a/src/lua.rs b/src/lua.rs index 212bb7aa..a8774f68 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -19,6 +19,7 @@ use crate::error::{Error, Result}; use crate::ffi; use crate::function::Function; use crate::hook::Debug; +use crate::memory::{MemoryState, ALLOCATOR}; use crate::scope::Scope; use crate::stdlib::StdLib; use crate::string::String; @@ -96,7 +97,7 @@ pub(crate) struct ExtraData { safe: bool, libs: StdLib, - mem_info: Option>, + mem_state: Option>, ref_thread: *mut ffi::lua_State, ref_stack_size: c_int, @@ -131,12 +132,6 @@ pub(crate) struct ExtraData { compiler: Option, } -#[derive(Default)] -struct MemoryInfo { - used_memory: isize, - memory_limit: isize, -} - /// Mode of the Lua garbage collector (GC). /// /// In Lua 5.4 GC can work in two modes: incremental and generational. @@ -269,8 +264,8 @@ impl Drop for ExtraData { }; *mlua_expect!(self.registry_unref_list.lock(), "unref list poisoned") = None; - if let Some(mem_info) = self.mem_info { - drop(unsafe { Box::from_raw(mem_info.as_ptr()) }); + if let Some(mem_state) = self.mem_state { + drop(unsafe { Box::from_raw(mem_state.as_ptr()) }); } } } @@ -368,77 +363,19 @@ impl Lua { /// Creates a new Lua state with required `libs` and `options` unsafe fn inner_new(libs: StdLib, options: LuaOptions) -> Lua { - unsafe extern "C" fn allocator( - extra_data: *mut c_void, - ptr: *mut c_void, - osize: usize, - nsize: usize, - ) -> *mut c_void { - use std::alloc::{self, Layout}; - - let mem_info = &mut *(extra_data as *mut MemoryInfo); - - if nsize == 0 { - // Free memory - if !ptr.is_null() { - let layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN); - alloc::dealloc(ptr as *mut u8, layout); - mem_info.used_memory -= osize as isize; - } - return ptr::null_mut(); - } - - // Do not allocate more than isize::MAX - if nsize > isize::MAX as usize { - return ptr::null_mut(); - } - - // Are we fit to the memory limits? - let mut mem_diff = nsize as isize; - if !ptr.is_null() { - mem_diff -= osize as isize; - } - let new_used_memory = mem_info.used_memory + mem_diff; - if mem_info.memory_limit > 0 && new_used_memory > mem_info.memory_limit { - return ptr::null_mut(); - } - mem_info.used_memory += mem_diff; - - if ptr.is_null() { - // Allocate new memory - let new_layout = match Layout::from_size_align(nsize, ffi::SYS_MIN_ALIGN) { - Ok(layout) => layout, - Err(_) => return ptr::null_mut(), - }; - let new_ptr = alloc::alloc(new_layout) as *mut c_void; - if new_ptr.is_null() { - alloc::handle_alloc_error(new_layout); - } - return new_ptr; - } - - // Reallocate memory - let old_layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN); - let new_ptr = alloc::realloc(ptr as *mut u8, old_layout, nsize) as *mut c_void; - if new_ptr.is_null() { - alloc::handle_alloc_error(old_layout); - } - new_ptr - } - // Skip Rust allocator for non-vendored LuaJIT (see https://github.com/khvzak/mlua/issues/176) let use_rust_allocator = !(cfg!(feature = "luajit") && cfg!(not(feature = "vendored"))); - let (state, mem_info) = if use_rust_allocator { - let mut mem_info: *mut MemoryInfo = Box::into_raw(Box::default()); - let mut state = ffi::lua_newstate(allocator, mem_info as *mut c_void); + let (state, mem_state) = if use_rust_allocator { + let mut mem_state: *mut MemoryState = Box::into_raw(Box::default()); + let mut state = ffi::lua_newstate(ALLOCATOR, mem_state as *mut c_void); // If state is null (it's possible for LuaJIT on non-x86 arch) then switch to Lua internal allocator if state.is_null() { - drop(Box::from_raw(mem_info)); - mem_info = ptr::null_mut(); + drop(Box::from_raw(mem_state)); + mem_state = ptr::null_mut(); state = ffi::luaL_newstate(); } - (state, mem_info) + (state, mem_state) } else { (ffi::luaL_newstate(), ptr::null_mut()) }; @@ -449,7 +386,7 @@ impl Lua { let lua = Lua::init_from_ptr(state); let extra = lua.extra.get(); - (*extra).mem_info = NonNull::new(mem_info); + (*extra).mem_state = NonNull::new(mem_state); mlua_expect!( load_from_std_lib(state, libs), @@ -559,7 +496,7 @@ impl Lua { app_data: RefCell::new(FxHashMap::default()), safe: false, libs: StdLib::NONE, - mem_info: None, + mem_state: None, ref_thread, // We need 1 extra stack space to move values in and out of the ref stack. ref_stack_size: ffi::LUA_MINSTACK - 1, @@ -1124,8 +1061,8 @@ impl Lua { /// Returns the amount of memory (in bytes) currently used inside this Lua state. pub fn used_memory(&self) -> usize { unsafe { - match (*self.extra.get()).mem_info.map(|x| x.as_ref()) { - Some(mem_info) => mem_info.used_memory as usize, + match (*self.extra.get()).mem_state.map(|x| x.as_ref()) { + Some(mem_state) => mem_state.used_memory(), None => { // Get data from the Lua GC let used_kbytes = ffi::lua_gc(self.main_state, ffi::LUA_GCCOUNT, 0); @@ -1143,17 +1080,10 @@ impl Lua { /// Returns previous limit (zero means no limit). /// /// Does not work on module mode where Lua state is managed externally. - /// - /// Requires `feature = "lua54/lua53/lua52"` - #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))] - pub fn set_memory_limit(&self, memory_limit: usize) -> Result { + pub fn set_memory_limit(&self, limit: usize) -> Result { unsafe { - match (*self.extra.get()).mem_info.map(|mut x| x.as_mut()) { - Some(mem_info) => { - let prev_limit = mem_info.memory_limit as usize; - mem_info.memory_limit = memory_limit as isize; - Ok(prev_limit) - } + match (*self.extra.get()).mem_state.map(|mut x| x.as_mut()) { + Some(mem_state) => Ok(mem_state.set_memory_limit(limit)), None => Err(Error::MemoryLimitNotAvailable), } } @@ -3022,8 +2952,8 @@ impl Lua { pub(crate) unsafe fn unlikely_memory_error(&self) -> bool { // MemoryInfo is empty in module mode so we cannot predict memory limits (*self.extra.get()) - .mem_info - .map(|x| x.as_ref().memory_limit == 0) + .mem_state + .map(|x| x.as_ref().memory_limit() == 0) .unwrap_or_default() } @@ -3063,6 +2993,14 @@ impl LuaInner { } } +impl ExtraData { + #[cfg(feature = "luau")] + #[inline] + pub(crate) fn mem_state(&self) -> NonNull { + self.mem_state.unwrap() + } +} + struct StateGuard<'a>(&'a LuaInner, *mut ffi::lua_State); impl<'a> StateGuard<'a> { diff --git a/src/memory.rs b/src/memory.rs new file mode 100644 index 00000000..ce7e1e00 --- /dev/null +++ b/src/memory.rs @@ -0,0 +1,153 @@ +use std::alloc::{self, Layout}; +use std::os::raw::c_void; +use std::ptr; + +use crate::ffi; +#[cfg(feature = "luau")] +use crate::lua::ExtraData; + +pub(crate) static ALLOCATOR: ffi::lua_Alloc = allocator; + +#[derive(Default)] +pub(crate) struct MemoryState { + used_memory: isize, + memory_limit: isize, + // Can be set to temporary ignore the memory limit. + // This is used when calling `lua_pushcfunction` for lua5.1/jit/luau. + ignore_limit: bool, + // Indicates that the memory limit was reached on the last allocation. + #[cfg(feature = "luau")] + limit_reached: bool, +} + +impl MemoryState { + #[inline] + pub(crate) fn used_memory(&self) -> usize { + self.used_memory as usize + } + + #[inline] + pub(crate) fn memory_limit(&self) -> usize { + self.memory_limit as usize + } + + #[inline] + pub(crate) fn set_memory_limit(&mut self, limit: usize) -> usize { + let prev_limit = self.memory_limit; + self.memory_limit = limit as isize; + prev_limit as usize + } + + // This function is used primarily for calling `lua_pushcfunction` in lua5.1/jit + // to bypass the memory limit (if set). + #[cfg(any(feature = "lua51", feature = "luajit"))] + #[inline] + pub(crate) unsafe fn relax_limit_with(state: *mut ffi::lua_State, f: impl FnOnce()) { + let mut mem_state: *mut c_void = ptr::null_mut(); + if ffi::lua_getallocf(state, &mut mem_state) == ALLOCATOR { + (*(mem_state as *mut MemoryState)).ignore_limit = true; + f(); + (*(mem_state as *mut MemoryState)).ignore_limit = false; + } else { + f(); + } + } + + // Same as the above but for Luau + // It does not have `lua_getallocf` function, so instead we use `lua_callbacks` + #[cfg(feature = "luau")] + #[inline] + pub(crate) unsafe fn relax_limit_with(state: *mut ffi::lua_State, f: impl FnOnce()) { + let extra = (*ffi::lua_callbacks(state)).userdata as *mut ExtraData; + if extra.is_null() { + return f(); + } + let mem_state = (*extra).mem_state(); + (*mem_state.as_ptr()).ignore_limit = true; + f(); + (*mem_state.as_ptr()).ignore_limit = false; + } + + // Does nothing apart from calling `f()`, we don't need to bypass any limits + #[cfg(any(feature = "lua52", feature = "lua53", feature = "lua54"))] + #[inline] + pub(crate) unsafe fn relax_limit_with(_state: *mut ffi::lua_State, f: impl FnOnce()) { + f(); + } + + // Returns `true` if the memory limit was reached on the last memory operation + #[cfg(feature = "luau")] + pub(crate) unsafe fn limit_reached(state: *mut ffi::lua_State) -> bool { + let extra = (*ffi::lua_callbacks(state)).userdata as *mut ExtraData; + if extra.is_null() { + return false; + } + (*(*extra).mem_state().as_ptr()).limit_reached + } +} + +unsafe extern "C" fn allocator( + extra: *mut c_void, + ptr: *mut c_void, + osize: usize, + nsize: usize, +) -> *mut c_void { + let mem_state = &mut *(extra as *mut MemoryState); + #[cfg(feature = "luau")] + { + // Reset the flag + mem_state.limit_reached = false; + } + + if nsize == 0 { + // Free memory + if !ptr.is_null() { + let layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN); + alloc::dealloc(ptr as *mut u8, layout); + mem_state.used_memory -= osize as isize; + } + return ptr::null_mut(); + } + + // Do not allocate more than isize::MAX + if nsize > isize::MAX as usize { + return ptr::null_mut(); + } + + // Are we fit to the memory limits? + let mut mem_diff = nsize as isize; + if !ptr.is_null() { + mem_diff -= osize as isize; + } + let mem_limit = mem_state.memory_limit; + let new_used_memory = mem_state.used_memory + mem_diff; + if mem_limit > 0 && new_used_memory > mem_limit && !mem_state.ignore_limit { + #[cfg(feature = "luau")] + { + mem_state.limit_reached = true; + } + return ptr::null_mut(); + } + mem_state.used_memory += mem_diff; + + if ptr.is_null() { + // Allocate new memory + let new_layout = match Layout::from_size_align(nsize, ffi::SYS_MIN_ALIGN) { + Ok(layout) => layout, + Err(_) => return ptr::null_mut(), + }; + let new_ptr = alloc::alloc(new_layout) as *mut c_void; + if new_ptr.is_null() { + alloc::handle_alloc_error(new_layout); + } + return new_ptr; + } + + // Reallocate memory + let old_layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN); + let new_ptr = alloc::realloc(ptr as *mut u8, old_layout, nsize) as *mut c_void; + if new_ptr.is_null() { + alloc::handle_alloc_error(old_layout); + } + new_ptr +} diff --git a/src/util.rs b/src/util.rs index a34c2b2e..43825194 100644 --- a/src/util.rs +++ b/src/util.rs @@ -12,6 +12,7 @@ use rustc_hash::FxHashMap; use crate::error::{Error, Result}; use crate::ffi; +use crate::memory::MemoryState; static METATABLE_CACHE: Lazy> = Lazy::new(|| { let mut map = FxHashMap::with_capacity_and_hasher(32, Default::default()); @@ -89,8 +90,10 @@ pub unsafe fn protect_lua_call( ) -> Result<()> { let stack_start = ffi::lua_gettop(state) - nargs; - ffi::lua_pushcfunction(state, error_traceback); - ffi::lua_pushcfunction(state, f); + MemoryState::relax_limit_with(state, || { + ffi::lua_pushcfunction(state, error_traceback); + ffi::lua_pushcfunction(state, f); + }); if nargs > 0 { ffi::lua_rotate(state, stack_start + 1, 2); } @@ -147,8 +150,10 @@ where let stack_start = ffi::lua_gettop(state) - nargs; - ffi::lua_pushcfunction(state, error_traceback); - ffi::lua_pushcfunction(state, do_call::); + MemoryState::relax_limit_with(state, || { + ffi::lua_pushcfunction(state, error_traceback); + ffi::lua_pushcfunction(state, do_call::); + }); if nargs > 0 { ffi::lua_rotate(state, stack_start + 1, 2); } @@ -662,6 +667,13 @@ where } pub unsafe extern "C" fn error_traceback(state: *mut ffi::lua_State) -> c_int { + // This is a workaround for bug in Luau, when it calls error handler for memory allocation error + // See https://github.com/Roblox/luau/issues/880 + #[cfg(feature = "luau")] + if MemoryState::limit_reached(state) { + return 0; + } + if ffi::lua_checkstack(state, 2) == 0 { // If we don't have enough stack space to even check the error type, do // nothing so we don't risk shadowing a rust panic. diff --git a/tests/memory.rs b/tests/memory.rs index e359a5da..49f53b71 100644 --- a/tests/memory.rs +++ b/tests/memory.rs @@ -1,11 +1,7 @@ use std::sync::Arc; -use mlua::{GCMode, Lua, Result, UserData}; +use mlua::{Error, GCMode, Lua, Result, UserData}; -#[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))] -use mlua::Error; - -#[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))] #[test] fn test_memory_limit() -> Result<()> { let lua = Lua::new(); @@ -21,6 +17,15 @@ fn test_memory_limit() -> Result<()> { .into_function()?; f.call::<_, ()>(()).expect("should trigger no memory limit"); + if cfg!(feature = "luajit") && cfg!(not(feature = "vendored")) { + // we don't support setting memory limit for non-vendored luajit + assert!(matches!( + lua.set_memory_limit(0), + Err(Error::MemoryLimitNotAvailable) + )); + return Ok(()); + } + lua.set_memory_limit(initial_memory + 10000)?; match f.call::<_, ()>(()) { Err(Error::MemoryError(_)) => {}