Skip to main content

assay/lua/
mod.rs

1pub mod async_bridge;
2pub mod builtins;
3pub mod file_source;
4
5use anyhow::Result;
6use include_dir::{Dir, include_dir};
7use mlua::{Lua, LuaOptions, StdLib};
8
9static STDLIB_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/stdlib");
10
11/// Environment variable to override the global module search path.
12pub const MODULES_PATH_ENV: &str = "ASSAY_MODULES_PATH";
13
14/// Comma-separated list of additional globals to nil out at VM
15/// creation time (defense-in-depth knob for hardened deployments).
16/// Names support dotted paths into stdlib tables — e.g.
17/// `ASSAY_BLOCK_GLOBALS=dofile,os.execute,debug.getinfo`.
18pub const BLOCK_GLOBALS_ENV: &str = "ASSAY_BLOCK_GLOBALS";
19
20fn lua_err(e: mlua::Error) -> anyhow::Error {
21    anyhow::anyhow!("{e}")
22}
23
24pub fn create_vm(client: reqwest::Client) -> Result<Lua> {
25    create_vm_with_paths(client, None)
26}
27
28#[allow(dead_code)]
29pub fn create_vm_with_lib_path(client: reqwest::Client, lib_path: String) -> Result<Lua> {
30    create_vm_with_paths(client, Some(lib_path))
31}
32
33pub fn create_vm_with_paths(
34    client: reqwest::Client,
35    global_modules_path: Option<String>,
36) -> Result<Lua> {
37    let libs = StdLib::ALL_SAFE;
38    let lua = Lua::new_with(libs, LuaOptions::default()).map_err(lua_err)?;
39    lua.set_memory_limit(64 * 1024 * 1024).map_err(lua_err)?;
40    sandbox(&lua).map_err(lua_err)?;
41    register_fs_loader(&lua, global_modules_path).map_err(lua_err)?;
42    register_stdlib_loader(&lua).map_err(lua_err)?;
43    builtins::register_all(&lua, client).map_err(lua_err)?;
44    Ok(lua)
45}
46
47fn sandbox(lua: &Lua) -> mlua::Result<()> {
48    // Block bytecode-level escape hatches only. Source-level loaders
49    // (`load` / `loadfile` / `dofile`) stay available — operator scripts
50    // are trusted to compose themselves out of multiple files (seed +
51    // init bootstraps, shared helpers, etc.). `string.dump` stays
52    // blocked because it produces native bytecode that defeats the
53    // memory/CPU caps the runtime relies on.
54    let globals = lua.globals();
55    let string_lib: mlua::Table = globals.get("string")?;
56    string_lib.set("dump", mlua::Value::Nil)?;
57
58    if let Ok(extra) = std::env::var(BLOCK_GLOBALS_ENV) {
59        for raw in extra.split(',') {
60            let name = raw.trim();
61            if name.is_empty() {
62                continue;
63            }
64            nil_dotted_path(lua, name)?;
65        }
66    }
67
68    Ok(())
69}
70
71/// Resolve a dotted Lua path (e.g. `"os.execute"` or `"debug.getinfo"`)
72/// against globals and set the leaf to nil. A bare name (e.g.
73/// `"dofile"`) clears it from `_G`. Missing intermediate tables are
74/// silently skipped so a typo in `ASSAY_BLOCK_GLOBALS` doesn't fail
75/// VM creation.
76fn nil_dotted_path(lua: &Lua, path: &str) -> mlua::Result<()> {
77    let parts: Vec<&str> = path.split('.').filter(|s| !s.is_empty()).collect();
78    if parts.is_empty() {
79        return Ok(());
80    }
81    let mut current: mlua::Table = lua.globals();
82    for segment in &parts[..parts.len() - 1] {
83        let next: mlua::Value = current.get(*segment)?;
84        match next {
85            mlua::Value::Table(t) => current = t,
86            _ => return Ok(()),
87        }
88    }
89    current.set(parts[parts.len() - 1], mlua::Value::Nil)
90}
91
92fn register_stdlib_loader(lua: &Lua) -> mlua::Result<()> {
93    let package: mlua::Table = lua.globals().get("package")?;
94    let searchers: mlua::Table = package.get("searchers")?;
95
96    // Resolves `require("assay.ory.kratos")` -> "ory/kratos.lua" by replacing
97    // dots with slashes, matching standard Lua package loading convention.
98    // Tries "<path>.lua" first, then falls back to "<path>/init.lua" so
99    // both `stdlib/ory.lua` (flat convenience wrapper) and
100    // `stdlib/ory/kratos.lua` (nested submodule) resolve correctly.
101    let stdlib_searcher = lua.create_function(|lua, module_name: String| {
102        let rest = match module_name.strip_prefix("assay.") {
103            Some(r) => r,
104            None => {
105                return Ok(mlua::Value::String(
106                    lua.create_string(format!("not an assay.* module: {module_name}"))?,
107                ));
108            }
109        };
110
111        let base = rest.replace('.', "/");
112        let candidates = [format!("{base}.lua"), format!("{base}/init.lua")];
113
114        for path in &candidates {
115            if let Some(file) = STDLIB_DIR.get_file(path) {
116                let source = file.contents_utf8().ok_or_else(|| {
117                    mlua::Error::runtime(format!("stdlib {path}: invalid UTF-8"))
118                })?;
119                let loader = lua
120                    .load(source)
121                    .set_name(format!("@assay/{path}"))
122                    .into_function()?;
123                return Ok(mlua::Value::Function(loader));
124            }
125        }
126
127        Ok(mlua::Value::String(
128            lua.create_string(format!("no embedded stdlib file: {}", candidates[0]))?,
129        ))
130    })?;
131
132    let len = searchers.len()?;
133    searchers.set(len + 1, stdlib_searcher)?;
134
135    Ok(())
136}
137
138fn register_fs_loader(lua: &Lua, global_modules_path: Option<String>) -> mlua::Result<()> {
139    let package: mlua::Table = lua.globals().get("package")?;
140    let searchers: mlua::Table = package.get("searchers")?;
141
142    // Same dotted-path resolution as the stdlib loader: `assay.ory.kratos`
143    // -> "ory/kratos.lua", falling back to "ory/kratos/init.lua".
144    let fs_searcher = lua.create_function(move |lua, module_name: String| {
145        let rest = match module_name.strip_prefix("assay.") {
146            Some(r) => r,
147            None => {
148                return Ok(mlua::Value::String(
149                    lua.create_string(format!("not an assay.* module: {module_name}"))?,
150                ));
151            }
152        };
153        let base = rest.replace('.', "/");
154        let candidates = [format!("{base}.lua"), format!("{base}/init.lua")];
155
156        let try_load = |dir: &std::path::Path| -> Option<(std::path::PathBuf, String)> {
157            for rel in &candidates {
158                let full = dir.join(rel);
159                if let Ok(source) = std::fs::read_to_string(&full) {
160                    return Some((full, source));
161                }
162            }
163            None
164        };
165
166        // Priority 1: ./modules/<path>.lua (per-project)
167        if let Some((full, source)) = try_load(std::path::Path::new("./modules")) {
168            let loader = lua
169                .load(source)
170                .set_name(format!("@{}", full.display()))
171                .into_function()?;
172            return Ok(mlua::Value::Function(loader));
173        }
174
175        // Priority 2: $ASSAY_MODULES_PATH or ~/.assay/modules/<path>.lua
176        let global_path = if let Some(ref custom_path) = global_modules_path {
177            std::path::PathBuf::from(custom_path)
178        } else if let Ok(modules_env) = std::env::var(MODULES_PATH_ENV) {
179            std::path::PathBuf::from(modules_env)
180        } else if let Ok(home) = std::env::var("HOME") {
181            std::path::Path::new(&home).join(".assay/modules")
182        } else {
183            std::path::PathBuf::new()
184        };
185
186        if !global_path.as_os_str().is_empty()
187            && let Some((full, source)) = try_load(&global_path)
188        {
189            let loader = lua
190                .load(source)
191                .set_name(format!("@{}", full.display()))
192                .into_function()?;
193            return Ok(mlua::Value::Function(loader));
194        }
195
196        // Priority 3: Built-in modules are handled by register_stdlib_loader
197        // Return nil to fall through to the next searcher
198        Ok(mlua::Value::Nil)
199    })?;
200
201    let len = searchers.len()?;
202    searchers.set(len + 1, fs_searcher)?;
203
204    Ok(())
205}
206
207pub fn inject_env(lua: &Lua, env: &std::collections::HashMap<String, String>) -> Result<()> {
208    if env.is_empty() {
209        return Ok(());
210    }
211    let globals = lua.globals();
212    let env_table: mlua::Table = globals.get("env").map_err(lua_err)?;
213    let check_env: mlua::Table = env_table.get("_check_env").map_err(lua_err)?;
214    for (k, v) in env {
215        check_env.set(k.as_str(), v.as_str()).map_err(lua_err)?;
216    }
217    Ok(())
218}