subprocess 1.0.3

use std::env;
use std::ffi::{CString, OsStr, OsString};
use std::fs::File;
use std::io::{Error, Result};
use std::iter;
use std::marker::PhantomData;
use std::mem;
use std::os::unix::ffi::OsStrExt;
use std::os::unix::io::{AsRawFd, BorrowedFd, FromRawFd, RawFd};
use std::ptr;
use std::sync::Arc;
use std::time::{Duration, Instant};

use libc::{c_char, c_int};

use crate::exec::Redirection;
use crate::process::ExitStatus;
use crate::spawn::StandardStream;

pub use libc::ECHILD;

fn check_err<T: Ord + Default>(num: T) -> Result<T> {
    if num < T::default() {
        return Err(Error::last_os_error());
    }
    Ok(num)
}

pub fn set_nonblocking(f: &File) -> Result<()> {
    unsafe {
        let flags = check_err(libc::fcntl(f.as_raw_fd(), libc::F_GETFL))?;
        check_err(libc::fcntl(
            f.as_raw_fd(),
            libc::F_SETFL,
            flags | libc::O_NONBLOCK,
        ))?;
    }
    Ok(())
}

/// Create a pipe with both ends having `O_CLOEXEC` set.
///
/// Uses `pipe2(O_CLOEXEC)` where available to atomically set the flag, avoiding a race
/// where another thread's fork()+exec() could inherit the fds between pipe() and
/// fcntl(). Falls back to pipe()+fcntl() on platforms without pipe2.
#[cfg(not(any(target_os = "aix", target_vendor = "apple", target_os = "haiku")))]
pub fn pipe() -> Result<(File, File)> {
    let mut fds = [0; 2];
    check_err(unsafe { libc::pipe2(fds.as_mut_ptr(), libc::O_CLOEXEC) })?;
    Ok(unsafe { (File::from_raw_fd(fds[0]), File::from_raw_fd(fds[1])) })
}

#[cfg(any(target_os = "aix", target_vendor = "apple", target_os = "haiku"))]
pub fn pipe() -> Result<(File, File)> {
    let mut fds = [0; 2];
    check_err(unsafe { libc::pipe(fds.as_mut_ptr()) })?;
    // Set CLOEXEC on both ends. There is a small race window between pipe() and fcntl()
    // where another thread's fork()+exec() could inherit these fds - this matches what
    // Rust stdlib does on these platforms that lack pipe2().
    unsafe {
        check_err(libc::fcntl(fds[0], libc::F_SETFD, libc::FD_CLOEXEC))?;
        check_err(libc::fcntl(fds[1], libc::F_SETFD, libc::FD_CLOEXEC))?;
    }
    Ok(unsafe { (File::from_raw_fd(fds[0]), File::from_raw_fd(fds[1])) })
}

// marked unsafe because the child must not allocate before exec-ing
pub unsafe fn fork() -> Result<Option<u32>> {
    let pid = check_err(unsafe { libc::fork() })?;
    if pid == 0 {
        Ok(None) // child
    } else {
        Ok(Some(pid as u32)) // parent
    }
}

pub fn setuid(uid: u32) -> Result<()> {
    check_err(unsafe { libc::setuid(uid as libc::uid_t) })?;
    Ok(())
}

pub fn setgid(gid: u32) -> Result<()> {
    check_err(unsafe { libc::setgid(gid as libc::gid_t) })?;
    Ok(())
}

pub fn setpgid(pid: u32, pgid: u32) -> Result<()> {
    check_err(unsafe { libc::setpgid(pid as _, pgid as _) })?;
    Ok(())
}

fn os_to_cstring(s: &OsStr) -> Result<CString> {
    // Like CString::new, but returns an io::Result for consistency with everything else.
    CString::new(s.as_bytes()).map_err(|_| Error::from_raw_os_error(libc::EINVAL))
}

#[derive(Debug)]
struct CVec {
    // Individual C strings.  Each element self.ptrs[i] points to the data of
    // self.strings[i].as_bytes_with_nul().as_ptr().
    #[allow(dead_code)]
    strings: Vec<CString>,

    // nullptr-terminated vector of pointers to data inside self.strings.
    ptrs: Vec<*const c_char>,
}

