luau 0.732.0

Safe lifetime-bound Rust embedding API for the Luau runtime
use luau_vm::Thread as VmThread;
use luau_vm::VmResult;
use luau_vm::thread::{LUA_GLOBALS_INDEX, StackGuard};
use luau_vm::types::LUA_TTABLE;

use super::Lua;
use crate::error::Error;
use crate::lua::runtime::RuntimeData;
use crate::thread::Thread;

#[derive(Clone, Copy)]
pub(crate) enum ChunkLoad {
    Main,
    Dynamic,
    FreshSandbox,
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub(super) enum LuaSandboxState {
    Disabled,
    Enabled {
        original_globals_ref: i32,
        has_loaded: bool,
    },
}

impl Lua {
    /// Enables or disables the main thread's writable sandbox environment.
    ///
    /// The first successfully loaded chunk can retain Luau's safe-environment
    /// optimizations. Loading additional chunks or mutating reachable values
    /// through the safe API disables those optimizations so cached imports
    /// cannot become stale. The full lookup graph reachable through the
    /// immutable base, including nested tables and host state observed through
    /// callbacks or metamethods, must remain stable while the environment is
    /// safe. Mutations performed by Luau code or outside the safe API cannot be
    /// detected; disable the optimization on the sandbox or function
    /// environment—normally [`Lua::globals`]—with
    /// [`crate::Table::set_safe_env`] before they can occur.
    ///
    /// Disabling the sandbox restores the globals captured when it was enabled,
    /// discards writes made in the writable proxy, and clears the readonly and
    /// safe-environment flags applied to globals, top-level library tables, and
    /// the string metatable. Pre-existing flags are not snapshotted.
    ///
    /// Use [`crate::Chunk::into_sandboxed`] to give each script an isolated
    /// environment.
    pub fn sandbox(&self, enabled: bool) -> Result<(), Error> {
        let result = if enabled {
            unsafe { self.enable_sandbox() }
        } else {
            unsafe { self.disable_sandbox() }
        };
        result.map_err(|exit| Error::from_thread_exit(self.state.main_thread(), exit))
    }

    unsafe fn enable_sandbox(&self) -> VmResult {
        if matches!(
            self.runtime.sandbox_state(),
            LuaSandboxState::Enabled { .. }
        ) {
            return Ok(());
        }

        let thread = self.state.main_thread();
        let _stack = unsafe { StackGuard::new(thread) };
        unsafe {
            thread.push_value(LUA_GLOBALS_INDEX)?;
        }
        let original_globals_ref = unsafe { thread.ref_value(-1)? };

        let result = (|| unsafe {
            thread.sandbox()?;
            thread.sandbox_thread()
        })();
        if let Err(error) = result {
            unsafe {
                if thread.get_ref(original_globals_ref).is_ok() {
                    thread.replace(LUA_GLOBALS_INDEX);
                }
                let _ = Self::clear_sandbox_flags(thread);
                thread.unref_value(original_globals_ref);
            }
            return Err(error.into());
        }

        self.runtime.set_sandbox_state(LuaSandboxState::Enabled {
            original_globals_ref,
            has_loaded: false,
        });
        Ok(())
    }

    unsafe fn disable_sandbox(&self) -> VmResult {
        let LuaSandboxState::Enabled {
            original_globals_ref,
            has_loaded: _,
        } = self.runtime.sandbox_state()
        else {
            return Ok(());
        };

        let thread = self.state.main_thread();
        unsafe {
            thread.set_safe_env(LUA_GLOBALS_INDEX, 0);
            thread.get_ref(original_globals_ref)?;
            thread.replace(LUA_GLOBALS_INDEX);
            Self::clear_sandbox_flags(thread)?;
            thread.unref_value(original_globals_ref);
        }

        self.runtime.set_sandbox_state(LuaSandboxState::Disabled);
        Ok(())
    }

