Sandboxing support

5 files changed, 171 insertions, 8 deletions

diff --git a/src/ffi/luau/lauxlib.rs b/src/ffi/luau/lauxlib.rs
index 669a08f..af7d1c2 100644
--- a/src/ffi/luau/lauxlib.rs
+++ b/src/ffi/luau/lauxlib.rs
@@ -72,6 +72,11 @@ extern "C" {
     // TODO: luaL_findtable
 
     pub fn luaL_typename(L: *mut lua_State, idx: c_int) -> *const c_char;
+
+    // sandbox libraries and globals
+    #[link_name = "luaL_sandbox"]
+    pub fn luaL_sandbox_(L: *mut lua_State);
+    pub fn luaL_sandboxthread(L: *mut lua_State);
 }
 
 //
@@ -123,6 +128,29 @@ pub unsafe fn luaL_unref(L: *mut lua_State, t: c_int, r#ref: c_int) {
     lua::lua_unref(L, r#ref)
 }
 
+pub unsafe fn luaL_sandbox(L: *mut lua_State, enabled: c_int) {
+    use super::lua::*;
+
+    // set all libraries to read-only
+    lua_pushnil(L);
+    while lua_next(L, LUA_GLOBALSINDEX) != 0 {
+        if lua_istable(L, -1) != 0 {
+            lua_setreadonly(L, -1, enabled);
+        }
+        lua_pop(L, 1);
+    }
+
+    // set all builtin metatables to read-only
+    lua_pushliteral(L, "");
+    lua_getmetatable(L, -1);
+    lua_setreadonly(L, -1, enabled);
+    lua_pop(L, 2);
+
+    // set globals to readonly and activate safeenv since the env is immutable
+    lua_setreadonly(L, LUA_GLOBALSINDEX, enabled);
+    lua_setsafeenv(L, LUA_GLOBALSINDEX, enabled);
+}
+
 //
 // TODO: Generic Buffer Manipulation
 //
diff --git a/src/ffi/luau/lualib.rs b/src/ffi/luau/lualib.rs
index 2f1ba7e..92b8642 100644
--- a/src/ffi/luau/lualib.rs
+++ b/src/ffi/luau/lualib.rs
@@ -26,8 +26,4 @@ extern "C" {
 
     // open all builtin libraries
     pub fn luaL_openlibs(L: *mut lua_State);
-
-    // sandbox libraries and globals
-    pub fn luaL_sandbox(L: *mut lua_State);
-    pub fn luaL_sandboxthread(L: *mut lua_State);
 }
diff --git a/src/lua.rs b/src/lua.rs
index 9687e7a..9bfbebb 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -108,6 +108,9 @@ struct ExtraData {
     hook_callback: Option<HookCallback>,
     #[cfg(feature = "lua54")]
     warn_callback: Option<WarnCallback>,
+
+    #[cfg(feature = "luau")]
+    sandboxed: bool,
 }
 
 #[cfg_attr(any(feature = "lua51", feature = "luajit"), allow(dead_code))]
@@ -548,6 +551,8 @@ impl Lua {
             hook_callback: None,
             #[cfg(feature = "lua54")]
             warn_callback: None,
+            #[cfg(feature = "luau")]
+            sandboxed: false,
         }));
 
         mlua_expect!(
@@ -750,6 +755,60 @@ impl Lua {
         self.entrypoint(move |lua, _: ()| func(lua))
     }
 
+    /// Enables (or disables) sandbox mode on this Lua instance.
+    ///
+    /// This method, in particular:
+    /// - Set all libraries to read-only
+    /// - Set all builtin metatables to read-only
+    /// - Set globals to read-only (and activates safeenv)
+    /// - Setup local environment table that performs writes locally and proxies reads
+    ///   to the global environment.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use mlua::{Lua, Result};
+    /// # fn main() -> Result<()> {
+    /// let lua = Lua::new();
+    ///
+    /// lua.sandbox(true)?;
+    /// lua.load("var = 123").exec()?;
+    /// assert_eq!(lua.globals().get::<_, u32>("var")?, 123);
+    ///
+    /// // Restore the global environment (clear changes made in sandbox)
+    /// lua.sandbox(false)?;
+    /// assert_eq!(lua.globals().get::<_, Option<u32>>("var")?, None);
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
+    /// Requires `feature = "luau"`
+    #[cfg(feature = "luau")]
+    pub fn sandbox(&self, enabled: bool) -> Result<()> {
+        unsafe {
+            let extra = &mut *self.extra.get();
+            if extra.sandboxed != enabled {
+                let state = self.main_state.ok_or(Error::MainThreadNotAvailable)?;
+                check_stack(state, 3)?;
+                protect_lua!(state, 0, 0, |state| {
+                    if enabled {
+                        ffi::luaL_sandbox(state, 1);
+                        ffi::luaL_sandboxthread(state);
+                    } else {
+                        // Restore original `LUA_GLOBALSINDEX`
+                        self.ref_thread_exec(|ref_thread| {
+                            ffi::lua_xpush(ref_thread, state, ffi::LUA_GLOBALSINDEX);
+                            ffi::lua_replace(state, ffi::LUA_GLOBALSINDEX);
+                        });
+                        ffi::luaL_sandbox(state, 0);
+                    }
+                })?;
+                extra.sandboxed = enabled;
+            }
+            Ok(())
+        }
+    }
+
     /// Sets a 'hook' function that will periodically be called as Lua code executes.
     ///
     /// When exactly the hook function is called depends on the contents of the `triggers`
@@ -1424,7 +1483,11 @@ impl Lua {
         &'lua self,
         func: Function<'lua>,
     ) -> Result<Thread<'lua>> {
-        #[cfg(any(feature = "lua54", all(feature = "luajit", feature = "vendored")))]
+        #[cfg(any(
+            feature = "lua54",
+            all(feature = "luajit", feature = "vendored"),
+            feature = "luau",
+        ))]
         unsafe {
             let _sg = StackGuard::new(self.state);
             check_stack(self.state, 1)?;
@@ -1434,6 +1497,14 @@ impl Lua {
                 let thread_state = ffi::lua_tothread(extra.ref_thread, index);
                 self.push_ref(&func.0);
                 ffi::lua_xmove(self.state, thread_state, 1);
+
+                #[cfg(feature = "luau")]
+                {
+                    // Inherit `LUA_GLOBALSINDEX` from the caller
+                    ffi::lua_xpush(self.state, thread_state, ffi::LUA_GLOBALSINDEX);
+                    ffi::lua_replace(thread_state, ffi::LUA_GLOBALSINDEX);
+                }
+
                 return Ok(Thread(LuaRef { lua: self, index }));
             }
         };
