Skip to main content

assay/lua/builtins/
core.rs

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    // A key outside the policy allowlist reads as absent rather than
38    // erroring: presence itself is information a caller shouldn't get.
39    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    // env.set(key, val) — set env var (nil val = unset)
65    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    // env.list() — returns table of {key, value} for all env vars
77    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    // fs.read_bytes(path) → string (binary-safe; Lua strings can hold any bytes)
139    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    // fs.write_bytes(path, data) → write binary data (Lua string with arbitrary bytes)
169    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        // Use symlink_metadata to detect symlinks without following them.
187        // A symlink to a directory should be removed as a file (unlink),
188        // not recursively delete the target directory.
189        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        // Use symlink_metadata separately to correctly detect symlinks,
239        // since std::fs::metadata follows symlinks (is_symlink always false).
240        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    // fs.copy(src, dst) — copy file, returns bytes copied
273    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    // fs.rename(src, dst) — atomic rename
282    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    // fs.glob(pattern) — glob pattern matching, returns array of path strings
292    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    // fs.tempdir() — create a temporary directory, returns path string
307    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    // fs.chmod(path, mode) — set file permissions (octal integer, e.g. 493 = 0o755)
323    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    // fs.readdir(path, opts?) — recursive directory listing
331    // opts: { depth = N } — max recursion depth (nil = unlimited)
332    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    // fs.lines(path) — stateful iterator yielding one line per call.
398    // Designed for `for line in fs.lines(path) do ... end`. Streams via
399    // BufReader so multi-GB files don't land in memory. Lines are
400    // stripped of their trailing `\n` (and `\r\n` on Windows files).
401    // Returns an iterator function; Lua's for-loop calls it until nil.
402    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    // fs.sub_in_file(path, pattern, repl) — `sed -i` equivalent.
422    // In-place search-and-replace using Lua's native pattern engine
423    // (same semantics as string.gsub, including %0-%9 backreferences
424    // and function replacements). Reads the file, substitutes, and
425    // only writes back if at least one match was made — so repeated
426    // calls on an already-substituted file are a no-op on disk.
427    // Returns the count of substitutions.
428    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    // string.split(s, sep?) — awk-style field split.
454    // When `sep` is nil or empty: splits on any run of whitespace and
455    // skips leading/trailing empty fields (matches awk default FS and
456    // Python's str.split() with no arg). When `sep` is provided: splits
457    // on the literal string (not a Lua pattern — use string.gmatch if
458    // you need pattern semantics). Returns a 1-indexed array table.
459    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}