asonix
5aa97212b3
Some checks failed
/ check (armv7-unknown-linux-musleabihf) (push) Successful in 2m9s
/ clippy (push) Failing after 1m23s
/ tests (push) Successful in 1m39s
/ check (aarch64-unknown-linux-musl) (push) Successful in 2m8s
/ check (x86_64-unknown-linux-musl) (push) Successful in 2m4s
59 lines
1.4 KiB
Rust
59 lines
1.4 KiB
Rust
use crate::error::{Error, ErrorKind};
|
|
use actix_web::web::{Bytes, BytesMut};
|
|
use futures_core::Stream;
|
|
use streem::IntoStreamer;
|
|
|
|
pub(crate) fn limit_stream<'a, S>(
|
|
input: S,
|
|
limit: usize,
|
|
) -> impl Stream<Item = Result<Bytes, Error>> + Send + 'a
|
|
where
|
|
S: Stream<Item = reqwest::Result<Bytes>> + Send + 'a,
|
|
{
|
|
streem::try_from_fn(move |yielder| async move {
|
|
let stream = std::pin::pin!(input);
|
|
let mut stream = stream.into_streamer();
|
|
|
|
let mut count = 0;
|
|
|
|
while let Some(bytes) = stream.try_next().await? {
|
|
count += bytes.len();
|
|
|
|
if count > limit {
|
|
return Err(ErrorKind::BodyTooLarge.into());
|
|
}
|
|
|
|
yielder.yield_ok(bytes).await;
|
|
}
|
|
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
pub(crate) async fn aggregate<S>(input: S) -> Result<Bytes, Error>
|
|
where
|
|
S: Stream<Item = Result<Bytes, Error>>,
|
|
{
|
|
let stream = std::pin::pin!(input);
|
|
let mut streamer = stream.into_streamer();
|
|
|
|
let mut buf = Vec::new();
|
|
|
|
while let Some(bytes) = streamer.try_next().await? {
|
|
buf.push(bytes);
|
|
}
|
|
|
|
if buf.len() == 1 {
|
|
return Ok(buf.pop().expect("buf has exactly one element"));
|
|
}
|
|
|
|
let total_size: usize = buf.iter().map(|b| b.len()).sum();
|
|
|
|
let mut bytes_mut = BytesMut::with_capacity(total_size);
|
|
|
|
for bytes in &buf {
|
|
bytes_mut.extend_from_slice(&bytes);
|
|
}
|
|
|
|
Ok(bytes_mut.freeze())
|
|
}
|