use futures_core::Stream;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use super::BodyStream;
use crate::Error;
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct ReadIntoBytes {
buffer: Vec<u8>,
stream: BodyStream,
}
impl ReadIntoBytes {
pub(crate) fn new(buffer: Vec<u8>, stream: BodyStream) -> Self {
Self { buffer, stream }
}
}
impl ReadIntoBytes {
fn project(self: Pin<&mut Self>) -> (Pin<&mut BodyStream>, &mut Vec<u8>) {
let this = self.get_mut();
let ptr = &mut this.stream;
(Pin::new(ptr), &mut this.buffer)
}
}
impl Future for ReadIntoBytes {
type Output = Result<Vec<u8>, Error>;
fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
let (mut stream, buffer) = self.project();
loop {
return match stream.as_mut().poll_next(context) {
Poll::Ready(Some(Ok(frame))) => {
let data = match frame.into_data() {
Ok(bytes) => bytes,
Err(_) => continue,
};
let len = data.len();
if let Err(error) = buffer.try_reserve(len) {
buffer.clear();
Poll::Ready(Err(error.into()))
} else {
buffer.extend_from_slice(&data);
continue;
}
}
Poll::Ready(Some(Err(error))) => {
buffer.clear();
Poll::Ready(Err(error))
}
Poll::Ready(None) => {
let bytes = buffer.split_off(0);
Poll::Ready(Ok(bytes))
}
Poll::Pending => {
Poll::Pending
}
};
}
}
}