fix some significant undefined behavior caused by double dropping userdata

src/error.rs

@@ -10,7 +10,7 @@ pub enum Error {
     #[error("Lua runtime error: {0}")]
     Runtime(String),
     #[error("Error when running a __gc metamethod")]
-    GcFailure(String),
+    UserdataGc(String),
     /// Ran into a not enough memory error when trying to grow the lua stack.
     #[error("Failed to allocate memory")]
     MemoryAlloc,

src/function.rs

@@ -59,9 +59,6 @@ impl<'a> Function<'a> {
             };
 
             self.push_to_lua_stack(self.state)?;
-
-            /* let args_val = args.into_lua_value_vec(self.state)?;
-            args_val.push_to_lua_stack(self.state)?; */
             args_val.push_to_lua_stack(self.state)?;
 
             match lua::lua_pcall(s, args_len, lua::LUA_MULTRET, handler_idx) {
@@ -84,8 +81,10 @@ impl<'a> Function<'a> {
             let val = if ret_count > 1 {
                 let vals = ValueVec::from_lua_stack(self.state)?;
                 Value::Variable(vals)
-            } else {
+            } else if ret_count == 1 {
                 Value::from_lua_stack(self.state)?
+            } else {
+                Value::None
             };
 
             R::from_lua(self.state, val)

src/main.rs

@@ -1,4 +1,6 @@
-use std::{cell::{Ref, RefCell}, marker::PhantomData, mem, sync::Arc};
+use std::{
+    cell::{Ref, RefCell}, ffi::CStr, marker::PhantomData, mem, ops::Deref, sync::Arc
+};
 
 use mlua_sys as lua;
 
@@ -35,43 +37,32 @@ fn main() -> Result<()> {
 
     let globals = lua.globals()?;
 
-    let a = |_lua: &State, (num,): (i32,)| -> Result<i32> {
+    let v1 = RefCell::new(Vec2 { x: 50.0, y: 5.0 });
-        println!("Rust got number from lua: {}", num);
+    let ud = lua.create_userdata(UserdataRef::from(v1.borrow()))?;
-        Ok(999)
+    globals.set("v1", ud)?;
-    };
 
-    let f = lua.create_function(a)?;
+    let v2 = Vec2 { x: 10.0, y: 15.0 };
-    globals.set("native_test", f)?;
+    let ud = lua.create_userdata(UserdataRef::from(&v2))?;
+    globals.set("v2", ud)?;
 
-    /* let ud = lua.create_userdata(UserdataProxy::<Vec2>::new())?;
+    //lua.gc_stop()?;
-    globals.set("Vec2", ud)?; */
+    let chunk = lua.load(
-
+        "text.lua",
-    let cell = RefCell::new(Vec2 {
+        r#"
-        x: 50.0,
-        y: 5.0
-    });
-
-    {
-        let ud = lua.create_userdata(UserdataRef {
-            ud_ptr: unsafe {
-                mem::transmute::<Ref<Vec2>, Ref<Vec2>>(cell.borrow())
-            },
-        })?;
-        globals.set("vec2", ud)?;
-    }
-
-    let chunk = lua.load("text.lua", r#"
         require "util"
 
-        if vec2 ~= nil then
+        --if vec2 ~= nil then
-            print("vec2: (" .. vec2.x .. ", " .. vec2.y .. ")")
+        print("v1: (" .. v1.x .. ", " .. v1.y .. ")")
-        end
+        --print("v2: (" .. v2.x .. ", " .. v2.y .. ")")
+        --end
 
         --print("vec2.x = " .. vec2.x)
         --print("vec2.y = " .. vec2.y)
-    "#)?;
+    "#,
+    )?;
 
-    for _ in 0..2 {
+    const MAX_RUNS: i32 = 400;
+    for i in 0..MAX_RUNS {
         // I don't care about the result of this execution, so I set the result as a
         // Value which can be anything
         //
@@ -84,22 +75,18 @@ fn main() -> Result<()> {
             panic!("{}", e);
         }
         
-        let vec2 = globals.get::<_, AnyUserdata>("vec2")?;
+        //print_refs(&lua).unwrap();
-        vec2.garbage_collect()?;
+        println!("i = {}", i);
-        globals.set("vec2", Value::Nil)?;
+        globals.set("v1", Value::Nil)?;
+        lua.gc_collect()?;
 
-        let mut t = cell.borrow_mut();
+        let mut t = v1.borrow_mut();
         t.x += 50.0;
         t.y += 5.0;
         drop(t); 
 
-        let ud = lua.create_userdata(UserdataRef {
+        let ud = lua.create_userdata(UserdataRef::from(v1.borrow()))?;
-            // TODO: avoid unsafe here
+        globals.raw_set("v1", ud)?;
-            ud_ptr: unsafe {
-                mem::transmute::<Ref<Vec2>, Ref<Vec2>>(cell.borrow())
-            },
-        })?;
-        globals.raw_set("vec2", ud)?;
     }
 
     unsafe {
@@ -109,21 +96,59 @@ fn main() -> Result<()> {
     Ok(())
 }
 
+/// Prints the types of the registry
+pub fn print_refs(lua: &State) -> Result<()> {
+    unsafe {
+        let _g = StackGuard::new(&lua);
+        let s = lua.state_ptr();
+
+        Value::String(String::from("Hello World")).push_to_lua_stack(&lua)?;
+        let r = lua::luaL_ref(s, lua::LUA_REGISTRYINDEX);
+        
+        for i in 0..r {
+            let ty = lua::lua_rawgeti(s, lua::LUA_REGISTRYINDEX, i as i64);
+            
+            match ty {
+                lua::LUA_TUSERDATA => {
+                    lua::lua_getmetatable(s, -1);
+                    lua::lua_pushliteral(s, "__name");
+                    lua::lua_rawget(s, -2);
+                    
+                    let v = Value::from_lua_stack(&lua)?;
+                    println!("{}: {}", i, v.as_string().unwrap());
+                },
+                _ => {
+                    let tyname = CStr::from_ptr(lua::lua_typename(s, ty));
+                    let tyname = tyname.to_str().unwrap();
+                    println!("{}: {}", i, tyname);
+                }
+            }
+            
+            lua::lua_pop(s, 1);
+        }
+        lua::luaL_unref(s, lua::LUA_REGISTRYINDEX, r);
+    }
+    Ok(())
+}
+
 /// A LuaRef is a a handle to something in Lua.
 ///
 /// These are created with `luaL_ref`, they can be created from anything on the top of the stack
 /// with [`LuaRef::from_stack`]. The references are automatically freed with `luaL_unref` when
 /// the inner Arc detects a single strong count.
 #[derive(Clone)]
-pub struct LuaRef<'a>(Arc<i32>, &'a State);
+pub struct LuaRef<'a> {
+    lref: Arc<i32>,
+    state: &'a State,
+}
 
 impl<'a> Drop for LuaRef<'a> {
     fn drop(&mut self) {
         unsafe {
-            let s = self.1.state_ptr();
+            let s = self.state.state_ptr();
 
-            if Arc::strong_count(&self.0) == 1 {
+            if Arc::strong_count(&self.lref) == 1 {
-                lua::luaL_unref(s, lua::LUA_REGISTRYINDEX, *self.0);
+                lua::luaL_unref(s, lua::LUA_REGISTRYINDEX, *self.lref);
             }
         }
     }
@@ -131,13 +156,15 @@ impl<'a> Drop for LuaRef<'a> {
 
 impl<'a> LuaRef<'a> {
     pub fn new(lua_ref: i32, state: &'a State) -> Self {
-        Self(Arc::new(lua_ref), state)
+        Self {
+            lref: Arc::new(lua_ref),
+            state,
+        }
     }
 
     /// Creates a reference to what is at the top of the stack.
     pub unsafe fn from_stack(state: &'a State) -> Result<Self> {
         let s = state.state_ptr();
-
         let r = lua::luaL_ref(s, lua::LUA_REGISTRYINDEX);
         if r == lua::LUA_REFNIL {
             Err(Error::Nil)
@@ -153,7 +180,7 @@ impl<'a> PushToLuaStack<'a> for LuaRef<'a> {
 
         state.ensure_stack(1)?;
         let top = lua::lua_gettop(s);
-        let ty = lua::lua_rawgeti(s, lua::LUA_REGISTRYINDEX, *self.0 as i64);
+        let ty = lua::lua_rawgeti(s, lua::LUA_REGISTRYINDEX, *self.lref as i64);
         let new_top = lua::lua_gettop(s);
 
         if ty == lua::LUA_TNIL || ty == lua::LUA_TNONE || top == new_top {
@@ -205,9 +232,10 @@ macro_rules! impl_as_lua_number {
             fn from_lua(_lua: &State, val: Value) -> crate::Result<Self> {
                 match val {
                     Value::Number(v) => Ok(v as $ty),
-                    _ => {
+                    _ => Err(Error::UnexpectedType(
-                        Err(Error::UnexpectedType("Number".to_string(), val.type_name().to_string()))
+                        "Number".to_string(),
-                    },
+                        val.type_name().to_string(),
+                    )),
                 }
             }
         }
@@ -218,7 +246,6 @@ macro_rules! impl_as_lua_number {
                 v.push_to_lua_stack(state)
             }
         }
-        
     };
 }
 
@@ -292,14 +319,9 @@ impl Userdata for Vec2 {
         builder
             .field_getter("x", |_, this| Ok(this.x))
             .field_getter("y", |_, this| Ok(this.y))
-            
             .field_setter("x", |_, this, x: f32| this.x = x)
             .field_setter("y", |_, this, y: f32| this.y = y)
-
+            .function("new", |lua, (x, y)| lua.create_userdata(Vec2 { x, y }))
-            .function("new", |lua, (x, y)| {
-                lua.create_userdata(Vec2 { x, y, })
-            })
-
             // method test
             .method("add", |lua, lhs: &Vec2, (rhs,): (Ref<Vec3>,)| {
                 let lx = lhs.x;
@@ -308,9 +330,11 @@ impl Userdata for Vec2 {
                 let rx = rhs.x;
                 let ry = rhs.y;
 
-                lua.create_userdata(Vec2 { x: lx + rx, y: ly + ry, })
+                lua.create_userdata(Vec2 {
+                    x: lx + rx,
+                    y: ly + ry,
+                })
             })
-
             .meta_method(MetaMethod::Add, |lua, lhs: &Vec2, (rhs,): (Ref<Vec2>,)| {
                 let lx = lhs.x;
                 let ly = lhs.y;
@@ -318,7 +342,10 @@ impl Userdata for Vec2 {
                 let rx = rhs.x;
                 let ry = rhs.y;
 
-                lua.create_userdata(Vec2 { x: lx + rx, y: ly + ry, })
+                lua.create_userdata(Vec2 {
+                    x: lx + rx,
+                    y: ly + ry,
+                })
             });
 
         Ok(())
@@ -360,8 +387,48 @@ impl<T: Userdata> Userdata for UserdataProxy<T> {
     }
 }
 
+pub enum Borrow<'a, T> {
+    Wrapped(Ref<'a, T>),
+    Raw(&'a T),
+}
+
+impl<'a, T> Deref for Borrow<'a, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Borrow::Wrapped(w) => w,
+            Borrow::Raw(w) => w,
+        }
+    }
+}
+
 pub struct UserdataRef<'a, T: Userdata> {
-    ud_ptr: Ref<'a, T>
+    ud_ptr: Borrow<'a, T>,
+}
+
+impl<'a, T: Userdata> From<&'a T> for UserdataRef<'static, T> {
+    fn from(value: &'a T) -> Self {
+        let ud = Borrow::Raw(value);
+
+        Self {
+            ud_ptr: unsafe {
+                mem::transmute(ud)
+            }
+        }
+    }
+}
+
+impl<'a, T: Userdata> From<Ref<'a, T>> for UserdataRef<'static, T> {
+    fn from(value: Ref<'a, T>) -> Self {
+        let ud = Borrow::Wrapped(value);
+
+        Self {
+            ud_ptr: unsafe {
+                mem::transmute(ud)
+            },
+        }
+    }
 }
 
 impl<'a, T: Userdata + 'static> Userdata for UserdataRef<'a, T> {

src/state.rs

@@ -1,6 +1,7 @@
 use core::ffi;
 use std::{alloc::{self, Layout}, any::TypeId, cell::RefCell, collections::HashMap, ffi::{CStr, CString}, mem, ptr::{self, NonNull}, str::Utf8Error, sync::Arc};
 
+use lua::lua_gc;
 use mlua_sys as lua;
 
 use crate::{lua_error_guard, AnyUserdata, AsLua, Chunk, Error, FromLua, FromLuaStack, FromLuaVec, Function, IntoChunkData, IntoLuaVec, LuaRef, MetaMethod, PushToLuaStack, Result, StackGuard, Table, Userdata, UserdataBuilder, Value, ValueVec};
@@ -88,8 +89,8 @@ struct ClosureData<'a> {
 }
 
 #[derive(Default)]
-struct ExtraSpace<'a> {
+pub struct ExtraSpace<'a> {
-    userdata_metatables: HashMap<TypeId, LuaRef<'a>>,
+    pub userdata_metatables: HashMap<TypeId, LuaRef<'a>>,
 }
 
 pub struct State {
@@ -101,10 +102,17 @@ impl Drop for State {
         unsafe { 
             let extra = self.get_extra_space_ptr();
         
+            {
+                // clear the refs to anything in lua before we close it and
+                // attempt to drop extra after
+                let extra = &mut *extra;
+                extra.userdata_metatables.clear();
+            }
+
             lua::lua_close(self.lua.as_ptr()); 
             extra.drop_in_place();
 
-            // must be dealloced since it wasn't memory created from lua (like userdata)
+            // must be dealloced since it wasn't memory created from lua (i.e. userdata)
             alloc::dealloc(extra.cast(), Layout::new::<ExtraSpace>());
         }
     }
@@ -152,7 +160,7 @@ impl State {
         }
     }
 
-    fn get_extra_space(&self) -> &mut ExtraSpace {
+    pub fn get_extra_space(&self) -> &mut ExtraSpace {
         unsafe {
             self.get_extra_space_ptr().as_mut()
                 .expect("Somehow the Lua state's extra data got deleted!")
@@ -496,23 +504,16 @@ impl State {
 
         // if a Gc metamethod was not manually defined, set it
         if !mt.raw_has_key(MetaMethod::Gc)? {
-            let gc_func = self.create_function(move |lua: &State, ud: AnyUserdata| {
+            let gc_func = self.create_function(|_lua, ud: AnyUserdata| {
                 unsafe {
-                    // if this panics, there's a weird bug.
+                    let ud_ptr = ud.as_ptr_unchecked::<T>()?;
-                    let ud_ptr = ud.as_ptr_unchecked::<T>().unwrap();
                     ud_ptr.drop_in_place();
-
-                    lua::luaL_unref(lua.state_ptr(), lua::LUA_REGISTRYINDEX, *ud.unsafe_ud.lref.0);
-
-                    let extra = lua.get_extra_space();
-                    extra.userdata_metatables.remove(&TypeId::of::<T>());
-                    //todo!()
                 }
 
                 Ok(())
             })?;
 
-            mt.set_meta(MetaMethod::Gc, gc_func)?;
+            mt.set(MetaMethod::Gc, gc_func)?;
         }
 
         let extra = self.get_extra_space();
@@ -577,4 +578,35 @@ impl State {
             Ok(s)
         }
     }
+
+    /// Returns a boolean indiciating if the Lua garbage collector is running.
+    pub fn is_gc_running(&self) -> Result<bool> {
+        unsafe {
+            let s = self.state_ptr();
+
+            Ok(lua_gc(s, lua::LUA_GCISRUNNING) != 0)
+        }
+    }
+
+    /// Triggers a full garbage-collection cycle.
+    pub fn gc_collect(&self) -> Result<()> {
+        unsafe {
+            let s = self.state_ptr();
+
+            lua_gc(s, lua::LUA_GCCOLLECT);
+        }
+
+        Ok(())
+    }
+
+    /// Disable the garbage collector
+    pub fn gc_stop(&self) -> Result<()> {
+        unsafe {
+            let s = self.state_ptr();
+
+            lua_gc(s, lua::LUA_GCSTOP);
+        }
+
+        Ok(())
+    }
 }

src/table.rs

@@ -8,7 +8,7 @@ pub struct Table<'a> {
     state: &'a State,
     pub(crate) lref: LuaRef<'a>,
     /// a boolean indicating if this Table is a Metatable
-    mt_marker: bool,
+    pub(crate) mt_marker: bool,
 }
 
 impl<'a> Table<'a> {

src/userdata/mod.rs

@@ -388,18 +388,22 @@ impl<'a> AnyUserdata<'a> {
     }
 
     /// Returns the metatable of this userdata
-    pub fn get_metatable(&self) -> crate::Result<Table<'a>> {
+    pub fn get_metatable(&'a self) -> crate::Result<Table<'a>> {
         unsafe {
             self.state.ensure_stack(2)?;
             let _g = StackGuard::new(self.state);
 
-            self.unsafe_ud.lref.push_to_lua_stack(self.state)?;
+            //self.unsafe_ud.lref.push_to_lua_stack(self.state)?;
+            self.push_to_lua_stack(self.state)?;
             let s = self.state.state_ptr();
 
             if lua::lua_getmetatable(s, -1) == 0 {
                 Err(crate::Error::MissingMetatable)
             } else {
-                Ok(Table::from_lua_stack(self.state)?)
+                let mut t = Table::from_lua_stack(self.state)?;
+                t.mt_marker = true;
+
+                Ok(t)
             }
         }
     }
@@ -429,12 +433,11 @@ impl<'a> AnyUserdata<'a> {
 
     /// Execute a method on this userdata. This finds a function with `name` and executes it
     /// with the userdata as the first argument.
-    pub fn execute_method<A, R>(&self, name: &str, args: A) -> crate::Result<R>
+    pub fn execute_method<A, R>(&'a self, name: &str, args: A) -> crate::Result<R>
     where
         A: IntoLuaVec<'a>,
         R: FromLua<'a>,
     {
-        //let s: &Self = unsafe { mem::transmute(self) };
         let name = name.to_string();
         let mt = self.get_metatable()?;
         let f = mt.get::<_, Function>(name)?;
@@ -444,11 +447,6 @@ impl<'a> AnyUserdata<'a> {
 
         f.exec(args)
     }
-
-    /// Trigger garbage collection on the userdata.
-    pub fn garbage_collect(self) -> crate::Result<()> {
-        self.execute_method("__gc", ()).into()
-    }
 }
 
 impl<'a> FromLuaStack<'a> for AnyUserdata<'a> {