// Copyright 2024 RustFS Team // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // Default encryption block size - aligned with system default read buffer size (1MB) pub const DEFAULT_ENCRYPTION_BLOCK_SIZE: usize = 1024 * 1024; use std::future::Future; use std::pin::Pin; use tokio::io::AsyncReadExt; macro_rules! delegate_reader_capabilities_generic { ($name:ident<$inner_ty:ident>, $inner:ident) => { impl<$inner_ty> crate::EtagResolvable for $name<$inner_ty> where $inner_ty: crate::EtagResolvable, { fn try_resolve_etag(&mut self) -> Option { self.$inner.try_resolve_etag() } } impl<$inner_ty> crate::HashReaderDetector for $name<$inner_ty> where $inner_ty: crate::HashReaderDetector, { fn is_hash_reader(&self) -> bool { self.$inner.is_hash_reader() } fn as_hash_reader_mut(&mut self) -> Option<&mut dyn crate::HashReaderMut> { self.$inner.as_hash_reader_mut() } } impl<$inner_ty> crate::TryGetIndex for $name<$inner_ty> where $inner_ty: crate::TryGetIndex, { fn try_get_index(&self) -> Option<&crate::compress_index::Index> { self.$inner.try_get_index() } } }; } macro_rules! delegate_reader_capabilities_generic_no_index { ($name:ident<$inner_ty:ident>, $inner:ident) => { impl<$inner_ty> crate::EtagResolvable for $name<$inner_ty> where $inner_ty: crate::EtagResolvable, { fn try_resolve_etag(&mut self) -> Option { self.$inner.try_resolve_etag() } } impl<$inner_ty> crate::HashReaderDetector for $name<$inner_ty> where $inner_ty: crate::HashReaderDetector, { fn is_hash_reader(&self) -> bool { self.$inner.is_hash_reader() } fn as_hash_reader_mut(&mut self) -> Option<&mut dyn crate::HashReaderMut> { self.$inner.as_hash_reader_mut() } } }; } mod limit_reader; pub use limit_reader::LimitReader; mod etag_reader; pub use etag_reader::EtagReader; mod compress_index; mod compress_reader; pub use compress_reader::{CompressReader, DecompressReader}; mod encrypt_reader; pub use encrypt_reader::{DecryptReader, EncryptReader}; mod hardlimit_reader; pub use hardlimit_reader::HardLimitReader; mod hash_reader; pub use hash_reader::*; mod checksum; pub use checksum::*; mod errors; pub use errors::*; pub mod reader; pub use reader::WarpReader; mod writer; pub use writer::*; mod http_reader; pub use http_reader::*; pub use compress_index::{Index, TryGetIndex}; mod etag; pub type BoxReadBlockFuture<'a> = Pin> + Send + 'a>>; pub trait BlockReadable { fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a>; } fn read_block_via_async_read<'a, R>(reader: &'a mut R, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> where R: tokio::io::AsyncRead + Unpin + Send + Sync + 'a, { Box::pin(async move { let mut total = 0; while total < buf.len() { match reader.read(&mut buf[total..]).await { Ok(0) => return Ok(total), Ok(n) => total += n, Err(err) => return Err(err), } } Ok(total) }) } pub trait ReadStream: tokio::io::AsyncRead + Unpin + Send + Sync {} impl ReadStream for T where T: tokio::io::AsyncRead + Unpin + Send + Sync {} pub trait ReaderCapabilities: EtagResolvable + HashReaderDetector + TryGetIndex {} impl ReaderCapabilities for T where T: EtagResolvable + HashReaderDetector + TryGetIndex {} pub trait Reader: ReadStream + ReaderCapabilities + BlockReadable {} impl Reader for T where T: ReadStream + ReaderCapabilities + BlockReadable {} pub type DynReader = Box; // Trait for types that can be recursively searched for etag capability pub trait EtagResolvable { fn is_etag_reader(&self) -> bool { false } fn try_resolve_etag(&mut self) -> Option { None } } // Generic function that can work with any EtagResolvable type pub fn resolve_etag_generic(reader: &mut R) -> Option where R: EtagResolvable, { reader.try_resolve_etag() } /// Trait to detect and manipulate HashReader instances pub trait HashReaderDetector { fn is_hash_reader(&self) -> bool { false } fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> { None } } impl BlockReadable for crate::WarpReader where R: tokio::io::AsyncRead + Unpin + Send + Sync, { fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> { read_block_via_async_read(self, buf) } } impl BlockReadable for tokio::io::BufReader where R: tokio::io::AsyncRead + Unpin + Send + Sync, { fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> { read_block_via_async_read(self, buf) } } impl BlockReadable for crate::LimitReader where R: Reader, { fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> { read_block_via_async_read(self, buf) } } impl BlockReadable for crate::DecryptReader where R: Reader, { fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> { read_block_via_async_read(self, buf) } } pub fn boxed_reader(reader: R) -> DynReader where R: Reader + 'static, { Box::new(reader) } pub fn wrap_reader(reader: R) -> DynReader where R: ReadStream + 'static, { boxed_reader(WarpReader::new(reader)) } impl EtagResolvable for Box where T: EtagResolvable + ?Sized, { fn try_resolve_etag(&mut self) -> Option { self.as_mut().try_resolve_etag() } } impl HashReaderDetector for Box where T: HashReaderDetector + ?Sized, { fn is_hash_reader(&self) -> bool { self.as_ref().is_hash_reader() } fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> { self.as_mut().as_hash_reader_mut() } } impl TryGetIndex for Box where T: TryGetIndex + ?Sized, { fn try_get_index(&self) -> Option<&compress_index::Index> { self.as_ref().try_get_index() } } impl BlockReadable for Box where T: BlockReadable + ?Sized, { fn read_block<'a>(&'a mut self, buf: &'a mut [u8]) -> BoxReadBlockFuture<'a> { self.as_mut().read_block(buf) } } #[cfg(test)] mod tests { use super::*; use std::collections::VecDeque; use std::io::{self, ErrorKind}; use std::pin::Pin; use std::task::{Context, Poll}; use tokio::io::{AsyncRead, ReadBuf}; enum ReadStep { Data(Vec), Error(ErrorKind), Eof, } struct StepReader { steps: VecDeque, } impl StepReader { fn new(steps: impl IntoIterator) -> Self { Self { steps: steps.into_iter().collect(), } } } impl AsyncRead for StepReader { fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { match self.steps.pop_front().unwrap_or(ReadStep::Eof) { ReadStep::Data(data) => { buf.put_slice(&data); Poll::Ready(Ok(())) } ReadStep::Error(kind) => Poll::Ready(Err(io::Error::new(kind, "synthetic read failure"))), ReadStep::Eof => Poll::Ready(Ok(())), } } } #[tokio::test] async fn test_read_block_via_async_read_preserves_midstream_error_kind() { let reader = StepReader::new([ReadStep::Data(b"ab".to_vec()), ReadStep::Error(ErrorKind::ConnectionReset)]); let mut reader = WarpReader::new(reader); let mut buf = [0_u8; 4]; let err = reader.read_block(&mut buf).await.unwrap_err(); assert_eq!(err.kind(), ErrorKind::ConnectionReset); } #[tokio::test] async fn test_read_block_via_async_read_returns_zero_on_initial_eof() { let mut reader = WarpReader::new(StepReader::new([ReadStep::Eof])); let mut buf = [0_u8; 4]; let n = reader.read_block(&mut buf).await.unwrap(); assert_eq!(n, 0); } }