mirror of
https://git.asonix.dog/asonix/pict-rs
synced 2024-12-25 12:01:24 +00:00
281 lines
8.2 KiB
Rust
281 lines
8.2 KiB
Rust
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 {}
|