Serialize only known (registered) userdata.

khvzak · khvzak · commit 0c7db4916c76 · 2021-04-16T22:01:55.000+01:00
This reverts commit 7332c6a. Non-static userdata is a special case and can cause segfault if try to serialize it. Now it should be safe, plus I added non-static userdata destructor to generate better error messages in case of accessing destructed userdata.
diff --git a/src/lua.rs b/src/lua.rs
@@ -1,6 +1,6 @@
 use std::any::TypeId;
 use std::cell::{RefCell, UnsafeCell};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::ffi::CString;
 use std::marker::PhantomData;
 use std::os::raw::{c_char, c_int, c_void};
@@ -58,6 +58,7 @@ pub struct Lua {
 // Data associated with the lua_State.
 struct ExtraData {
     registered_userdata: HashMap<TypeId, c_int>,
+    registered_userdata_mt: HashSet<isize>,
     registry_unref_list: Arc<Mutex<Option<Vec<c_int>>>>,
 
     libs: StdLib,
@@ -322,6 +323,7 @@ impl Lua {
 
         let extra = Arc::new(Mutex::new(ExtraData {
             registered_userdata: HashMap::new(),
+            registered_userdata_mt: HashSet::new(),
             registry_unref_list: Arc::new(Mutex::new(Some(Vec::new()))),
             ref_thread,
             libs: StdLib::NONE,
@@ -1569,34 +1571,52 @@ impl Lua {
             ffi::lua_pop(self.state, 1);
         }
 
+        let ptr = ffi::lua_topointer(self.state, -1);
         let id = protect_lua_closure(self.state, 1, 0, |state| {
             ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX)
         })?;
 
         let mut extra = mlua_expect!(self.extra.lock(), "extra is poisoned");
         extra.registered_userdata.insert(type_id, id);
+        extra.registered_userdata_mt.insert(ptr as isize);
 
         Ok(id)
     }
 
-    // Pushes a LuaRef value onto the stack, checking that it's not destructed
+    pub(crate) fn register_userdata_metatable(&self, id: isize) {
+        let mut extra = mlua_expect!(self.extra.lock(), "extra is poisoned");
+        extra.registered_userdata_mt.insert(id);
+    }
+
+    pub(crate) fn deregister_userdata_metatable(&self, id: isize) {
+        let mut extra = mlua_expect!(self.extra.lock(), "extra is poisoned");
+        extra.registered_userdata_mt.remove(&id);
+    }
+
+    // Pushes a LuaRef value onto the stack, checking that it's a registered
+    // and not destructed UserData.
     // Uses 2 stack spaces, does not call checkstack
     #[cfg(feature = "serialize")]
     pub(crate) unsafe fn push_userdata_ref(&self, lref: &LuaRef) -> Result<()> {
         self.push_ref(lref);
         if ffi::lua_getmetatable(self.state, -1) == 0 {
-            Err(Error::UserDataTypeMismatch)
-        } else {
-            // Check that userdata is not destructed
-            get_destructed_userdata_metatable(self.state);
-            let eq = ffi::lua_rawequal(self.state, -1, -2) == 1;
+            return Err(Error::UserDataTypeMismatch);
+        }
+        // Check that userdata is registered
+        let ptr = ffi::lua_topointer(self.state, -1);
+        let extra = mlua_expect!(self.extra.lock(), "extra is poisoned");
+        if extra.registered_userdata_mt.contains(&(ptr as isize)) {
+            ffi::lua_pop(self.state, 1);
+            return Ok(());
+        }
+        // Maybe userdata was destructed?
+        get_destructed_userdata_metatable(self.state);
+        if ffi::lua_rawequal(self.state, -1, -2) != 0 {
             ffi::lua_pop(self.state, 2);
-            if eq {
-                Err(Error::UserDataDestructed)
-            } else {
-                Ok(())
-            }
+            return Err(Error::UserDataDestructed);
         }
+        ffi::lua_pop(self.state, 2);
+        Err(Error::UserDataTypeMismatch)
     }
 
     // Creates a Function out of a Callback containing a 'static Fn.  This is safe ONLY because the
diff --git a/src/scope.rs b/src/scope.rs
@@ -312,7 +312,8 @@ impl<'lua, 'scope> Scope<'lua, 'scope> {
             let _sg = StackGuard::new(lua.state);
             assert_stack(lua.state, 6);
 
-            push_userdata(lua.state, ())?;
+            // We need to wrap dummy userdata because their memory can be accessed by serializer
+            push_userdata(lua.state, UserDataCell::new(UserDataWrapped::new(())))?;
             #[cfg(any(feature = "lua54", feature = "lua53"))]
             ffi::lua_pushlightuserdata(lua.state, data.as_ptr() as *mut c_void);
             #[cfg(any(feature = "lua52", feature = "lua51", feature = "luajit"))]
@@ -355,9 +356,22 @@ impl<'lua, 'scope> Scope<'lua, 'scope> {
                 ffi::lua_pop(lua.state, 1);
             }
 
+            let mt_id = ffi::lua_topointer(lua.state, -1);
             ffi::lua_setmetatable(lua.state, -2);
 
-            Ok(AnyUserData(lua.pop_ref()))
+            let ud = AnyUserData(lua.pop_ref());
+            lua.register_userdata_metatable(mt_id as isize);
+            self.destructors.borrow_mut().push((ud.0.clone(), |ud| {
+                let state = ud.lua.state;
+                assert_stack(state, 2);
+                ud.lua.push_ref(&ud);
+                ffi::lua_getmetatable(state, -1);
+                let mt_id = ffi::lua_topointer(state, -1);
+                ffi::lua_pop(state, 1);
+                ud.lua.deregister_userdata_metatable(mt_id as isize);
+                vec![Box::new(take_userdata::<UserDataCell<()>>(state))]
+            }));
+            Ok(ud)
         }
     }
 
diff --git a/tests/scope.rs b/tests/scope.rs
@@ -13,7 +13,9 @@ extern "system" {}
 use std::cell::Cell;
 use std::rc::Rc;
 
-use mlua::{Error, Function, Lua, MetaMethod, Result, String, UserData, UserDataMethods};
+use mlua::{
+    AnyUserData, Error, Function, Lua, MetaMethod, Result, String, UserData, UserDataMethods,
+};
 
 #[test]
 fn scope_func() -> Result<()> {
@@ -57,17 +59,62 @@ fn scope_drop() -> Result<()> {
 
     lua.scope(|scope| {
         lua.globals()
-            .set("test", scope.create_userdata(MyUserdata(rc.clone()))?)?;
+            .set("static_ud", scope.create_userdata(MyUserdata(rc.clone()))?)?;
         assert_eq!(Rc::strong_count(&rc), 2);
         Ok(())
     })?;
     assert_eq!(Rc::strong_count(&rc), 1);
 
-    match lua.load("test:method()").exec() {
-        Err(Error::CallbackError { .. }) => {}
+    match lua.load("static_ud:method()").exec() {
+        Err(Error::CallbackError { ref cause, .. }) => match cause.as_ref() {
+            Error::CallbackDestructed => {}
+            e => panic!("expected CallbackDestructed, got {:?}", e),
+        },
         r => panic!("improper return for destructed userdata: {:?}", r),
     };
 
+    let static_ud = lua.globals().get::<_, AnyUserData>("static_ud")?;
+    match static_ud.borrow::<MyUserdata>() {
+        Ok(_) => panic!("borrowed destructed userdata"),
+        Err(Error::UserDataDestructed) => {}
+        Err(e) => panic!("expected UserDataDestructed, got {:?}", e),
+    }
+
+    // Check non-static UserData drop
+    struct MyUserDataRef<'a>(&'a Cell<i64>);
+
+    impl<'a> UserData for MyUserDataRef<'a> {
+        fn add_methods<'lua, M: UserDataMethods<'lua, Self>>(methods: &mut M) {
+            methods.add_method("inc", |_, data, ()| {
+                data.0.set(data.0.get() + 1);
+                Ok(())
+            });
+        }
+    }
+
+    let i = Cell::new(1);
+    lua.scope(|scope| {
+        lua.globals().set(
+            "nonstatic_ud",
+            scope.create_nonstatic_userdata(MyUserDataRef(&i))?,
+        )
+    })?;
+
+    match lua.load("nonstatic_ud:inc(1)").exec() {
+        Err(Error::CallbackError { ref cause, .. }) => match cause.as_ref() {
+            Error::CallbackDestructed => {}
+            e => panic!("expected CallbackDestructed, got {:?}", e),
+        },
+        r => panic!("improper return for destructed userdata: {:?}", r),
+    };
+
+    let nonstatic_ud = lua.globals().get::<_, AnyUserData>("nonstatic_ud")?;
+    match nonstatic_ud.borrow::<MyUserDataRef>() {
+        Ok(_) => panic!("borrowed destructed userdata"),
+        Err(Error::UserDataDestructed) => {}
+        Err(e) => panic!("expected UserDataDestructed, got {:?}", e),
+    }
+
     Ok(())
 }
 
diff --git a/tests/serde.rs b/tests/serde.rs
@@ -104,6 +104,19 @@ fn test_serialize_in_scope() -> LuaResult<()> {
         Err(e) => panic!("expected destructed error, got {}", e),
     }
 
+    struct MyUserDataRef<'a>(&'a ());
+
+    impl<'a> UserData for MyUserDataRef<'a> {}
+
+    lua.scope(|scope| {
+        let ud = scope.create_nonstatic_userdata(MyUserDataRef(&()))?;
+        match serde_json::to_value(&ud) {
+            Ok(v) => panic!("expected serialization error, got {}", v),
+            Err(serde_json::Error { .. }) => {}
+        };
+        Ok(())
+    })?;
+
     Ok(())
 }
 
