Create C closures under lua_cpcall so an allocation failure doesn't escape (Lua 5.1)

Only applies to module mode, where external Lua state memory allocator cannot be controlled.

Fixes #725

Author: khvzak

src/state.rs

@@ -316,7 +316,7 @@ impl Clone for Lua {
 
 impl fmt::Debug for Lua {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "Lua({:p})", self.lock().state())
+        write!(f, "Lua({:p})", self.state())
     }
 }
 
@@ -2400,6 +2400,12 @@ impl Lua {
         .await
     }
 
+    /// Returns a pointer to the underlying Lua state.
+    #[doc(hidden)]
+    pub fn state(&self) -> *mut ffi::lua_State {
+        self.lock().state()
+    }
+
     /// Returns a weak reference to the Lua instance.
     ///
     /// This is useful for creating a reference to the Lua instance that does not prevent it from

src/util/error.rs

@@ -149,6 +149,67 @@ pub(crate) unsafe fn pop_error(state: *mut ffi::lua_State, err_code: c_int) -> E
     }
 }
 
+// Create C closures under `lua_cpcall` so an allocation failure doesn't escape
+#[cfg(any(feature = "lua51", feature = "luajit"))]
+unsafe fn push_protected_cfunctions(state: *mut ffi::lua_State, f: ffi::lua_CFunction) -> Result<()> {
+    if !MemoryState::get(state).is_null() {
+        MemoryState::relax_limit_with(state, || {
+            ffi::lua_pushcfunction(state, error_traceback);
+            ffi::lua_pushcfunction(state, f);
+        });
+        return Ok(());
+    }
+
+    static ERROR_TRACEBACK_KEY: u8 = 0;
+    static FUNCTION_KEY: u8 = 0;
+
+    unsafe extern "C-unwind" fn do_push(state: *mut ffi::lua_State) -> c_int {
+        let f = ffi::lua_tolightuserdata(state, -1) as *const ffi::lua_CFunction;
+        ffi::lua_pop(state, 1);
+
+        ffi::lua_pushcfunction(state, error_traceback);
+        ffi::lua_rawsetp(
+            state,
+            ffi::LUA_REGISTRYINDEX,
+            &ERROR_TRACEBACK_KEY as *const u8 as *const c_void,
+        );
+        ffi::lua_pushcfunction(state, *f);
+        ffi::lua_rawsetp(
+            state,
+            ffi::LUA_REGISTRYINDEX,
+            &FUNCTION_KEY as *const u8 as *const c_void,
+        );
+        0
+    }
+
+    let ret = ffi::lua_cpcall(state, do_push, &f as *const ffi::lua_CFunction as *mut c_void);
+    if ret != ffi::LUA_OK {
+        return Err(pop_error(state, ret));
+    }
+
+    ffi::lua_rawgetp(
+        state,
+        ffi::LUA_REGISTRYINDEX,
+        &ERROR_TRACEBACK_KEY as *const u8 as *const c_void,
+    );
+    ffi::lua_rawgetp(
+        state,
+        ffi::LUA_REGISTRYINDEX,
+        &FUNCTION_KEY as *const u8 as *const c_void,
+    );
+    Ok(())
+}
+
+#[cfg(not(any(feature = "lua51", feature = "luajit")))]
+#[inline]
+unsafe fn push_protected_cfunctions(state: *mut ffi::lua_State, f: ffi::lua_CFunction) -> Result<()> {
+    MemoryState::relax_limit_with(state, || {
+        ffi::lua_pushcfunction(state, error_traceback);
+        ffi::lua_pushcfunction(state, f);
+    });
+    Ok(())
+}
+
 // Call a function that calls into the Lua API and may trigger a Lua error (longjmp) in a safe way.
 // Wraps the inner function in a call to `lua_pcall`, so the inner function only has access to a
 // limited lua stack. `nargs` is the same as the the parameter to `lua_pcall`, and `nresults` is
@@ -162,10 +223,7 @@ pub(crate) unsafe fn protect_lua_call(
 ) -> Result<()> {
     let stack_start = ffi::lua_gettop(state) - nargs;
 
-    MemoryState::relax_limit_with(state, || {
-        ffi::lua_pushcfunction(state, error_traceback);
-        ffi::lua_pushcfunction(state, f);
-    });
+    push_protected_cfunctions(state, f)?;
     if nargs > 0 {
         ffi::lua_rotate(state, stack_start + 1, 2);
     }
@@ -223,10 +281,7 @@ where
 
     let stack_start = ffi::lua_gettop(state) - nargs;
 
-    MemoryState::relax_limit_with(state, || {
-        ffi::lua_pushcfunction(state, error_traceback);
-        ffi::lua_pushcfunction(state, do_call::<F, R>);
-    });
+    push_protected_cfunctions(state, do_call::<F, R>)?;
     if nargs > 0 {
         ffi::lua_rotate(state, stack_start + 1, 2);
     }

tests/module/loader/tests/load.rs

@@ -41,6 +41,30 @@ fn test_module_error() -> Result<()> {
     .exec()
 }
 
+#[cfg(any(feature = "lua51", feature = "luajit"))]
+#[test]
+fn test_module_protected_call_setup() -> Result<()> {
+    let lua = make_lua()?;
+    lua.load("test_protected_call_setup = require('test_module').test_protected_call_setup")
+        .exec()?;
+
+    lua.set_memory_limit(1)?;
+
+    let state = lua.state();
+    let check_passed = unsafe {
+        mlua::ffi::lua_getglobal(state, c"test_protected_call_setup".as_ptr());
+        mlua::ffi::lua_call(state, 0, 1);
+        let ok = mlua::ffi::lua_toboolean(state, -1) != 0;
+        mlua::ffi::lua_pop(state, 1);
+        ok
+    };
+    lua.set_memory_limit(0)?;
+
+    assert!(check_passed);
+
+    Ok(())
+}
+
 #[cfg(any(
     feature = "lua55",
     feature = "lua54",

tests/module/src/lib.rs

@@ -12,12 +12,23 @@ fn check_userdata(_: &Lua, ud: LuaAnyUserData) -> LuaResult<i32> {
     Ok(ud.borrow::<MyUserData>()?.0)
 }
 
+#[cfg(any(feature = "lua51", feature = "luajit"))]
+fn test_protected_call_setup(lua: &Lua, _: ()) -> LuaResult<bool> {
+    let res = lua.create_function(|_, ()| Ok(()));
+    Ok(matches!(res, Err(LuaError::MemoryError(_))))
+}
+
 #[mlua::lua_module]
 fn test_module(lua: &Lua) -> LuaResult<LuaTable> {
     let exports = lua.create_table()?;
     exports.set("sum", lua.create_function(sum)?)?;
     exports.set("used_memory", lua.create_function(used_memory)?)?;
     exports.set("check_userdata", lua.create_function(check_userdata)?)?;
+    #[cfg(any(feature = "lua51", feature = "luajit"))]
+    exports.set(
+        "test_protected_call_setup",
+        lua.create_function(test_protected_call_setup)?,
+    )?;
     Ok(exports)
 }