use std::{
cmp, fmt,
future::Future,
io,
pin::Pin,
task::{Context, Poll},
};
use tracing::trace;
use actix_http::ContentEncoding;
use actix_web::error::Error;
use futures_core::{Stream, ready};
use pin_project_lite::pin_project;
use super::named::File;
pin_project! {
#[doc(hidden)]
pub struct ChunkedReadFile<F, Fut> {
size: u64,
offset: u64,
#[pin]
state: ChunkedReadFileState<Fut>,
counter: u64,
callback: F,
}
}
pin_project! {
#[project = ChunkedReadFileStateProj]
#[project_replace = ChunkedReadFileStateProjReplace]
enum ChunkedReadFileState<Fut> {
File { file: Option<File>, },
Future { #[pin] fut: Fut },
}
}
impl<F, Fut> fmt::Debug for ChunkedReadFile<F, Fut> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("ChunkedReadFile")
}
}
pub(crate) fn new_chunked_read(size: u64, offset: u64, file: File) -> impl Stream<Item = Result<Bytes, Error>> {
ChunkedReadFile {
size,
offset,
state: ChunkedReadFileState::File { file: Some(file) },
counter: 0,
callback: chunked_read_file_callback,
}
}
async fn chunked_read_file_callback(mut file: File, offset: u64, max_bytes: usize) -> Result<(File, Bytes), Error> {
use io::{Read as _, Seek as _};
let res = actix_web::web::block(move || {
let mut buf = Vec::with_capacity(max_bytes);
file.seek(io::SeekFrom::Start(offset))?;
let n_bytes = file.by_ref().take(max_bytes as u64).read_to_end(&mut buf)?;
if n_bytes == 0 {
Err(io::Error::from(io::ErrorKind::UnexpectedEof))
} else {
Ok((file, Bytes::from(buf)))
}
})
.await??;
Ok(res)
}
impl<F, Fut> Stream for ChunkedReadFile<F, Fut>
where
F: Fn(File, u64, usize) -> Fut,
Fut: Future<Output = Result<(File, Bytes), Error>>,
{
type Item = Result<Bytes, Error>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut this = self.as_mut().project();
match this.state.as_mut().project() {
ChunkedReadFileStateProj::File { file } => {
let size = *this.size;
let offset = *this.offset;
let counter = *this.counter;
if size == counter {
Poll::Ready(None)
} else {
let max_bytes = cmp::min(size.saturating_sub(counter), 65_536) as usize;
let file = file.take().expect("ChunkedReadFile polled after completion");
let fut = (this.callback)(file, offset, max_bytes);
this.state.project_replace(ChunkedReadFileState::Future { fut });
self.poll_next(cx)
}
}
ChunkedReadFileStateProj::Future { fut } => {
let (file, bytes) = ready!(fut.poll(cx))?;
this.state
.project_replace(ChunkedReadFileState::File { file: Some(file) });
*this.offset += bytes.len() as u64;
*this.counter += bytes.len() as u64;
Poll::Ready(Some(Ok(bytes)))
}
}
}
}
pin_project! {
#[doc(hidden)]
pub struct ChunkedReadCompressedFile<F, Fut> {
size: u64,
offset: u64,
#[pin]
state: ChunkedReadCompressedFileState<Fut>,
counter: u64,
callback: F,
}
}
pin_project! {
#[project = ChunkedReadCompressedFileStateProj]
#[project_replace = ChunkedReadCompressedFileStateProjReplace]
enum ChunkedReadCompressedFileState<Fut> {
File { file: Option<ContentDecoder>, },
Future { #[pin] fut: Fut },
}
}
impl<F, Fut> fmt::Debug for ChunkedReadCompressedFile<F, Fut> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("ChunkedReadFile")
}
}
pub(crate) fn new_chunked_compressed_read(
size: u64,
offset: u64,
file: ContentDecoder,
) -> impl Stream<Item = Result<Bytes, Error>> {
ChunkedReadCompressedFile {
size,
offset,
state: ChunkedReadCompressedFileState::File { file: Some(file) },
counter: 0,
callback: chunked_read_compressed_file_callback,
}
}
pub(crate) fn compressed_size(mut f: File, encoding: ContentEncoding) -> Result<usize, Error> {
use std::io::{Read, Seek};
match encoding {
ContentEncoding::Brotli => {
f.seek(io::SeekFrom::Start(0))?;
let decoder = brotli::Decompressor::new(f, 8_096);
let bytes = decoder.bytes();
let (_, hint) = bytes.size_hint();
Ok(hint.unwrap_or_else(|| bytes.count()))
}
ContentEncoding::Gzip => {
f.seek(io::SeekFrom::Start(0))?;
let decoder = GzDecoder::new(f);
let bytes = decoder.bytes();
let (_, hint) = bytes.size_hint();
Ok(hint.unwrap_or_else(|| bytes.count()))
}
ContentEncoding::Deflate => {
f.seek(io::SeekFrom::Start(0))?;
let decoder = ZlibDecoder::new(f);
let bytes = decoder.bytes();
let (_, hint) = bytes.size_hint();
Ok(hint.unwrap_or_else(|| bytes.count()))
}
ContentEncoding::Zstd => {
let decoder = ZstdDecoder::new(f)
.expect("Failed to create zstd decoder. This is a bug. Please report it to the vacuna repository.");
let bytes = decoder.bytes();
let (_, hint) = bytes.size_hint();
Ok(hint.unwrap_or_else(|| bytes.count()))
}
_ => {
unreachable!();
}
}
}
async fn chunked_read_compressed_file_callback(
mut file: ContentDecoder,
offset: u64,
max_bytes: usize,
) -> Result<(ContentDecoder, Bytes), Error> {
trace!("pre-read content uncompressing it for pre-compressed");
let res = actix_web::web::block(move || {
trace!("read content uncompressing it for pre-compressed");
let buf = file.read_at(offset, max_bytes as u64)?;
if let Some(b) = buf {
trace!(
"after-read content uncompressing it for pre-compressed read:{}",
b.len()
);
trace!("after-read content read:{:?}", String::from_utf8(b.to_vec()));
Ok((file, b))
} else {
Err(io::Error::from(io::ErrorKind::UnexpectedEof))
}
})
.await??;
Ok(res)
}
impl<F, Fut> Stream for ChunkedReadCompressedFile<F, Fut>
where
F: Fn(ContentDecoder, u64, usize) -> Fut,
Fut: Future<Output = Result<(ContentDecoder, Bytes), Error>>,
{
type Item = Result<Bytes, Error>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut this = self.as_mut().project();
match this.state.as_mut().project() {
ChunkedReadCompressedFileStateProj::File { file } => {
let size = *this.size;
let offset = *this.offset;
let counter = *this.counter;
if size == counter {
Poll::Ready(None)
} else {
let max_bytes = cmp::min(size.saturating_sub(counter), 65_536) as usize;
let file = file.take().expect("ChunkedReadFile polled after completion");
let fut = (this.callback)(file, offset, max_bytes);
this.state
.project_replace(ChunkedReadCompressedFileState::Future { fut });
self.poll_next(cx)
}
}
ChunkedReadCompressedFileStateProj::Future { fut } => {
let (file, bytes) = ready!(fut.poll(cx))?;
this.state
.project_replace(ChunkedReadCompressedFileState::File { file: Some(file) });
*this.offset += bytes.len() as u64;
*this.counter += bytes.len() as u64;
Poll::Ready(Some(Ok(bytes)))
}
}
}
}
use flate2::read::{GzDecoder, ZlibDecoder};
use zstd::stream::Decoder as ZstdDecoder;
use bytes::Bytes;
pub struct ContentDecoder {
file: File,
encoding: ContentEncoding,
}
impl ContentDecoder {
pub fn new(file: File, encoding: ContentEncoding) -> Self {
trace!("decoding {file:?} from encoding {encoding:?}");
ContentDecoder { file, encoding }
}
fn read_at(&mut self, offset: u64, max_bytes: u64) -> io::Result<Option<Bytes>> {
trace!("decode reading from {offset} bytes {max_bytes} ");
use std::io::{Read, Seek};
let data = match self.encoding {
ContentEncoding::Brotli => {
let mut f = self.file.try_clone()?;
f.seek(io::SeekFrom::Start(0))?;
let decoder = brotli::Decompressor::new(f, 8_096);
decoder
.bytes()
.skip(offset as usize)
.take(max_bytes as usize)
.filter_map(Result::ok)
.collect::<Vec<u8>>()
}
ContentEncoding::Gzip => {
let mut f = self.file.try_clone()?;
f.seek(io::SeekFrom::Start(0))?;
let decoder = GzDecoder::new(f);
decoder
.bytes()
.skip(offset as usize)
.take(max_bytes as usize)
.filter_map(Result::ok)
.collect::<Vec<u8>>()
}
ContentEncoding::Deflate => {
let mut f = self.file.try_clone()?;
f.seek(io::SeekFrom::Start(0))?;
let decoder = ZlibDecoder::new(f);
decoder
.bytes()
.skip(offset as usize)
.take(max_bytes as usize)
.filter_map(Result::ok)
.collect::<Vec<u8>>()
}
ContentEncoding::Zstd => {
let mut f = self.file.try_clone()?;
f.seek(io::SeekFrom::Start(0))?;
let decoder = ZstdDecoder::new(f)
.expect("Failed to create zstd decoder. This is a bug. Please report it to the vacuna repository.");
decoder
.bytes()
.skip(offset as usize)
.take(max_bytes as usize)
.filter_map(Result::ok)
.collect::<Vec<u8>>()
}
_ => {
unreachable!();
}
};
if !data.is_empty() {
Ok(Some(data.into()))
} else {
Ok(None)
}
}
}