use std::ffi::{OsStr, OsString};
use std::io::{self, Read, Write};
use std::mem;
use std::os::fd::{AsRawFd, OwnedFd, RawFd};
use std::os::unix::ffi::OsStrExt;
use std::os::unix::process::ExitStatusExt;
pub use std::process::{ExitStatus, Output};
use rustix::fs::{Mode, OFlags};
use rustix::pipe::pipe;
use rustix::process::{Pid, Signal};
pub struct Command {
program: OsString,
args: Vec<OsString>,
pub(crate) stdin: Stdio,
pub(crate) stdout: Stdio,
pub(crate) stderr: Stdio,
}
impl Command {
pub fn new<S: AsRef<OsStr>>(program: S) -> Self {
let program = program.as_ref().to_os_string();
Self {
program,
stdout: Default::default(),
stderr: Default::default(),
stdin: Default::default(),
args: Default::default(),
}
}
pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Self {
let arg = arg.as_ref().to_os_string();
self.args.push(arg);
self
}
pub fn args<I, S>(&mut self, args: I) -> &mut Self
where
I: IntoIterator<Item = S>,
S: AsRef<OsStr>,
{
for arg in args {
let arg = arg.as_ref().to_os_string();
self.args.push(arg);
}
self
}
pub fn stdin<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Self {
self.stdin = cfg.into();
self
}
pub fn stdout<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Self {
self.stdout = cfg.into();
self
}
pub fn stderr<T: Into<Stdio>>(&mut self, cfg: T) -> &mut Self {
self.stderr = cfg.into();
self
}
pub fn get_program(&self) -> &OsStr {
OsStr::from_bytes(self.program.as_bytes())
}
}
impl From<Command> for std::process::Command {
fn from(command: Command) -> Self {
let mut std_command = std::process::Command::new(command.program);
std_command.args(command.args);
let stdin: Option<std::process::Stdio> = command.stdin.into();
if let Some(stdin) = stdin {
std_command.stdin(stdin);
}
let stdout: Option<std::process::Stdio> = command.stdout.into();
if let Some(stdout) = stdout {
std_command.stdout(stdout);
}
let stderr: Option<std::process::Stdio> = command.stderr.into();
if let Some(stderr) = stderr {
std_command.stderr(stderr);
}
std_command
}
}
pub struct Child {
pub stdin: Option<ChildStdin>,
pub stdout: Option<ChildStdout>,
pub stderr: Option<ChildStderr>,
exit_signal: OwnedFd,
pid: u32,
}
impl Child {
pub(crate) fn new(
pid: i32,
exit_signal: OwnedFd,
stdin: Option<OwnedFd>,
stdout: Option<OwnedFd>,
stderr: Option<OwnedFd>,
) -> io::Result<Self> {
Ok(Self {
exit_signal,
pid: pid as u32,
stdin: stdin.map(ChildStdin::new).transpose()?,
stdout: stdout.map(ChildStdout::new).transpose()?,
stderr: stderr.map(ChildStderr::new).transpose()?,
})
}
pub fn kill(&mut self) -> io::Result<()> {
let pid = Pid::from_raw(self.pid as i32).unwrap();
rustix::process::kill_process(pid, Signal::Kill)?;
Ok(())
}
pub fn id(&self) -> u32 {
self.pid
}
pub fn wait(&mut self) -> io::Result<ExitStatus> {
unsafe {
let mut status: libc::c_int = 0;
if libc::waitpid(self.pid as i32, &mut status, 0) == -1 {
Err(io::Error::last_os_error())
} else {
match self.exit_signal()? {
Some(exit_signal) => Ok(exit_signal),
None => Ok(ExitStatus::from_raw(status)),
}
}
}
}
pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
unsafe {
let mut status: libc::c_int = 0;
let pid = libc::waitpid(self.pid as i32, &mut status, libc::WNOHANG);
if pid == -1 {
Err(io::Error::last_os_error())
} else if pid == 0 {
Ok(None)
} else {
match self.exit_signal()? {
Some(exit_signal) => Ok(Some(exit_signal)),
None => Ok(Some(ExitStatus::from_raw(status))),
}
}
}
}
pub fn wait_with_output(mut self) -> io::Result<Output> {
let _ = self.stdin.take();
let reader = ChildReader::new(self.stdout.take(), self.stderr.take())?;
let (stdout, stderr) = reader.read()?;
let status = self.wait()?;
Ok(Output { status, stdout, stderr })
}
fn exit_signal(&self) -> io::Result<Option<ExitStatus>> {
rustix::fs::fcntl_setfl(&self.exit_signal, OFlags::NONBLOCK)?;
let mut bytes = [0; mem::size_of::<u32>()];
let read = match rustix::io::read(&self.exit_signal, &mut bytes) {
Ok(read) => read,
Err(err) if err.kind() == io::ErrorKind::WouldBlock => return Ok(None),
Err(err) => return Err(err.into()),
};
if read == mem::size_of::<u32>() {
let signal = u32::from_le_bytes(bytes);
assert!(signal <= 0x7f);
Ok(Some(ExitStatus::from_raw(signal as i32)))
} else {
Ok(None)
}
}
}
#[derive(Default)]
pub struct Stdio {
pub(crate) ty: StdioType,
}
impl Stdio {
pub fn piped() -> Self {
Self { ty: StdioType::Piped }
}
pub fn inherit() -> Self {
Self { ty: StdioType::Inherit }
}
pub fn null() -> Self {
Self { ty: StdioType::Null }
}
pub(crate) fn make_pipe(&self, stdin: bool) -> io::Result<(Option<OwnedFd>, Option<OwnedFd>)> {
match self.ty {
StdioType::Inherit | StdioType::Default => Ok((None, None)),
StdioType::Piped => {
let (rx, tx) = pipe()?;
Ok((Some(rx), Some(tx)))
},
StdioType::Null => {
let null_fd = rustix::fs::open("/dev/null", OFlags::RDWR, Mode::empty())?;
if stdin {
Ok((Some(null_fd), None))
} else {
Ok((None, Some(null_fd)))
}
},
}
}
}
impl From<Stdio> for Option<std::process::Stdio> {
fn from(stdio: Stdio) -> Option<std::process::Stdio> {
match stdio.ty {
StdioType::Default => None,
StdioType::Inherit => Some(std::process::Stdio::inherit()),
StdioType::Piped => Some(std::process::Stdio::piped()),
StdioType::Null => Some(std::process::Stdio::null()),
}
}
}
#[derive(Default, Copy, Clone)]
pub(crate) enum StdioType {
#[default]
Default,
Piped,
Inherit,
Null,
}
pub struct ChildStdin {
fd: OwnedFd,
}
impl ChildStdin {
fn new(fd: OwnedFd) -> io::Result<Self> {
Ok(Self { fd })
}
}
impl Write for ChildStdin {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
rustix::io::write(&self.fd, buf).map_err(io::Error::from)
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
pub struct ChildStdout {
fd: OwnedFd,
}
impl ChildStdout {
fn new(fd: OwnedFd) -> io::Result<Self> {
Ok(Self { fd })
}
}
impl Read for ChildStdout {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
rustix::io::read(&self.fd, buf).map_err(io::Error::from)
}
}
pub type ChildStderr = ChildStdout;
struct ChildReader {
poll_fds: Vec<libc::pollfd>,
stdout: Option<ChildStdout>,
stderr: Option<ChildStderr>,
stdout_buffer: Vec<u8>,
stderr_buffer: Vec<u8>,
}
impl ChildReader {
fn new(stdout: Option<ChildStdout>, stderr: Option<ChildStderr>) -> io::Result<Self> {
let mut poll_fds = Vec::new();
if let Some(stdout) = &stdout {
rustix::fs::fcntl_setfl(&stdout.fd, OFlags::NONBLOCK)?;
let fd = stdout.fd.as_raw_fd();
poll_fds.push(libc::pollfd { fd, events: libc::POLLIN, revents: 0 });
}
if let Some(stderr) = &stderr {
rustix::fs::fcntl_setfl(&stderr.fd, OFlags::NONBLOCK)?;
let fd = stderr.fd.as_raw_fd();
poll_fds.push(libc::pollfd { fd, events: libc::POLLIN, revents: 0 });
}
Ok(Self {
poll_fds,
stdout,
stderr,
stdout_buffer: Default::default(),
stderr_buffer: Default::default(),
})
}
fn read(mut self) -> io::Result<(Vec<u8>, Vec<u8>)> {
while !self.poll_fds.is_empty() {
let result = unsafe { libc::poll(self.poll_fds.as_mut_ptr(), 2, -1) };
if result == -1 {
return Err(io::Error::last_os_error());
}
for i in (0..self.poll_fds.len()).rev() {
let poll_fd = &self.poll_fds[i];
if poll_fd.revents == 0 {
continue;
}
let (stdio, buffer) = self.stdio_from_fd(poll_fd.fd);
match stdio.read_to_end(buffer) {
Ok(_) => {
self.poll_fds.remove(i);
},
Err(err) if err.kind() == io::ErrorKind::WouldBlock => (),
Err(err) => return Err(err),
}
}
}
Ok((self.stdout_buffer, self.stderr_buffer))
}
fn stdio_from_fd(&mut self, fd: RawFd) -> (&mut ChildStdout, &mut Vec<u8>) {
match (self.stdout.as_mut(), self.stderr.as_mut()) {
(Some(stdout), _) if stdout.fd.as_raw_fd() == fd => (stdout, &mut self.stdout_buffer),
(_, Some(stderr)) if stderr.fd.as_raw_fd() == fd => (stderr, &mut self.stderr_buffer),
_ => unreachable!(),
}
}
}