use std::{any::TypeId, ffi::OsStr, io, process::Stdio}; use mlua::{AnyUserData, ExternalError, IntoLua, IntoLuaMulti, Lua, MetaMethod, Table, UserData, UserDataMethods, Value}; use tokio::process::{ChildStderr, ChildStdin, ChildStdout}; use yazi_shared::wtf8::FromWtf8; use super::{Child, output::Output}; use crate::{Error, process::Status}; pub struct Command { inner: tokio::process::Command, memory: Option, } const NULL: u8 = 0; const PIPED: u8 = 1; const INHERIT: u8 = 2; impl Command { pub fn install(lua: &Lua) -> mlua::Result<()> { let new = lua.create_function(|_, (_, program): (Table, String)| { let mut inner = tokio::process::Command::new(program); inner.kill_on_drop(true).stdin(Stdio::null()).stdout(Stdio::null()).stderr(Stdio::null()); Ok(Self { inner, memory: None }) })?; let command = lua.create_table_from([ // Stdio ("NULL", NULL), ("PIPED", PIPED), ("INHERIT", INHERIT), ])?; command.set_metatable(Some(lua.create_table_from([(MetaMethod::Call.name(), new)])?))?; lua.globals().raw_set("Command", command) } #[cfg(unix)] fn spawn(&mut self) -> io::Result { if let Some(max) = self.memory { unsafe { self.inner.pre_exec(move || { let rlp = libc::rlimit { rlim_cur: max as _, rlim_max: max as _ }; libc::setrlimit(libc::RLIMIT_AS, &rlp); Ok(()) }); } } self.inner.spawn().map(Child::new) } #[cfg(windows)] fn spawn(&mut self) -> io::Result { use std::os::windows::io::RawHandle; use windows_sys::Win32::{Foundation::CloseHandle, System::JobObjects::{AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, JOB_OBJECT_LIMIT_PROCESS_MEMORY, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, SetInformationJobObject}}; fn create_job(handle: RawHandle, memory: Option) -> io::Result { unsafe { let job = CreateJobObjectW(std::ptr::null_mut(), std::ptr::null()); if job.is_null() { return Err(io::Error::last_os_error()); } let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed(); info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; if let Some(m) = memory { info.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_PROCESS_MEMORY; info.ProcessMemoryLimit = m; } if SetInformationJobObject( job, JobObjectExtendedLimitInformation, &mut info as *mut _ as *mut _, std::mem::size_of_val(&info) as u32, ) == 0 { CloseHandle(job); return Err(io::Error::last_os_error()); } if AssignProcessToJobObject(job, handle) == 0 { CloseHandle(job); return Err(io::Error::last_os_error()); } Ok(job) } } let child = self.inner.spawn()?; if let Some(handle) = child.raw_handle() && let Ok(job) = create_job(handle, self.memory) { return Ok(Child::new(child, Some(job))); } Ok(Child::new(child, None)) } async fn output(&mut self) -> io::Result { self.inner.stdout(Stdio::piped()); self.inner.stderr(Stdio::piped()); self.spawn()?.wait_with_output().await } async fn status(&mut self) -> io::Result { self.spawn()?.status().await } } impl UserData for Command { fn add_methods>(methods: &mut M) { #[inline] fn make_stdio(v: Value) -> mlua::Result { match v { Value::Integer(n) => { return Ok(match n as u8 { PIPED => Stdio::piped(), INHERIT => Stdio::inherit(), _ => Stdio::null(), }); } Value::UserData(ud) => match ud.type_id() { Some(t) if t == TypeId::of::() => { return Ok(ud.take::()?.try_into()?); } Some(t) if t == TypeId::of::() => { return Ok(ud.take::()?.try_into()?); } Some(t) if t == TypeId::of::() => { return Ok(ud.take::()?.try_into()?); } _ => {} }, _ => {} } Err( "must be one of Command.NULL, Command.PIPED, Command.INHERIT, or a ChildStdin, ChildStdout, or ChildStderr".into_lua_err(), ) } methods.add_function("arg", |lua, (ud, arg): (AnyUserData, Value)| { let mut me = ud.borrow_mut::()?; match arg { Value::Nil => return lua.create_sequence_from(me.inner.as_std().get_args())?.into_lua(lua), Value::String(s) => { me.inner.arg(OsStr::from_wtf8(&s.as_bytes())?); } Value::Table(t) => { for s in t.sequence_values::() { me.inner.arg(OsStr::from_wtf8(&s?.as_bytes())?); } } _ => Err("arg must be a string or table of strings".into_lua_err())?, } ud.into_lua(lua) }); methods.add_function("cwd", |_, (ud, dir): (AnyUserData, mlua::String)| { ud.borrow_mut::()?.inner.current_dir(dir.to_str()?.as_ref()); Ok(ud) }); methods.add_function( "env", |_, (ud, key, value): (AnyUserData, mlua::String, mlua::String)| { ud.borrow_mut::()? .inner .env(OsStr::from_wtf8(&key.as_bytes())?, OsStr::from_wtf8(&value.as_bytes())?); Ok(ud) }, ); methods.add_function("stdin", |_, (ud, stdio): (AnyUserData, Value)| { ud.borrow_mut::()?.inner.stdin(make_stdio(stdio)?); Ok(ud) }); methods.add_function("stdout", |_, (ud, stdio): (AnyUserData, Value)| { ud.borrow_mut::()?.inner.stdout(make_stdio(stdio)?); Ok(ud) }); methods.add_function("stderr", |_, (ud, stdio): (AnyUserData, Value)| { ud.borrow_mut::()?.inner.stderr(make_stdio(stdio)?); Ok(ud) }); methods.add_function("memory", |_, (ud, max): (AnyUserData, usize)| { ud.borrow_mut::()?.memory = Some(max); Ok(ud) }); methods.add_method_mut("spawn", |lua, me, ()| match me.spawn() { Ok(child) => child.into_lua_multi(lua), Err(e) => (Value::Nil, Error::Io(e)).into_lua_multi(lua), }); methods.add_async_method_mut("output", |lua, mut me, ()| async move { match me.output().await { Ok(output) => Output::new(output).into_lua_multi(&lua), Err(e) => (Value::Nil, Error::Io(e)).into_lua_multi(&lua), } }); methods.add_async_method_mut("status", |lua, mut me, ()| async move { match me.status().await { Ok(status) => Status::new(status).into_lua_multi(&lua), Err(e) => (Value::Nil, Error::Io(e)).into_lua_multi(&lua), } }); } }