pict-rs/src/process.rs

use crate::store::Store;
use actix_rt::task::JoinHandle;
use actix_web::web::Bytes;
use std::{
    future::Future,
    pin::Pin,
    process::{ExitStatus, Stdio},
    task::{Context, Poll},
    time::Duration,
};
use tokio::{
    io::{AsyncRead, AsyncWriteExt, ReadBuf},
    process::{Child, ChildStdin, ChildStdout, Command},
    sync::oneshot::{channel, Receiver},
};
use tracing::{Instrument, Span};

#[derive(Debug)]
struct StatusError(ExitStatus);

pub(crate) struct Process {
    child: Child,
    timeout: Duration,
}

impl std::fmt::Debug for Process {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Process").field("child", &"Child").finish()
    }
}

struct DropHandle {
    inner: JoinHandle<()>,
}

pub(crate) struct ProcessRead<I> {
    inner: I,
    err_recv: Receiver<std::io::Error>,
    err_closed: bool,
    #[allow(dead_code)]
    handle: DropHandle,
    eof: bool,
    sleep: Pin<Box<actix_rt::time::Sleep>>,
}

#[derive(Debug, thiserror::Error)]
pub(crate) enum ProcessError {
    #[error("Required command {0} not found, make sure it exists in pict-rs' $PATH")]
    NotFound(String),

    #[error("Cannot run command {0} due to invalid permissions on binary, make sure the pict-rs user has permission to run it")]
    PermissionDenied(String),

    #[error("Reached process spawn limit")]
    LimitReached,

    #[error("Process timed out")]
    Timeout,

    #[error("Failed with status {0}")]
    Status(ExitStatus),

    #[error("Unknown process error")]
    Other(#[source] std::io::Error),
}

impl Process {
    pub(crate) fn run(command: &str, args: &[&str], timeout: u64) -> Result<Self, ProcessError> {
        let res = tracing::trace_span!(parent: None, "Create command")
            .in_scope(|| Self::spawn(Command::new(command).args(args), timeout));

        match res {
            Ok(this) => Ok(this),
            Err(e) => match e.kind() {
                std::io::ErrorKind::NotFound => Err(ProcessError::NotFound(command.to_string())),
                std::io::ErrorKind::PermissionDenied => {
                    Err(ProcessError::PermissionDenied(command.to_string()))
                }
                std::io::ErrorKind::WouldBlock => Err(ProcessError::LimitReached),
                _ => Err(ProcessError::Other(e)),
            },
        }
    }

    fn spawn(cmd: &mut Command, timeout: u64) -> std::io::Result<Self> {
        let timeout = Duration::from_secs(timeout);

        tracing::trace_span!(parent: None, "Spawn command").in_scope(|| {
            let cmd = cmd
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .kill_on_drop(true);

            cmd.spawn().map(|child| Process { child, timeout })
        })
    }

    #[tracing::instrument(skip(self))]
    pub(crate) async fn wait(self) -> Result<(), ProcessError> {
        let Process { mut child, timeout } = self;

        match actix_rt::time::timeout(timeout, child.wait()).await {
            Ok(Ok(status)) if status.success() => Ok(()),
            Ok(Ok(status)) => Err(ProcessError::Status(status)),
            Ok(Err(e)) => Err(ProcessError::Other(e)),
            Err(_) => {
                child.kill().await.map_err(ProcessError::Other)?;
                Err(ProcessError::Timeout)
            }
        }
    }

    pub(crate) fn bytes_read(self, input: Bytes) -> ProcessRead<ChildStdout> {
        self.spawn_fn(move |mut stdin| {
            let mut input = input;
            async move { stdin.write_all_buf(&mut input).await }
        })
    }

    pub(crate) fn read(self) -> ProcessRead<ChildStdout> {
        self.spawn_fn(|_| async { Ok(()) })
    }

