|
1 | 1 | use std::any::TypeId; |
2 | 2 | use std::cell::{RefCell, UnsafeCell}; |
3 | | -use std::collections::HashMap; |
| 3 | +use std::collections::{HashMap, HashSet}; |
4 | 4 | use std::ffi::CString; |
5 | 5 | use std::marker::PhantomData; |
6 | 6 | use std::os::raw::{c_char, c_int, c_void}; |
@@ -58,6 +58,7 @@ pub struct Lua { |
58 | 58 | // Data associated with the lua_State. |
59 | 59 | struct ExtraData { |
60 | 60 | registered_userdata: HashMap<TypeId, c_int>, |
| 61 | + registered_userdata_mt: HashSet<isize>, |
61 | 62 | registry_unref_list: Arc<Mutex<Option<Vec<c_int>>>>, |
62 | 63 |
|
63 | 64 | libs: StdLib, |
@@ -322,6 +323,7 @@ impl Lua { |
322 | 323 |
|
323 | 324 | let extra = Arc::new(Mutex::new(ExtraData { |
324 | 325 | registered_userdata: HashMap::new(), |
| 326 | + registered_userdata_mt: HashSet::new(), |
325 | 327 | registry_unref_list: Arc::new(Mutex::new(Some(Vec::new()))), |
326 | 328 | ref_thread, |
327 | 329 | libs: StdLib::NONE, |
@@ -1569,34 +1571,52 @@ impl Lua { |
1569 | 1571 | ffi::lua_pop(self.state, 1); |
1570 | 1572 | } |
1571 | 1573 |
|
| 1574 | + let ptr = ffi::lua_topointer(self.state, -1); |
1572 | 1575 | let id = protect_lua_closure(self.state, 1, 0, |state| { |
1573 | 1576 | ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX) |
1574 | 1577 | })?; |
1575 | 1578 |
|
1576 | 1579 | let mut extra = mlua_expect!(self.extra.lock(), "extra is poisoned"); |
1577 | 1580 | extra.registered_userdata.insert(type_id, id); |
| 1581 | + extra.registered_userdata_mt.insert(ptr as isize); |
1578 | 1582 |
|
1579 | 1583 | Ok(id) |
1580 | 1584 | } |
1581 | 1585 |
|
1582 | | - // Pushes a LuaRef value onto the stack, checking that it's not destructed |
| 1586 | + pub(crate) fn register_userdata_metatable(&self, id: isize) { |
| 1587 | + let mut extra = mlua_expect!(self.extra.lock(), "extra is poisoned"); |
| 1588 | + extra.registered_userdata_mt.insert(id); |
| 1589 | + } |
| 1590 | + |
| 1591 | + pub(crate) fn deregister_userdata_metatable(&self, id: isize) { |
| 1592 | + let mut extra = mlua_expect!(self.extra.lock(), "extra is poisoned"); |
| 1593 | + extra.registered_userdata_mt.remove(&id); |
| 1594 | + } |
| 1595 | + |
| 1596 | + // Pushes a LuaRef value onto the stack, checking that it's a registered |
| 1597 | + // and not destructed UserData. |
1583 | 1598 | // Uses 2 stack spaces, does not call checkstack |
1584 | 1599 | #[cfg(feature = "serialize")] |
1585 | 1600 | pub(crate) unsafe fn push_userdata_ref(&self, lref: &LuaRef) -> Result<()> { |
1586 | 1601 | self.push_ref(lref); |
1587 | 1602 | if ffi::lua_getmetatable(self.state, -1) == 0 { |
1588 | | - Err(Error::UserDataTypeMismatch) |
1589 | | - } else { |
1590 | | - // Check that userdata is not destructed |
1591 | | - get_destructed_userdata_metatable(self.state); |
1592 | | - let eq = ffi::lua_rawequal(self.state, -1, -2) == 1; |
| 1603 | + return Err(Error::UserDataTypeMismatch); |
| 1604 | + } |
| 1605 | + // Check that userdata is registered |
| 1606 | + let ptr = ffi::lua_topointer(self.state, -1); |
| 1607 | + let extra = mlua_expect!(self.extra.lock(), "extra is poisoned"); |
| 1608 | + if extra.registered_userdata_mt.contains(&(ptr as isize)) { |
| 1609 | + ffi::lua_pop(self.state, 1); |
| 1610 | + return Ok(()); |
| 1611 | + } |
| 1612 | + // Maybe userdata was destructed? |
| 1613 | + get_destructed_userdata_metatable(self.state); |
| 1614 | + if ffi::lua_rawequal(self.state, -1, -2) != 0 { |
1593 | 1615 | ffi::lua_pop(self.state, 2); |
1594 | | - if eq { |
1595 | | - Err(Error::UserDataDestructed) |
1596 | | - } else { |
1597 | | - Ok(()) |
1598 | | - } |
| 1616 | + return Err(Error::UserDataDestructed); |
1599 | 1617 | } |
| 1618 | + ffi::lua_pop(self.state, 2); |
| 1619 | + Err(Error::UserDataTypeMismatch) |
1600 | 1620 | } |
1601 | 1621 |
|
1602 | 1622 | // Creates a Function out of a Callback containing a 'static Fn. This is safe ONLY because the |
|
0 commit comments