#![allow(clippy::cargo_common_metadata)] use std::{ cell::RefCell, env::{ self, consts::{ARCH, OS}, }, path::MAIN_SEPARATOR, process::Stdio, rc::Rc, sync::Arc, }; use mlua::prelude::*; use lune_utils::TableBuilder; use mlua_luau_scheduler::{Functions, LuaSpawnExt}; use options::ProcessSpawnOptionsStdio; use os_str_bytes::RawOsString; use stream::{ChildProcessReader, ChildProcessWriter}; use tokio::{io::AsyncWriteExt, process::Child, sync::RwLock}; mod options; mod stream; mod tee_writer; mod wait_for_child; use self::options::ProcessSpawnOptions; use self::wait_for_child::wait_for_child; use lune_utils::path::get_current_dir; /** Creates the `process` standard library module. # Errors Errors when out of memory. */ #[allow(clippy::missing_panics_doc)] pub fn module(lua: &Lua) -> LuaResult { let mut cwd_str = get_current_dir() .to_str() .expect("cwd should be valid UTF-8") .to_string(); if !cwd_str.ends_with(MAIN_SEPARATOR) { cwd_str.push(MAIN_SEPARATOR); } // Create constants for OS & processor architecture let os = lua.create_string(OS.to_lowercase())?; let arch = lua.create_string(ARCH.to_lowercase())?; let endianness = lua.create_string(if cfg!(target_endian = "big") { "big" } else { "little" })?; // Create readonly args array let args_vec = lua .app_data_ref::>() .ok_or_else(|| LuaError::runtime("Missing args vec in Lua app data"))? .clone(); let args_tab = TableBuilder::new(lua)? .with_sequential_values(args_vec)? .build_readonly()?; // Create proxied table for env that gets & sets real env vars let env_tab = TableBuilder::new(lua)? .with_metatable( TableBuilder::new(lua)? .with_function(LuaMetaMethod::Index.name(), process_env_get)? .with_function(LuaMetaMethod::NewIndex.name(), process_env_set)? .with_function(LuaMetaMethod::Iter.name(), process_env_iter)? .build_readonly()?, )? .build_readonly()?; // Create our process exit function, the scheduler crate provides this let fns = Functions::new(lua)?; let process_exit = fns.exit; // Create the full process table TableBuilder::new(lua)? .with_value("os", os)? .with_value("arch", arch)? .with_value("endianness", endianness)? .with_value("args", args_tab)? .with_value("cwd", cwd_str)? .with_value("env", env_tab)? .with_value("exit", process_exit)? .with_async_function("exec", process_exec)? .with_function("create", process_create)? .build_readonly() } fn process_env_get<'lua>( lua: &'lua Lua, (_, key): (LuaValue<'lua>, String), ) -> LuaResult> { match env::var_os(key) { Some(value) => { let raw_value = RawOsString::new(value); Ok(LuaValue::String( lua.create_string(raw_value.to_raw_bytes())?, )) } None => Ok(LuaValue::Nil), } } fn process_env_set<'lua>( _: &'lua Lua, (_, key, value): (LuaValue<'lua>, String, Option), ) -> LuaResult<()> { // Make sure key is valid, otherwise set_var will panic if key.is_empty() { Err(LuaError::RuntimeError("Key must not be empty".to_string())) } else if key.contains('=') { Err(LuaError::RuntimeError( "Key must not contain the equals character '='".to_string(), )) } else if key.contains('\0') { Err(LuaError::RuntimeError( "Key must not contain the NUL character".to_string(), )) } else if let Some(value) = value { // Make sure value is valid, otherwise set_var will panic if value.contains('\0') { Err(LuaError::RuntimeError( "Value must not contain the NUL character".to_string(), )) } else { env::set_var(&key, &value); Ok(()) } } else { env::remove_var(&key); Ok(()) } } fn process_env_iter<'lua>( lua: &'lua Lua, (_, ()): (LuaValue<'lua>, ()), ) -> LuaResult> { let mut vars = env::vars_os().collect::>().into_iter(); lua.create_function_mut(move |lua, (): ()| match vars.next() { Some((key, value)) => { let raw_key = RawOsString::new(key); let raw_value = RawOsString::new(value); Ok(( LuaValue::String(lua.create_string(raw_key.to_raw_bytes())?), LuaValue::String(lua.create_string(raw_value.to_raw_bytes())?), )) } None => Ok((LuaValue::Nil, LuaValue::Nil)), }) } async fn process_exec( lua: &Lua, (program, args, options): (String, Option>, ProcessSpawnOptions), ) -> LuaResult { let res = lua .spawn(async move { let cmd = spawn_command_with_stdin(program, args, options.clone()).await?; wait_for_child(cmd, options.stdio.stdout, options.stdio.stderr).await }) .await?; /* NOTE: If an exit code was not given by the child process, we default to 1 if it yielded any error output, otherwise 0 An exit code may be missing if the process was terminated by some external signal, which is the only time we use this default */ let code = res .status .code() .unwrap_or(i32::from(!res.stderr.is_empty())); // Construct and return a readonly lua table with results TableBuilder::new(lua)? .with_value("ok", code == 0)? .with_value("code", code)? .with_value("stdout", lua.create_string(&res.stdout)?)? .with_value("stderr", lua.create_string(&res.stderr)?)? .build_readonly() } #[allow(clippy::await_holding_refcell_ref)] fn process_create( lua: &Lua, (program, args, options): (String, Option>, ProcessSpawnOptions), ) -> LuaResult { // We do not want the user to provide stdio options for process.create, // so we reset the options, regardless of what the user provides us let mut spawn_options = options.clone(); spawn_options.stdio = ProcessSpawnOptionsStdio::default(); let (code_tx, code_rx) = tokio::sync::broadcast::channel(4); let code_rx_rc = Rc::new(RefCell::new(code_rx)); let child = spawn_command(program, args, spawn_options)?; let child_arc = Arc::new(RwLock::new(child)); let child_arc_clone = Arc::clone(&child_arc); let mut child_lock = tokio::task::block_in_place(|| child_arc_clone.blocking_write()); let stdin = child_lock.stdin.take().unwrap(); let stdout = child_lock.stdout.take().unwrap(); let stderr = child_lock.stderr.take().unwrap(); let child_arc_inner = Arc::clone(&child_arc); // Spawn a background task to wait for the child to exit and send the exit code let status_handle = tokio::spawn(async move { let res = child_arc_inner.write().await.wait().await; if let Ok(output) = res { let code = output.code().unwrap_or_default(); code_tx .send(code) .expect("ExitCode receiver was unexpectedly dropped"); } }); TableBuilder::new(lua)? .with_value("stdout", ChildProcessReader(stdout))? .with_value("stderr", ChildProcessReader(stderr))? .with_value("stdin", ChildProcessWriter(stdin))? .with_async_function("kill", move |_, ()| { // First, stop the status task so the RwLock is dropped status_handle.abort(); let child_arc_clone = Arc::clone(&child_arc); // Then get another RwLock to write to the child process and kill it async move { Ok(child_arc_clone.write().await.kill().await?) } })? .with_async_function("status", move |lua, ()| { let code_rx_rc_clone = Rc::clone(&code_rx_rc); async move { // Exit code of 9 corresponds to SIGKILL, which should be the only case where // the receiver gets suddenly dropped let code = code_rx_rc_clone.borrow_mut().recv().await.unwrap_or(9); TableBuilder::new(lua)? .with_value("code", code)? .with_value("ok", code == 0)? .build_readonly() } })? .build_readonly() } async fn spawn_command_with_stdin( program: String, args: Option>, mut options: ProcessSpawnOptions, ) -> LuaResult { let stdin = options.stdio.stdin.take(); let mut child = spawn_command(program, args, options)?; if let Some(stdin) = stdin { let mut child_stdin = child.stdin.take().unwrap(); child_stdin.write_all(&stdin).await.into_lua_err()?; } Ok(child) } fn spawn_command( program: String, args: Option>, options: ProcessSpawnOptions, ) -> LuaResult { let stdout = options.stdio.stdout; let stderr = options.stdio.stderr; let child = options .into_command(program, args) .stdin(Stdio::piped()) .stdout(stdout.as_stdio()) .stderr(stderr.as_stdio()) .spawn()?; Ok(child) }