From 8a2c5f65bbbe2c36282bfef58a3dcd3f15763363 Mon Sep 17 00:00:00 2001 From: Filip Tibell Date: Sat, 19 Aug 2023 19:43:08 -0500 Subject: [PATCH] Add back process builtin --- src/lune/builtins/mod.rs | 5 + src/lune/builtins/process/mod.rs | 305 ++++++++++++++++++++++ src/lune/builtins/process/pipe_inherit.rs | 46 ++++ src/lune/builtins/process/tee_writer.rs | 64 +++++ src/lune/mod.rs | 1 + src/lune/scheduler/mod.rs | 8 + src/lune/util/table_builder.rs | 29 ++ 7 files changed, 458 insertions(+) create mode 100644 src/lune/builtins/process/mod.rs create mode 100644 src/lune/builtins/process/pipe_inherit.rs create mode 100644 src/lune/builtins/process/tee_writer.rs diff --git a/src/lune/builtins/mod.rs b/src/lune/builtins/mod.rs index b44dd86..8c312b9 100644 --- a/src/lune/builtins/mod.rs +++ b/src/lune/builtins/mod.rs @@ -4,6 +4,7 @@ use mlua::prelude::*; mod fs; mod luau; +mod process; mod serde; mod stdio; mod task; @@ -13,6 +14,7 @@ pub enum LuneBuiltin { Fs, Luau, Task, + Process, Serde, Stdio, } @@ -26,6 +28,7 @@ where Self::Fs => "fs", Self::Luau => "luau", Self::Task => "task", + Self::Process => "process", Self::Serde => "serde", Self::Stdio => "stdio", } @@ -36,6 +39,7 @@ where Self::Fs => fs::create(lua), Self::Luau => luau::create(lua), Self::Task => task::create(lua), + Self::Process => process::create(lua), Self::Serde => serde::create(lua), Self::Stdio => stdio::create(lua), }; @@ -56,6 +60,7 @@ impl FromStr for LuneBuiltin { "fs" => Ok(Self::Fs), "luau" => Ok(Self::Luau), "task" => Ok(Self::Task), + "process" => Ok(Self::Process), "serde" => Ok(Self::Serde), "stdio" => Ok(Self::Stdio), _ => Err(format!("Unknown builtin library '{s}'")), diff --git a/src/lune/builtins/process/mod.rs b/src/lune/builtins/process/mod.rs new file mode 100644 index 0000000..895db4e --- /dev/null +++ b/src/lune/builtins/process/mod.rs @@ -0,0 +1,305 @@ +use std::{ + collections::HashMap, + env::{self, consts}, + path::{self, PathBuf}, + process::Stdio, +}; + +use directories::UserDirs; +use dunce::canonicalize; +use mlua::prelude::*; +use os_str_bytes::RawOsString; +use tokio::process::Command; + +use crate::lune::{scheduler::Scheduler, util::TableBuilder}; + +mod tee_writer; + +mod pipe_inherit; +use pipe_inherit::pipe_and_inherit_child_process_stdio; + +const PROCESS_EXIT_IMPL_LUA: &str = r#" +exit(...) +yield() +"#; + +pub fn create(lua: &'static Lua) -> LuaResult { + let cwd_str = { + let cwd = canonicalize(env::current_dir()?)?; + let cwd_str = cwd.to_string_lossy().to_string(); + if !cwd_str.ends_with(path::MAIN_SEPARATOR) { + format!("{cwd_str}{}", path::MAIN_SEPARATOR) + } else { + cwd_str + } + }; + // Create constants for OS & processor architecture + let os = lua.create_string(&consts::OS.to_lowercase())?; + let arch = lua.create_string(&consts::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, this is a bit involved since + // we have no way to yield from c / rust, we need to load a lua + // chunk that will set the exit code and yield for us instead + let coroutine_yield = lua + .globals() + .get::<_, LuaTable>("coroutine")? + .get::<_, LuaFunction>("yield")?; + let set_scheduler_exit_code = lua.create_function(|lua, code: Option| { + let sched = lua + .app_data_ref::<&Scheduler>() + .expect("Lua struct is missing scheduler"); + sched.set_exit_code(code.unwrap_or_default()); + Ok(()) + })?; + let process_exit = lua + .load(PROCESS_EXIT_IMPL_LUA) + .set_name("=process.exit") + .set_environment( + TableBuilder::new(lua)? + .with_value("yield", coroutine_yield)? + .with_value("exit", set_scheduler_exit_code)? + .build_readonly()?, + ) + .into_function()?; + // 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.as_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 { + match value { + Some(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(()) + } + } + None => { + 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.as_raw_bytes())?), + LuaValue::String(lua.create_string(raw_value.as_raw_bytes())?), + )) + } + None => Ok((LuaValue::Nil, LuaValue::Nil)), + }) +} + +async fn process_spawn<'lua>( + lua: &'static 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()); + // Substitute leading tilde (~) for the actual home dir + if cwd.starts_with("~") { + if let Some(user_dirs) = UserDirs::new() { + cwd = user_dirs.home_dir().join(cwd.strip_prefix("~").unwrap()) + } + }; + 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 + } + }; + // Set dir to run in and env variables + 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()) + .spawn()?; + // Inherit the output and stderr if wanted + let result = if child_stdio_inherit { + pipe_and_inherit_child_process_stdio(child).await + } else { + let output = child.wait_with_output().await?; + 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 = 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(&stdout)?)? + .with_value("stderr", lua.create_string(&stderr)?)? + .build_readonly() +} diff --git a/src/lune/builtins/process/pipe_inherit.rs b/src/lune/builtins/process/pipe_inherit.rs new file mode 100644 index 0000000..0e4b9a3 --- /dev/null +++ b/src/lune/builtins/process/pipe_inherit.rs @@ -0,0 +1,46 @@ +use std::process::ExitStatus; + +use mlua::prelude::*; +use tokio::{io, process::Child, task}; + +use super::tee_writer::AsyncTeeWriter; + +pub async fn pipe_and_inherit_child_process_stdio( + mut child: Child, +) -> LuaResult<(ExitStatus, Vec, Vec)> { + let mut child_stdout = child.stdout.take().unwrap(); + let mut child_stderr = child.stderr.take().unwrap(); + + /* + NOTE: We do not need to register these + independent tasks spawning in the scheduler + + This function is only used by `process.spawn` which in + turn registers a task with the scheduler that awaits this + */ + + let stdout_thread = task::spawn(async move { + let mut stdout = io::stdout(); + let mut tee = AsyncTeeWriter::new(&mut stdout); + + io::copy(&mut child_stdout, &mut tee).await.into_lua_err()?; + + Ok::<_, LuaError>(tee.into_vec()) + }); + + let stderr_thread = task::spawn(async move { + let mut stderr = io::stderr(); + let mut tee = AsyncTeeWriter::new(&mut stderr); + + io::copy(&mut child_stderr, &mut tee).await.into_lua_err()?; + + Ok::<_, LuaError>(tee.into_vec()) + }); + + let status = child.wait().await.expect("Child process failed to start"); + + let stdout_buffer = stdout_thread.await.expect("Tee writer for stdout errored"); + let stderr_buffer = stderr_thread.await.expect("Tee writer for stderr errored"); + + Ok::<_, LuaError>((status, stdout_buffer?, stderr_buffer?)) +} diff --git a/src/lune/builtins/process/tee_writer.rs b/src/lune/builtins/process/tee_writer.rs new file mode 100644 index 0000000..fee7776 --- /dev/null +++ b/src/lune/builtins/process/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/src/lune/mod.rs b/src/lune/mod.rs index 8ac8a88..3aaa11c 100644 --- a/src/lune/mod.rs +++ b/src/lune/mod.rs @@ -49,6 +49,7 @@ impl Lune { V: Into>, { self.args = args.into(); + self.lua.set_app_data(self.args.clone()); self } diff --git a/src/lune/scheduler/mod.rs b/src/lune/scheduler/mod.rs index 6f457d7..3e6a58e 100644 --- a/src/lune/scheduler/mod.rs +++ b/src/lune/scheduler/mod.rs @@ -69,6 +69,14 @@ impl<'lua, 'fut> Scheduler<'lua, 'fut> { this } + pub fn set_exit_code(&self, code: impl Into) { + assert!( + self.state.exit_code().is_none(), + "Exit code may only be set exactly once" + ); + self.state.set_exit_code(code.into()) + } + #[doc(hidden)] pub fn into_static(self) -> &'static Self { Box::leak(Box::new(self)) diff --git a/src/lune/util/table_builder.rs b/src/lune/util/table_builder.rs index 9f022d4..25ded8e 100644 --- a/src/lune/util/table_builder.rs +++ b/src/lune/util/table_builder.rs @@ -26,6 +26,35 @@ impl<'lua> TableBuilder<'lua> { Ok(self) } + pub fn with_values(self, values: Vec<(K, V)>) -> LuaResult + where + K: IntoLua<'lua>, + V: IntoLua<'lua>, + { + for (key, value) in values { + self.tab.raw_set(key, value)?; + } + Ok(self) + } + + pub fn with_sequential_value(self, value: V) -> LuaResult + where + V: IntoLua<'lua>, + { + self.tab.raw_push(value)?; + Ok(self) + } + + pub fn with_sequential_values(self, values: Vec) -> LuaResult + where + V: IntoLua<'lua>, + { + for value in values { + self.tab.raw_push(value)?; + } + Ok(self) + } + pub fn with_function(self, key: K, func: F) -> LuaResult where K: IntoLua<'lua>,