impl CVec {
    fn new(slice: &[impl AsRef<OsStr>]) -> Result<CVec> {
        let maybe_strings: Result<Vec<CString>> =
            slice.iter().map(|x| os_to_cstring(x.as_ref())).collect();
        let strings = maybe_strings?;
        let ptrs: Vec<_> = strings
            .iter()
            .map(|s| s.as_bytes_with_nul().as_ptr() as _)
            .chain(iter::once(ptr::null()))
            .collect();
        Ok(CVec { strings, ptrs })
    }

    pub fn as_c_vec(&self) -> *const *const c_char {
        self.ptrs.as_ptr()
    }
}

fn split_path(mut path: &OsStr) -> impl Iterator<Item = &OsStr> {
    // Can't use `env::split`_path because it allocates OsString objects, and we need to
    // iterate over PATH after fork() when allocations are strictly verboten.  We can't
    // use `str::split()` either because PATH is an `OsStr`, and there is no
    // `OsStr::split()`.
    std::iter::from_fn(move || {
        while let Some(pos) = path.as_bytes().iter().position(|&c| c == b':') {
            let piece = OsStr::from_bytes(&path.as_bytes()[..pos]);
            path = OsStr::from_bytes(&path.as_bytes()[pos + 1..]);
            if !piece.is_empty() {
                return Some(piece);
            }
        }
        let piece = path;
        path = OsStr::new("");
        if !piece.is_empty() {
            return Some(piece);
        }
        None
    })
}

struct PrepExec {
    cmd: OsString,
    argvec: CVec,
    envvec: Option<CVec>,
    search_path: Option<OsString>,
    prealloc_exe: Vec<u8>,
}

impl PrepExec {
    fn new(
        cmd: OsString,
        argvec: CVec,
        envvec: Option<CVec>,
        search_path: Option<OsString>,
    ) -> PrepExec {
        // Avoid allocation after fork() by pre-allocating the buffer that will be used
        // for constructing the executable C string.

        // Allocate enough room for "<pathdir>/<command>\0", pathdir being the longest
        // component of PATH.
        let mut max_exe_len = cmd.len() + 1;
        if let Some(ref search_path) = search_path {
            // make sure enough room is present for the largest of the PATH components,
            // plus 1 for the intervening '/'.
            max_exe_len += 1 + split_path(search_path).map(OsStr::len).max().unwrap_or(0);
        }

        PrepExec {
            cmd,
            argvec,
            envvec,
            search_path,
            prealloc_exe: Vec::with_capacity(max_exe_len),
        }
    }

    fn exec(self) -> Result<()> {
        // Invoked after fork() - no heap allocation or deallocation allowed.
        // Wrap in ManuallyDrop so that ? doesn't deallocate on exec failure.
        let mut this = std::mem::ManuallyDrop::new(self);
        let mut exe = std::mem::ManuallyDrop::new(std::mem::take(&mut this.prealloc_exe));

        if let Some(ref search_path) = this.search_path {
            let mut err = Ok(());
            // POSIX requires execvp and execve, but not execvpe (although glibc provides
            // one), so we have to iterate over PATH ourselves
            for dir in split_path(search_path.as_os_str()) {
                err = this.libc_exec(PrepExec::assemble_exe(
                    &mut exe,
                    &[dir.as_bytes(), b"/", this.cmd.as_bytes()],
                ));
                // if exec succeeds, we won't run anymore; if we're here, it failed
                assert!(err.is_err());
            }
            // we haven't found the command anywhere on the path, just return the last
            // error
            return err;
        }

        this.libc_exec(PrepExec::assemble_exe(&mut exe, &[this.cmd.as_bytes()]))?;

        // failed exec can only return Err(..)
        unreachable!();
    }

    fn assemble_exe<'a>(storage: &'a mut Vec<u8>, components: &[&[u8]]) -> &'a [u8] {
        storage.truncate(0);
        for comp in components {
            storage.extend_from_slice(comp);
        }
        // `storage` will be passed to libc::execve so it must end with \0.
        storage.push(0u8);
        storage.as_slice()
    }

    fn libc_exec(&self, exe: &[u8]) -> Result<()> {
        unsafe {
            match self.envvec.as_ref() {
                Some(envvec) => {
                    libc::execve(exe.as_ptr() as _, self.argvec.as_c_vec(), envvec.as_c_vec())
                }
                None => libc::execv(exe.as_ptr() as _, self.argvec.as_c_vec()),
            }
        };
        Err(Error::last_os_error())
    }
}

