From f27c49f931c07b26513284324cf20b2941b6db4b Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Mon, 7 Nov 2022 00:10:57 +0000 Subject: [PATCH] Fix bug when recycled Registry slot can be set to Nil. This can result in allocating the same slot twice and rewriting old value. Lua uses (registry) table length to find next free slot and having Nil in the middle of the table can impact length calculation. With this fix we ensure that Nil values uses a special LUA_REFNIL slot. --- src/lua.rs | 55 +++++++++++++++++++++++++++++++++----------------- src/types.rs | 37 +++++++++++++++++++++++++++++---- tests/tests.rs | 33 ++++++++++++++++++++++++++++++ 3 files changed, 103 insertions(+), 22 deletions(-) diff --git a/src/lua.rs b/src/lua.rs index 38d7c42..d1001de 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -88,6 +88,8 @@ pub(crate) struct ExtraData { registered_userdata: FxHashMap, registered_userdata_mt: FxHashMap<*const c_void, Option>, + + // When Lua instance dropped, setting `None` would prevent collecting `RegistryKey`s registry_unref_list: Arc>>>, #[cfg(not(feature = "send"))] @@ -2060,6 +2062,12 @@ impl Lua { /// [`RegistryKey`]: crate::RegistryKey pub fn create_registry_value<'lua, T: ToLua<'lua>>(&'lua self, t: T) -> Result { let t = t.to_lua(self)?; + if t == Value::Nil { + // Special case to skip calling `luaL_ref` and use `LUA_REFNIL` instead + let unref_list = unsafe { (*self.extra.get()).registry_unref_list.clone() }; + return Ok(RegistryKey::new(ffi::LUA_REFNIL, unref_list)); + } + unsafe { let _sg = StackGuard::new(self.state); check_stack(self.state, 4)?; @@ -2067,27 +2075,20 @@ impl Lua { let unref_list = (*self.extra.get()).registry_unref_list.clone(); self.push_value(t)?; - // Try to reuse previously allocated RegistryKey + // Try to reuse previously allocated slot let unref_list2 = unref_list.clone(); let mut unref_list2 = mlua_expect!(unref_list2.lock(), "unref list poisoned"); if let Some(registry_id) = unref_list2.as_mut().and_then(|x| x.pop()) { // It must be safe to replace the value without triggering memory error ffi::lua_rawseti(self.state, ffi::LUA_REGISTRYINDEX, registry_id as Integer); - return Ok(RegistryKey { - registry_id, - unref_list, - }); + return Ok(RegistryKey::new(registry_id, unref_list)); } // Allocate a new RegistryKey let registry_id = protect_lua!(self.state, 1, 0, |state| { ffi::luaL_ref(state, ffi::LUA_REGISTRYINDEX) })?; - - Ok(RegistryKey { - registry_id, - unref_list, - }) + Ok(RegistryKey::new(registry_id, unref_list)) } } @@ -2102,13 +2103,16 @@ impl Lua { return Err(Error::MismatchedRegistryKey); } - let value = unsafe { - let _sg = StackGuard::new(self.state); - check_stack(self.state, 1)?; + let value = match key.is_nil() { + true => Value::Nil, + false => unsafe { + let _sg = StackGuard::new(self.state); + check_stack(self.state, 1)?; - let id = key.registry_id as Integer; - ffi::lua_rawgeti(self.state, ffi::LUA_REGISTRYINDEX, id); - self.pop_value() + let id = key.registry_id as Integer; + ffi::lua_rawgeti(self.state, ffi::LUA_REGISTRYINDEX, id); + self.pop_value() + }, }; T::from_lua(value, self) } @@ -2147,13 +2151,28 @@ impl Lua { } let t = t.to_lua(self)?; + if t == Value::Nil && key.is_nil() { + // Nothing to replace + return Ok(()); + } else if t != Value::Nil && key.registry_id == ffi::LUA_REFNIL { + // We cannot update `LUA_REFNIL` slot + let err = "cannot replace nil value with non-nil".to_string(); + return Err(Error::RuntimeError(err)); + } + unsafe { let _sg = StackGuard::new(self.state); check_stack(self.state, 2)?; - self.push_value(t)?; - // It must be safe to replace the value without triggering memory error let id = key.registry_id as Integer; + if t == Value::Nil { + self.push_value(Value::Integer(id))?; + key.set_nil(true); + } else { + self.push_value(t)?; + key.set_nil(false); + } + // It must be safe to replace the value without triggering memory error ffi::lua_rawseti(self.state, ffi::LUA_REGISTRYINDEX, id); } Ok(()) diff --git a/src/types.rs b/src/types.rs index 7b1f653..9404d10 100644 --- a/src/types.rs +++ b/src/types.rs @@ -1,6 +1,7 @@ use std::cell::UnsafeCell; use std::hash::{Hash, Hasher}; use std::os::raw::{c_int, c_void}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::{fmt, mem, ptr}; @@ -104,6 +105,7 @@ pub(crate) struct DestructedUserdata; /// [`AnyUserData::get_user_value`]: crate::AnyUserData::get_user_value pub struct RegistryKey { pub(crate) registry_id: c_int, + pub(crate) is_nil: AtomicBool, pub(crate) unref_list: Arc>>>, } @@ -129,15 +131,27 @@ impl Eq for RegistryKey {} impl Drop for RegistryKey { fn drop(&mut self) { - let mut unref_list = mlua_expect!(self.unref_list.lock(), "unref list poisoned"); - if let Some(list) = unref_list.as_mut() { - list.push(self.registry_id); + // We don't need to collect nil slot + if self.registry_id > ffi::LUA_REFNIL { + let mut unref_list = mlua_expect!(self.unref_list.lock(), "unref list poisoned"); + if let Some(list) = unref_list.as_mut() { + list.push(self.registry_id); + } } } } impl RegistryKey { - // Destroys the RegistryKey without adding to the drop list + // Creates a new instance of `RegistryKey` + pub(crate) const fn new(id: c_int, unref_list: Arc>>>) -> Self { + RegistryKey { + registry_id: id, + is_nil: AtomicBool::new(id == ffi::LUA_REFNIL), + unref_list, + } + } + + // Destroys the `RegistryKey` without adding to the unref list pub(crate) fn take(self) -> c_int { let registry_id = self.registry_id; unsafe { @@ -146,6 +160,21 @@ impl RegistryKey { } registry_id } + + // Returns true if this `RegistryKey` holds a nil value + #[inline(always)] + pub(crate) fn is_nil(&self) -> bool { + self.is_nil.load(Ordering::Relaxed) + } + + // Marks value of this `RegistryKey` as `Nil` + #[inline(always)] + pub(crate) fn set_nil(&self, enabled: bool) { + // We cannot replace previous value with nil in as this will break + // Lua mechanism to find free keys. + // Instead, we set a special flag to mark value as nil. + self.is_nil.store(enabled, Ordering::Relaxed); + } } pub(crate) struct LuaRef<'lua> { diff --git a/tests/tests.rs b/tests/tests.rs index 59fb28f..a71691e 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -793,6 +793,17 @@ fn test_replace_registry_value() -> Result<()> { let key = lua.create_registry_value::(42)?; lua.replace_registry_value(&key, "new value")?; assert_eq!(lua.registry_value::(&key)?, "new value"); + lua.replace_registry_value(&key, Value::Nil)?; + assert_eq!(lua.registry_value::(&key)?, Value::Nil); + lua.replace_registry_value(&key, 123)?; + assert_eq!(lua.registry_value::(&key)?, 123); + + // It should be impossible to replace (initial) nil value with non-nil + let key2 = lua.create_registry_value(Value::Nil)?; + match lua.replace_registry_value(&key2, "abc") { + Err(Error::RuntimeError(_)) => {} + r => panic!("expected RuntimeError, got {r:?}"), + } Ok(()) } @@ -844,6 +855,28 @@ fn test_mismatched_registry_key() -> Result<()> { Ok(()) } +#[test] +fn test_registry_value_reuse() -> Result<()> { + let lua = Lua::new(); + + let r1 = lua.create_registry_value("value1")?; + let r1_slot = format!("{r1:?}"); + drop(r1); + + // Previous slot must not be reused by nil value + let r2 = lua.create_registry_value(Value::Nil)?; + let r2_slot = format!("{r2:?}"); + assert_ne!(r1_slot, r2_slot); + drop(r2); + + // But should be reused by non-nil value + let r3 = lua.create_registry_value("value3")?; + let r3_slot = format!("{r3:?}"); + assert_eq!(r1_slot, r3_slot); + + Ok(()) +} + #[test] fn test_application_data() -> Result<()> { let lua = Lua::new();