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