From 41a1a0d15af476dcc82272e86b12ab333ad5d4c4 Mon Sep 17 00:00:00 2001
From: Alex Orlenko <zxteam@protonmail.com>
Date: Sun, 11 Apr 2021 00:07:04 +0100
Subject: [PATCH] Add more checks for destructed userdata in AnyUserData

---
 src/lua.rs      | 11 +++++------
 src/userdata.rs | 12 ++++++------
 tests/scope.rs  | 17 +++++++++++++++++
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/src/lua.rs b/src/lua.rs
index 3d8fc09..111c3a8 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -25,10 +25,10 @@ use crate::userdata::{
     AnyUserData, MetaMethod, UserData, UserDataFields, UserDataMethods, UserDataWrapped,
 };
 use crate::util::{
-    assert_stack, callback_error, check_stack, get_gc_userdata, get_main_state, get_userdata,
-    get_wrapped_error, init_error_registry, init_gc_metatable_for, init_userdata_metatable,
-    pop_error, push_gc_userdata, push_userdata, push_wrapped_error, StackGuard, WrappedError,
-    WrappedPanic,
+    assert_stack, callback_error, check_stack, get_destructed_userdata_metatable, get_gc_userdata,
+    get_main_state, get_userdata, get_wrapped_error, init_error_registry, init_gc_metatable_for,
+    init_userdata_metatable, pop_error, push_gc_userdata, push_userdata, push_wrapped_error,
+    StackGuard, WrappedError, WrappedPanic,
 };
 use crate::value::{FromLua, FromLuaMulti, MultiValue, Nil, ToLua, ToLuaMulti, Value};
 
@@ -44,7 +44,7 @@ use {
 };
 
 #[cfg(feature = "serialize")]
-use {crate::util::get_destructed_userdata_metatable, serde::Serialize};
+use serde::Serialize;
 
 /// Top level Lua struct which holds the Lua state itself.
 pub struct Lua {
@@ -1588,7 +1588,6 @@ impl Lua {
     // Pushes a LuaRef value onto the stack, checking that it's a registered
     // and not destructed UserData.
     // Uses 3 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 {
diff --git a/src/userdata.rs b/src/userdata.rs
index a2082e7..41cdd29 100644
--- a/src/userdata.rs
+++ b/src/userdata.rs
@@ -685,8 +685,8 @@ impl<'lua> AnyUserData<'lua> {
         let v = v.to_lua(lua)?;
         unsafe {
             let _sg = StackGuard::new(lua.state);
-            assert_stack(lua.state, 2);
-            lua.push_ref(&self.0);
+            assert_stack(lua.state, 3);
+            lua.push_userdata_ref(&self.0)?;
             lua.push_value(v)?;
             ffi::lua_setuservalue(lua.state, -2);
             Ok(())
@@ -702,8 +702,8 @@ impl<'lua> AnyUserData<'lua> {
         let lua = self.0.lua;
         let res = unsafe {
             let _sg = StackGuard::new(lua.state);
-            assert_stack(lua.state, 2);
-            lua.push_ref(&self.0);
+            assert_stack(lua.state, 3);
+            lua.push_userdata_ref(&self.0)?;
             ffi::lua_getuservalue(lua.state, -1);
             lua.pop_value()
         };
@@ -754,9 +754,9 @@ impl<'lua> AnyUserData<'lua> {
         unsafe {
             let lua = self.0.lua;
             let _sg = StackGuard::new(lua.state);
-            assert_stack(lua.state, 2);
+            assert_stack(lua.state, 3);
 
-            lua.push_ref(&self.0);
+            lua.push_userdata_ref(&self.0)?;
             if ffi::lua_getmetatable(lua.state, -1) == 0 {
                 return Err(Error::UserDataTypeMismatch);
             }
diff --git a/tests/scope.rs b/tests/scope.rs
index 7c29bff..92f150a 100644
--- a/tests/scope.rs
+++ b/tests/scope.rs
@@ -286,6 +286,15 @@ fn scope_userdata_drop() -> Result<()> {
         Err(err) => panic!("improper borrow error for destructed userdata: {:?}", err),
     }
 
+    match ud.get_metatable() {
+        Ok(_) => panic!("successful metatable retrieval of destructed userdata"),
+        Err(Error::UserDataDestructed) => {}
+        Err(err) => panic!(
+            "improper metatable error for destructed userdata: {:?}",
+            err
+        ),
+    }
+
     Ok(())
 }
 
@@ -336,6 +345,14 @@ fn scope_nonstatic_userdata_drop() -> Result<()> {
         Err(Error::UserDataDestructed) => {}
         Err(err) => panic!("improper borrow error for destructed userdata: {:?}", err),
     }
+    match ud.get_metatable() {
+        Ok(_) => panic!("successful metatable retrieval of destructed userdata"),
+        Err(Error::UserDataDestructed) => {}
+        Err(err) => panic!(
+            "improper metatable error for destructed userdata: {:?}",
+            err
+        ),
+    }
 
     Ok(())
 }