use bytes::Bytes;
use futures_util::Stream;
use http_body::{Body, Frame};
use pin_project_lite::pin_project;
use std::{
io,
pin::Pin,
task::{Context, Poll},
};
use tokio::io::{AsyncRead, AsyncReadExt, Take};
use tokio_util::io::ReaderStream;
mod serve_dir;
mod serve_file;
pub use self::{
serve_dir::{
future::ResponseFuture as ServeFileSystemResponseFuture,
DefaultServeDirFallback,
ResponseBody as ServeFileSystemResponseBody,
ServeDir,
},
serve_file::ServeFile,
};
pin_project! {
#[derive(Debug)]
pub struct AsyncReadBody<T> {
#[pin]
reader: ReaderStream<T>,
}
}
impl<T> AsyncReadBody<T>
where
T: AsyncRead,
{
fn with_capacity(read: T, capacity: usize) -> Self {
Self {
reader: ReaderStream::with_capacity(read, capacity),
}
}
fn with_capacity_limited(
read: T,
capacity: usize,
max_read_bytes: u64,
) -> AsyncReadBody<Take<T>> {
AsyncReadBody {
reader: ReaderStream::with_capacity(read.take(max_read_bytes), capacity),
}
}
}
impl<T> Body for AsyncReadBody<T>
where
T: AsyncRead,
{
type Data = Bytes;
type Error = io::Error;
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
match std::task::ready!(self.project().reader.poll_next(cx)) {
Some(Ok(chunk)) => Poll::Ready(Some(Ok(Frame::data(chunk)))),
Some(Err(err)) => Poll::Ready(Some(Err(err))),
None => Poll::Ready(None),
}
}
}