/// Prepare everything needed to `exec()` the provided `cmd` after `fork()`.
///
/// Since code executed in the child after a `fork()` is not allowed to
/// allocate (because the lock might be held), this allocates everything
/// beforehand.
pub fn prep_exec(
    cmd: impl AsRef<OsStr>,
    args: &[impl AsRef<OsStr>],
    env: Option<&[impl AsRef<OsStr>]>,
) -> Result<impl FnOnce() -> Result<()>> {
    let cmd = cmd.as_ref().to_owned();
    let argvec = CVec::new(args)?;
    let envvec = if let Some(env) = env {
        Some(CVec::new(env)?)
    } else {
        None
    };

    let search_path = if !cmd.as_bytes().contains(&b'/') {
        env::var_os("PATH")
            // treat empty path as non-existent
            .and_then(|p| if p.is_empty() { None } else { Some(p) })
    } else {
        None
    };

    // Allocate now and return a closure that just does the exec.
    let prep = PrepExec::new(cmd, argvec, envvec, search_path);
    Ok(move || prep.exec())
}

/// Prepare a `chdir()` call for use after `fork()`.
///
/// Like `prep_exec`, this pre-allocates the CString before fork so the
/// returned closure only performs the async-signal-safe `chdir` syscall.
pub fn prep_chdir(dir: &OsStr) -> Result<impl FnOnce() -> Result<()>> {
    let dir = os_to_cstring(dir)?;
    Ok(move || {
        let ret = check_err(unsafe { libc::chdir(dir.as_ptr()) });
        // Prevent deallocation after fork - only async-signal-safe
        // operations are allowed.
        std::mem::forget(dir);
        ret.map(|_| ())
    })
}

pub fn _exit(status: u8) -> ! {
    unsafe { libc::_exit(status as c_int) }
}

pub const WNOHANG: i32 = libc::WNOHANG;

pub fn waitpid(pid: u32, flags: i32) -> Result<(u32, ExitStatus)> {
    let mut status = 0 as c_int;
    let pid = check_err(unsafe {
        libc::waitpid(
            pid as libc::pid_t,
            &mut status as *mut c_int,
            flags as c_int,
        )
    })?;
    Ok((pid as u32, ExitStatus::from_raw(status)))
}

#[cfg(target_os = "linux")]
pub fn pidfd_open(pid: u32) -> Result<std::os::unix::io::OwnedFd> {
    let fd =
        check_err(unsafe { libc::syscall(libc::SYS_pidfd_open, pid as libc::pid_t, 0) })? as RawFd;
    Ok(unsafe { std::os::unix::io::OwnedFd::from_raw_fd(fd) })
}

pub use libc::{SIGKILL, SIGTERM};

pub fn kill(pid: u32, signal: i32) -> Result<()> {
    check_err(unsafe { libc::kill(pid as c_int, signal) })?;
    Ok(())
}

pub fn killpg(pgid: u32, signal: i32) -> Result<()> {
    check_err(unsafe { libc::killpg(pgid as c_int, signal) })?;
    Ok(())
}

pub const F_GETFD: i32 = libc::F_GETFD;
pub const F_SETFD: i32 = libc::F_SETFD;
pub const FD_CLOEXEC: i32 = libc::FD_CLOEXEC;

pub fn fcntl(fd: i32, cmd: i32, arg1: Option<i32>) -> Result<i32> {
    check_err(unsafe {
        match arg1 {
            Some(arg1) => libc::fcntl(fd, cmd, arg1),
            None => libc::fcntl(fd, cmd),
        }
    })
}

pub fn dup2(oldfd: i32, newfd: i32) -> Result<()> {
    check_err(unsafe { libc::dup2(oldfd, newfd) })?;
    Ok(())
}

