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