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
121 .contents_utf8()
122 .ok_or_else(|| mlua::Error::runtime(format!("stdlib {path}: invalid UTF-8")))?;
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(lua.create_string(format!(
132 "no embedded stdlib file: {}",
133 candidates[0]
134 ))?))
135 })?;
136
137 let len = searchers.len()?;
138 searchers.set(len + 1, stdlib_searcher)?;
139
140 Ok(())
141}
142
143fn register_fs_loader(lua: &Lua, global_modules_path: Option<String>) -> mlua::Result<()> {
144 let package: mlua::Table = lua.globals().get("package")?;
145 let searchers: mlua::Table = package.get("searchers")?;
146
147 let fs_searcher = lua.create_function(move |lua, module_name: String| {
150 let rest = match module_name.strip_prefix("assay.") {
151 Some(r) => r,
152 None => {
153 return Ok(mlua::Value::String(
154 lua.create_string(format!("not an assay.* module: {module_name}"))?,
155 ));
156 }
157 };
158 let base = rest.replace('.', "/");
159 let candidates = [format!("{base}.lua"), format!("{base}/init.lua")];
160
161 let try_load = |dir: &std::path::Path| -> Option<(std::path::PathBuf, String)> {
162 for rel in &candidates {
163 let full = dir.join(rel);
164 if let Ok(source) = std::fs::read_to_string(&full) {
165 return Some((full, source));
166 }
167 }
168 None
169 };
170
171 if let Some((full, source)) = try_load(std::path::Path::new("./modules")) {
173 let loader = lua
174 .load(source)
175 .set_name(format!("@{}", full.display()))
176 .into_function()?;
177 return Ok(mlua::Value::Function(loader));
178 }
179
180 let global_path = if let Some(ref custom_path) = global_modules_path {
182 std::path::PathBuf::from(custom_path)
183 } else if let Ok(modules_env) = std::env::var(MODULES_PATH_ENV) {
184 std::path::PathBuf::from(modules_env)
185 } else if let Ok(home) = std::env::var("HOME") {
186 std::path::Path::new(&home).join(".assay/modules")
187 } else {
188 std::path::PathBuf::new()
189 };
190
191 if !global_path.as_os_str().is_empty()
192 && let Some((full, source)) = try_load(&global_path)
193 {
194 let loader = lua
195 .load(source)
196 .set_name(format!("@{}", full.display()))
197 .into_function()?;
198 return Ok(mlua::Value::Function(loader));
199 }
200
201 Ok(mlua::Value::Nil)
204 })?;
205
206 let len = searchers.len()?;
207 searchers.set(len + 1, fs_searcher)?;
208
209 Ok(())
210}
211
212pub fn inject_env(lua: &Lua, env: &std::collections::HashMap<String, String>) -> Result<()> {
213 if env.is_empty() {
214 return Ok(());
215 }
216 let globals = lua.globals();
217 let env_table: mlua::Table = globals.get("env").map_err(lua_err)?;
218 let check_env: mlua::Table = env_table.get("_check_env").map_err(lua_err)?;
219 for (k, v) in env {
220 check_env.set(k.as_str(), v.as_str()).map_err(lua_err)?;
221 }
222 Ok(())
223}