ACTUALLY expose `RegistryKey` API

Also fixes a safety issue with RegistryKey, where you could use RegistryKeys
with mismatching Lua instances.

Cargo.toml

@@ -16,9 +16,11 @@ travis-ci = { repository = "chucklefish/rlua", branch = "master" }
 default = ["builtin-lua"]
 # Builds the correct version of Lua 5.3 inside the crate. If you want to link a
 # specialized version of lua into your binary, you can disable this feature to
-# do that, but care must be taken.  `rlua` expects the linked lua to define
-# LUA_INTEGER as long long, and LUA_NUMBER as double, and may make other
-# assumptions about how lua is built.
+# do that, but care must be taken.  `rlua` makes at least the following
+# assumptions about the linked lua library:
+#   * LUA_INTEGER is long long
+#   * LUA_NUMBER as double
+#   * LUA_EXTRASPACE is sizeof(void*)
 builtin-lua = ["gcc"]
 
 [dependencies]

src/conversion.rs

@@ -270,14 +270,16 @@ impl<'lua, T: FromLua<'lua>> FromLua<'lua> for Vec<T> {
 }
 
 impl<'lua, K: Eq + Hash + ToLua<'lua>, V: ToLua<'lua>, S: BuildHasher> ToLua<'lua>
-    for HashMap<K, V, S> {
+    for HashMap<K, V, S>
+{
     fn to_lua(self, lua: &'lua Lua) -> Result<Value<'lua>> {
         Ok(Value::Table(lua.create_table_from(self)?))
     }
 }
 
 impl<'lua, K: Eq + Hash + FromLua<'lua>, V: FromLua<'lua>, S: BuildHasher + Default> FromLua<'lua>
-    for HashMap<K, V, S> {
+    for HashMap<K, V, S>
+{
     fn from_lua(value: Value<'lua>, _: &'lua Lua) -> Result<Self> {
         if let Value::Table(table) = value {
             table.pairs().collect()

src/ffi.rs

@@ -3,24 +3,19 @@
 #![allow(unused)]
 
 use std::ptr;
+use std::mem;
 use std::os::raw::{c_char, c_double, c_int, c_longlong, c_void};
 
 pub type lua_Integer = c_longlong;
 pub type lua_Number = c_double;
 
 pub enum lua_State {}
-pub type lua_Alloc = unsafe extern "C" fn(
-    ud: *mut c_void,
-    ptr: *mut c_void,
-    osize: usize,
-    nsize: usize,
-) -> *mut c_void;
+pub type lua_Alloc =
+    unsafe extern "C" fn(ud: *mut c_void, ptr: *mut c_void, osize: usize, nsize: usize)
+        -> *mut c_void;
 pub type lua_KContext = *mut c_void;
-pub type lua_KFunction = unsafe extern "C" fn(
-    state: *mut lua_State,
-    status: c_int,
-    ctx: lua_KContext,
-) -> c_int;
+pub type lua_KFunction =
+    unsafe extern "C" fn(state: *mut lua_State, status: c_int, ctx: lua_KContext) -> c_int;
 pub type lua_CFunction = unsafe extern "C" fn(state: *mut lua_State) -> c_int;
 
 pub const LUA_OK: c_int = 0;
@@ -177,6 +172,10 @@ extern "C" {
     pub fn luaL_len(push_state: *mut lua_State, index: c_int) -> lua_Integer;
 }
 
+pub unsafe fn lua_getextraspace(state: *mut lua_State) -> *mut c_void {
+    (state as *mut c_void).offset(-(mem::size_of::<*mut c_void>() as isize))
+}
+
 pub unsafe fn lua_pop(state: *mut lua_State, n: c_int) {
     lua_settop(state, -n - 1);
 }

src/lib.rs

@@ -64,7 +64,7 @@ mod userdata;
 mod tests;
 
 pub use error::{Error, ExternalError, ExternalResult, Result};
-pub use types::{Integer, LightUserData, Number};
+pub use types::{Integer, LightUserData, Number, RegistryKey};
 pub use multi::Variadic;
 pub use string::String;
 pub use table::{Table, TablePairs, TableSequence};

src/lua.rs

@@ -14,7 +14,7 @@ use ffi;
 use error::*;
 use util::*;
 use value::{FromLua, FromLuaMulti, MultiValue, Nil, ToLua, ToLuaMulti, Value};
-use types::{Callback, Integer, LightUserData, LuaRef, Number, RegistryKey};
+use types::{Callback, Integer, LightUserData, LuaRef, Number, RegistryKey, SharedId};
 use string::String;
 use table::Table;
 use function::Function;
@@ -28,6 +28,12 @@ pub struct Lua {
     ephemeral: bool,
 }
 
+// Data associated with the main lua_State via lua_getextraspace.
+struct ExtraData {
+    lua_id: SharedId,
+    registered_userdata: HashMap<TypeId, c_int>,
+}
+
 impl Drop for Lua {
     fn drop(&mut self) {
         unsafe {
@@ -40,6 +46,9 @@ impl Drop for Lua {
                     }
                 }
 
+                let extra_data = *(ffi::lua_getextraspace(self.state) as *mut *mut ExtraData);
+                Box::from_raw(extra_data);
+
                 ffi::lua_close(self.state);
             }
         }
@@ -496,7 +505,10 @@ impl Lua {
 
                 ffi::lua_pop(self.state, 1);
 
-                Ok(RegistryKey(id))
+                Ok(RegistryKey {
+                    lua_id: (*self.extra()).lua_id.clone(),
+                    registry_id: id,
+                })
             })
         }
     }
@@ -507,6 +519,12 @@ impl Lua {
     /// value previously placed by `create_registry_value`.
     pub fn registry_value<'lua, T: FromLua<'lua>>(&'lua self, key: &RegistryKey) -> Result<T> {
         unsafe {
+            lua_assert!(
+                self.state,
+                key.lua_id == (*self.extra()).lua_id,
+                "Lua instance passed RegistryKey created from a different Lua"
+            );
+
             stack_err_guard(self.state, 0, || {
                 check_stack(self.state, 2);
 
@@ -516,7 +534,7 @@ impl Lua {
                 );
                 ffi::lua_rawget(self.state, ffi::LUA_REGISTRYINDEX);
 
-                ffi::lua_rawgeti(self.state, -1, key.0 as ffi::lua_Integer);
+                ffi::lua_rawgeti(self.state, -1, key.registry_id as ffi::lua_Integer);
 
                 let t = T::from_lua(self.pop_value(self.state), self)?;
 
@@ -533,6 +551,12 @@ impl Lua {
     /// `create_registry_value`
     pub fn remove_registry_value<'lua>(&'lua self, key: RegistryKey) {
         unsafe {
+            lua_assert!(
+                self.state,
+                key.lua_id == (*self.extra()).lua_id,
+                "Lua instance passed RegistryKey created from a different Lua"
+            );
+
             stack_guard(self.state, 0, || {
                 check_stack(self.state, 2);
 
@@ -543,7 +567,7 @@ impl Lua {
                 ffi::lua_rawget(self.state, ffi::LUA_REGISTRYINDEX);
 
                 ffi::lua_pushnil(self.state);
-                ffi::lua_rawseti(self.state, -2, key.0 as ffi::lua_Integer);
+                ffi::lua_rawseti(self.state, -2, key.registry_id as ffi::lua_Integer);
 
                 ffi::lua_pop(self.state, 1);
             })
@@ -707,15 +731,7 @@ impl Lua {
         stack_err_guard(self.state, 0, move || {
             check_stack(self.state, 5);
 
-            ffi::lua_pushlightuserdata(
-                self.state,
-                &LUA_USERDATA_REGISTRY_KEY as *const u8 as *mut c_void,
-            );
-            ffi::lua_gettable(self.state, ffi::LUA_REGISTRYINDEX);
-            let registered_userdata = get_userdata::<HashMap<TypeId, c_int>>(self.state, -1)?;
-            ffi::lua_pop(self.state, 1);
-
-            if let Some(table_id) = (*registered_userdata).get(&TypeId::of::<T>()) {
+            if let Some(table_id) = (*self.extra()).registered_userdata.get(&TypeId::of::<T>()) {
                 return Ok(*table_id);
             }
 
@@ -822,7 +838,9 @@ impl Lua {
             let id = gc_guard(self.state, || {
                 ffi::luaL_ref(self.state, ffi::LUA_REGISTRYINDEX)
             });
-            (*registered_userdata).insert(TypeId::of::<T>(), id);
+            (*self.extra())
+                .registered_userdata
+                .insert(TypeId::of::<T>(), id);
             Ok(id)
         })
     }
@@ -875,25 +893,6 @@ impl Lua {
                 ffi::lua_pop(state, 1);
             }
 
-            // Create the userdata registry table
-
-            ffi::lua_pushlightuserdata(
-                state,
-                &LUA_USERDATA_REGISTRY_KEY as *const u8 as *mut c_void,
-            );
-
-            push_userdata::<HashMap<TypeId, c_int>>(state, HashMap::new()).unwrap();
-
-            ffi::lua_newtable(state);
-
-            push_string(state, "__gc").unwrap();
-            ffi::lua_pushcfunction(state, userdata_destructor::<HashMap<TypeId, c_int>>);
-            ffi::lua_rawset(state, -3);
-
-            ffi::lua_setmetatable(state, -2);
-
-            ffi::lua_rawset(state, ffi::LUA_REGISTRYINDEX);
-
             // Create the function metatable
 
             ffi::lua_pushlightuserdata(
@@ -932,6 +931,14 @@ impl Lua {
             ffi::lua_rawset(state, -3);
 
             ffi::lua_pop(state, 1);
+
+            // Create ExtraData, and place it in the lua_State "extra space"
+
+            let extra_data = Box::into_raw(Box::new(ExtraData {
+                lua_id: SharedId::new(),
+                registered_userdata: HashMap::new(),
+            }));
+            *(ffi::lua_getextraspace(state) as *mut *mut ExtraData) = extra_data;
         });
 
         Lua {
@@ -996,8 +1003,11 @@ impl Lua {
             })
         }
     }
+
+    unsafe fn extra(&self) -> *mut ExtraData {
+        *(ffi::lua_getextraspace(self.main_state) as *mut *mut ExtraData)
+    }
 }
 
-static LUA_USERDATA_REGISTRY_KEY: u8 = 0;
 static FUNCTION_METATABLE_REGISTRY_KEY: u8 = 0;
 static USER_REGISTRY_KEY: u8 = 0;

src/prelude.rs

@@ -4,7 +4,8 @@ pub use {AnyUserData as LuaAnyUserData, Error as LuaError, ExternalError as LuaE
          ExternalResult as LuaExternalResult, FromLua, FromLuaMulti, Function as LuaFunction,
          Integer as LuaInteger, LightUserData as LuaLightUserData, Lua,
          MetaMethod as LuaMetaMethod, MultiValue as LuaMultiValue, Nil as LuaNil,
-         Number as LuaNumber, Result as LuaResult, String as LuaString, Table as LuaTable,
-         TablePairs as LuaTablePairs, TableSequence as LuaTableSequence, Thread as LuaThread,
-         ThreadStatus as LuaThreadStatus, ToLua, ToLuaMulti, UserData as LuaUserData,
-         UserDataMethods as LuaUserDataMethods, Value as LuaValue};
+         Number as LuaNumber, RegistryKey as LuaRegistryKey, Result as LuaResult,
+         String as LuaString, Table as LuaTable, TablePairs as LuaTablePairs,
+         TableSequence as LuaTableSequence, Thread as LuaThread, ThreadStatus as LuaThreadStatus,
+         ToLua, ToLuaMulti, UserData as LuaUserData, UserDataMethods as LuaUserDataMethods,
+         Value as LuaValue};

src/tests.rs

@@ -536,6 +536,28 @@ fn test_registry_value() {
     f.call::<_, ()>(()).unwrap();
 }
 
+#[test]
+#[should_panic]
+fn test_mismatched_lua_ref() {
+    let lua1 = Lua::new();
+    let lua2 = Lua::new();
+
+    let s = lua1.create_string("hello").unwrap();
+    let f = lua2.create_function(|_, _: String| Ok(())).unwrap();
+
+    f.call::<_, ()>(s).unwrap();
+}
+
+#[test]
+#[should_panic]
+fn test_mismatched_registry_key() {
+    let lua1 = Lua::new();
+    let lua2 = Lua::new();
+
+    let r = lua1.create_registry_value("hello").unwrap();
+    lua2.remove_registry_value(r);
+}
+
 // TODO: Need to use compiletest-rs or similar to make sure these don't compile.
 /*
 #[test]

src/types.rs

@@ -1,5 +1,6 @@
 use std::fmt;
 use std::os::raw::{c_int, c_void};
+use std::rc::Rc;
 
 use ffi;
 use error::Result;
@@ -15,12 +16,32 @@ pub type Number = ffi::lua_Number;
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub struct LightUserData(pub *mut c_void);
 
-/// An auto generated key into the Lua registry.
-pub struct RegistryKey(pub(crate) c_int);
+// Clone-able id where every clone of the same id compares equal, and are guaranteed unique.
+#[derive(Debug, Clone)]
+pub(crate) struct SharedId(Rc<()>);
 
-pub(crate) type Callback<'lua> = Box<
-    FnMut(&'lua Lua, MultiValue<'lua>) -> Result<MultiValue<'lua>> + 'lua,
->;
+impl SharedId {
+    pub(crate) fn new() -> SharedId {
+        SharedId(Rc::new(()))
+    }
+}
+
+impl PartialEq for SharedId {
+    fn eq(&self, other: &SharedId) -> bool {
+        Rc::ptr_eq(&self.0, &other.0)
+    }
+}
+
+impl Eq for SharedId {}
+
+/// An auto generated key into the Lua registry.
+pub struct RegistryKey {
+    pub(crate) lua_id: SharedId,
+    pub(crate) registry_id: c_int,
+}
+
+pub(crate) type Callback<'lua> =
+    Box<FnMut(&'lua Lua, MultiValue<'lua>) -> Result<MultiValue<'lua>> + 'lua>;
 
 pub(crate) struct LuaRef<'lua> {
     pub lua: &'lua Lua,