From 7a3f19b857b9720f6b282f05039af7ee4144e0ad Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Fri, 29 Nov 2024 13:36:51 +0000
Subject: [PATCH] Ensure that buffer with Luau compiled code is always freed

---
 mlua-sys/src/luau/compat.rs | 15 +++++++++++++--
 src/state/raw.rs            |  4 ++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/mlua-sys/src/luau/compat.rs b/mlua-sys/src/luau/compat.rs
index a738dae3..3d189cec 100644
--- a/mlua-sys/src/luau/compat.rs
+++ b/mlua-sys/src/luau/compat.rs
@@ -326,12 +326,16 @@ pub unsafe fn luaL_loadbufferenv(
     mut size: usize,
     name: *const c_char,
     mode: *const c_char,
-    env: c_int,
+    mut env: c_int,
 ) -> c_int {
     extern "C" {
         fn free(p: *mut c_void);
     }
 
+    unsafe extern "C-unwind" fn data_dtor(data: *mut c_void) {
+        free(*(data as *mut *mut c_char) as *mut c_void);
+    }
+
     let chunk_is_text = size == 0 || (*data as u8) >= b'\t';
     if !mode.is_null() {
         let modeb = CStr::from_ptr(mode).to_bytes();
@@ -345,9 +349,16 @@ pub unsafe fn luaL_loadbufferenv(
     }
 
     if chunk_is_text {
+        if env < 0 {
+            env -= 1;
+        }
+        let data_ud = lua_newuserdatadtor(L, mem::size_of::<*mut c_char>(), data_dtor) as *mut *mut c_char;
         let data = luau_compile_(data, size, ptr::null_mut(), &mut size);
+        ptr::write(data_ud, data);
+        // By deferring the `free(data)` to the userdata destructor, we ensure that
+        // even if `luau_load` throws an error, the `data` is still released.
         let ok = luau_load(L, name, data, size, env) == 0;
-        free(data as *mut c_void);
+        lua_replace(L, -2); // replace data with the result
         if !ok {
             return LUA_ERRSYNTAX;
         }
diff --git a/src/state/raw.rs b/src/state/raw.rs
index 0af877f5..0add603d 100644
--- a/src/state/raw.rs
+++ b/src/state/raw.rs
@@ -327,10 +327,10 @@ impl RawLua {
                 Some(ChunkMode::Text) => cstr!("t"),
                 None => cstr!("bt"),
             };
-            let status = if cfg!(not(feature = "luau")) || self.unlikely_memory_error() {
+            let status = if self.unlikely_memory_error() {
                 self.load_chunk_inner(state, name, env, mode, source)
             } else {
-                // Only Luau can trigger an exception during chunk loading
+                // Luau and Lua 5.2 can trigger an exception during chunk loading
                 protect_lua!(state, 0, 1, |state| {
                     self.load_chunk_inner(state, name, env, mode, source)
                 })?