pub fn make_redirection_to_standard_stream(which: StandardStream) -> Result<Arc<Redirection>> {
    let file = unsafe { File::from_raw_fd(which as RawFd) };
    let stream = Arc::new(Redirection::File(file));
    // Leak the Arc so the object we return doesn't close the
    // underlying file descriptor. We didn't open it, and it is shared
    // by everything else, so we are not allowed to close it either.
    mem::forget(Arc::clone(&stream));
    Ok(stream)
}

pub fn reset_sigpipe() -> Result<()> {
    // This is called after forking to reset SIGPIPE handling to the defaults that Unix
    // programs expect.  Quoting std::process::Command::do_exec:
    //
    // """
    // libstd ignores SIGPIPE, and signal-handling libraries often set a mask. Child
    // processes inherit ignored signals and the signal mask from their parent, but most
    // UNIX programs do not reset these things on their own, so we need to clean things up
    // now to avoid confusing the program we're about to run.
    // """
    unsafe {
        let mut set: mem::MaybeUninit<libc::sigset_t> = mem::MaybeUninit::uninit();
        check_err(libc::sigemptyset(set.as_mut_ptr()))?;
        let set = set.assume_init();
        check_err(libc::pthread_sigmask(
            libc::SIG_SETMASK,
            &set,
            ptr::null_mut(),
        ))?;
        if libc::signal(libc::SIGPIPE, libc::SIG_DFL) == libc::SIG_ERR {
            return Err(Error::last_os_error());
        }
    }
    Ok(())
}

#[repr(C)]
pub struct PollFd<'a>(libc::pollfd, PhantomData<&'a ()>);

impl PollFd<'_> {
    pub fn new(fd: Option<BorrowedFd<'_>>, events: i16) -> PollFd<'_> {
        PollFd(
            libc::pollfd {
                fd: fd.map(|f| f.as_raw_fd()).unwrap_or(-1),
                events,
                revents: 0,
            },
            PhantomData,
        )
    }

    pub fn test(&self, mask: i16) -> bool {
        self.0.revents & mask != 0
    }
}

pub use libc::{POLLHUP, POLLIN, POLLOUT};

pub fn poll(fds: &mut [PollFd<'_>], mut timeout: Option<Duration>) -> Result<usize> {
    let deadline = timeout.map(|timeout| Instant::now() + timeout);
    loop {
        // poll() accepts a maximum timeout of 2**31-1 ms, which is less than 25 days.
        // The caller can specify Durations much larger than that, so support them by
        // waiting in a loop.
        let (timeout_ms, overflow) = timeout
            .map(|timeout| {
                let timeout = timeout.as_millis();
                if timeout <= i32::MAX as u128 {
                    (timeout as i32, false)
                } else {
                    (i32::MAX, true)
                }
            })
            .unwrap_or((-1, false));
        let fds_ptr = fds.as_ptr() as *mut libc::pollfd;
        let cnt = unsafe { check_err(libc::poll(fds_ptr, fds.len() as libc::nfds_t, timeout_ms))? };
        if cnt != 0 || !overflow {
            return Ok(cnt as usize);
        }

        let deadline = deadline.unwrap();
        let now = Instant::now();
        if now >= deadline {
            return Ok(0);
        }
        timeout = Some(deadline - now);
    }
}

#[cfg(test)]
mod tests {
    use super::split_path;
    use std::ffi::OsStr;
    use std::os::unix::ffi::OsStrExt;

    fn s(s: &str) -> Vec<&str> {
        split_path(OsStr::new(s))
            .map(|osstr| std::str::from_utf8(osstr.as_bytes()).unwrap())
            .collect()
    }

    #[test]
    fn test_split_path() {
        let empty = Vec::<&OsStr>::new();

        assert_eq!(s("a:b"), vec!["a", "b"]);
        assert_eq!(s("one:twothree"), vec!["one", "twothree"]);
        assert_eq!(s("a:"), vec!["a"]);
        assert_eq!(s(""), empty);
        assert_eq!(s(":"), empty);
        assert_eq!(s("::"), empty);
        assert_eq!(s(":::"), empty);
        assert_eq!(s("a::b"), vec!["a", "b"]);
        assert_eq!(s(":a::::b:"), vec!["a", "b"]);
    }
}