mirror of
https://github.com/rustfs/rustfs.git
synced 2026-09-06 12:09:12 +00:00
cxymds
655f6ae452
* fix(s3): harden Snowball extract error boundaries * fix(s3): close Snowball extract compatibility gaps * fix(s3): verify Snowball request body completion * test(s3): reject forged Snowball streaming signatures * build(deps): pin Snowball archive parser limits * fix(s3): preserve Snowball trailer and member errors * docs(architecture): register Snowball tar fork cleanup * refactor(s3): route Snowball errors through object boundary * ci(deps): allow pinned tokio-tar source * fix: align Snowball archive codec detection * fix(s3): harden Snowball codec compatibility * fix(s3): preserve Snowball codec compatibility * test(zip): align yield wake assertion with Tokio * fix(rio): preserve legacy large-block reads * fix(zip): accept blank tar numeric fields
739 lines
26 KiB
Rust
739 lines
26 KiB
Rust
// 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.
|
|
|
|
use minlz::{Encoder as MinlzEncoder, crc::crc};
|
|
use pin_project_lite::pin_project;
|
|
use rand::RngExt;
|
|
use rustfs_rio::{
|
|
EtagResolvable, HashReaderDetector, HashReaderMut, Index, MAX_S2_DECOMPRESSED_BLOCK_SIZE, S2Decoder, TryGetIndex,
|
|
};
|
|
use rustfs_utils::CompressionAlgorithm;
|
|
use std::cmp::min;
|
|
use std::fmt;
|
|
use std::io;
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
use tokio::io::{AsyncRead, ReadBuf};
|
|
|
|
const MAGIC_CHUNK: &[u8] = b"\xff\x06\x00\x00S2sTwO";
|
|
const CHUNK_TYPE_COMPRESSED_DATA: u8 = 0x00;
|
|
const CHUNK_TYPE_UNCOMPRESSED_DATA: u8 = 0x01;
|
|
const CHUNK_TYPE_PADDING: u8 = 0xfe;
|
|
const DEFAULT_BLOCK_SIZE: usize = 1 << 20;
|
|
const MAX_CHUNK_SIZE: usize = (1 << 24) - 1;
|
|
const CHECKSUM_SIZE: usize = 4;
|
|
const CHUNK_HEADER_LEN: usize = 4;
|
|
const ENCRYPTED_PADDING_MULTIPLE: usize = 256;
|
|
const MIN_INDEX_SIZE: usize = 8 << 20;
|
|
const MAX_READY_READS_PER_POLL: usize = 64;
|
|
|
|
pin_project! {
|
|
#[derive(Debug)]
|
|
pub struct CompressReader<R> {
|
|
#[pin]
|
|
inner: R,
|
|
buffer: Vec<u8>,
|
|
pos: usize,
|
|
done: bool,
|
|
block_size: usize,
|
|
index: Index,
|
|
written: usize,
|
|
uncompressed_written: usize,
|
|
temp_buffer: Vec<u8>,
|
|
read_buffer: Vec<u8>,
|
|
wrote_stream_header: bool,
|
|
padding_multiple: Option<usize>,
|
|
block_encoder: S2BlockEncoder,
|
|
}
|
|
}
|
|
|
|
struct S2BlockEncoder {
|
|
inner: MinlzEncoder,
|
|
}
|
|
|
|
impl S2BlockEncoder {
|
|
fn new() -> Self {
|
|
Self {
|
|
inner: MinlzEncoder::new(),
|
|
}
|
|
}
|
|
|
|
fn encode(&mut self, uncompressed: &[u8]) -> Vec<u8> {
|
|
self.inner.encode(uncompressed)
|
|
}
|
|
}
|
|
|
|
impl fmt::Debug for S2BlockEncoder {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.debug_struct("S2BlockEncoder").finish_non_exhaustive()
|
|
}
|
|
}
|
|
|
|
impl<R> CompressReader<R>
|
|
where
|
|
R: AsyncRead + Unpin + Send + Sync,
|
|
{
|
|
pub fn new(inner: R, _compression_algorithm: CompressionAlgorithm) -> Self {
|
|
Self::with_block_size(inner, DEFAULT_BLOCK_SIZE, CompressionAlgorithm::default())
|
|
}
|
|
|
|
/// Create an encoder with a caller-selected S2 block size.
|
|
///
|
|
/// Zero selects the default. Larger values are capped at the maximum block
|
|
/// accepted by the paired decoder, preserving this infallible API without
|
|
/// allowing it to emit a stream that RustFS cannot read back.
|
|
pub fn with_block_size(inner: R, block_size: usize, _compression_algorithm: CompressionAlgorithm) -> Self {
|
|
let block_size = if block_size == 0 {
|
|
DEFAULT_BLOCK_SIZE
|
|
} else {
|
|
block_size.min(MAX_S2_DECOMPRESSED_BLOCK_SIZE)
|
|
};
|
|
Self {
|
|
inner,
|
|
buffer: Vec::new(),
|
|
pos: 0,
|
|
done: false,
|
|
block_size,
|
|
index: Index::new(),
|
|
written: 0,
|
|
uncompressed_written: 0,
|
|
temp_buffer: Vec::with_capacity(block_size),
|
|
read_buffer: vec![0u8; block_size],
|
|
wrote_stream_header: false,
|
|
padding_multiple: None,
|
|
block_encoder: S2BlockEncoder::new(),
|
|
}
|
|
}
|
|
|
|
pub fn with_encrypted_padding(inner: R, _compression_algorithm: CompressionAlgorithm) -> Self {
|
|
let mut reader = Self::new(inner, CompressionAlgorithm::default());
|
|
reader.padding_multiple = Some(ENCRYPTED_PADDING_MULTIPLE);
|
|
reader
|
|
}
|
|
}
|
|
|
|
impl<R> TryGetIndex for CompressReader<R> {
|
|
fn try_get_index(&self) -> Option<&Index> {
|
|
(self.uncompressed_written > MIN_INDEX_SIZE).then_some(&self.index)
|
|
}
|
|
}
|
|
|
|
impl<R> AsyncRead for CompressReader<R>
|
|
where
|
|
R: AsyncRead + Unpin + Send + Sync,
|
|
{
|
|
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
let mut this = self.project();
|
|
let mut ready_reads = 0usize;
|
|
|
|
if *this.pos < this.buffer.len() {
|
|
let to_copy = min(buf.remaining(), this.buffer.len() - *this.pos);
|
|
buf.put_slice(&this.buffer[*this.pos..*this.pos + to_copy]);
|
|
*this.pos += to_copy;
|
|
if *this.pos == this.buffer.len() {
|
|
this.buffer.clear();
|
|
*this.pos = 0;
|
|
}
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
if *this.done {
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
while this.temp_buffer.len() < *this.block_size {
|
|
if ready_reads >= MAX_READY_READS_PER_POLL {
|
|
cx.waker().wake_by_ref();
|
|
return Poll::Pending;
|
|
}
|
|
|
|
let remaining = *this.block_size - this.temp_buffer.len();
|
|
let mut read_buf = ReadBuf::new(&mut this.read_buffer[..remaining]);
|
|
match this.inner.as_mut().poll_read(cx, &mut read_buf) {
|
|
Poll::Pending => {
|
|
return Poll::Pending;
|
|
}
|
|
Poll::Ready(Ok(())) => {
|
|
ready_reads += 1;
|
|
let n = read_buf.filled().len();
|
|
if n == 0 {
|
|
break;
|
|
}
|
|
this.temp_buffer.extend_from_slice(read_buf.filled());
|
|
}
|
|
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
|
|
}
|
|
}
|
|
|
|
if this.temp_buffer.is_empty() {
|
|
if let Some(padding_multiple) = *this.padding_multiple
|
|
&& let Some(padding_chunk) = build_padding_chunk(*this.written, padding_multiple)?
|
|
{
|
|
*this.written += padding_chunk.len();
|
|
this.index.total_compressed = *this.written as i64;
|
|
*this.buffer = padding_chunk;
|
|
*this.pos = 0;
|
|
*this.done = true;
|
|
|
|
let to_copy = min(buf.remaining(), this.buffer.len());
|
|
buf.put_slice(&this.buffer[..to_copy]);
|
|
*this.pos += to_copy;
|
|
if *this.pos == this.buffer.len() {
|
|
this.buffer.clear();
|
|
*this.pos = 0;
|
|
}
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
*this.done = true;
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
let mut out = Vec::new();
|
|
if !*this.wrote_stream_header {
|
|
out.extend_from_slice(MAGIC_CHUNK);
|
|
*this.written += MAGIC_CHUNK.len();
|
|
*this.wrote_stream_header = true;
|
|
}
|
|
|
|
if let Err(err) = this.index.add(*this.written as i64, *this.uncompressed_written as i64) {
|
|
return Poll::Ready(Err(err));
|
|
}
|
|
|
|
let block = build_s2_chunk(this.temp_buffer.as_slice(), this.block_encoder)?;
|
|
*this.uncompressed_written += this.temp_buffer.len();
|
|
*this.written += block.len();
|
|
this.index.total_uncompressed = *this.uncompressed_written as i64;
|
|
this.index.total_compressed = *this.written as i64;
|
|
|
|
out.extend_from_slice(&block);
|
|
this.temp_buffer.clear();
|
|
*this.buffer = out;
|
|
*this.pos = 0;
|
|
|
|
let to_copy = min(buf.remaining(), this.buffer.len());
|
|
buf.put_slice(&this.buffer[..to_copy]);
|
|
*this.pos += to_copy;
|
|
if *this.pos == this.buffer.len() {
|
|
this.buffer.clear();
|
|
*this.pos = 0;
|
|
}
|
|
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
|
|
impl<R> EtagResolvable for CompressReader<R>
|
|
where
|
|
R: EtagResolvable,
|
|
{
|
|
fn try_resolve_etag(&mut self) -> Option<String> {
|
|
self.inner.try_resolve_etag()
|
|
}
|
|
}
|
|
|
|
impl<R> HashReaderDetector for CompressReader<R>
|
|
where
|
|
R: HashReaderDetector,
|
|
{
|
|
fn is_hash_reader(&self) -> bool {
|
|
self.inner.is_hash_reader()
|
|
}
|
|
|
|
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
|
|
self.inner.as_hash_reader_mut()
|
|
}
|
|
}
|
|
|
|
pin_project! {
|
|
#[derive(Debug)]
|
|
pub struct DecompressReader<R> {
|
|
#[pin]
|
|
inner: S2Decoder<R>,
|
|
}
|
|
}
|
|
|
|
impl<R> DecompressReader<R>
|
|
where
|
|
R: AsyncRead + Unpin + Send + Sync,
|
|
{
|
|
pub fn new(inner: R, _compression_algorithm: CompressionAlgorithm) -> Self {
|
|
Self {
|
|
inner: S2Decoder::new_at_legacy_chunk_boundary(inner),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<R> AsyncRead for DecompressReader<R>
|
|
where
|
|
R: AsyncRead + Unpin + Send + Sync,
|
|
{
|
|
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
self.project().inner.poll_read(cx, buf)
|
|
}
|
|
}
|
|
|
|
impl<R> EtagResolvable for DecompressReader<R>
|
|
where
|
|
R: EtagResolvable,
|
|
{
|
|
fn try_resolve_etag(&mut self) -> Option<String> {
|
|
self.inner.get_mut().try_resolve_etag()
|
|
}
|
|
}
|
|
|
|
impl<R> HashReaderDetector for DecompressReader<R>
|
|
where
|
|
R: HashReaderDetector,
|
|
{
|
|
fn is_hash_reader(&self) -> bool {
|
|
self.inner.get_ref().is_hash_reader()
|
|
}
|
|
|
|
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
|
|
self.inner.get_mut().as_hash_reader_mut()
|
|
}
|
|
}
|
|
|
|
fn build_s2_chunk(uncompressed: &[u8], encoder: &mut S2BlockEncoder) -> io::Result<Vec<u8>> {
|
|
let compressed = encode_block(uncompressed, encoder);
|
|
let checksum = crc(uncompressed);
|
|
let dst_limit = uncompressed.len().saturating_sub(uncompressed.len() / 32).saturating_sub(5);
|
|
let (chunk_type, payload) = if compressed.len() <= dst_limit {
|
|
(CHUNK_TYPE_COMPRESSED_DATA, compressed)
|
|
} else {
|
|
(CHUNK_TYPE_UNCOMPRESSED_DATA, uncompressed.to_vec())
|
|
};
|
|
|
|
let chunk_len = payload.len() + CHECKSUM_SIZE;
|
|
if chunk_len > MAX_CHUNK_SIZE {
|
|
return Err(io::Error::new(io::ErrorKind::InvalidData, "S2 chunk exceeds 24-bit framing limit"));
|
|
}
|
|
|
|
let mut out = Vec::with_capacity(CHUNK_HEADER_LEN + chunk_len);
|
|
out.push(chunk_type);
|
|
out.push((chunk_len & 0xff) as u8);
|
|
out.push(((chunk_len >> 8) & 0xff) as u8);
|
|
out.push(((chunk_len >> 16) & 0xff) as u8);
|
|
out.extend_from_slice(&checksum.to_le_bytes());
|
|
out.extend_from_slice(&payload);
|
|
Ok(out)
|
|
}
|
|
|
|
fn encode_block(uncompressed: &[u8], encoder: &mut S2BlockEncoder) -> Vec<u8> {
|
|
encoder.encode(uncompressed)
|
|
}
|
|
|
|
fn build_padding_chunk(current_size: usize, padding_multiple: usize) -> io::Result<Option<Vec<u8>>> {
|
|
if padding_multiple == 0 || current_size.is_multiple_of(padding_multiple) {
|
|
return Ok(None);
|
|
}
|
|
|
|
let padding_len = (padding_multiple - ((current_size + CHUNK_HEADER_LEN) % padding_multiple)) % padding_multiple;
|
|
if padding_len > MAX_CHUNK_SIZE {
|
|
return Err(io::Error::new(io::ErrorKind::InvalidData, "S2 padding exceeds 24-bit framing limit"));
|
|
}
|
|
|
|
let mut out = Vec::with_capacity(CHUNK_HEADER_LEN + padding_len);
|
|
out.push(CHUNK_TYPE_PADDING);
|
|
out.push((padding_len & 0xff) as u8);
|
|
out.push(((padding_len >> 8) & 0xff) as u8);
|
|
out.push(((padding_len >> 16) & 0xff) as u8);
|
|
|
|
if padding_len > 0 {
|
|
let mut padding = vec![0u8; padding_len];
|
|
rand::rng().fill(padding.as_mut_slice());
|
|
out.extend_from_slice(&padding);
|
|
}
|
|
|
|
Ok(Some(out))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::io::Cursor;
|
|
use std::pin::Pin;
|
|
use std::sync::{
|
|
Arc,
|
|
atomic::{AtomicUsize, Ordering},
|
|
};
|
|
use std::task::{Context, Poll, Wake, Waker};
|
|
use tokio::io::AsyncReadExt;
|
|
|
|
#[derive(Default)]
|
|
struct WakeCounter(AtomicUsize);
|
|
|
|
impl Wake for WakeCounter {
|
|
fn wake(self: Arc<Self>) {
|
|
self.0.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
}
|
|
|
|
struct AlwaysReadyOneByte {
|
|
bytes: Vec<u8>,
|
|
position: usize,
|
|
read_calls: Arc<AtomicUsize>,
|
|
}
|
|
|
|
impl AlwaysReadyOneByte {
|
|
fn new(bytes: Vec<u8>, read_calls: Arc<AtomicUsize>) -> Self {
|
|
Self {
|
|
bytes,
|
|
position: 0,
|
|
read_calls,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl AsyncRead for AlwaysReadyOneByte {
|
|
fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
self.read_calls.fetch_add(1, Ordering::Relaxed);
|
|
if self.position == self.bytes.len() || buf.remaining() == 0 {
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
let byte = self.bytes[self.position];
|
|
self.position += 1;
|
|
buf.put_slice(&[byte]);
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
|
|
struct PendingAfterBytes<R> {
|
|
inner: R,
|
|
max_chunk: usize,
|
|
pending_next: bool,
|
|
}
|
|
|
|
impl<R> PendingAfterBytes<R> {
|
|
fn new(inner: R, max_chunk: usize) -> Self {
|
|
Self {
|
|
inner,
|
|
max_chunk,
|
|
pending_next: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<R: AsyncRead + Unpin> AsyncRead for PendingAfterBytes<R> {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
if self.pending_next {
|
|
self.pending_next = false;
|
|
cx.waker().wake_by_ref();
|
|
return Poll::Pending;
|
|
}
|
|
|
|
let allowed = self.max_chunk.min(buf.remaining());
|
|
if allowed == 0 {
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
let mut scratch = vec![0u8; allowed];
|
|
let mut limited = ReadBuf::new(&mut scratch);
|
|
match Pin::new(&mut self.inner).poll_read(cx, &mut limited) {
|
|
Poll::Pending => Poll::Pending,
|
|
Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
|
|
Poll::Ready(Ok(())) => {
|
|
let filled = limited.filled();
|
|
if !filled.is_empty() {
|
|
buf.put_slice(filled);
|
|
self.pending_next = true;
|
|
}
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn s2_chunk_types(stream: &[u8]) -> Vec<u8> {
|
|
let mut chunk_types = Vec::new();
|
|
let mut offset = 0usize;
|
|
while offset + CHUNK_HEADER_LEN <= stream.len() {
|
|
let chunk_type = stream[offset];
|
|
let chunk_len =
|
|
(stream[offset + 1] as usize) | ((stream[offset + 2] as usize) << 8) | ((stream[offset + 3] as usize) << 16);
|
|
chunk_types.push(chunk_type);
|
|
offset += CHUNK_HEADER_LEN + chunk_len;
|
|
}
|
|
chunk_types
|
|
}
|
|
|
|
#[test]
|
|
fn minlz_encoded_payload_decodes_with_minlz() {
|
|
let plaintext = b"compressible-rio-v2-block-".repeat(4096);
|
|
let mut encoder = S2BlockEncoder::new();
|
|
let compressed = encode_block(&plaintext, &mut encoder);
|
|
let decoded = minlz::decode(&compressed).expect("decode payload");
|
|
|
|
assert_eq!(decoded, plaintext);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_compress_reader_roundtrip() {
|
|
let plaintext = b"hello-rio-v2-s2-".repeat(32_768);
|
|
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
|
|
let mut compressed = Vec::new();
|
|
reader.read_to_end(&mut compressed).await.expect("read compressed data");
|
|
|
|
assert!(compressed.starts_with(MAGIC_CHUNK));
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
#[test]
|
|
fn s2_compress_reader_yields_after_ready_read_budget() {
|
|
let read_calls = Arc::new(AtomicUsize::new(0));
|
|
let source = AlwaysReadyOneByte::new(vec![b'x'; MAX_READY_READS_PER_POLL + 1], read_calls.clone());
|
|
let mut reader = CompressReader::new(source, CompressionAlgorithm::default());
|
|
let wake_counter = Arc::new(WakeCounter::default());
|
|
let waker = Waker::from(wake_counter.clone());
|
|
let mut cx = Context::from_waker(&waker);
|
|
let mut output = [0u8; 1];
|
|
let mut read_buf = ReadBuf::new(&mut output);
|
|
|
|
assert!(Pin::new(&mut reader).poll_read(&mut cx, &mut read_buf).is_pending());
|
|
assert!(read_buf.filled().is_empty());
|
|
assert_eq!(read_calls.load(Ordering::Relaxed), MAX_READY_READS_PER_POLL);
|
|
assert_eq!(reader.temp_buffer.len(), MAX_READY_READS_PER_POLL);
|
|
assert_eq!(wake_counter.0.load(Ordering::Relaxed), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn s2_compress_reader_normalizes_non_decodable_block_sizes() {
|
|
let zero = CompressReader::with_block_size(Cursor::new(Vec::<u8>::new()), 0, CompressionAlgorithm::default());
|
|
assert_eq!(zero.block_size, DEFAULT_BLOCK_SIZE);
|
|
|
|
let oversized = CompressReader::with_block_size(
|
|
Cursor::new(Vec::<u8>::new()),
|
|
MAX_S2_DECOMPRESSED_BLOCK_SIZE + 1,
|
|
CompressionAlgorithm::default(),
|
|
);
|
|
assert_eq!(oversized.block_size, MAX_S2_DECOMPRESSED_BLOCK_SIZE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_compress_reader_max_block_roundtrips_with_paired_decoder() {
|
|
let plaintext = pseudo_random_bytes(MAX_S2_DECOMPRESSED_BLOCK_SIZE);
|
|
let mut reader = CompressReader::with_block_size(
|
|
Cursor::new(plaintext.clone()),
|
|
MAX_S2_DECOMPRESSED_BLOCK_SIZE,
|
|
CompressionAlgorithm::default(),
|
|
);
|
|
let mut compressed = Vec::new();
|
|
reader.read_to_end(&mut compressed).await.expect("read maximum S2 block");
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor
|
|
.read_to_end(&mut actual)
|
|
.await
|
|
.expect("paired decoder should accept maximum S2 block");
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_decompress_reader_accepts_legacy_block_above_16_mib() {
|
|
const PRE_CAP_LEGACY_BLOCK_SIZE: usize = (16 << 20) + 1;
|
|
let plaintext = vec![b'x'; PRE_CAP_LEGACY_BLOCK_SIZE];
|
|
let mut encoder = S2BlockEncoder::new();
|
|
let mut fixture = MAGIC_CHUNK.to_vec();
|
|
fixture.extend_from_slice(
|
|
&build_s2_chunk(&plaintext, &mut encoder).expect("the pre-cap writer format should encode a block above 16 MiB"),
|
|
);
|
|
assert_eq!(fixture[MAGIC_CHUNK.len()], CHUNK_TYPE_COMPRESSED_DATA);
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(fixture), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor
|
|
.read_to_end(&mut actual)
|
|
.await
|
|
.expect("legacy rio-v2 blocks above the current writer limit should remain readable");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_decompress_reader_accepts_an_indexed_headerless_tail() {
|
|
let plaintext = b"indexed-rio-v2-s2-tail-".repeat(32_768);
|
|
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
|
|
let mut compressed = Vec::new();
|
|
reader.read_to_end(&mut compressed).await.expect("read compressed data");
|
|
assert!(compressed.starts_with(MAGIC_CHUNK));
|
|
|
|
let mut decompressor =
|
|
DecompressReader::new(Cursor::new(compressed[MAGIC_CHUNK.len()..].to_vec()), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor
|
|
.read_to_end(&mut actual)
|
|
.await
|
|
.expect("indexed tail should decode without the stream header");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_compress_reader_roundtrip_near_erasure_boundary() {
|
|
let size = 4 * 1024 * 1024 - 97;
|
|
let plaintext = pseudo_random_bytes(size);
|
|
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
|
|
let mut compressed = Vec::new();
|
|
reader.read_to_end(&mut compressed).await.expect("read compressed data");
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
fn pseudo_random_bytes(size: usize) -> Vec<u8> {
|
|
(0..size)
|
|
.scan(0x9e37_79b9_7f4a_7c15u64, |state, _| {
|
|
*state ^= *state << 7;
|
|
*state ^= *state >> 9;
|
|
*state = state.wrapping_mul(0xbf58_476d_1ce4_e5b9);
|
|
Some((*state >> 32) as u8)
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_compress_reader_roundtrip_large_random() {
|
|
let plaintext = pseudo_random_bytes(8 * 1024 * 1024 + 123);
|
|
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
|
|
let mut compressed = Vec::new();
|
|
reader.read_to_end(&mut compressed).await.expect("read compressed data");
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_compress_reader_roundtrip_with_pending_source() {
|
|
let plaintext = pseudo_random_bytes(2 * 1024 * 1024 + 17);
|
|
let pending_reader = PendingAfterBytes::new(Cursor::new(plaintext.clone()), 257);
|
|
let mut reader = CompressReader::new(pending_reader, CompressionAlgorithm::default());
|
|
let mut compressed = Vec::new();
|
|
reader.read_to_end(&mut compressed).await.expect("read compressed data");
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_decompress_reader_returns_bytes_on_first_read() {
|
|
let plaintext = b"abcdefghijklmnopqrstuvwxyz".to_vec();
|
|
let mut compressed = Vec::new();
|
|
CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default())
|
|
.read_to_end(&mut compressed)
|
|
.await
|
|
.expect("compress plaintext");
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
|
|
let mut buf = [0u8; 64];
|
|
let n = decompressor.read(&mut buf).await.expect("read first decompressed chunk");
|
|
|
|
assert!(n > 0);
|
|
assert_eq!(&buf[..n], plaintext.as_slice());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_decompress_reader_resumes_chunk_body_after_pending() {
|
|
let plaintext = pseudo_random_bytes(1024 * 1024 + 123);
|
|
let mut compressed = Vec::new();
|
|
CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default())
|
|
.read_to_end(&mut compressed)
|
|
.await
|
|
.expect("compress plaintext");
|
|
|
|
let pending_reader = PendingAfterBytes::new(Cursor::new(compressed), 257);
|
|
let mut decompressor = DecompressReader::new(pending_reader, CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_decompress_reader_handles_concatenated_streams_after_pending() {
|
|
let first = pseudo_random_bytes(16 * 1024 * 1024);
|
|
let second = pseudo_random_bytes(12 * 1024 * 1024 + 123);
|
|
let mut first_compressed = Vec::new();
|
|
CompressReader::new(Cursor::new(first.clone()), CompressionAlgorithm::default())
|
|
.read_to_end(&mut first_compressed)
|
|
.await
|
|
.expect("compress first stream");
|
|
let mut second_compressed = Vec::new();
|
|
CompressReader::new(Cursor::new(second.clone()), CompressionAlgorithm::default())
|
|
.read_to_end(&mut second_compressed)
|
|
.await
|
|
.expect("compress second stream");
|
|
|
|
first_compressed.extend_from_slice(&second_compressed);
|
|
|
|
let pending_reader = PendingAfterBytes::new(Cursor::new(first_compressed), 4096);
|
|
let mut decompressor = DecompressReader::new(pending_reader, CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
|
|
|
|
let mut expected = first;
|
|
expected.extend_from_slice(&second);
|
|
assert_eq!(actual, expected);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_compress_reader_with_encrypted_padding_emits_padding_frame() {
|
|
let plaintext = b"encrypted-padding-check-".repeat(8192);
|
|
let mut reader = CompressReader::with_encrypted_padding(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
|
|
let mut compressed = Vec::new();
|
|
reader.read_to_end(&mut compressed).await.expect("read compressed data");
|
|
|
|
assert_eq!(compressed.len() % ENCRYPTED_PADDING_MULTIPLE, 0);
|
|
assert!(s2_chunk_types(&compressed).contains(&CHUNK_TYPE_PADDING));
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn s2_compress_reader_skips_index_for_small_streams() {
|
|
let plaintext = b"index-threshold-check-".repeat(16_384);
|
|
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
|
|
let mut compressed = Vec::new();
|
|
reader.read_to_end(&mut compressed).await.expect("read compressed data");
|
|
|
|
assert!(reader.try_get_index().is_none());
|
|
|
|
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
|
|
let mut actual = Vec::new();
|
|
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
|
|
|
|
assert_eq!(actual, plaintext);
|
|
}
|
|
}
|