oysh/crates/oyster_runtime/src/pipeline.rs

use std::{
    borrow::Cow,
    ffi::{OsStr, OsString},
    fs::File,
    io,
    os::unix::{
        io::FromRawFd,
        process::{CommandExt, ExitStatusExt},
    },
    process::{Child, Command, Stdio},
    slice::Iter,
};

use nix::{
    fcntl::OFlag,
    unistd::{self, Pid},
};
use oyster_parser::ast::{self, Redirect};

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

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

impl<'a> PreparedCommand<'a> {
    /// Run this command with the given context.
    fn spawn(self, pgid: &mut Pid) -> Result<Child, RuntimeError> {
        let args = self.args.map(|w| match w {
            Cow::Borrowed(s) => Cow::Borrowed(OsStr::new(s)),
            Cow::Owned(s) => Cow::Owned(OsString::from(s)),
        });

        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));
        cmd.process_group(pgid.as_raw());

        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),
                })
            }
        };

        let child_id = Pid::from_raw(child.id() as i32);
        if *pgid == Pid::from_raw(0) {
            *pgid = child_id;
        };

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

        Ok(child)
    }
}

impl<'a> From<&'a ast::Command<'a>> for PreparedCommand<'a> {
    fn from(command: &'a ast::Command<'a>) -> Self {
        let mut words = WordBuilder(command.0.iter());
        let cmd = words.next().expect("words need to have >1 parts");
        let args = words;
        let redirect = command.1;

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

impl Shell {
    pub(crate) fn run_pipeline(
        &mut self,
        pipeline: &ast::Pipeline,
    ) -> Result<Status, RuntimeError> {
        let mut children = Vec::with_capacity(pipeline.0.len());
        let mut cmds = pipeline.0.iter().map(PreparedCommand::from);
        let mut pgid = Pid::from_raw(0);

        let mut last_cmd = cmds.next().expect("pipelines need to have >1 commands");
        for mut cmd in cmds {
            let (output, input) = create_pipe()?;
            cmd.stdin = Some(output);

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

            children.push(last_cmd.spawn(&mut pgid)?);

            last_cmd = cmd;
        }

        children.push(last_cmd.spawn(&mut pgid)?);

        // TODO: kill children if error occured
        // TODO: set foreground group then wait for foreground group

        let mut last_status = Status::SUCCESS;
        for mut c in children {
            // TODO: handle error
            let status = c.wait().unwrap();
            last_status = Status(
                status
                    .code()
                    .unwrap_or_else(|| Status::SIG_BASE.0 + status.signal().unwrap()),
            );
        }

        Ok(last_status)
    }
}

/// 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)) })
}

/// Build words from WordParts.
struct WordBuilder<'a>(Iter<'a, ast::Word<'a>>);

impl<'a> Iterator for WordBuilder<'a> {
    type Item = Cow<'a, str>;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|word| {
            let mut words = word.0.iter();

            let mut s = match words.next().expect("words need to have >1 parts") {
                ast::WordPart::Text(text) => Cow::from(*text),
            };

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

            s
        })
    }
}