oysh/crates/oyster_runtime/src/pipeline.rs

use std::{
    borrow::Cow,
    ffi::{OsStr, OsString},
    fs::File,
    io,
    os::unix::{
        io::{AsRawFd, FromRawFd, RawFd},
        process::CommandExt,
    },
    process::{self, Command, Stdio},
};

use nix::{
    errno::Errno,
    fcntl::OFlag,
    libc,
    sys::{
        signal::{self, SaFlags, SigAction, SigHandler},
        signalfd::SigSet,
        wait::{self, WaitPidFlag, WaitStatus},
    },
    unistd::{self, ForkResult, Pid},
};
use oyster_parser::ast::{self, Redirect};

use crate::{builtins::Builtin, RuntimeError, Shell, Status};

/// The specific kind of the command.
enum CommandKind {
    External,
    Builtin(Builtin),
}

/// A command that's ready to be run.
struct PreparedCommand<'a> {
    kind: CommandKind,
    cmd: Cow<'a, str>,
    args: Vec<Cow<'a, str>>,
    redirect: Redirect,
    stdin: Option<File>,
    stdout: Option<File>,
    stderr: Option<File>,
}

impl<'a> PreparedCommand<'a> {
    /// Create a new PreparedCommand.
    fn new(command: &'a ast::Command, shell: &Shell) -> Self {
        let mut words = command.0.iter().map(|word| {
            let mut words = word.0.iter();

            let mut s = match words.next().unwrap() {
                ast::WordPart::Text(text) => Cow::from(*text),
            };

            for part in words {
                match part {
                    ast::WordPart::Text(text) => s.to_mut().push_str(text),
                }
            }

            s
        });
        let cmd = words.next().unwrap();
        let args = words.collect();
        let redirect = command.1;

        let kind = match shell.builtins().get(&cmd) {
            Some(builtin) => CommandKind::Builtin(*builtin),
            None => CommandKind::External,
        };

        PreparedCommand {
            kind,
            cmd,
            args,
            redirect,
            stdin: None,
            stdout: None,
            stderr: None,
        }
    }

    /// Run this command with the given context.
    fn spawn(self, shell: &mut Shell, pgid: &mut Pid) -> Result<(), RuntimeError> {
        let pre_exec = {
            let pgid = *pgid;
            move || {
                let _ = unistd::setpgid(Pid::from_raw(0), pgid);
                let _ = unistd::tcsetpgrp(libc::STDIN_FILENO, unistd::getpid());

                let default = SigAction::new(SigHandler::SigDfl, SaFlags::empty(), SigSet::empty());
                unsafe {
                    let _ = signal::sigaction(signal::Signal::SIGTSTP, &default);
                    let _ = signal::sigaction(signal::Signal::SIGTTOU, &default);
                }

                Ok(())
            }
        };

        let args = self.args.iter().map(|w| match w {
            Cow::Borrowed(s) => Cow::Borrowed(OsStr::new(s)),
            Cow::Owned(s) => Cow::Owned(OsString::from(s)),
        });

        let child = match self.kind {
            CommandKind::External => {
                let mut cmd = Command::new(self.cmd.as_ref());
                cmd.args(args);
                cmd.stdin(self.stdin.map_or(Stdio::inherit(), Stdio::from));
                cmd.stdout(self.stdout.map_or(Stdio::inherit(), Stdio::from));
                cmd.stderr(self.stderr.map_or(Stdio::inherit(), Stdio::from));

                unsafe {
                    cmd.pre_exec(pre_exec);
                }

                let child = match cmd.spawn() {
                    Ok(child) => child,
                    Err(err) => {
                        return Err(match err.kind() {
                            io::ErrorKind::NotFound => {
                                RuntimeError::CommandNotFound(self.cmd.into_owned())
                            }
                            io::ErrorKind::PermissionDenied => {
                                RuntimeError::PermissionDenied(self.cmd.into_owned())
                            }
                            _ => RuntimeError::SpawnFailed(err),
                        })
                    }
                };

                Pid::from_raw(child.id() as i32)
            }
            CommandKind::Builtin(builtin) => {
                match unsafe { unistd::fork() }.map_err(RuntimeError::ForkFailed)? {
                    ForkResult::Parent { child } => child,
                    ForkResult::Child => {
                        fn redirect(old: Option<File>, new: RawFd) {
                            if let Some(old) = old {
                                let _ = unistd::dup2(old.as_raw_fd(), new);
                            }
                        }

                        redirect(self.stdin, libc::STDIN_FILENO);
                        redirect(self.stdout, libc::STDOUT_FILENO);
                        redirect(self.stderr, libc::STDERR_FILENO);

                        let _ = pre_exec();
                        (builtin.fun)(shell, &self.args);
                        process::exit(0);
                    }
                }
            }
        };

        if *pgid == Pid::from_raw(0) {
            *pgid = child;
        }

        // prevent race conditions
        let _ = unistd::setpgid(child, *pgid);

        Ok(())
    }
}

impl Shell {
    pub(crate) fn run_pipeline(
        &mut self,
        pipeline: &ast::Pipeline,
    ) -> Result<Status, RuntimeError> {
        let mut pgid = Pid::from_raw(0);
        let status = (|| {
            let mut cmds = pipeline.0.iter();

            let mut last_cmd = cmds
                .next()
                .map(|cmd| PreparedCommand::new(cmd, self))
                .unwrap();
            for cmd in cmds {
                let mut cmd = PreparedCommand::new(cmd, self);
                let (output, input) = create_pipe()?;
                cmd.stdin = Some(output);

                match last_cmd.redirect {
                    Redirect::None => (),
                    Redirect::Stdout => last_cmd.stdout = Some(input),
                }

                last_cmd.spawn(self, &mut pgid)?;

                last_cmd = cmd;
            }

            last_cmd.spawn(self, &mut pgid)?;

            wait_pgid(pgid)
        })();

        if status.is_err() && pgid != Pid::from_raw(0) {
            let _ = signal::killpg(pgid, signal::Signal::SIGTERM);
            let _ = signal::killpg(pgid, signal::Signal::SIGCONT);
        }

        let _ = unistd::tcsetpgrp(libc::STDIN_FILENO, unistd::getpgid(None).unwrap());

        status
    }
}

/// Wait for all processes of the specified job to terminate.
fn wait_pgid(pgid: Pid) -> Result<Status, RuntimeError> {
    let pgid = Pid::from_raw(-pgid.as_raw());
    let mut last_status = Status::SUCCESS;

    loop {
        match wait::waitpid(pgid, Some(WaitPidFlag::WUNTRACED)) {
            Ok(WaitStatus::Exited(_, code)) => last_status = Status(code),
            Ok(WaitStatus::Signaled(_, signal, _)) => {
                last_status = Status(Status::SIG_BASE.0 + signal as i32)
            }
            Ok(WaitStatus::Stopped(_, _)) => todo!("put job into background job list"),
            Ok(_) => (),
            Err(err) => match err {
                Errno::ECHILD => break Ok(last_status), // no more children
                _ => break Err(RuntimeError::WaidPid(err)),
            },
        }
    }
}

/// Create a new unix pipe.
fn create_pipe() -> Result<(File, File), RuntimeError> {
    let (output, input) =
        unistd::pipe2(OFlag::O_CLOEXEC).map_err(RuntimeError::PipeCreationFailed)?;

    Ok(unsafe { (File::from_raw_fd(output), File::from_raw_fd(input)) })
}