use std::{
collections::BTreeMap,
env,
ffi::{OsStr, OsString},
io,
path::{Path, PathBuf},
process::{Output, Stdio},
};
use smol::{io::AsyncReadExt as _, process::Command, unblock};
use crate::utils::{CommandError, format_failure_stream, std_output_enabled};
mod detached;
#[derive(Debug, Clone)]
pub struct Host {
env: BTreeMap<OsString, OsString>,
cwd: PathBuf,
home: Option<PathBuf>,
app_dirs: Vec<PathBuf>,
}
impl Host {
#[must_use]
pub fn current() -> Self {
let env = env::vars_os().collect();
Self {
env,
cwd: env::current_dir().expect("process must have a working directory"),
home: dirs::home_dir(),
app_dirs: default_app_dirs(),
}
}
pub fn new<P, K, V>(
path_dirs: impl IntoIterator<Item = P>,
vars: impl IntoIterator<Item = (K, V)>,
) -> Self
where
P: AsRef<Path>,
K: AsRef<OsStr>,
V: AsRef<OsStr>,
{
let mut env = BTreeMap::new();
seed_process_plumbing(&mut env);
let path = env::join_paths(
path_dirs
.into_iter()
.map(|dir| dir.as_ref().as_os_str().to_os_string()),
)
.expect("Host::new PATH entries must join into a valid PATH string");
env.insert(OsString::from("PATH"), path);
for (key, value) in vars {
env.insert(key.as_ref().to_os_string(), value.as_ref().to_os_string());
}
let home = home_dir_from_env(&env);
Self {
env,
cwd: env::current_dir().expect("process must have a working directory"),
home,
app_dirs: Vec::new(),
}
}
#[must_use]
pub fn with_cwd(mut self, cwd: impl Into<PathBuf>) -> Self {
self.cwd = cwd.into();
self
}
#[must_use]
pub fn with_app_dirs(mut self, app_dirs: impl IntoIterator<Item = PathBuf>) -> Self {
self.app_dirs = app_dirs.into_iter().collect();
self
}
#[must_use]
pub fn env(&self, key: impl AsRef<OsStr>) -> Option<&OsStr> {
env_get(&self.env, key.as_ref())
}
#[must_use]
pub fn env_string(&self, key: impl AsRef<OsStr>) -> Option<String> {
self.env(key)
.and_then(|value| value.to_str().map(ToOwned::to_owned))
}
#[must_use]
pub fn path_entries(&self) -> Vec<PathBuf> {
self.env("PATH")
.map(|paths| {
env::split_paths(paths)
.filter(|entry| !entry.as_os_str().is_empty())
.collect()
})
.unwrap_or_default()
}
#[must_use]
pub fn cwd(&self) -> &Path {
&self.cwd
}
#[must_use]
pub fn home_dir(&self) -> Option<&Path> {
self.home.as_deref()
}
pub fn current_exe() -> io::Result<PathBuf> {
env::current_exe()
}
#[must_use]
pub fn app_dirs(&self) -> &[PathBuf] {
&self.app_dirs
}
pub async fn which(&self, name: impl AsRef<OsStr>) -> Result<PathBuf, which::Error> {
let name = name.as_ref().to_os_string();
let paths = self.joined_path();
let cwd = self.cwd.clone();
unblock(move || which::which_in(name, paths, cwd)).await
}
#[must_use]
pub fn with_env(&self, key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) -> Self {
let mut host = self.clone();
host.env
.insert(key.as_ref().to_os_string(), value.as_ref().to_os_string());
host
}
#[must_use]
pub fn command(&self, program: impl AsRef<OsStr>) -> Command {
withhold_std_handles_from_children();
let mut command = Command::new(self.resolve_program(program.as_ref()));
command.env_clear().envs(&self.env).current_dir(&self.cwd);
command
}
#[must_use]
pub fn std_command(&self, program: impl AsRef<OsStr>) -> std::process::Command {
withhold_std_handles_from_children();
let mut command = std::process::Command::new(self.resolve_program(program.as_ref()));
command.env_clear().envs(&self.env).current_dir(&self.cwd);
command
}
pub async fn output(
&self,
program: impl AsRef<OsStr>,
args: impl IntoIterator<Item = impl AsRef<OsStr>>,
) -> Result<Output, CommandError> {
let program = program.as_ref();
let program_name = program.to_string_lossy().into_owned();
let args = args
.into_iter()
.map(|argument| argument.as_ref().to_os_string())
.collect::<Vec<_>>();
tracing::debug!(program = %program_name, ?args, "spawning");
let started = std::time::Instant::now();
let mut command = self.command(program);
command
.args(&args)
.kill_on_drop(true)
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let mut child = command.spawn().map_err(|source| CommandError::Spawn {
program: program_name.clone(),
source,
})?;
let echo = std_output_enabled();
let stdout_task = smol::spawn(drain_child_pipe(
child.stdout.take().expect("stdout is piped"),
io::stdout(),
echo,
));
let stderr_task = smol::spawn(drain_child_pipe(
child.stderr.take().expect("stderr is piped"),
io::stderr(),
echo,
));
let status = child.status().await.map_err(|source| CommandError::Spawn {
program: program_name.clone(),
source,
})?;
let stdout = stdout_task.await.map_err(|source| CommandError::Spawn {
program: program_name.clone(),
source,
})?;
let stderr = stderr_task.await.map_err(|source| CommandError::Spawn {
program: program_name.clone(),
source,
})?;
tracing::debug!(
program = %program_name,
%status,
elapsed_ms = started.elapsed().as_millis(),
"exited"
);
Ok(Output {
status,
stdout,
stderr,
})
}
pub async fn run(
&self,
program: impl AsRef<OsStr>,
args: impl IntoIterator<Item = impl AsRef<OsStr>>,
) -> Result<String, CommandError> {
let program = program.as_ref();
let output = self.output(program, args).await?;
if output.status.success() {
Ok(String::from_utf8_lossy(&output.stdout).to_string())
} else {
Err(CommandError::Failed {
program: program.to_string_lossy().into_owned(),
status: output.status,
report: format!(
"{}{}",
format_failure_stream("stderr", &output.stderr),
format_failure_stream("stdout", &output.stdout),
),
})
}
}
pub async fn run_detached(
&self,
program: impl AsRef<OsStr>,
args: impl IntoIterator<Item = impl AsRef<OsStr>>,
) -> Result<std::process::ExitStatus, CommandError> {
let program_name = program.as_ref().to_string_lossy().into_owned();
let args = args
.into_iter()
.map(|argument| argument.as_ref().to_os_string())
.collect::<Vec<_>>();
tracing::debug!(program = %program_name, ?args, "spawning detached");
let resolved = self.resolve_program(program.as_ref());
let env = self.env.clone();
let cwd = self.cwd.clone();
let status = unblock(move || detached::run(&resolved, &args, &env, &cwd))
.await
.map_err(|source| CommandError::Spawn {
program: program_name.clone(),
source,
})?;
tracing::debug!(program = %program_name, %status, "detached launcher exited");
Ok(status)
}
fn resolve_program(&self, program: &OsStr) -> OsString {
let path = Path::new(program);
if path.components().count() > 1 {
return program.to_os_string();
}
let paths = self.joined_path();
which::which_in(program, paths, &self.cwd)
.map_or_else(|_| program.to_os_string(), PathBuf::into_os_string)
}
fn joined_path(&self) -> Option<OsString> {
let entries = self.path_entries();
if entries.is_empty() {
return None;
}
Some(
env::join_paths(entries)
.expect("PATH entries produced by split_paths re-join into a PATH string"),
)
}
}
#[cfg(windows)]
fn withhold_std_handles_from_children() {
use windows_sys::Win32::{
Foundation::{HANDLE_FLAG_INHERIT, INVALID_HANDLE_VALUE, SetHandleInformation},
System::Console::{GetStdHandle, STD_ERROR_HANDLE, STD_INPUT_HANDLE, STD_OUTPUT_HANDLE},
};
for (name, id) in [
("stdin", STD_INPUT_HANDLE),
("stdout", STD_OUTPUT_HANDLE),
("stderr", STD_ERROR_HANDLE),
] {
let handle = unsafe { GetStdHandle(id) };
if handle.is_null() || handle == INVALID_HANDLE_VALUE {
continue;
}
let cleared = unsafe { SetHandleInformation(handle, HANDLE_FLAG_INHERIT, 0) };
assert!(
cleared != 0,
"failed to make {name} non-inheritable: {}",
io::Error::last_os_error()
);
}
}
#[cfg(not(windows))]
const fn withhold_std_handles_from_children() {
}
async fn drain_child_pipe(
mut reader: impl smol::io::AsyncRead + Unpin,
mut sink: impl io::Write,
echo: bool,
) -> io::Result<Vec<u8>> {
let mut collected = Vec::new();
let mut chunk = [0u8; 8192];
loop {
let read = reader.read(&mut chunk).await?;
if read == 0 {
break;
}
if echo {
let _ = sink.write_all(&chunk[..read]);
let _ = sink.flush();
}
collected.extend_from_slice(&chunk[..read]);
}
Ok(collected)
}
fn env_get<'a>(env: &'a BTreeMap<OsString, OsString>, key: &OsStr) -> Option<&'a OsStr> {
if cfg!(target_os = "windows") {
env.iter()
.find(|(existing, _)| existing.as_os_str().eq_ignore_ascii_case(key))
.map(|(_, value)| value.as_os_str())
} else {
env.get(key).map(OsString::as_os_str)
}
}
fn home_dir_from_env(env: &BTreeMap<OsString, OsString>) -> Option<PathBuf> {
if cfg!(target_os = "windows") {
env_get(env, "USERPROFILE".as_ref())
.or_else(|| env_get(env, "HOME".as_ref()))
.map(PathBuf::from)
} else {
env_get(env, "HOME".as_ref())
.or_else(|| env_get(env, "USERPROFILE".as_ref()))
.map(PathBuf::from)
}
}
fn default_app_dirs() -> Vec<PathBuf> {
if cfg!(target_os = "macos") {
vec![PathBuf::from("/Applications")]
} else {
Vec::new()
}
}
#[cfg(target_os = "windows")]
fn seed_process_plumbing(env: &mut BTreeMap<OsString, OsString>) {
for key in ["SystemRoot", "SystemDrive", "windir", "ComSpec", "PATHEXT"] {
if let Some(value) = env::var_os(key) {
env.entry(OsString::from(key)).or_insert(value);
}
}
}
#[cfg(not(target_os = "windows"))]
const fn seed_process_plumbing(_env: &mut BTreeMap<OsString, OsString>) {}
#[cfg(test)]
mod tests {
use super::Host;
use crate::toolchain::testing::TestMachine;
const UNDECLARED: &str = "WATERUI_TEST_NEVER_DECLARED";
#[test]
fn declared_host_env_contains_only_what_was_declared() {
let host = Host::new(
Vec::<std::path::PathBuf>::new(),
[(String::from("WATERUI_TEST_DECLARED"), String::from("yes"))],
);
assert_eq!(
host.env_string("WATERUI_TEST_DECLARED").as_deref(),
Some("yes")
);
assert!(
host.env(UNDECLARED).is_none(),
"declared hosts must not see ambient environment variables"
);
assert!(host.path_entries().is_empty());
}
#[test]
fn declared_host_home_comes_from_declared_env() {
let machine = TestMachine::new();
let host = machine.host(Vec::<(String, String)>::new());
assert_eq!(host.home_dir(), Some(machine.home().as_path()));
assert_eq!(host.cwd(), machine.root());
assert!(
host.app_dirs().is_empty(),
"declared hosts never see installed application bundles"
);
}
#[test]
fn which_resolves_only_the_host_path() {
let machine = TestMachine::new();
let host = machine.host(Vec::<(String, String)>::new());
smol::block_on(async {
assert!(host.which("waterui-test-missing-tool").await.is_err());
assert!(
host.which("cargo").await.is_err(),
"real cargo must not leak"
);
machine.install("cargo");
let resolved = host
.which("cargo")
.await
.expect("installed fake tool must resolve");
assert_eq!(resolved.parent(), Some(machine.bin().as_path()));
});
}
#[test]
fn spawned_children_see_the_declared_environment() {
let machine = TestMachine::new();
machine.install("cargo");
let host = machine.host([(
String::from("WATERUI_FAKE_CARGO_VERSION"),
String::from("9.9.9-waterui-test"),
)]);
let output = smol::block_on(host.run("cargo", ["--version"]))
.expect("fake cargo must run under the declared host");
assert!(output.contains("9.9.9-waterui-test"));
}
#[test]
fn run_reports_nonzero_exit_with_output() {
let machine = TestMachine::new();
machine.install("rustup");
let host = machine.host(Vec::<(String, String)>::new());
let error = smol::block_on(host.run("rustup", ["frobnicate"]))
.expect_err("a failing tool must surface as an error");
assert!(error.to_string().contains("rustup"));
}
}