1pub mod async_bridge;
2pub mod builtins;
3pub mod file_source;
4
5#[cfg(feature = "server")]
6#[allow(unused_imports)]
7pub use builtins::LuaAxumRouter;
8
9use anyhow::Result;
10use include_dir::{Dir, include_dir};
11use mlua::{Lua, LuaOptions, StdLib};
12
13static STDLIB_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/stdlib");
14
15pub const MODULES_PATH_ENV: &str = "ASSAY_MODULES_PATH";
17
18pub const BLOCK_GLOBALS_ENV: &str = "ASSAY_BLOCK_GLOBALS";
23
24fn lua_err(e: mlua::Error) -> anyhow::Error {
25 anyhow::anyhow!("{e}")
26}
27
28pub fn create_vm(client: reqwest::Client) -> Result<Lua> {
29 create_vm_with_paths(client, None)
30}
31
32#[allow(dead_code)]
33pub fn create_vm_with_lib_path(client: reqwest::Client, lib_path: String) -> Result<Lua> {
34 create_vm_with_paths(client, Some(lib_path))
35}
36
37pub fn create_vm_with_paths(
38 client: reqwest::Client,
39 global_modules_path: Option<String>,
40) -> Result<Lua> {
41 let libs = StdLib::ALL_SAFE;
42 let lua = Lua::new_with(libs, LuaOptions::default()).map_err(lua_err)?;
43 lua.set_memory_limit(64 * 1024 * 1024).map_err(lua_err)?;
44 sandbox(&lua).map_err(lua_err)?;
45 register_fs_loader(&lua, global_modules_path).map_err(lua_err)?;
46 register_stdlib_loader(&lua).map_err(lua_err)?;
47 builtins::register_all(&lua, client).map_err(lua_err)?;
48 Ok(lua)
49}
50
51fn sandbox(lua: &Lua) -> mlua::Result<()> {
52 let globals = lua.globals();
59 let string_lib: mlua::Table = globals.get("string")?;
60 string_lib.set("dump", mlua::Value::Nil)?;
61
62 if let Ok(extra) = std::env::var(BLOCK_GLOBALS_ENV) {
63 for raw in extra.split(',') {
64 let name = raw.trim();
65 if name.is_empty() {
66 continue;
67 }
68 nil_dotted_path(lua, name)?;
69 }
70 }
71
72 Ok(())
73}
74
75fn nil_dotted_path(lua: &Lua, path: &str) -> mlua::Result<()> {
81 let parts: Vec<&str> = path.split('.').filter(|s| !s.is_empty()).collect();
82 if parts.is_empty() {
83 return Ok(());
84 }
85 let mut current: mlua::Table = lua.globals();
86 for segment in &parts[..parts.len() - 1] {
87 let next: mlua::Value = current.get(*segment)?;
88 match next {
89 mlua::Value::Table(t) => current = t,
90 _ => return Ok(()),
91 }
92 }
93 current.set(parts[parts.len() - 1], mlua::Value::Nil)
94}
95
96fn register_stdlib_loader(lua: &Lua) -> mlua::Result<()> {
97 let package: mlua::Table = lua.globals().get("package")?;
98 let searchers: mlua::Table = package.get("searchers")?;
99
100 let stdlib_searcher = lua.create_function(|lua, module_name: String| {
106 let rest = match module_name.strip_prefix("assay.") {
107 Some(r) => r,
108 None => {
109 return Ok(mlua::Value::String(
110 lua.create_string(format!("not an assay.* module: {module_name}"))?,
111 ));
112 }
113 };
114
115 let base = rest.replace('.', "/");
116 let candidates = [format!("{base}.lua"), format!("{base}/init.lua")];
117
118 for path in &candidates {
119 if let Some(file) = STDLIB_DIR.get_file(path) {
120 let source = file.contents_utf8().ok_or_else(|| {
121 mlua::Error::runtime(format!("stdlib {path}: invalid UTF-8"))
122 })?;
123 let loader = lua
124 .load(source)
125 .set_name(format!("@assay/{path}"))
126 .into_function()?;
127 return Ok(mlua::Value::Function(loader));
128 }
129 }
130
131 Ok(mlua::Value::String(
132 lua.create_string(format!("no embedded stdlib file: {}", candidates[0]))?,
133 ))
134 })?;
135
136 let len = searchers.len()?;
137 searchers.set(len + 1, stdlib_searcher)?;
138
139 Ok(())
140}
141
142fn register_fs_loader(lua: &Lua, global_modules_path: Option<String>) -> mlua::Result<()> {
143 let package: mlua::Table = lua.globals().get("package")?;
144 let searchers: mlua::Table = package.get("searchers")?;
145
146 let fs_searcher = lua.create_function(move |lua, module_name: String| {
149 let rest = match module_name.strip_prefix("assay.") {
150 Some(r) => r,
151 None => {
152 return Ok(mlua::Value::String(
153 lua.create_string(format!("not an assay.* module: {module_name}"))?,
154 ));
155 }
156 };
157 let base = rest.replace('.', "/");
158 let candidates = [format!("{base}.lua"), format!("{base}/init.lua")];
159
160 let try_load = |dir: &std::path::Path| -> Option<(std::path::PathBuf, String)> {
161 for rel in &candidates {
162 let full = dir.join(rel);
163 if let Ok(source) = std::fs::read_to_string(&full) {
164 return Some((full, source));
165 }
166 }
167 None
168 };
169
170 if let Some((full, source)) = try_load(std::path::Path::new("./modules")) {
172 let loader = lua
173 .load(source)
174 .set_name(format!("@{}", full.display()))
175 .into_function()?;
176 return Ok(mlua::Value::Function(loader));
177 }
178
179 let global_path = if let Some(ref custom_path) = global_modules_path {
181 std::path::PathBuf::from(custom_path)
182 } else if let Ok(modules_env) = std::env::var(MODULES_PATH_ENV) {
183 std::path::PathBuf::from(modules_env)
184 } else if let Ok(home) = std::env::var("HOME") {
185 std::path::Path::new(&home).join(".assay/modules")
186 } else {
187 std::path::PathBuf::new()
188 };
189
190 if !global_path.as_os_str().is_empty()
191 && let Some((full, source)) = try_load(&global_path)
192 {
193 let loader = lua
194 .load(source)
195 .set_name(format!("@{}", full.display()))
196 .into_function()?;
197 return Ok(mlua::Value::Function(loader));
198 }
199
200 Ok(mlua::Value::Nil)
203 })?;
204
205 let len = searchers.len()?;
206 searchers.set(len + 1, fs_searcher)?;
207
208 Ok(())
209}
210
211pub fn inject_env(lua: &Lua, env: &std::collections::HashMap<String, String>) -> Result<()> {
212 if env.is_empty() {
213 return Ok(());
214 }
215 let globals = lua.globals();
216 let env_table: mlua::Table = globals.get("env").map_err(lua_err)?;
217 let check_env: mlua::Table = env_table.get("_check_env").map_err(lua_err)?;
218 for (k, v) in env {
219 check_env.set(k.as_str(), v.as_str()).map_err(lua_err)?;
220 }
221 Ok(())
222}