diff --git a/luneTypes.d.luau b/luneTypes.d.luau index 44bc1dc..132cbc8 100644 --- a/luneTypes.d.luau +++ b/luneTypes.d.luau @@ -322,7 +322,7 @@ declare process: { options: { cwd: string?, env: { [string]: string }?, - shell: boolean | string, + shell: (boolean | string)?, stdio: ("inherit" | "default")?, }? ) -> { diff --git a/src/lib/globals/process.rs b/src/lib/globals/process.rs index c9520e8..9c901e8 100644 --- a/src/lib/globals/process.rs +++ b/src/lib/globals/process.rs @@ -1,12 +1,26 @@ -use std::{env, process::Stdio, sync::Weak}; +use std::{ + collections::HashMap, + env, + path::PathBuf, + process::{Command, Stdio}, + sync::Weak, +}; use mlua::prelude::*; use os_str_bytes::RawOsString; -use smol::{channel::Sender, process::Command}; +use smol::channel::Sender; -use crate::{utils::table::TableBuilder, LuneMessage}; +use crate::{ + utils::{process::pipe_and_inherit_child_process_stdio, table::TableBuilder}, + LuneMessage, +}; pub fn create(lua: &Lua, args_vec: Vec) -> LuaResult<()> { + let cwd = env::current_dir()?.canonicalize()?; + let mut cwd_str = cwd.to_string_lossy().to_string(); + if !cwd_str.ends_with('/') { + cwd_str = format!("{}/", cwd_str); + } // Create readonly args array let args_tab = TableBuilder::new(lua)? .with_sequential_values(args_vec)? @@ -26,9 +40,10 @@ pub fn create(lua: &Lua, args_vec: Vec) -> LuaResult<()> { "process", TableBuilder::new(lua)? .with_value("args", args_tab)? + .with_value("cwd", cwd_str)? .with_value("env", env_tab)? .with_async_function("exit", process_exit)? - .with_async_function("spawn", process_spawn)? + .with_function("spawn", process_spawn)? .build_readonly()?, ) } @@ -112,37 +127,141 @@ async fn process_exit(lua: &Lua, exit_code: Option) -> LuaResult<()> { Ok(()) } -async fn process_spawn( - lua: &Lua, - (program, args): (String, Option>), -) -> LuaResult { - // Create and spawn our child process - let pwd = env::current_dir()?; - let mut cmd = Command::new(program); - if let Some(args) = args { - cmd.args(args); - } - let output = cmd - .current_dir(pwd) +fn process_spawn<'a>( + lua: &'a Lua, + (mut program, args, options): (String, Option>, Option>), +) -> LuaResult> { + // Parse any given options or create defaults + let (child_cwd, child_envs, child_shell, child_stdio_inherit) = match options { + Some(options) => { + let mut cwd = env::current_dir()?; + let mut envs = HashMap::new(); + let mut shell = None; + let mut inherit = false; + match options.raw_get("cwd")? { + LuaValue::Nil => {} + LuaValue::String(s) => { + cwd = PathBuf::from(s.to_string_lossy().to_string()); + if !cwd.exists() { + return Err(LuaError::RuntimeError( + "Invalid value for option 'cwd' - path does not exist".to_string(), + )); + } + } + value => { + return Err(LuaError::RuntimeError(format!( + "Invalid type for option 'cwd' - expected 'string', got '{}'", + value.type_name() + ))) + } + } + match options.raw_get("env")? { + LuaValue::Nil => {} + LuaValue::Table(t) => { + for pair in t.pairs::() { + let (k, v) = pair?; + envs.insert(k, v); + } + } + value => { + return Err(LuaError::RuntimeError(format!( + "Invalid type for option 'env' - expected 'table', got '{}'", + value.type_name() + ))) + } + } + match options.raw_get("shell")? { + LuaValue::Nil => {} + LuaValue::String(s) => shell = Some(s.to_string_lossy().to_string()), + LuaValue::Boolean(true) => { + shell = match env::consts::FAMILY { + "unix" => Some("/bin/sh".to_string()), + "windows" => Some("/bin/sh".to_string()), + _ => None, + }; + } + value => { + return Err(LuaError::RuntimeError(format!( + "Invalid type for option 'shell' - expected 'true' or 'string', got '{}'", + value.type_name() + ))) + } + } + match options.raw_get("stdio")? { + LuaValue::Nil => {} + LuaValue::String(s) => { + match s.to_str()? { + "inherit" => { + inherit = true; + }, + "default" => { + inherit = false; + } + _ => return Err(LuaError::RuntimeError( + format!("Invalid value for option 'stdio' - expected 'inherit' or 'default', got '{}'", s.to_string_lossy()), + )) + } + } + value => { + return Err(LuaError::RuntimeError(format!( + "Invalid type for option 'stdio' - expected 'string', got '{}'", + value.type_name() + ))) + } + } + Ok::<_, LuaError>((cwd, envs, shell, inherit)) + } + None => Ok((env::current_dir()?, HashMap::new(), None, false)), + }?; + // Run a shell using the command param if wanted + let child_args = if let Some(shell) = child_shell { + let shell_args = match args { + Some(args) => vec!["-c".to_string(), format!("{} {}", program, args.join(" "))], + None => vec!["-c".to_string(), program], + }; + program = shell; + Some(shell_args) + } else { + args + }; + // Create command with the wanted options + let mut cmd = match child_args { + None => Command::new(program), + Some(args) => { + let mut cmd = Command::new(program); + cmd.args(args); + cmd + } + }; + // FUTURE: Implement and test for tilde (~) to home dir substitution in child_cwd + cmd.current_dir(child_cwd); + cmd.envs(child_envs); + // Spawn the child process + let child = cmd .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) - .output() - .await?; + .spawn()?; + // Inherit the output and stderr if wanted + let result = if child_stdio_inherit { + pipe_and_inherit_child_process_stdio(child) + } else { + let output = child.wait_with_output()?; + Ok((output.status, output.stdout, output.stderr)) + }; + // Extract result + let (status, stdout, stderr) = result?; // NOTE: If an exit code was not given by the child process, // we default to 1 if it yielded any error output, otherwise 0 - let code = output - .status - .code() - .unwrap_or(match output.stderr.is_empty() { - true => 0, - false => 1, - }); + let code = status.code().unwrap_or(match stderr.is_empty() { + true => 0, + false => 1, + }); // 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(&output.stdout)?)? - .with_value("stderr", lua.create_string(&output.stderr)?)? + .with_value("stdout", lua.create_string(&stdout)?)? + .with_value("stderr", lua.create_string(&stderr)?)? .build_readonly() } diff --git a/src/lib/lib.rs b/src/lib/lib.rs index 6236635..937c197 100644 --- a/src/lib/lib.rs +++ b/src/lib/lib.rs @@ -219,6 +219,7 @@ mod tests { net_json_decode: "net/json/decode", net_json_encode: "net/json/encode", process_args: "process/args", + process_cwd: "process/cwd", process_env: "process/env", process_exit: "process/exit", process_spawn: "process/spawn", diff --git a/src/lib/utils/mod.rs b/src/lib/utils/mod.rs index 8775dcd..a43306a 100644 --- a/src/lib/utils/mod.rs +++ b/src/lib/utils/mod.rs @@ -1,4 +1,5 @@ pub mod formatting; pub mod net; +pub mod process; pub mod table; pub mod task; diff --git a/src/lib/utils/process.rs b/src/lib/utils/process.rs new file mode 100644 index 0000000..49ad233 --- /dev/null +++ b/src/lib/utils/process.rs @@ -0,0 +1,74 @@ +// https://stackoverflow.com/questions/71141122/- + +use std::{ + io, + io::Write, + process::{Child, ExitStatus}, +}; + +use mlua::prelude::*; + +pub struct TeeWriter<'a, W0: Write, W1: Write> { + w0: &'a mut W0, + w1: &'a mut W1, +} + +impl<'a, W0: Write, W1: Write> TeeWriter<'a, W0, W1> { + pub fn new(w0: &'a mut W0, w1: &'a mut W1) -> Self { + Self { w0, w1 } + } +} + +impl<'a, W0: Write, W1: Write> Write for TeeWriter<'a, W0, W1> { + fn write(&mut self, buf: &[u8]) -> io::Result { + // We have to use write_all() otherwise what + // happens if different amounts are written? + self.w0.write_all(buf)?; + self.w1.write_all(buf)?; + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.w0.flush()?; + self.w1.flush()?; + Ok(()) + } +} + +pub fn pipe_and_inherit_child_process_stdio( + mut child: Child, +) -> LuaResult<(ExitStatus, Vec, Vec)> { + // https://stackoverflow.com/questions/71141122/- + let mut child_stdout = child.stdout.take().unwrap(); + let mut child_stderr = child.stderr.take().unwrap(); + std::thread::scope(|s| { + let stdout_thread = s.spawn(|| { + let stdout = io::stdout(); + let mut log = Vec::new(); + let mut stdout = stdout.lock(); + let mut tee = TeeWriter::new(&mut stdout, &mut log); + + io::copy(&mut child_stdout, &mut tee).map_err(LuaError::external)?; + + Ok(log) + }); + + let stderr_thread = s.spawn(|| { + let stderr = io::stderr(); + let mut log = Vec::new(); + let mut stderr = stderr.lock(); + let mut tee = TeeWriter::new(&mut stderr, &mut log); + + io::copy(&mut child_stderr, &mut tee).map_err(LuaError::external)?; + + Ok(log) + }); + + let status = child.wait().expect("child wasn't running"); + + let stdout_log: Result<_, LuaError> = stdout_thread.join().expect("stdout thread panicked"); + let stderr_log: Result<_, LuaError> = stderr_thread.join().expect("stderr thread panicked"); + + Ok::<_, LuaError>((status, stdout_log?, stderr_log?)) + }) +} diff --git a/src/tests/process/cwd.luau b/src/tests/process/cwd.luau new file mode 100644 index 0000000..a41d42c --- /dev/null +++ b/src/tests/process/cwd.luau @@ -0,0 +1,7 @@ +assert(process.cwd ~= nil, "Process cwd is missing") + +assert(type(process.cwd) == "string", "Process cwd is not a string") + +assert(#process.cwd > 0, "Process cwd is an empty string") + +assert(string.sub(process.cwd, -1) == "/", "Process cwd does not end with '/'") diff --git a/src/tests/process/spawn.luau b/src/tests/process/spawn.luau index cd6bda8..e7bcc55 100644 --- a/src/tests/process/spawn.luau +++ b/src/tests/process/spawn.luau @@ -1,3 +1,5 @@ +-- Spawning a child process should work with options + local result = process.spawn("ls", { "-a", }) @@ -9,3 +11,62 @@ assert(result.stdout ~= "", "Stdout was empty") assert(string.find(result.stdout, "Cargo.toml") ~= nil, "Missing Cargo.toml in output") assert(string.find(result.stdout, ".gitignore") ~= nil, "Missing .gitignore in output") + +-- It should also work the same when spawned using a shell + +local shellResult = process.spawn("ls", { + "-a", +}, { + shell = true, +}) + +assert(shellResult.ok, "Failed to spawn child process (shell)") + +assert(shellResult.stderr == "", "Stderr was not empty (shell)") +assert(shellResult.stdout ~= "", "Stdout was empty (shell)") + +assert(string.find(shellResult.stdout, "Cargo.toml") ~= nil, "Missing Cargo.toml in output (shell)") +assert(string.find(shellResult.stdout, ".gitignore") ~= nil, "Missing .gitignore in output (shell)") + +-- Make sure the cwd option actually uses the directory we want +local rootPwd = process.spawn("pwd", {}, { + cwd = "/", +}).stdout +rootPwd = string.gsub(rootPwd, "^%s+", "") +rootPwd = string.gsub(rootPwd, "%s+$", "") +if rootPwd ~= "/" then + error( + string.format( + "Current working directory for child process was not set correctly!" + .. "\nExpected '/', got '%s'", + rootPwd + ) + ) +end + +-- Setting cwd should not change the cwd of this process + +local before = process.spawn("pwd").stdout +process.spawn("ls", {}, { + cwd = "/", + shell = true, +}) +local after = process.spawn("pwd").stdout +assert(before == after, "Current working directory changed after running child process") + +-- Inheriting stdio & environment variables should work + +task.delay(2, function() + local message = "Hello from child process!" + local result = process.spawn("echo", { + '"$TEST_VAR"', + }, { + env = { TEST_VAR = message }, + shell = "bash", + stdio = "inherit", + }) + assert( + result.stdout == (message .. "\n"), -- Note that echo adds a newline + "Inheriting stdio did not return proper output" + ) +end)