From 63493e78de43a525a6531a572a653ac7b678b5bd Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Wed, 5 Jun 2024 18:33:21 +0200
Subject: [PATCH] Get sandbox working

---
 crates/lune-std-task/src/lib.rs |  6 -----
 crates/lune/src/rt/runtime.rs   | 47 ++++++++++++++++++++++++++++++---
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/crates/lune-std-task/src/lib.rs b/crates/lune-std-task/src/lib.rs
index 47a78d5..dce0873 100644
--- a/crates/lune-std-task/src/lib.rs
+++ b/crates/lune-std-task/src/lib.rs
@@ -33,12 +33,6 @@ pub fn module(lua: &Lua) -> LuaResult<LuaTable> {
         .set_environment(task_delay_env)
         .into_function()?;
 
-    // Overwrite resume & wrap functions on the coroutine global
-    // with ones that are compatible with our scheduler
-    let co = lua.globals().get::<_, LuaTable>("coroutine")?;
-    co.set("resume", fns.resume.clone())?;
-    co.set("wrap", fns.wrap.clone())?;
-
     TableBuilder::new(lua)?
         .with_value("cancel", fns.cancel)?
         .with_value("defer", fns.defer)?
diff --git a/crates/lune/src/rt/runtime.rs b/crates/lune/src/rt/runtime.rs
index a5501c6..a302f53 100644
--- a/crates/lune/src/rt/runtime.rs
+++ b/crates/lune/src/rt/runtime.rs
@@ -9,8 +9,8 @@ use std::{
     },
 };
 
-use mlua::Lua;
-use mlua_luau_scheduler::Scheduler;
+use mlua::prelude::*;
+use mlua_luau_scheduler::{Functions, Scheduler};
 
 use super::{RuntimeError, RuntimeResult};
 
@@ -94,6 +94,12 @@ impl Runtime {
             eprintln!("{}", RuntimeError::from(e));
         });
 
+        // Overwrite resume & wrap functions on the coroutine global
+        // with ones that are compatible with our scheduler
+        // We also sandbox the VM, preventing further modifications
+        // to the global environment, and enabling optimizations
+        inject_scheduler_functions_and_sandbox(&self.lua)?;
+
         // Load our "main" thread
         let main = self
             .lua
@@ -105,12 +111,45 @@ impl Runtime {
         sched.run().await;
 
         // Return the exit code - default to FAILURE if we got any errors
-        Ok(sched.get_exit_code().unwrap_or({
+        let exit_code = sched.get_exit_code().unwrap_or({
             if got_any_error.load(Ordering::SeqCst) {
                 ExitCode::FAILURE
             } else {
                 ExitCode::SUCCESS
             }
-        }))
+        });
+
+        Ok(exit_code)
     }
 }
+
+fn inject_scheduler_functions_and_sandbox(lua: &Lua) -> LuaResult<()> {
+    let fns = Functions::new(lua)?;
+
+    let co = lua.globals().get::<_, LuaTable>("coroutine")?;
+    co.set("resume", fns.resume.clone())?;
+    co.set("wrap", fns.wrap.clone())?;
+
+    lua.sandbox(true)?;
+
+    // NOTE: We need to create the _G table after
+    // sandboxing, otherwise it will be read-only
+    #[cfg(any(
+        feature = "std-datetime",
+        feature = "std-fs",
+        feature = "std-luau",
+        feature = "std-net",
+        feature = "std-process",
+        feature = "std-regex",
+        feature = "std-roblox",
+        feature = "std-serde",
+        feature = "std-stdio",
+        feature = "std-task",
+    ))]
+    {
+        let g_global = lune_std::LuneStandardGlobal::GTable;
+        lua.globals().set(g_global.name(), g_global.create(lua)?)?;
+    }
+
+    Ok(())
+}