diff --git a/src/state/raw.rs b/src/state/raw.rs index a691367..0af877f 100644 --- a/src/state/raw.rs +++ b/src/state/raw.rs @@ -319,41 +319,60 @@ impl RawLua { let state = self.state(); unsafe { let _sg = StackGuard::new(state); - check_stack(state, 2)?; + check_stack(state, 3)?; - let mode_str = match mode { + let name = name.map(CStr::as_ptr).unwrap_or(ptr::null()); + let mode = match mode { Some(ChunkMode::Binary) => cstr!("b"), Some(ChunkMode::Text) => cstr!("t"), None => cstr!("bt"), }; - - match ffi::luaL_loadbufferenv( - state, - source.as_ptr() as *const c_char, - source.len(), - name.map(|n| n.as_ptr()).unwrap_or_else(ptr::null), - mode_str, - match env { - Some(env) => { - self.push_ref(&env.0); - -1 - } - _ => 0, - }, - ) { - ffi::LUA_OK => { - #[cfg(feature = "luau-jit")] - if (*self.extra.get()).enable_jit && ffi::luau_codegen_supported() != 0 { - ffi::luau_codegen_compile(state, -1); - } - - Ok(Function(self.pop_ref())) - } + let status = if cfg!(not(feature = "luau")) || self.unlikely_memory_error() { + self.load_chunk_inner(state, name, env, mode, source) + } else { + // Only Luau can trigger an exception during chunk loading + protect_lua!(state, 0, 1, |state| { + self.load_chunk_inner(state, name, env, mode, source) + })? + }; + match status { + ffi::LUA_OK => Ok(Function(self.pop_ref())), err => Err(pop_error(state, err)), } } } + pub(crate) unsafe fn load_chunk_inner( + &self, + state: *mut ffi::lua_State, + name: *const c_char, + env: Option<&Table>, + mode: *const c_char, + source: &[u8], + ) -> c_int { + let status = ffi::luaL_loadbufferenv( + state, + source.as_ptr() as *const c_char, + source.len(), + name, + mode, + match env { + Some(env) => { + self.push_ref(&env.0); + -1 + } + _ => 0, + }, + ); + #[cfg(feature = "luau-jit")] + if status == ffi::LUA_OK { + if (*self.extra.get()).enable_jit && ffi::luau_codegen_supported() != 0 { + ffi::luau_codegen_compile(state, -1); + } + } + status + } + /// Sets a 'hook' function for a thread (coroutine). #[cfg(not(feature = "luau"))] pub(crate) unsafe fn set_thread_hook( diff --git a/tests/memory.rs b/tests/memory.rs index e765e06..ba30761 100644 --- a/tests/memory.rs +++ b/tests/memory.rs @@ -31,6 +31,16 @@ fn test_memory_limit() -> Result<()> { lua.set_memory_limit(0)?; f.call::<()>(()).expect("should trigger no memory limit"); + // Test memory limit during chunk loading + lua.set_memory_limit(1024)?; + match lua + .load("local t = {}; for i = 1,10000 do t[i] = i end") + .into_function() + { + Err(Error::MemoryError(_)) => {} + _ => panic!("did not trigger memory error"), + }; + Ok(()) }