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
11pub const MODULES_PATH_ENV: &str = "ASSAY_MODULES_PATH";
13
14pub 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 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
71fn 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 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 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 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 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 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}