    pub(crate) fn pipe_async_read<A: AsyncRead + Unpin + 'static>(
        self,
        mut async_read: A,
    ) -> ProcessRead<ChildStdout> {
        self.spawn_fn(move |mut stdin| async move {
            tokio::io::copy(&mut async_read, &mut stdin)
                .await
                .map(|_| ())
        })
    }

    pub(crate) fn store_read<S: Store + 'static>(
        self,
        store: S,
        identifier: S::Identifier,
    ) -> ProcessRead<ChildStdout> {
        self.spawn_fn(move |mut stdin| {
            let store = store;
            let identifier = identifier;

            async move { store.read_into(&identifier, &mut stdin).await }
        })
    }

    #[allow(unknown_lints)]
    #[allow(clippy::let_with_type_underscore)]
    #[tracing::instrument(level = "trace", skip_all)]
    fn spawn_fn<F, Fut>(self, f: F) -> ProcessRead<ChildStdout>
    where
        F: FnOnce(ChildStdin) -> Fut + 'static,
        Fut: Future<Output = std::io::Result<()>>,
    {
        let Process { mut child, timeout } = self;

        let stdin = child.stdin.take().expect("stdin exists");
        let stdout = child.stdout.take().expect("stdout exists");

        let (tx, rx) = tracing::trace_span!(parent: None, "Create channel")
            .in_scope(channel::<std::io::Error>);

        let span = tracing::info_span!(parent: None, "Background process task");
        span.follows_from(Span::current());

        let handle = tracing::trace_span!(parent: None, "Spawn task").in_scope(|| {
            actix_rt::spawn(
                async move {
                    let child_fut = async {
                        (f)(stdin).await?;

                        child.wait().await
                    };

                    let err = match actix_rt::time::timeout(timeout, child_fut).await {
                        Ok(Ok(status)) if status.success() => return,
                        Ok(Ok(status)) => std::io::Error::new(
                            std::io::ErrorKind::Other,
                            ProcessError::Status(status),
                        ),
                        Ok(Err(e)) => e,
                        Err(_) => std::io::ErrorKind::TimedOut.into(),
                    };

                    let _ = tx.send(err);

                    let _ = child.kill().await;
                }
                .instrument(span),
            )
        });

        let sleep = Box::pin(actix_rt::time::sleep(timeout));

        ProcessRead {
            inner: stdout,
            err_recv: rx,
            err_closed: false,
            handle: DropHandle { inner: handle },
            eof: false,
            sleep,
        }
    }
}

impl<I> AsyncRead for ProcessRead<I>
where
    I: AsyncRead + Unpin,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        if !self.err_closed {
            if let Poll::Ready(res) = Pin::new(&mut self.err_recv).poll(cx) {
                self.err_closed = true;

                if let Ok(err) = res {
                    return Poll::Ready(Err(err));
                }

                if self.eof {
                    return Poll::Ready(Ok(()));
                }
            }

            if let Poll::Ready(()) = self.sleep.as_mut().poll(cx) {
                self.err_closed = true;

                return Poll::Ready(Err(std::io::ErrorKind::TimedOut.into()));
            }
        }

        if !self.eof {
            let before_size = buf.filled().len();

            return match Pin::new(&mut self.inner).poll_read(cx, buf) {
                Poll::Ready(Ok(())) => {
                    if buf.filled().len() == before_size {
                        self.eof = true;

                        if !self.err_closed {
                            // reached end of stream & haven't received process signal
                            return Poll::Pending;
                        }
                    }

                    Poll::Ready(Ok(()))
                }
                Poll::Ready(Err(e)) => {
                    self.eof = true;

                    Poll::Ready(Err(e))
                }
                Poll::Pending => Poll::Pending,
            };
        }

        if self.err_closed && self.eof {
            return Poll::Ready(Ok(()));
        }

        Poll::Pending
    }
}

impl Drop for DropHandle {
    fn drop(&mut self) {
        self.inner.abort();
    }
}

impl std::fmt::Display for StatusError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "Command failed with bad status: {}", self.0)
    }
}

impl std::error::Error for StatusError {}