    unsafe fn clear_sandbox_flags(thread: &VmThread) -> VmResult {
        let _stack = unsafe { StackGuard::new(thread) };

        unsafe {
            thread.push_nil()?;
            while thread.next(LUA_GLOBALS_INDEX)? != 0 {
                if thread.type_of(-1) == LUA_TTABLE {
                    thread.set_safe_env(-1, 0);
                    thread.set_readonly(-1, 0);
                }

                thread.pop(1);
            }

            thread.set_safe_env(LUA_GLOBALS_INDEX, 0);
            thread.set_readonly(LUA_GLOBALS_INDEX, 0);

            thread.push_string("")?;
            if thread.get_metatable(-1)? != 0 {
                thread.set_safe_env(-1, 0);
                thread.set_readonly(-1, 0);
                thread.pop(2);
            } else {
                thread.pop(1);
            }
        }

        Ok(())
    }
}

impl Thread<'_> {
    /// Replaces this thread's globals with a writable isolated environment.
    ///
    /// Writes remain local while reads fall through to the thread's current
    /// globals. This is the safe-layer counterpart to Luau's
    /// `luaL_sandboxthread`. The new environment's safe-environment
    /// optimization remains disabled because its inherited globals may be
    /// mutable; embedders can opt in explicitly with
    /// [`crate::Table::set_safe_env`] when they can uphold that invariant.
    pub fn sandbox(&self) -> Result<(), Error> {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            let is_main_thread = self.runtime().is_main_thread(thread);
            if is_main_thread {
                thread.set_safe_env(LUA_GLOBALS_INDEX, 0);
            }
            thread
                .sandbox_thread()
                .map_err(|exit| Error::from_thread_exit(thread, exit))?;
            thread.set_safe_env(LUA_GLOBALS_INDEX, 0);
        }
        Ok(())
    }

    pub(crate) fn sandbox_with_immutable_base(&self) -> Result<(), Error> {
        let LuaSandboxState::Enabled {
            original_globals_ref,
            has_loaded: _,
        } = self.runtime().sandbox_state()
        else {
            return Err(Error::runtime(
                "Lua::sandbox(true) must be enabled before sandboxing a thread",
            ));
        };

        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);

            thread
                .push_value(LUA_GLOBALS_INDEX)
                .map_err(|error| Error::from_thread_exit(thread, error))?;
            let previous_globals = thread.get_top();

            thread
                .get_ref(original_globals_ref)
                .map_err(|error| Error::from_thread_exit(thread, error))?;
            thread.replace(LUA_GLOBALS_INDEX);
            if let Err(error) = thread.sandbox_thread() {
                thread
                    .push_value(previous_globals)
                    .map_err(|restore| Error::from_thread_exit(thread, restore))?;
                thread.replace(LUA_GLOBALS_INDEX);
                return Err(Error::from_thread_exit(thread, error));
            }
        }
        Ok(())
    }
}

impl ChunkLoad {
    pub(crate) unsafe fn before_load(
        self,
        runtime: &RuntimeData,
        thread: &VmThread,
        environment_index: i32,
        uses_current_environment: bool,
    ) {
        if !uses_current_environment {
            return;
        }

        match self {
            Self::Main => {
                let LuaSandboxState::Enabled {
                    has_loaded,
                    original_globals_ref: _,
                } = runtime.sandbox_state()
                else {
                    return;
                };

                if has_loaded {
                    unsafe {
                        thread.set_safe_env(environment_index, 0);
                    }
                }
            }
            Self::Dynamic => unsafe {
                thread.set_safe_env(environment_index, 0);
            },
            Self::FreshSandbox => {}
        }
    }

    pub(crate) fn loaded(self, runtime: &RuntimeData, uses_current_environment: bool) {
        if !uses_current_environment {
            return;
        }

        match self {
            Self::Main => {
                let LuaSandboxState::Enabled {
                    original_globals_ref,
                    has_loaded: false,
                } = runtime.sandbox_state()
                else {
                    return;
                };
                runtime.set_sandbox_state(LuaSandboxState::Enabled {
                    original_globals_ref,
                    has_loaded: true,
                });
            }
            Self::Dynamic | Self::FreshSandbox => {}
        }
    }
}