From 67597b7c4928580f1a824c9819a18f703237594d Mon Sep 17 00:00:00 2001 From: Filip Tibell Date: Mon, 22 Apr 2024 22:56:00 +0200 Subject: [PATCH] Migrate process builtin to lune-std-process crate --- Cargo.lock | 5 + crates/lune-std-process/Cargo.toml | 10 + crates/lune-std-process/src/lib.rs | 182 +++++++++++++++++- crates/lune-std-process/src/options/kind.rs | 80 ++++++++ crates/lune-std-process/src/options/mod.rs | 177 +++++++++++++++++ crates/lune-std-process/src/options/stdio.rs | 56 ++++++ crates/lune-std-process/src/tee_writer.rs | 64 ++++++ crates/lune-std-process/src/wait_for_child.rs | 73 +++++++ crates/lune-utils/src/path.rs | 26 ++- 9 files changed, 669 insertions(+), 4 deletions(-) create mode 100644 crates/lune-std-process/src/options/kind.rs create mode 100644 crates/lune-std-process/src/options/mod.rs create mode 100644 crates/lune-std-process/src/options/stdio.rs create mode 100644 crates/lune-std-process/src/tee_writer.rs create mode 100644 crates/lune-std-process/src/wait_for_child.rs diff --git a/Cargo.lock b/Cargo.lock index 91bbf9f..b86ff79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1599,8 +1599,13 @@ dependencies = [ name = "lune-std-process" version = "0.1.0" dependencies = [ + "directories", "lune-utils", "mlua", + "mlua-luau-scheduler 0.0.1", + "os_str_bytes", + "pin-project", + "tokio", ] [[package]] diff --git a/crates/lune-std-process/Cargo.toml b/crates/lune-std-process/Cargo.toml index e44b683..bb050b0 100644 --- a/crates/lune-std-process/Cargo.toml +++ b/crates/lune-std-process/Cargo.toml @@ -12,5 +12,15 @@ workspace = true [dependencies] mlua = { version = "0.9.7", features = ["luau"] } +mlua-luau-scheduler = "0.0.1" + +directories = "5.0" +pin-project = "1.0" +os_str_bytes = { version = "7.0", features = ["conversions"] } + +tokio = { version = "1", default-features = false, features = [ + "sync", + "process", +] } lune-utils = { version = "0.1.0", path = "../lune-utils" } diff --git a/crates/lune-std-process/src/lib.rs b/crates/lune-std-process/src/lib.rs index 4490966..bd0d27c 100644 --- a/crates/lune-std-process/src/lib.rs +++ b/crates/lune-std-process/src/lib.rs @@ -1,8 +1,28 @@ #![allow(clippy::cargo_common_metadata)] +use std::{ + env::{ + self, + consts::{ARCH, OS}, + }, + process::Stdio, +}; + use mlua::prelude::*; use lune_utils::TableBuilder; +use mlua_luau_scheduler::{Functions, LuaSpawnExt}; +use os_str_bytes::RawOsString; +use tokio::io::AsyncWriteExt; + +mod options; +mod tee_writer; +mod wait_for_child; + +use self::options::ProcessSpawnOptions; +use self::wait_for_child::{wait_for_child, WaitForChildResult}; + +use lune_utils::path::get_current_dir; /** Creates the `process` standard library module. @@ -11,6 +31,166 @@ use lune_utils::TableBuilder; Errors when out of memory. */ +#[allow(clippy::missing_panics_doc)] pub fn module(lua: &Lua) -> LuaResult { - TableBuilder::new(lua)?.build_readonly() + let cwd_str = get_current_dir() + .to_str() + .expect("cwd should be valid UTF-8") + .to_string(); + // Create constants for OS & processor architecture + let os = lua.create_string(&OS.to_lowercase())?; + let arch = lua.create_string(&ARCH.to_lowercase())?; + // 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("args", args_tab)? + .with_value("cwd", cwd_str)? + .with_value("env", env_tab)? + .with_value("exit", process_exit)? + .with_async_function("spawn", process_spawn)? + .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_spawn( + lua: &Lua, + (program, args, options): (String, Option>, ProcessSpawnOptions), +) -> LuaResult { + let res = lua + .spawn(spawn_command(program, args, options)) + .await + .expect("Failed to receive result of spawned process"); + + /* + 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() +} + +async fn spawn_command( + program: String, + args: Option>, + mut options: ProcessSpawnOptions, +) -> LuaResult { + let stdout = options.stdio.stdout; + let stderr = options.stdio.stderr; + let stdin = options.stdio.stdin.take(); + + let mut child = options + .into_command(program, args) + .stdin(if stdin.is_some() { + Stdio::piped() + } else { + Stdio::null() + }) + .stdout(stdout.as_stdio()) + .stderr(stderr.as_stdio()) + .spawn()?; + + if let Some(stdin) = stdin { + let mut child_stdin = child.stdin.take().unwrap(); + child_stdin.write_all(&stdin).await.into_lua_err()?; + } + + wait_for_child(child, stdout, stderr).await } diff --git a/crates/lune-std-process/src/options/kind.rs b/crates/lune-std-process/src/options/kind.rs new file mode 100644 index 0000000..8eff17d --- /dev/null +++ b/crates/lune-std-process/src/options/kind.rs @@ -0,0 +1,80 @@ +use std::{fmt, process::Stdio, str::FromStr}; + +use mlua::prelude::*; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum ProcessSpawnOptionsStdioKind { + // TODO: We need better more obvious names + // for these, but that is a breaking change + #[default] + Default, + Forward, + Inherit, + None, +} + +impl ProcessSpawnOptionsStdioKind { + pub fn all() -> &'static [Self] { + &[Self::Default, Self::Forward, Self::Inherit, Self::None] + } + + pub fn as_stdio(self) -> Stdio { + match self { + Self::None => Stdio::null(), + Self::Forward => Stdio::inherit(), + _ => Stdio::piped(), + } + } +} + +impl fmt::Display for ProcessSpawnOptionsStdioKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let s = match *self { + Self::Default => "default", + Self::Forward => "forward", + Self::Inherit => "inherit", + Self::None => "none", + }; + f.write_str(s) + } +} + +impl FromStr for ProcessSpawnOptionsStdioKind { + type Err = LuaError; + fn from_str(s: &str) -> Result { + Ok(match s.trim().to_ascii_lowercase().as_str() { + "default" => Self::Default, + "forward" => Self::Forward, + "inherit" => Self::Inherit, + "none" => Self::None, + _ => { + return Err(LuaError::RuntimeError(format!( + "Invalid spawn options stdio kind - got '{}', expected one of {}", + s, + ProcessSpawnOptionsStdioKind::all() + .iter() + .map(|k| format!("'{k}'")) + .collect::>() + .join(", ") + ))) + } + }) + } +} + +impl<'lua> FromLua<'lua> for ProcessSpawnOptionsStdioKind { + fn from_lua(value: LuaValue<'lua>, _: &'lua Lua) -> LuaResult { + match value { + LuaValue::Nil => Ok(Self::default()), + LuaValue::String(s) => s.to_str()?.parse(), + _ => Err(LuaError::FromLuaConversionError { + from: value.type_name(), + to: "ProcessSpawnOptionsStdioKind", + message: Some(format!( + "Invalid spawn options stdio kind - expected string, got {}", + value.type_name() + )), + }), + } + } +} diff --git a/crates/lune-std-process/src/options/mod.rs b/crates/lune-std-process/src/options/mod.rs new file mode 100644 index 0000000..8ef8be0 --- /dev/null +++ b/crates/lune-std-process/src/options/mod.rs @@ -0,0 +1,177 @@ +use std::{ + collections::HashMap, + env::{self}, + path::PathBuf, +}; + +use directories::UserDirs; +use mlua::prelude::*; +use tokio::process::Command; + +mod kind; +mod stdio; + +pub(super) use kind::*; +pub(super) use stdio::*; + +#[derive(Debug, Clone, Default)] +pub(super) struct ProcessSpawnOptions { + pub cwd: Option, + pub envs: HashMap, + pub shell: Option, + pub stdio: ProcessSpawnOptionsStdio, +} + +impl<'lua> FromLua<'lua> for ProcessSpawnOptions { + fn from_lua(value: LuaValue<'lua>, _: &'lua Lua) -> LuaResult { + let mut this = Self::default(); + let value = match value { + LuaValue::Nil => return Ok(this), + LuaValue::Table(t) => t, + _ => { + return Err(LuaError::FromLuaConversionError { + from: value.type_name(), + to: "ProcessSpawnOptions", + message: Some(format!( + "Invalid spawn options - expected table, got {}", + value.type_name() + )), + }) + } + }; + + /* + If we got a working directory to use: + + 1. Substitute leading tilde (~) for the users home dir + 2. Make sure it exists + */ + match value.get("cwd")? { + LuaValue::Nil => {} + LuaValue::String(s) => { + let mut cwd = PathBuf::from(s.to_str()?); + if let Ok(stripped) = cwd.strip_prefix("~") { + let user_dirs = UserDirs::new().ok_or_else(|| { + LuaError::runtime( + "Invalid value for option 'cwd' - failed to get home directory", + ) + })?; + cwd = user_dirs.home_dir().join(stripped); + } + if !cwd.exists() { + return Err(LuaError::runtime( + "Invalid value for option 'cwd' - path does not exist", + )); + }; + this.cwd = Some(cwd); + } + value => { + return Err(LuaError::RuntimeError(format!( + "Invalid type for option 'cwd' - expected string, got '{}'", + value.type_name() + ))) + } + } + + /* + If we got environment variables, make sure they are strings + */ + match value.get("env")? { + LuaValue::Nil => {} + LuaValue::Table(e) => { + for pair in e.pairs::() { + let (k, v) = pair.context("Environment variables must be strings")?; + this.envs.insert(k, v); + } + } + value => { + return Err(LuaError::RuntimeError(format!( + "Invalid type for option 'env' - expected table, got '{}'", + value.type_name() + ))) + } + } + + /* + If we got a shell to use: + + 1. When given as a string, use that literally + 2. When set to true, use a default shell for the platform + */ + match value.get("shell")? { + LuaValue::Nil => {} + LuaValue::String(s) => this.shell = Some(s.to_string_lossy().to_string()), + LuaValue::Boolean(true) => { + this.shell = match env::consts::FAMILY { + "unix" => Some("/bin/sh".to_string()), + "windows" => Some("powershell".to_string()), + _ => None, + }; + } + value => { + return Err(LuaError::RuntimeError(format!( + "Invalid type for option 'shell' - expected 'true' or 'string', got '{}'", + value.type_name() + ))) + } + } + + /* + If we got options for stdio handling, parse those as well - note that + we accept a separate "stdin" value here for compatibility with older + scripts, but the user should preferrably pass it in the stdio table + */ + this.stdio = value.get("stdio")?; + match value.get("stdin")? { + LuaValue::Nil => {} + LuaValue::String(s) => this.stdio.stdin = Some(s.as_bytes().to_vec()), + value => { + return Err(LuaError::RuntimeError(format!( + "Invalid type for option 'stdin' - expected 'string', got '{}'", + value.type_name() + ))) + } + } + + Ok(this) + } +} + +impl ProcessSpawnOptions { + pub fn into_command(self, program: impl Into, args: Option>) -> Command { + let mut program = program.into(); + + // Run a shell using the command param if wanted + let pargs = match self.shell { + None => args, + Some(shell) => { + let shell_args = match args { + Some(args) => vec!["-c".to_string(), format!("{} {}", program, args.join(" "))], + None => vec!["-c".to_string(), program.to_string()], + }; + program = shell.to_string(); + Some(shell_args) + } + }; + + // Create command with the wanted options + let mut cmd = match pargs { + None => Command::new(program), + Some(args) => { + let mut cmd = Command::new(program); + cmd.args(args); + cmd + } + }; + + // Set dir to run in and env variables + if let Some(cwd) = self.cwd { + cmd.current_dir(cwd); + } + if !self.envs.is_empty() { + cmd.envs(self.envs); + } + + cmd + } +} diff --git a/crates/lune-std-process/src/options/stdio.rs b/crates/lune-std-process/src/options/stdio.rs new file mode 100644 index 0000000..4c12ff4 --- /dev/null +++ b/crates/lune-std-process/src/options/stdio.rs @@ -0,0 +1,56 @@ +use mlua::prelude::*; + +use super::kind::ProcessSpawnOptionsStdioKind; + +#[derive(Debug, Clone, Default)] +pub struct ProcessSpawnOptionsStdio { + pub stdout: ProcessSpawnOptionsStdioKind, + pub stderr: ProcessSpawnOptionsStdioKind, + pub stdin: Option>, +} + +impl From for ProcessSpawnOptionsStdio { + fn from(value: ProcessSpawnOptionsStdioKind) -> Self { + Self { + stdout: value, + stderr: value, + ..Default::default() + } + } +} + +impl<'lua> FromLua<'lua> for ProcessSpawnOptionsStdio { + fn from_lua(value: LuaValue<'lua>, lua: &'lua Lua) -> LuaResult { + match value { + LuaValue::Nil => Ok(Self::default()), + LuaValue::String(s) => { + Ok(ProcessSpawnOptionsStdioKind::from_lua(LuaValue::String(s), lua)?.into()) + } + LuaValue::Table(t) => { + let mut this = Self::default(); + + if let Some(stdin) = t.get("stdin")? { + this.stdin = stdin; + } + + if let Some(stdout) = t.get("stdout")? { + this.stdout = stdout; + } + + if let Some(stderr) = t.get("stderr")? { + this.stderr = stderr; + } + + Ok(this) + } + _ => Err(LuaError::FromLuaConversionError { + from: value.type_name(), + to: "ProcessSpawnOptionsStdio", + message: Some(format!( + "Invalid spawn options stdio - expected string or table, got {}", + value.type_name() + )), + }), + } + } +} diff --git a/crates/lune-std-process/src/tee_writer.rs b/crates/lune-std-process/src/tee_writer.rs new file mode 100644 index 0000000..fee7776 --- /dev/null +++ b/crates/lune-std-process/src/tee_writer.rs @@ -0,0 +1,64 @@ +use std::{ + io::Write, + pin::Pin, + task::{Context, Poll}, +}; + +use pin_project::pin_project; +use tokio::io::{self, AsyncWrite}; + +#[pin_project] +pub struct AsyncTeeWriter<'a, W> +where + W: AsyncWrite + Unpin, +{ + #[pin] + writer: &'a mut W, + buffer: Vec, +} + +impl<'a, W> AsyncTeeWriter<'a, W> +where + W: AsyncWrite + Unpin, +{ + pub fn new(writer: &'a mut W) -> Self { + Self { + writer, + buffer: Vec::new(), + } + } + + pub fn into_vec(self) -> Vec { + self.buffer + } +} + +impl<'a, W> AsyncWrite for AsyncTeeWriter<'a, W> +where + W: AsyncWrite + Unpin, +{ + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let mut this = self.project(); + match this.writer.as_mut().poll_write(cx, buf) { + Poll::Ready(res) => { + this.buffer + .write_all(buf) + .expect("Failed to write to internal tee buffer"); + Poll::Ready(res) + } + Poll::Pending => Poll::Pending, + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().writer.as_mut().poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().writer.as_mut().poll_shutdown(cx) + } +} diff --git a/crates/lune-std-process/src/wait_for_child.rs b/crates/lune-std-process/src/wait_for_child.rs new file mode 100644 index 0000000..4343041 --- /dev/null +++ b/crates/lune-std-process/src/wait_for_child.rs @@ -0,0 +1,73 @@ +use std::process::ExitStatus; + +use mlua::prelude::*; +use tokio::{ + io::{self, AsyncRead, AsyncReadExt}, + process::Child, + task, +}; + +use super::{options::ProcessSpawnOptionsStdioKind, tee_writer::AsyncTeeWriter}; + +#[derive(Debug, Clone)] +pub(super) struct WaitForChildResult { + pub status: ExitStatus, + pub stdout: Vec, + pub stderr: Vec, +} + +async fn read_with_stdio_kind( + read_from: Option, + kind: ProcessSpawnOptionsStdioKind, +) -> LuaResult> +where + R: AsyncRead + Unpin, +{ + Ok(match kind { + ProcessSpawnOptionsStdioKind::None | ProcessSpawnOptionsStdioKind::Forward => Vec::new(), + ProcessSpawnOptionsStdioKind::Default => { + let mut read_from = + read_from.expect("read_from must be Some when stdio kind is Default"); + + let mut buffer = Vec::new(); + + read_from.read_to_end(&mut buffer).await.into_lua_err()?; + + buffer + } + ProcessSpawnOptionsStdioKind::Inherit => { + let mut read_from = + read_from.expect("read_from must be Some when stdio kind is Inherit"); + + let mut stdout = io::stdout(); + let mut tee = AsyncTeeWriter::new(&mut stdout); + + io::copy(&mut read_from, &mut tee).await.into_lua_err()?; + + tee.into_vec() + } + }) +} + +pub(super) async fn wait_for_child( + mut child: Child, + stdout_kind: ProcessSpawnOptionsStdioKind, + stderr_kind: ProcessSpawnOptionsStdioKind, +) -> LuaResult { + let stdout_opt = child.stdout.take(); + let stderr_opt = child.stderr.take(); + + let stdout_task = task::spawn(read_with_stdio_kind(stdout_opt, stdout_kind)); + let stderr_task = task::spawn(read_with_stdio_kind(stderr_opt, stderr_kind)); + + let status = child.wait().await.expect("Child process failed to start"); + + let stdout_buffer = stdout_task.await.into_lua_err()??; + let stderr_buffer = stderr_task.await.into_lua_err()??; + + Ok(WaitForChildResult { + status, + stdout: stdout_buffer, + stderr: stderr_buffer, + }) +} diff --git a/crates/lune-utils/src/path.rs b/crates/lune-utils/src/path.rs index a510e22..febaac6 100644 --- a/crates/lune-utils/src/path.rs +++ b/crates/lune-utils/src/path.rs @@ -1,6 +1,6 @@ use std::{ env::{current_dir, current_exe}, - path::{Path, PathBuf}, + path::{Path, PathBuf, MAIN_SEPARATOR}, sync::Arc, }; @@ -11,14 +11,25 @@ static CWD: Lazy> = Lazy::new(create_cwd); static EXE: Lazy> = Lazy::new(create_exe); fn create_cwd() -> Arc { - let cwd = current_dir().expect("failed to find current working directory"); + let mut cwd = current_dir() + .expect("failed to find current working directory") + .to_str() + .expect("current working directory is not valid UTF-8") + .to_string(); + if !cwd.ends_with(MAIN_SEPARATOR) { + cwd.push(MAIN_SEPARATOR); + } dunce::canonicalize(cwd) .expect("failed to canonicalize current working directory") .into() } fn create_exe() -> Arc { - let exe = current_exe().expect("failed to find current executable"); + let exe = current_exe() + .expect("failed to find current executable") + .to_str() + .expect("current executable is not valid UTF-8") + .to_string(); dunce::canonicalize(exe) .expect("failed to canonicalize current executable") .into() @@ -29,6 +40,11 @@ fn create_exe() -> Arc { This absolute path is canonicalized and does not contain any `.` or `..` components, and it is also in a friendly (non-UNC) format. + + This path is also guaranteed to: + + - Be valid UTF-8. + - End with the platform's main path separator. */ #[must_use] pub fn get_current_dir() -> Arc { @@ -40,6 +56,10 @@ pub fn get_current_dir() -> Arc { This absolute path is canonicalized and does not contain any `.` or `..` components, and it is also in a friendly (non-UNC) format. + + This path is also guaranteed to: + + - Be valid UTF-8. */ #[must_use] pub fn get_current_exe() -> Arc {