@@ -1442,7 +1513,11 @@ impl Lua {
 
     /// Resets thread (coroutine) and returns to the cache for later use.
     #[cfg(feature = "async")]
-    #[cfg(any(feature = "lua54", all(feature = "luajit", feature = "vendored")))]
+    #[cfg(any(
+        feature = "lua54",
+        all(feature = "luajit", feature = "vendored"),
+        feature = "luau",
+    ))]
     pub(crate) unsafe fn recycle_thread(&self, thread: &mut Thread) {
         let extra = &mut *self.extra.get();
         let thread_state = ffi::lua_tothread(extra.ref_thread, thread.0.index);
diff --git a/src/table.rs b/src/table.rs
index 4be837c..cb1e85c 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -356,7 +356,13 @@ impl<'lua> Table<'lua> {
     pub fn set_readonly(&self, enabled: bool) {
         let lua = self.0.lua;
         unsafe {
-            lua.ref_thread_exec(|refthr| ffi::lua_setreadonly(refthr, self.0.index, enabled as _));
+            lua.ref_thread_exec(|refthr| {
+                ffi::lua_setreadonly(refthr, self.0.index, enabled as _);
+                if !enabled {
+                    // Reset "safeenv" flag
+                    ffi::lua_setsafeenv(refthr, self.0.index, 0);
+                }
+            });
         }
     }
 
diff --git a/src/thread.rs b/src/thread.rs
index 0be0664..d710c92 100644
--- a/src/thread.rs
+++ b/src/thread.rs
@@ -214,6 +214,13 @@ impl<'lua> Thread<'lua> {
             lua.push_ref(&func.0);
             ffi::lua_xmove(lua.state, thread_state, 1);
 
+            #[cfg(feature = "luau")]
+            {
+                // Inherit `LUA_GLOBALSINDEX` from the caller
+                ffi::lua_xpush(lua.state, thread_state, ffi::LUA_GLOBALSINDEX);
+                ffi::lua_replace(thread_state, ffi::LUA_GLOBALSINDEX);
+            }
+
             Ok(())
         }
     }
@@ -278,6 +285,53 @@ impl<'lua> Thread<'lua> {
             recycle: false,
         }
     }
+
+    /// Enables sandbox mode on this thread.
+    ///
+    /// Under the hood replaces the global environment table with a new table,
+    /// that performs writes locally and proxies reads to caller's global environment.
+    ///
+    /// This mode ideally should be used together with the global sandbox mode [`Lua::sandbox()`].
+    ///
+    /// Please note that Luau links environment table with chunk when loading it into Lua state.
+    /// Therefore you need to load chunks into a thread to link with the thread environment.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use mlua::{Lua, Result};
+    /// # fn main() -> Result<()> {
+    /// let lua = Lua::new();
+    /// let thread = lua.create_thread(lua.create_function(|lua2, ()| {
+    ///     lua2.load("var = 123").exec()?;
+    ///     assert_eq!(lua2.globals().get::<_, u32>("var")?, 123);
+    ///     Ok(())
+    /// })?)?;
+    /// thread.sandbox()?;
+    /// thread.resume(())?;
+    ///
+    /// // The global environment should be unchanged
+    /// assert_eq!(lua.globals().get::<_, Option<u32>>("var")?, None);
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
+    /// Requires `feature = "luau"`
+    #[cfg(any(feature = "luau", doc))]
+    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
+    #[doc(hidden)]
+    pub fn sandbox(&self) -> Result<()> {
+        let lua = self.0.lua;
+        unsafe {
+            let thread = lua.ref_thread_exec(|t| ffi::lua_tothread(t, self.0.index));
+            check_stack(thread, 1)?;
+            check_stack(lua.state, 3)?;
+            // Inherit `LUA_GLOBALSINDEX` from the caller
+            ffi::lua_xpush(lua.state, thread, ffi::LUA_GLOBALSINDEX);
+            ffi::lua_replace(thread, ffi::LUA_GLOBALSINDEX);
+            protect_lua!(lua.state, 0, 0, |_| ffi::luaL_sandboxthread(thread))
+        }
+    }
 }
 
 impl<'lua> PartialEq for Thread<'lua> {
@@ -295,7 +349,11 @@ impl<'lua, R> AsyncThread<'lua, R> {
 }
 
 #[cfg(feature = "async")]
-#[cfg(any(feature = "lua54", all(feature = "luajit", feature = "vendored")))]
+#[cfg(any(
+    feature = "lua54",
+    all(feature = "luajit", feature = "vendored"),
+    feature = "luau",
+))]
 impl<'lua, R> Drop for AsyncThread<'lua, R> {
     fn drop(&mut self) {
         if self.recycle {