1use data_encoding::BASE64;
2use mlua::{Lua, Value};
3use std::os::unix::fs::PermissionsExt;
4use std::time::{SystemTime, UNIX_EPOCH};
5use tracing::{error, info, warn};
6
7static TEMPDIR_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
8
9pub fn register_log(lua: &Lua) -> mlua::Result<()> {
10 let log_table = lua.create_table()?;
11
12 let info_fn = lua.create_function(|_, msg: String| {
13 info!(target: "lua", "{}", msg);
14 Ok(())
15 })?;
16 log_table.set("info", info_fn)?;
17
18 let warn_fn = lua.create_function(|_, msg: String| {
19 warn!(target: "lua", "{}", msg);
20 Ok(())
21 })?;
22 log_table.set("warn", warn_fn)?;
23
24 let error_fn = lua.create_function(|_, msg: String| {
25 error!(target: "lua", "{}", msg);
26 Ok(())
27 })?;
28 log_table.set("error", error_fn)?;
29
30 lua.globals().set("log", log_table)?;
31 Ok(())
32}
33
34pub fn register_env(lua: &Lua) -> mlua::Result<()> {
35 let env_table = lua.create_table()?;
36
37 let process_get_fn = lua.create_function(|lua, name: String| {
40 if !crate::lua::policy::env_visible(lua, &name) {
41 return Ok(None);
42 }
43 match std::env::var(&name) {
44 Ok(val) => Ok(Some(val)),
45 Err(_) => Ok(None),
46 }
47 })?;
48 env_table.set("_process_get", process_get_fn)?;
49 env_table.set("_check_env", lua.create_table()?)?;
50
51 lua.globals().set("env", env_table)?;
52
53 lua.load(
54 r#"
55 function env.get(name)
56 local val = env._check_env[name]
57 if val ~= nil then return val end
58 return env._process_get(name)
59 end
60 "#,
61 )
62 .exec()?;
63
64 let set_fn = lua.create_function(|_, (key, val): (String, Option<String>)| {
66 match val {
67 Some(v) => unsafe { std::env::set_var(&key, &v) },
68 None => unsafe { std::env::remove_var(&key) },
69 }
70 Ok(())
71 })?;
72 lua.globals()
73 .get::<mlua::Table>("env")?
74 .set("set", set_fn)?;
75
76 let list_fn = lua.create_function(|lua, ()| {
78 let results = lua.create_table()?;
79 let visible = std::env::vars().filter(|(key, _)| crate::lua::policy::env_visible(lua, key));
80 for (i, (key, val)) in (1..).zip(visible) {
81 let entry = lua.create_table()?;
82 entry.set("key", key)?;
83 entry.set("value", val)?;
84 results.set(i, entry)?;
85 }
86 Ok(results)
87 })?;
88 lua.globals()
89 .get::<mlua::Table>("env")?
90 .set("list", list_fn)?;
91
92 Ok(())
93}
94
95pub fn register_sleep(lua: &Lua) -> mlua::Result<()> {
96 let sleep_fn = lua.create_async_function(|_, seconds: f64| async move {
97 let duration = std::time::Duration::from_secs_f64(seconds);
98 tokio::time::sleep(duration).await;
99 Ok(())
100 })?;
101 lua.globals().set("sleep", sleep_fn)?;
102 Ok(())
103}
104
105pub fn register_time(lua: &Lua) -> mlua::Result<()> {
106 let time_fn = lua.create_function(|_, ()| {
107 let secs = SystemTime::now()
108 .duration_since(UNIX_EPOCH)
109 .map_err(|e| mlua::Error::runtime(format!("time(): {e}")))?
110 .as_secs_f64();
111 Ok(secs)
112 })?;
113 lua.globals().set("time", time_fn)?;
114 Ok(())
115}
116
117pub fn register_fs(lua: &Lua) -> mlua::Result<()> {
118 use crate::lua::file_source::FileSourceHandle;
119
120 let fs_table = lua.create_table()?;
121
122 let read_fn = lua.create_function(|lua, path: String| -> mlua::Result<String> {
123 let bytes = match lua.app_data_ref::<FileSourceHandle>() {
124 Some(source) => source.read(&path).ok_or_else(|| {
125 mlua::Error::runtime(format!(
126 "fs.read: failed to read {path:?}: not found in file source"
127 ))
128 })?,
129 None => std::fs::read(&path).map_err(|e| {
130 mlua::Error::runtime(format!("fs.read: failed to read {path:?}: {e}"))
131 })?,
132 };
133 String::from_utf8(bytes)
134 .map_err(|e| mlua::Error::runtime(format!("fs.read: invalid UTF-8 in {path:?}: {e}")))
135 })?;
136 fs_table.set("read", read_fn)?;
137
138 let read_bytes_fn = lua.create_function(|lua, path: String| {
140 let bytes = match lua.app_data_ref::<FileSourceHandle>() {
141 Some(source) => source.read(&path).ok_or_else(|| {
142 mlua::Error::runtime(format!(
143 "fs.read_bytes: failed to read {path:?}: not found in file source"
144 ))
145 })?,
146 None => std::fs::read(&path).map_err(|e| {
147 mlua::Error::runtime(format!("fs.read_bytes: failed to read {path:?}: {e}"))
148 })?,
149 };
150 lua.create_string(&bytes)
151 })?;
152 fs_table.set("read_bytes", read_bytes_fn)?;
153
154 let write_fn = lua.create_function(|_, (path, content): (String, String)| {
155 let p = std::path::Path::new(&path);
156 if let Some(parent) = p.parent() {
157 std::fs::create_dir_all(parent).map_err(|e| {
158 mlua::Error::runtime(format!(
159 "fs.write: failed to create directories for {path:?}: {e}"
160 ))
161 })?;
162 }
163 std::fs::write(&path, &content)
164 .map_err(|e| mlua::Error::runtime(format!("fs.write: failed to write {path:?}: {e}")))
165 })?;
166 fs_table.set("write", write_fn)?;
167
168 let write_bytes_fn = lua.create_function(|_, (path, data): (String, mlua::String)| {
170 let p = std::path::Path::new(&path);
171 if let Some(parent) = p.parent() {
172 std::fs::create_dir_all(parent).map_err(|e| {
173 mlua::Error::runtime(format!(
174 "fs.write_bytes: failed to create directories for {path:?}: {e}"
175 ))
176 })?;
177 }
178 std::fs::write(&path, data.as_bytes()).map_err(|e| {
179 mlua::Error::runtime(format!("fs.write_bytes: failed to write {path:?}: {e}"))
180 })
181 })?;
182 fs_table.set("write_bytes", write_bytes_fn)?;
183
184 let remove_fn = lua.create_function(|_, path: String| {
185 let p = std::path::Path::new(&path);
186 let is_dir = match std::fs::symlink_metadata(&path) {
190 Ok(m) => m.file_type().is_dir(),
191 Err(_) => p.is_dir(),
192 };
193 if is_dir {
194 std::fs::remove_dir_all(&path).map_err(|e| {
195 mlua::Error::runtime(format!(
196 "fs.remove: failed to remove directory {path:?}: {e}"
197 ))
198 })
199 } else {
200 std::fs::remove_file(&path).map_err(|e| {
201 mlua::Error::runtime(format!("fs.remove: failed to remove {path:?}: {e}"))
202 })
203 }
204 })?;
205 fs_table.set("remove", remove_fn)?;
206
207 let list_fn =
208 lua.create_function(|lua, path: String| {
209 let entries = lua.create_table()?;
210 for (i, entry) in (1..).zip(std::fs::read_dir(&path).map_err(|e| {
211 mlua::Error::runtime(format!("fs.list: failed to list {path:?}: {e}"))
212 })?) {
213 let entry = entry.map_err(|e| {
214 mlua::Error::runtime(format!("fs.list: error reading entry in {path:?}: {e}"))
215 })?;
216 let info = lua.create_table()?;
217 let name = entry.file_name().to_string_lossy().to_string();
218 info.set("name", name)?;
219 let file_type = entry.file_type().map_err(|e| {
220 mlua::Error::runtime(format!("fs.list: failed to get file type: {e}"))
221 })?;
222 if file_type.is_dir() {
223 info.set("type", "directory")?;
224 } else if file_type.is_symlink() {
225 info.set("type", "symlink")?;
226 } else {
227 info.set("type", "file")?;
228 }
229 entries.set(i, info)?;
230 }
231 Ok(entries)
232 })?;
233 fs_table.set("list", list_fn)?;
234
235 let stat_fn = lua.create_function(|lua, path: String| {
236 let metadata = std::fs::metadata(&path)
237 .map_err(|e| mlua::Error::runtime(format!("fs.stat: failed to stat {path:?}: {e}")))?;
238 let is_symlink = std::fs::symlink_metadata(&path)
241 .map(|m| m.file_type().is_symlink())
242 .unwrap_or(false);
243 let info = lua.create_table()?;
244 info.set("size", metadata.len())?;
245 info.set("is_file", metadata.is_file())?;
246 info.set("is_dir", metadata.is_dir())?;
247 info.set("is_symlink", is_symlink)?;
248 if let Ok(modified) = metadata.modified()
249 && let Ok(duration) = modified.duration_since(std::time::UNIX_EPOCH)
250 {
251 info.set("modified", duration.as_secs_f64())?;
252 }
253 if let Ok(created) = metadata.created()
254 && let Ok(duration) = created.duration_since(std::time::UNIX_EPOCH)
255 {
256 info.set("created", duration.as_secs_f64())?;
257 }
258 Ok(info)
259 })?;
260 fs_table.set("stat", stat_fn)?;
261
262 let mkdir_fn = lua.create_function(|_, path: String| {
263 std::fs::create_dir_all(&path)
264 .map_err(|e| mlua::Error::runtime(format!("fs.mkdir: failed to create {path:?}: {e}")))
265 })?;
266 fs_table.set("mkdir", mkdir_fn)?;
267
268 let exists_fn =
269 lua.create_function(|_, path: String| Ok(std::path::Path::new(&path).exists()))?;
270 fs_table.set("exists", exists_fn)?;
271
272 let copy_fn = lua.create_function(|_, (src, dst): (String, String)| {
274 let bytes = std::fs::copy(&src, &dst).map_err(|e| {
275 mlua::Error::runtime(format!("fs.copy: failed to copy {src:?} to {dst:?}: {e}"))
276 })?;
277 Ok(bytes)
278 })?;
279 fs_table.set("copy", copy_fn)?;
280
281 let rename_fn = lua.create_function(|_, (src, dst): (String, String)| {
283 std::fs::rename(&src, &dst).map_err(|e| {
284 mlua::Error::runtime(format!(
285 "fs.rename: failed to rename {src:?} to {dst:?}: {e}"
286 ))
287 })
288 })?;
289 fs_table.set("rename", rename_fn)?;
290
291 let glob_fn = lua.create_function(|lua, pattern: String| {
293 let paths = glob::glob(&pattern).map_err(|e| {
294 mlua::Error::runtime(format!("fs.glob: invalid pattern {pattern:?}: {e}"))
295 })?;
296 let results = lua.create_table()?;
297 for (i, entry) in (1..).zip(paths) {
298 let path = entry
299 .map_err(|e| mlua::Error::runtime(format!("fs.glob: error reading entry: {e}")))?;
300 results.set(i, path.to_string_lossy().to_string())?;
301 }
302 Ok(results)
303 })?;
304 fs_table.set("glob", glob_fn)?;
305
306 let tempdir_fn = lua.create_function(|_, ()| {
308 let base = std::env::temp_dir();
309 let nanos: u64 = std::time::SystemTime::now()
310 .duration_since(std::time::UNIX_EPOCH)
311 .unwrap_or_default()
312 .as_nanos() as u64;
313 let seq = TEMPDIR_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
314 let dir = base.join(format!("assay-{nanos:x}-{seq}"));
315 std::fs::create_dir_all(&dir).map_err(|e| {
316 mlua::Error::runtime(format!("fs.tempdir: failed to create {dir:?}: {e}"))
317 })?;
318 Ok(dir.to_string_lossy().to_string())
319 })?;
320 fs_table.set("tempdir", tempdir_fn)?;
321
322 let chmod_fn = lua.create_function(|_, (path, mode): (String, u32)| {
324 let perms = std::fs::Permissions::from_mode(mode);
325 std::fs::set_permissions(&path, perms)
326 .map_err(|e| mlua::Error::runtime(format!("fs.chmod: failed to chmod {path:?}: {e}")))
327 })?;
328 fs_table.set("chmod", chmod_fn)?;
329
330 let readdir_fn = lua.create_function(|lua, args: mlua::MultiValue| {
333 let mut args_iter = args.into_iter();
334 let path: String = args_iter
335 .next()
336 .ok_or_else(|| mlua::Error::runtime("fs.readdir: path required"))
337 .and_then(|v| lua.unpack(v))?;
338
339 let max_depth: Option<usize> = if let Some(Value::Table(opts)) = args_iter.next() {
340 opts.get::<Option<usize>>("depth")?
341 } else {
342 None
343 };
344
345 let results = lua.create_table()?;
346 let mut i = 1u64;
347 let base = std::path::PathBuf::from(&path);
348
349 fn walk(
350 base: &std::path::Path,
351 dir: &std::path::Path,
352 results: &mlua::Table,
353 lua: &Lua,
354 i: &mut u64,
355 depth: usize,
356 max_depth: Option<usize>,
357 ) -> mlua::Result<()> {
358 let entries = std::fs::read_dir(dir).map_err(|e| {
359 mlua::Error::runtime(format!("fs.readdir: failed to read {dir:?}: {e}"))
360 })?;
361 for entry in entries {
362 let entry = entry.map_err(|e| {
363 mlua::Error::runtime(format!("fs.readdir: error reading entry: {e}"))
364 })?;
365 let file_type = entry.file_type().map_err(|e| {
366 mlua::Error::runtime(format!("fs.readdir: failed to get file type: {e}"))
367 })?;
368 let rel_path = entry
369 .path()
370 .strip_prefix(base)
371 .unwrap_or(&entry.path())
372 .to_string_lossy()
373 .to_string();
374 let info = lua.create_table()?;
375 info.set("path", rel_path)?;
376 if file_type.is_dir() {
377 info.set("type", "directory")?;
378 } else if file_type.is_symlink() {
379 info.set("type", "symlink")?;
380 } else {
381 info.set("type", "file")?;
382 }
383 results.set(*i, info)?;
384 *i += 1;
385 if file_type.is_dir() && (max_depth.is_none() || depth < max_depth.unwrap()) {
386 walk(base, &entry.path(), results, lua, i, depth + 1, max_depth)?;
387 }
388 }
389 Ok(())
390 }
391
392 walk(&base, &base, &results, lua, &mut i, 1, max_depth)?;
393 Ok(results)
394 })?;
395 fs_table.set("readdir", readdir_fn)?;
396
397 let lines_fn = lua.create_function(|lua, path: String| {
403 use std::io::BufRead;
404 let file = std::fs::File::open(&path)
405 .map_err(|e| mlua::Error::runtime(format!("fs.lines: failed to open {path:?}: {e}")))?;
406 let iter =
407 std::sync::Arc::new(std::sync::Mutex::new(std::io::BufReader::new(file).lines()));
408 lua.create_function(move |_, ()| {
409 let mut it = iter
410 .lock()
411 .map_err(|e| mlua::Error::runtime(format!("fs.lines: lock poisoned: {e}")))?;
412 match it.next() {
413 Some(Ok(line)) => Ok(Some(line)),
414 Some(Err(e)) => Err(mlua::Error::runtime(format!("fs.lines: read error: {e}"))),
415 None => Ok(None),
416 }
417 })
418 })?;
419 fs_table.set("lines", lines_fn)?;
420
421 let sub_in_file_fn = lua.create_function(
429 |lua, (path, pattern, repl): (String, String, mlua::Value)| {
430 let content = std::fs::read_to_string(&path).map_err(|e| {
431 mlua::Error::runtime(format!("fs.sub_in_file: failed to read {path:?}: {e}"))
432 })?;
433 let string_table: mlua::Table = lua.globals().get("string")?;
434 let gsub: mlua::Function = string_table.get("gsub")?;
435 let (new_content, count): (String, u64) = gsub.call((content, pattern, repl))?;
436 if count > 0 {
437 std::fs::write(&path, &new_content).map_err(|e| {
438 mlua::Error::runtime(format!("fs.sub_in_file: failed to write {path:?}: {e}"))
439 })?;
440 }
441 Ok(count)
442 },
443 )?;
444 fs_table.set("sub_in_file", sub_in_file_fn)?;
445
446 lua.globals().set("fs", fs_table)?;
447 Ok(())
448}
449
450pub fn register_string_helpers(lua: &Lua) -> mlua::Result<()> {
451 let string_table: mlua::Table = lua.globals().get("string")?;
452
453 let split_fn = lua.create_function(|lua, args: mlua::MultiValue| {
460 let mut args_iter = args.into_iter();
461 let s: String = args_iter
462 .next()
463 .ok_or_else(|| mlua::Error::runtime("string.split: string required"))
464 .and_then(|v| lua.unpack(v))?;
465 let sep: Option<String> = match args_iter.next() {
466 Some(mlua::Value::Nil) | None => None,
467 Some(v) => Some(lua.unpack(v)?),
468 };
469 let results = lua.create_table()?;
470 match sep {
471 Some(ref sep_str) if !sep_str.is_empty() => {
472 for (i, part) in (1..).zip(s.split(sep_str.as_str())) {
473 results.set(i, part)?;
474 }
475 }
476 _ => {
477 for (i, part) in (1..).zip(s.split_whitespace()) {
478 results.set(i, part)?;
479 }
480 }
481 }
482 Ok(results)
483 })?;
484 string_table.set("split", split_fn)?;
485
486 Ok(())
487}
488
489pub fn register_base64(lua: &Lua) -> mlua::Result<()> {
490 let b64_table = lua.create_table()?;
491
492 let encode_fn = lua.create_function(|_, input: String| Ok(BASE64.encode(input.as_bytes())))?;
493 b64_table.set("encode", encode_fn)?;
494
495 let decode_fn = lua.create_function(|_, input: String| {
496 let bytes = BASE64
497 .decode(input.as_bytes())
498 .map_err(|e| mlua::Error::runtime(format!("base64.decode: {e}")))?;
499 String::from_utf8(bytes)
500 .map_err(|e| mlua::Error::runtime(format!("base64.decode: invalid UTF-8: {e}")))
501 })?;
502 b64_table.set("decode", decode_fn)?;
503
504 lua.globals().set("base64", b64_table)?;
505 Ok(())
506}
507
508pub fn register_regex(lua: &Lua) -> mlua::Result<()> {
509 let regex_table = lua.create_table()?;
510
511 let match_fn = lua.create_function(|_, (text, pattern): (String, String)| {
512 let re = regex_lite::Regex::new(&pattern)
513 .map_err(|e| mlua::Error::runtime(format!("regex.match: invalid pattern: {e}")))?;
514 Ok(re.is_match(&text))
515 })?;
516 regex_table.set("match", match_fn)?;
517
518 let find_fn = lua.create_function(|lua, (text, pattern): (String, String)| {
519 let re = regex_lite::Regex::new(&pattern)
520 .map_err(|e| mlua::Error::runtime(format!("regex.find: invalid pattern: {e}")))?;
521 match re.captures(&text) {
522 Some(caps) => {
523 let result = lua.create_table()?;
524 let full_match = caps.get(0).map(|m| m.as_str()).unwrap_or("");
525 result.set("match", full_match.to_string())?;
526 let groups = lua.create_table()?;
527 for i in 1..caps.len() {
528 if let Some(m) = caps.get(i) {
529 groups.set(i, m.as_str().to_string())?;
530 }
531 }
532 result.set("groups", groups)?;
533 Ok(Value::Table(result))
534 }
535 None => Ok(Value::Nil),
536 }
537 })?;
538 regex_table.set("find", find_fn)?;
539
540 let find_all_fn = lua.create_function(|lua, (text, pattern): (String, String)| {
541 let re = regex_lite::Regex::new(&pattern)
542 .map_err(|e| mlua::Error::runtime(format!("regex.find_all: invalid pattern: {e}")))?;
543 let results = lua.create_table()?;
544 for (i, m) in re.find_iter(&text).enumerate() {
545 results.set(i + 1, m.as_str().to_string())?;
546 }
547 Ok(results)
548 })?;
549 regex_table.set("find_all", find_all_fn)?;
550
551 let replace_fn = lua.create_function(
552 |_, (text, pattern, replacement): (String, String, String)| {
553 let re = regex_lite::Regex::new(&pattern).map_err(|e| {
554 mlua::Error::runtime(format!("regex.replace: invalid pattern: {e}"))
555 })?;
556 Ok(re.replace_all(&text, replacement.as_str()).into_owned())
557 },
558 )?;
559 regex_table.set("replace", replace_fn)?;
560
561 lua.globals().set("regex", regex_table)?;
562 Ok(())
563}
564
565pub fn register_async(lua: &Lua) -> mlua::Result<()> {
566 let async_table = lua.create_table()?;
567
568 let spawn_fn = lua.create_async_function(|lua, func: mlua::Function| async move {
569 let thread = lua.create_thread(func)?;
570 let async_thread = thread.into_async::<mlua::MultiValue>(())?;
571 let join_handle: tokio::task::JoinHandle<Result<Vec<Value>, String>> =
572 tokio::task::spawn_local(async move {
573 let values = async_thread.await.map_err(|e| e.to_string())?;
574 Ok(values.into_vec())
575 });
576
577 let handle = lua.create_table()?;
578 let cell = std::rc::Rc::new(std::cell::RefCell::new(Some(join_handle)));
579 let cell_clone = cell.clone();
580
581 let await_fn = lua.create_async_function(move |lua, ()| {
582 let cell = cell_clone.clone();
583 async move {
584 let join_handle = cell
585 .borrow_mut()
586 .take()
587 .ok_or_else(|| mlua::Error::runtime("async handle already awaited"))?;
588 let result = join_handle.await.map_err(|e| {
589 mlua::Error::runtime(format!("async.spawn: task panicked: {e}"))
590 })?;
591 match result {
592 Ok(values) => {
593 let tbl = lua.create_table()?;
594 for (i, v) in values.into_iter().enumerate() {
595 tbl.set(i + 1, v)?;
596 }
597 Ok(Value::Table(tbl))
598 }
599 Err(msg) => Err(mlua::Error::runtime(msg)),
600 }
601 }
602 })?;
603 handle.set("await", await_fn)?;
604
605 Ok(handle)
606 })?;
607 async_table.set("spawn", spawn_fn)?;
608
609 let spawn_interval_fn =
610 lua.create_async_function(|lua, (seconds, func): (f64, mlua::Function)| async move {
611 if seconds <= 0.0 {
612 return Err(mlua::Error::runtime(
613 "async.spawn_interval: interval must be positive",
614 ));
615 }
616
617 let cancel = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
618 let cancel_clone = cancel.clone();
619
620 tokio::task::spawn_local({
621 let cancel = cancel_clone.clone();
622 async move {
623 let mut interval =
624 tokio::time::interval(std::time::Duration::from_secs_f64(seconds));
625 interval.tick().await;
626 loop {
627 interval.tick().await;
628 if cancel.load(std::sync::atomic::Ordering::Relaxed) {
629 break;
630 }
631 if let Err(e) = func.call_async::<()>(()).await {
632 error!("async.spawn_interval: callback error: {e}");
633 break;
634 }
635 }
636 }
637 });
638
639 let handle = lua.create_table()?;
640 let cancel_fn = lua.create_function(move |_, ()| {
641 cancel.store(true, std::sync::atomic::Ordering::Relaxed);
642 Ok(())
643 })?;
644 handle.set("cancel", cancel_fn)?;
645
646 Ok(handle)
647 })?;
648 async_table.set("spawn_interval", spawn_interval_fn)?;
649
650 lua.globals().set("async", async_table)?;
651 Ok(())
652}
653
654#[cfg(test)]
655mod tests {
656 use data_encoding::BASE64;
657
658 #[test]
659 fn test_base64_roundtrip() {
660 let input = "hello world";
661 let encoded = BASE64.encode(input.as_bytes());
662 assert_eq!(encoded, "aGVsbG8gd29ybGQ=");
663 let decoded = BASE64.decode(encoded.as_bytes()).unwrap();
664 assert_eq!(String::from_utf8(decoded).unwrap(), input);
665 }
666
667 #[test]
668 fn test_base64_empty() {
669 let encoded = BASE64.encode(b"");
670 assert_eq!(encoded, "");
671 let decoded = BASE64.decode(b"").unwrap();
672 assert!(decoded.is_empty());
673 }
674}