From 2a9664a90fecd242d8f3a46e6d97e0d57d39c0a5 Mon Sep 17 00:00:00 2001
From: qwreey <me@qwreey.moe>
Date: Thu, 31 Oct 2024 06:32:19 +0000
Subject: [PATCH] Add runtime flag for allowing unsafe libraries (#243)

---
 crates/lune-std/src/lib.rs            |  2 ++
 crates/lune-std/src/library.rs        | 32 ++++++++++++++++++++++++++-
 crates/lune-std/src/unsafe_library.rs | 22 ++++++++++++++++++
 crates/lune/src/cli/run.rs            |  7 +++++-
 crates/lune/src/rt/runtime.rs         | 11 +++++++++
 5 files changed, 72 insertions(+), 2 deletions(-)
 create mode 100644 crates/lune-std/src/unsafe_library.rs

diff --git a/crates/lune-std/src/lib.rs b/crates/lune-std/src/lib.rs
index a29bef0..9808912 100644
--- a/crates/lune-std/src/lib.rs
+++ b/crates/lune-std/src/lib.rs
@@ -6,10 +6,12 @@ mod global;
 mod globals;
 mod library;
 mod luaurc;
+mod unsafe_library;
 
 pub use self::global::LuneStandardGlobal;
 pub use self::globals::version::set_global_version;
 pub use self::library::LuneStandardLibrary;
+pub use self::unsafe_library::{get_unsafe_library_enabled, set_unsafe_library_enabled};
 
 /**
     Injects all standard globals into the given Lua state / VM.
diff --git a/crates/lune-std/src/library.rs b/crates/lune-std/src/library.rs
index 2ac783e..80e2b5c 100644
--- a/crates/lune-std/src/library.rs
+++ b/crates/lune-std/src/library.rs
@@ -2,6 +2,8 @@ use std::str::FromStr;
 
 use mlua::prelude::*;
 
+use crate::get_unsafe_library_enabled;
+
 /**
     A standard library probloxrovided by Lune.
 */
@@ -64,16 +66,44 @@ impl LuneStandardLibrary {
         }
     }
 
+    /**
+        Gets whether the library is unsafe.
+    */
+    #[must_use]
+    #[rustfmt::skip]
+    #[allow(unreachable_patterns)]
+    pub fn is_unsafe(&self) -> bool {
+        match self {
+            #[cfg(feature = "datetime")] Self::DateTime => false,
+            #[cfg(feature = "fs")]       Self::Fs       => false,
+            #[cfg(feature = "luau")]     Self::Luau     => false,
+            #[cfg(feature = "net")]      Self::Net      => false,
+            #[cfg(feature = "task")]     Self::Task     => false,
+            #[cfg(feature = "process")]  Self::Process  => false,
+            #[cfg(feature = "regex")]    Self::Regex    => false,
+            #[cfg(feature = "serde")]    Self::Serde    => false,
+            #[cfg(feature = "stdio")]    Self::Stdio    => false,
+            #[cfg(feature = "roblox")]   Self::Roblox   => false,
+            #[cfg(feature = "ffi")]      Self::Ffi      => true,
+
+            _ => unreachable!("no standard library enabled"),
+        }
+    }
+
     /**
         Creates the Lua module for the library.
 
         # Errors
 
-        If the library could not be created.
+        If the library could not be created, or if requiring an unsafe library without enabling the unsafe library.
     */
     #[rustfmt::skip]
     #[allow(unreachable_patterns)]
     pub fn module<'lua>(&self, lua: &'lua Lua) -> LuaResult<LuaMultiValue<'lua>> {
+        if self.is_unsafe() && !get_unsafe_library_enabled(lua) {
+            return Err(LuaError::external(format!("Standard library '{}' requires unsafe library enabled", self.name())));
+        }
+
         let res: LuaResult<LuaTable> = match self {
             #[cfg(feature = "datetime")] Self::DateTime => lune_std_datetime::module(lua),
             #[cfg(feature = "fs")]       Self::Fs       => lune_std_fs::module(lua),
diff --git a/crates/lune-std/src/unsafe_library.rs b/crates/lune-std/src/unsafe_library.rs
new file mode 100644
index 0000000..3c0cbf4
--- /dev/null
+++ b/crates/lune-std/src/unsafe_library.rs
@@ -0,0 +1,22 @@
+use mlua::prelude::*;
+
+struct UnsafeLibrary(bool);
+
+/**
+    Override unsafe library allowance
+*/
+pub fn set_unsafe_library_enabled(lua: &Lua, enabled: bool) {
+    lua.set_app_data(UnsafeLibrary(enabled));
+}
+
+/**
+    Returns whether unsafe libraries are allowed
+
+    # Panics
+
+    Panic if `UnsafeLib` app data doesn't exist.
+*/
+#[must_use]
+pub fn get_unsafe_library_enabled(lua: &Lua) -> bool {
+    lua.app_data_ref::<UnsafeLibrary>().unwrap().0
+}
diff --git a/crates/lune/src/cli/run.rs b/crates/lune/src/cli/run.rs
index 6267ed7..19eab0b 100644
--- a/crates/lune/src/cli/run.rs
+++ b/crates/lune/src/cli/run.rs
@@ -14,6 +14,9 @@ use super::utils::files::{discover_script_path_including_lune_dirs, strip_sheban
 /// Run a script
 #[derive(Debug, Clone, Parser)]
 pub struct RunCommand {
+    /// Allow unsafe libraries
+    #[clap(long, action)]
+    r#unsafe: bool,
     /// Script name or full path to the file to run
     script_path: String,
     /// Arguments to pass to the script, stored in process.args
@@ -41,7 +44,9 @@ impl RunCommand {
         };
 
         // Create a new lune runtime with all globals & run the script
-        let mut rt = Runtime::new().with_args(self.script_args);
+        let mut rt = Runtime::new()
+            .with_args(self.script_args)
+            .set_unsafe_lib_enabled(self.r#unsafe);
 
         let result = rt
             .run(&script_display_name, strip_shebang(script_contents))
diff --git a/crates/lune/src/rt/runtime.rs b/crates/lune/src/rt/runtime.rs
index 31e5b03..b6028dd 100644
--- a/crates/lune/src/rt/runtime.rs
+++ b/crates/lune/src/rt/runtime.rs
@@ -54,6 +54,7 @@ impl RuntimeInner {
                 feature = "std-serde",
                 feature = "std-stdio",
                 feature = "std-task",
+                feature = "std-ffi",
             ))]
             {
                 lune_std::set_global_version(lua, env!("CARGO_PKG_VERSION"));
@@ -76,6 +77,7 @@ impl RuntimeInner {
                 feature = "std-serde",
                 feature = "std-stdio",
                 feature = "std-task",
+                feature = "std-ffi",
             ))]
             {
                 let g_table = lune_std::LuneStandardGlobal::GTable;
@@ -130,6 +132,15 @@ impl Runtime {
         self
     }
 
+    /**
+        Sets arguments to give in `process.args` for Lune scripts.
+    */
+    #[must_use]
+    pub fn set_unsafe_lib_enabled(self, enabled: bool) -> Self {
+        lune_std::set_unsafe_library_enabled(self.inner.lua(), enabled);
+        self
+    }
+
     /**
         Runs a Lune script inside of the current runtime.