2021-09-04 19:20:31 +00:00
|
|
|
use crate::error::UploadError;
|
|
|
|
use actix_web::web::{Bytes, BytesMut};
|
|
|
|
use futures_core::stream::Stream;
|
2021-08-31 02:19:47 +00:00
|
|
|
use std::{
|
2021-09-04 00:53:53 +00:00
|
|
|
future::Future,
|
2021-08-31 02:19:47 +00:00
|
|
|
pin::Pin,
|
|
|
|
task::{Context, Poll},
|
|
|
|
};
|
2021-09-04 17:42:40 +00:00
|
|
|
use tokio::io::{AsyncRead, AsyncWriteExt, ReadBuf};
|
2021-09-04 00:53:53 +00:00
|
|
|
|
2021-09-04 19:20:31 +00:00
|
|
|
pub(crate) type LocalBoxStream<'a, T> = Pin<Box<dyn Stream<Item = T> + 'a>>;
|
|
|
|
pub(crate) type LocalBoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 'a>>;
|
2021-08-31 02:19:47 +00:00
|
|
|
|
|
|
|
pub(crate) struct Process {
|
|
|
|
child: tokio::process::Child,
|
|
|
|
}
|
|
|
|
|
2021-09-04 00:53:53 +00:00
|
|
|
pub(crate) struct ProcessRead<I> {
|
|
|
|
inner: I,
|
|
|
|
err_recv: tokio::sync::oneshot::Receiver<std::io::Error>,
|
|
|
|
err_closed: bool,
|
|
|
|
}
|
|
|
|
|
2021-09-04 19:20:31 +00:00
|
|
|
struct BytesFreezer<S>(S);
|
|
|
|
|
|
|
|
pub(crate) struct Once<T> {
|
|
|
|
inner: Option<T>,
|
2021-08-31 02:19:47 +00:00
|
|
|
}
|
|
|
|
|
2021-09-04 19:20:31 +00:00
|
|
|
pub(crate) struct Next<'a, S>(&'a mut S);
|
|
|
|
|
2021-08-31 02:19:47 +00:00
|
|
|
impl Process {
|
|
|
|
fn new(child: tokio::process::Child) -> Self {
|
|
|
|
Process { child }
|
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) fn spawn(cmd: &mut tokio::process::Command) -> std::io::Result<Self> {
|
|
|
|
cmd.stdin(std::process::Stdio::piped())
|
|
|
|
.stdout(std::process::Stdio::piped())
|
|
|
|
.spawn()
|
|
|
|
.map(Process::new)
|
|
|
|
}
|
|
|
|
|
2021-09-04 00:53:53 +00:00
|
|
|
pub(crate) fn bytes_read(mut self, mut input: Bytes) -> Option<impl AsyncRead + Unpin> {
|
|
|
|
let mut stdin = self.child.stdin.take()?;
|
|
|
|
let stdout = self.child.stdout.take()?;
|
|
|
|
|
|
|
|
let (tx, rx) = tokio::sync::oneshot::channel();
|
|
|
|
|
|
|
|
actix_rt::spawn(async move {
|
|
|
|
if let Err(e) = stdin.write_all_buf(&mut input).await {
|
|
|
|
let _ = tx.send(e);
|
|
|
|
}
|
|
|
|
});
|
|
|
|
|
|
|
|
Some(Box::pin(ProcessRead {
|
|
|
|
inner: stdout,
|
|
|
|
err_recv: rx,
|
|
|
|
err_closed: false,
|
|
|
|
}))
|
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) fn write_read(
|
|
|
|
mut self,
|
|
|
|
mut input_reader: impl AsyncRead + Unpin + 'static,
|
|
|
|
) -> Option<impl AsyncRead + Unpin> {
|
|
|
|
let mut stdin = self.child.stdin.take()?;
|
|
|
|
let stdout = self.child.stdout.take()?;
|
|
|
|
|
|
|
|
let (tx, rx) = tokio::sync::oneshot::channel();
|
|
|
|
|
|
|
|
actix_rt::spawn(async move {
|
|
|
|
if let Err(e) = tokio::io::copy(&mut input_reader, &mut stdin).await {
|
|
|
|
let _ = tx.send(e);
|
|
|
|
}
|
|
|
|
});
|
|
|
|
|
|
|
|
Some(Box::pin(ProcessRead {
|
|
|
|
inner: stdout,
|
|
|
|
err_recv: rx,
|
|
|
|
err_closed: false,
|
|
|
|
}))
|
|
|
|
}
|
2021-08-31 02:19:47 +00:00
|
|
|
}
|
|
|
|
|
2021-09-04 19:20:31 +00:00
|
|
|
pub(crate) fn bytes_stream(
|
|
|
|
input: impl AsyncRead + Unpin,
|
|
|
|
) -> impl Stream<Item = Result<Bytes, UploadError>> + Unpin {
|
|
|
|
BytesFreezer(tokio_util::codec::FramedRead::new(
|
|
|
|
input,
|
|
|
|
tokio_util::codec::BytesCodec::new(),
|
|
|
|
))
|
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) fn once<T>(input: T) -> Once<T> {
|
|
|
|
Once { inner: Some(input) }
|
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) fn next<'a, S>(stream: &'a mut S) -> Next<'a, S> {
|
|
|
|
Next(stream)
|
2021-09-04 00:53:53 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
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 let Poll::Ready(res) = Pin::new(&mut self.inner).poll_read(cx, buf) {
|
|
|
|
return Poll::Ready(res);
|
|
|
|
}
|
|
|
|
|
|
|
|
Poll::Pending
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-09-04 19:20:31 +00:00
|
|
|
impl<S> Stream for BytesFreezer<S>
|
|
|
|
where
|
|
|
|
S: Stream<Item = std::io::Result<BytesMut>> + Unpin,
|
|
|
|
{
|
|
|
|
type Item = Result<Bytes, UploadError>;
|
|
|
|
|
|
|
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
|
|
Pin::new(&mut self.0)
|
|
|
|
.poll_next(cx)
|
|
|
|
.map(|opt| opt.map(|res| res.map(|bytes_mut| bytes_mut.freeze())))
|
|
|
|
.map_err(UploadError::from)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<T> Stream for Once<T>
|
|
|
|
where
|
|
|
|
T: Future + Unpin,
|
|
|
|
{
|
|
|
|
type Item = <T as Future>::Output;
|
2021-08-31 02:19:47 +00:00
|
|
|
|
|
|
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
2021-09-04 19:20:31 +00:00
|
|
|
if let Some(mut fut) = self.inner.take() {
|
|
|
|
match Pin::new(&mut fut).poll(cx) {
|
|
|
|
Poll::Ready(item) => Poll::Ready(Some(item)),
|
|
|
|
Poll::Pending => {
|
|
|
|
self.inner = Some(fut);
|
|
|
|
Poll::Pending
|
|
|
|
}
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
Poll::Ready(None)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'a, S> Future for Next<'a, S>
|
|
|
|
where
|
|
|
|
S: Stream + Unpin,
|
|
|
|
{
|
|
|
|
type Output = Option<<S as Stream>::Item>;
|
|
|
|
|
|
|
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
|
|
Pin::new(&mut self.0).poll_next(cx)
|
2021-08-31 02:19:47 +00:00
|
|
|
}
|
|
|
|
}
|