From 655f6ae4523786bce596a8cf8dab607b9a37de21 Mon Sep 17 00:00:00 2001 From: cxymds Date: Mon, 31 Aug 2026 21:09:51 +0800 Subject: [PATCH] fix(s3): align Snowball codec compatibility (#6943) * 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 --- Cargo.lock | 4 + crates/e2e_test/Cargo.toml | 3 +- crates/e2e_test/src/multipart_auth_test.rs | 184 +++- crates/rio-v2/src/compress_reader.rs | 345 ++++---- crates/rio/Cargo.toml | 1 + crates/rio/src/lib.rs | 3 + crates/rio/src/s2_decoder.rs | 858 ++++++++++++++++++ crates/zip/Cargo.toml | 8 +- crates/zip/src/lib.rs | 954 ++++++++++++++++++++- rustfs/src/app/object/extract.rs | 128 ++- rustfs/src/app/object/mod.rs | 2 +- 11 files changed, 2272 insertions(+), 218 deletions(-) create mode 100644 crates/rio/src/s2_decoder.rs diff --git a/Cargo.lock b/Cargo.lock index ab5e7893c..82964573b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2162,6 +2162,7 @@ dependencies = [ "compression-core", "flate2", "liblzma", + "lz4", "memchr", "zstd", "zstd-safe", @@ -3943,6 +3944,7 @@ dependencies = [ "hyper-util", "local-ip-address", "md-5 0.11.0", + "minlz", "opentelemetry-proto", "prost 0.14.4", "rand 0.10.2", @@ -10448,6 +10450,7 @@ dependencies = [ "hyper", "hyper-util", "md-5 0.11.0", + "minlz", "pin-project-lite", "rand 0.10.2", "reqwest", @@ -10882,6 +10885,7 @@ version = "1.0.0-rc.4" dependencies = [ "async-compression", "hotpath", + "rustfs-rio", "thiserror 2.0.20", "tokio", ] diff --git a/crates/e2e_test/Cargo.toml b/crates/e2e_test/Cargo.toml index 0a0a70872..ac68cc646 100644 --- a/crates/e2e_test/Cargo.toml +++ b/crates/e2e_test/Cargo.toml @@ -101,7 +101,7 @@ aws-sdk-sts = { workspace = true, default-features = false, features = ["default aws-config = { workspace = true } aws-smithy-http-client = { workspace = true, default-features = false, features = ["rustls-aws-lc"] } aws-smithy-types.workspace = true -async-compression = { workspace = true, features = ["tokio", "bzip2", "xz"] } +async-compression = { workspace = true, features = ["tokio", "bzip2", "lz4", "xz"] } async-trait = { workspace = true } flate2.workspace = true http.workspace = true @@ -115,6 +115,7 @@ rustfs-signer.workspace = true # server's implementation: a shared helper could agree with a bug on both sides. data-encoding = { workspace = true } hmac = { workspace = true } +minlz.workspace = true sha1 = { workspace = true } serde_urlencoded = { workspace = true } tracing = { workspace = true } diff --git a/crates/e2e_test/src/multipart_auth_test.rs b/crates/e2e_test/src/multipart_auth_test.rs index b15998512..d2d45600a 100644 --- a/crates/e2e_test/src/multipart_auth_test.rs +++ b/crates/e2e_test/src/multipart_auth_test.rs @@ -15,7 +15,7 @@ //! Regression coverage for anonymous access on multipart control APIs. use crate::common::{RustFSTestEnvironment, init_logging, local_http_client}; -use async_compression::tokio::write::{BzEncoder, XzEncoder}; +use async_compression::tokio::write::{BzEncoder, Lz4Encoder, XzEncoder}; use aws_sdk_s3::error::{ProvideErrorMetadata, SdkError}; use aws_sdk_s3::operation::head_object::HeadObjectOutput; use aws_sdk_s3::primitives::ByteStream; @@ -23,7 +23,10 @@ use aws_sdk_s3::types::{ ServerSideEncryption, ServerSideEncryptionByDefault, ServerSideEncryptionConfiguration, ServerSideEncryptionRule, }; use chrono::{Duration as ChronoDuration, Utc}; -use flate2::{Compression, write::GzEncoder}; +use flate2::{ + Compression, + write::{GzEncoder, ZlibEncoder}, +}; use http::HeaderValue; use http::header::{CONTENT_TYPE, HOST}; use md5::{Digest as Md5Digest, Md5}; @@ -187,6 +190,12 @@ fn gzip_bytes(data: &[u8]) -> Vec { encoder.finish().expect("gzip encoder should finish") } +fn zlib_bytes(data: &[u8]) -> Vec { + let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(data).expect("zlib encoder should accept input"); + encoder.finish().expect("zlib encoder should finish") +} + fn zstd_bytes(data: &[u8]) -> Vec { let mut encoder = zstd::Encoder::new(Vec::new(), 0).expect("zstd encoder should initialize"); encoder.write_all(data).expect("zstd encoder should accept input"); @@ -209,6 +218,45 @@ async fn xz_bytes(data: &[u8]) -> Vec { encoder.into_inner().into_inner() } +async fn lz4_bytes(data: &[u8]) -> Vec { + let cursor = Cursor::new(Vec::new()); + let mut encoder = Lz4Encoder::new(cursor); + encoder.write_all(data).await.expect("LZ4 encoder should accept input"); + encoder.shutdown().await.expect("LZ4 encoder should finish"); + encoder.into_inner().into_inner() +} + +/// Encode the S2 framed stream shape emitted by minio-go PutObjectsSnowball +/// with `Compress: true`: 1 MiB independent blocks, better compression, +/// masked CRC-32C, and the `S2sTwO` stream identifier. +fn minio_go_snowball_s2_bytes(data: &[u8]) -> Vec { + const BLOCK_SIZE: usize = 1 << 20; + const CHECKSUM_SIZE: usize = 4; + + let mut output = b"\xff\x06\x00\x00S2sTwO".to_vec(); + let mut encoder = minlz::Encoder::new(); + for block in data.chunks(BLOCK_SIZE) { + let compressed = encoder.encode_better(block); + let compressed_limit = block.len().saturating_sub(block.len() / 32).saturating_sub(5); + let (chunk_type, payload) = if compressed.len() <= compressed_limit { + (0x00, compressed.as_slice()) + } else { + (0x01, block) + }; + let chunk_len = payload.len() + CHECKSUM_SIZE; + assert!(chunk_len < 1 << 24, "S2 fixture chunk must fit the 24-bit frame length"); + output.extend_from_slice(&[ + chunk_type, + (chunk_len & 0xff) as u8, + ((chunk_len >> 8) & 0xff) as u8, + ((chunk_len >> 16) & 0xff) as u8, + ]); + output.extend_from_slice(&minlz::crc::crc(block).to_le_bytes()); + output.extend_from_slice(payload); + } + output +} + fn assert_s3_error_code(result: Result>, code: &str) where T: std::fmt::Debug, @@ -4241,6 +4289,60 @@ async fn test_signed_put_object_extract_returns_archive_etag() -> Result<(), Box Ok(()) } +#[tokio::test] +async fn test_signed_put_object_extract_expands_s2_and_lz4_by_magic_with_raw_etags() +-> Result<(), Box> { + init_logging(); + + let mut env = RustFSTestEnvironment::new().await?; + env.start_rustfs_server(vec![]).await?; + + let bucket = "signed-extract-magic-codecs"; + let client = env.create_s3_client(); + client.create_bucket().bucket(bucket).send().await?; + + let s2_tar = make_tar(&[("s2/object.txt", b"s2-body")], &[]).await; + let s2_archive = minio_go_snowball_s2_bytes(&s2_tar); + let expected_s2_etag = format!("\"{}\"", md5_hex(&s2_archive)); + let s2_response = client + .put_object() + .bucket(bucket) + // minio-go intentionally uploads a compressed S2 stream with a .tar key. + .key("snowball-upload-0123456789abcdef.tar") + .body(ByteStream::from(s2_archive)) + .customize() + .mutate_request(|req| { + req.headers_mut().insert("x-amz-meta-snowball-auto-extract", "true"); + }) + .send() + .await?; + assert_eq!(s2_response.e_tag(), Some(expected_s2_etag.as_str())); + + let s2_object = client.get_object().bucket(bucket).key("s2/object.txt").send().await?; + assert_eq!(s2_object.body.collect().await?.into_bytes().as_ref(), b"s2-body"); + + let lz4_tar = make_tar(&[("lz4/object.txt", b"lz4-body")], &[]).await; + let lz4_archive = lz4_bytes(&lz4_tar).await; + let expected_lz4_etag = format!("\"{}\"", md5_hex(&lz4_archive)); + let lz4_response = client + .put_object() + .bucket(bucket) + .key("also-looks-like-a-plain.tar") + .body(ByteStream::from(lz4_archive)) + .customize() + .mutate_request(|req| { + req.headers_mut().insert("x-amz-meta-snowball-auto-extract", "true"); + }) + .send() + .await?; + assert_eq!(lz4_response.e_tag(), Some(expected_lz4_etag.as_str())); + + let lz4_object = client.get_object().bucket(bucket).key("lz4/object.txt").send().await?; + assert_eq!(lz4_object.body.collect().await?.into_bytes().as_ref(), b"lz4-body"); + + Ok(()) +} + #[tokio::test] async fn test_signed_put_object_extract_preserves_entry_mtime() -> Result<(), Box> { init_logging(); @@ -5106,8 +5208,8 @@ async fn test_signed_put_object_extract_expands_tzst_archive() -> Result<(), Box } #[tokio::test] -async fn test_signed_put_object_extract_rejects_missing_archive_extension() -> Result<(), Box> -{ +async fn test_signed_put_object_extract_uses_magic_without_requiring_or_trusting_extension() +-> Result<(), Box> { init_logging(); let mut env = RustFSTestEnvironment::new().await?; @@ -5120,8 +5222,7 @@ async fn test_signed_put_object_extract_rejects_missing_archive_extension() -> R admin_client.create_bucket().bucket(bucket).send().await?; let tar_bytes = make_tar(&[("plain.txt", b"plain-body")], &[]).await; - - let result = admin_client + admin_client .put_object() .bucket(bucket) .key(archive_key) @@ -5131,15 +5232,80 @@ async fn test_signed_put_object_extract_rejects_missing_archive_extension() -> R req.headers_mut().insert("x-amz-meta-snowball-auto-extract", "true"); }) .send() - .await; + .await?; - assert_s3_error_code(result, "InvalidArgument"); + let plain = admin_client.get_object().bucket(bucket).key("plain.txt").send().await?; + assert_eq!(plain.body.collect().await?.into_bytes().as_ref(), b"plain-body"); + + let raw_with_gzip_suffix = make_tar(&[("raw-with-wrong-suffix.txt", b"raw-body")], &[]).await; + admin_client + .put_object() + .bucket(bucket) + .key("raw-but-named.tar.gz") + .body(ByteStream::from(raw_with_gzip_suffix)) + .customize() + .mutate_request(|req| { + req.headers_mut().insert("x-amz-meta-snowball-auto-extract", "true"); + }) + .send() + .await?; + + let raw = admin_client + .get_object() + .bucket(bucket) + .key("raw-with-wrong-suffix.txt") + .send() + .await?; + assert_eq!(raw.body.collect().await?.into_bytes().as_ref(), b"raw-body"); + + let gzip_with_tar_suffix = gzip_bytes(&make_tar(&[("gzip-with-wrong-suffix.txt", b"gzip-body")], &[]).await); + admin_client + .put_object() + .bucket(bucket) + .key("gzip-but-named.tar") + .body(ByteStream::from(gzip_with_tar_suffix)) + .customize() + .mutate_request(|req| { + req.headers_mut().insert("x-amz-meta-snowball-auto-extract", "true"); + }) + .send() + .await?; + + let gzip = admin_client + .get_object() + .bucket(bucket) + .key("gzip-with-wrong-suffix.txt") + .send() + .await?; + assert_eq!(gzip.body.collect().await?.into_bytes().as_ref(), b"gzip-body"); + + let zlib_archive = zlib_bytes(&make_tar(&[("zlib-extension.txt", b"zlib-body")], &[]).await); + admin_client + .put_object() + .bucket(bucket) + .key("bundle.zlib") + .body(ByteStream::from(zlib_archive)) + .customize() + .mutate_request(|req| { + req.headers_mut().insert("x-amz-meta-snowball-auto-extract", "true"); + }) + .send() + .await?; + + let zlib = admin_client + .get_object() + .bucket(bucket) + .key("zlib-extension.txt") + .send() + .await?; + assert_eq!(zlib.body.collect().await?.into_bytes().as_ref(), b"zlib-body"); Ok(()) } #[tokio::test] -async fn test_signed_put_object_extract_rejects_invalid_tar_gz_payload() -> Result<(), Box> { +async fn test_signed_put_object_extract_rejects_invalid_archive_payload() -> Result<(), Box> +{ init_logging(); let mut env = RustFSTestEnvironment::new().await?; diff --git a/crates/rio-v2/src/compress_reader.rs b/crates/rio-v2/src/compress_reader.rs index caafd156f..d50ae09d8 100644 --- a/crates/rio-v2/src/compress_reader.rs +++ b/crates/rio-v2/src/compress_reader.rs @@ -12,10 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -use minlz::{Encoder as MinlzEncoder, crc::crc, decode}; +use minlz::{Encoder as MinlzEncoder, crc::crc}; use pin_project_lite::pin_project; use rand::RngExt; -use rustfs_rio::{EtagResolvable, HashReaderDetector, HashReaderMut, Index, TryGetIndex}; +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; @@ -25,18 +27,16 @@ use std::task::{Context, Poll}; use tokio::io::{AsyncRead, ReadBuf}; const MAGIC_CHUNK: &[u8] = b"\xff\x06\x00\x00S2sTwO"; -const MAGIC_CHUNK_SNAPPY: &[u8] = b"\xff\x06\x00\x00sNaPpY"; const CHUNK_TYPE_COMPRESSED_DATA: u8 = 0x00; const CHUNK_TYPE_UNCOMPRESSED_DATA: u8 = 0x01; -const CHUNK_TYPE_INDEX: u8 = 0x99; const CHUNK_TYPE_PADDING: u8 = 0xfe; -const CHUNK_TYPE_STREAM_IDENTIFIER: u8 = 0xff; 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)] @@ -88,7 +88,17 @@ where 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(), @@ -125,6 +135,7 @@ where { fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { 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); @@ -142,6 +153,11 @@ where } 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) { @@ -149,6 +165,7 @@ where return Poll::Pending; } Poll::Ready(Ok(())) => { + ready_reads += 1; let n = read_buf.filled().len(); if n == 0 { break; @@ -243,18 +260,7 @@ pin_project! { #[derive(Debug)] pub struct DecompressReader { #[pin] - inner: R, - buffer: Vec, - buffer_pos: usize, - finished: bool, - header_buf: [u8; CHUNK_HEADER_LEN], - header_read: usize, - chunk_type: u8, - chunk_buf: Vec, - chunk_len: usize, - chunk_read: usize, - reading_chunk: bool, - stream_initialized: bool, + inner: S2Decoder, } } @@ -264,18 +270,7 @@ where { pub fn new(inner: R, _compression_algorithm: CompressionAlgorithm) -> Self { Self { - inner, - buffer: Vec::new(), - buffer_pos: 0, - finished: false, - header_buf: [0u8; CHUNK_HEADER_LEN], - header_read: 0, - chunk_type: 0, - chunk_buf: Vec::new(), - chunk_len: 0, - chunk_read: 0, - reading_chunk: false, - stream_initialized: false, + inner: S2Decoder::new_at_legacy_chunk_boundary(inner), } } } @@ -285,125 +280,7 @@ where R: AsyncRead + Unpin + Send + Sync, { fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { - let mut this = self.project(); - - if *this.buffer_pos < this.buffer.len() { - let to_copy = min(buf.remaining(), this.buffer.len() - *this.buffer_pos); - buf.put_slice(&this.buffer[*this.buffer_pos..*this.buffer_pos + to_copy]); - *this.buffer_pos += to_copy; - if *this.buffer_pos == this.buffer.len() { - this.buffer.clear(); - *this.buffer_pos = 0; - } - return Poll::Ready(Ok(())); - } - - loop { - if *this.finished { - return Poll::Ready(Ok(())); - } - - if !*this.reading_chunk { - while *this.header_read < CHUNK_HEADER_LEN { - let mut read_buf = ReadBuf::new(&mut this.header_buf[*this.header_read..]); - match this.inner.as_mut().poll_read(cx, &mut read_buf) { - Poll::Pending => return Poll::Pending, - Poll::Ready(Ok(())) => { - let n = read_buf.filled().len(); - if n == 0 { - if *this.header_read == 0 { - *this.finished = true; - return Poll::Ready(Ok(())); - } - return Poll::Ready(Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "unexpected EOF while reading S2 chunk header", - ))); - } - *this.header_read += n; - } - Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), - } - } - - *this.chunk_type = this.header_buf[0]; - *this.chunk_len = - (this.header_buf[1] as usize) | ((this.header_buf[2] as usize) << 8) | ((this.header_buf[3] as usize) << 16); - *this.header_read = 0; - - if this.chunk_buf.len() < *this.chunk_len { - this.chunk_buf.resize(*this.chunk_len, 0); - } - *this.chunk_read = 0; - *this.reading_chunk = true; - } - - while *this.chunk_read < *this.chunk_len { - let mut read_buf = ReadBuf::new(&mut this.chunk_buf[*this.chunk_read..*this.chunk_len]); - match this.inner.as_mut().poll_read(cx, &mut read_buf) { - Poll::Pending => return Poll::Pending, - Poll::Ready(Ok(())) => { - let n = read_buf.filled().len(); - if n == 0 { - return Poll::Ready(Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "unexpected EOF while reading S2 chunk body", - ))); - } - *this.chunk_read += n; - } - Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), - } - } - - let chunk = &this.chunk_buf[..*this.chunk_len]; - *this.reading_chunk = false; - match *this.chunk_type { - CHUNK_TYPE_STREAM_IDENTIFIER => { - if chunk != &MAGIC_CHUNK[CHUNK_HEADER_LEN..] && chunk != &MAGIC_CHUNK_SNAPPY[CHUNK_HEADER_LEN..] { - return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "invalid S2 stream identifier"))); - } - *this.stream_initialized = true; - continue; - } - CHUNK_TYPE_COMPRESSED_DATA => { - *this.stream_initialized = true; - let decompressed = decode_chunk(chunk, true)?; - *this.buffer = decompressed; - } - CHUNK_TYPE_UNCOMPRESSED_DATA => { - *this.stream_initialized = true; - let decompressed = decode_chunk(chunk, false)?; - *this.buffer = decompressed; - } - CHUNK_TYPE_INDEX | CHUNK_TYPE_PADDING | 0x80..=0xfd => { - *this.stream_initialized = true; - continue; - } - _ => { - if !*this.stream_initialized && *this.chunk_type != CHUNK_TYPE_COMPRESSED_DATA { - return Poll::Ready(Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("unknown S2 chunk type: 0x{:02x}", *this.chunk_type), - ))); - } - return Poll::Ready(Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("unknown S2 chunk type: 0x{:02x}", *this.chunk_type), - ))); - } - } - - *this.buffer_pos = 0; - let to_copy = min(buf.remaining(), this.buffer.len()); - buf.put_slice(&this.buffer[..to_copy]); - *this.buffer_pos += to_copy; - if *this.buffer_pos == this.buffer.len() { - this.buffer.clear(); - *this.buffer_pos = 0; - } - return Poll::Ready(Ok(())); - } + self.project().inner.poll_read(cx, buf) } } @@ -412,7 +289,7 @@ where R: EtagResolvable, { fn try_resolve_etag(&mut self) -> Option { - self.inner.try_resolve_etag() + self.inner.get_mut().try_resolve_etag() } } @@ -421,11 +298,11 @@ where R: HashReaderDetector, { fn is_hash_reader(&self) -> bool { - self.inner.is_hash_reader() + self.inner.get_ref().is_hash_reader() } fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> { - self.inner.as_hash_reader_mut() + self.inner.get_mut().as_hash_reader_mut() } } @@ -483,42 +360,57 @@ fn build_padding_chunk(current_size: usize, padding_multiple: usize) -> io::Resu Ok(Some(out)) } -fn decode_chunk(chunk: &[u8], compressed: bool) -> io::Result> { - if chunk.len() < CHECKSUM_SIZE { - return Err(io::Error::new(io::ErrorKind::InvalidData, "S2 chunk smaller than checksum header")); - } - - let expected_crc = u32::from_le_bytes(chunk[..CHECKSUM_SIZE].try_into().expect("checksum header")); - let payload = &chunk[CHECKSUM_SIZE..]; - let decompressed = if compressed { - decode(payload).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, format!("S2 decode error: {err}")))? - } else { - payload.to_vec() - }; - - let actual_crc = crc(&decompressed); - if actual_crc != expected_crc { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!( - "S2 CRC mismatch: expected={expected_crc:08x} actual={actual_crc:08x} compressed={compressed} payload_len={} decompressed_len={}", - payload.len(), - decompressed.len() - ), - )); - } - - Ok(decompressed) -} - #[cfg(test)] mod tests { use super::*; use std::io::Cursor; use std::pin::Pin; - use std::task::{Context, Poll}; + 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.0.fetch_add(1, Ordering::Relaxed); + } + } + + struct AlwaysReadyOneByte { + bytes: Vec, + position: usize, + read_calls: Arc, + } + + impl AlwaysReadyOneByte { + fn new(bytes: Vec, read_calls: Arc) -> 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> { + 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 { inner: R, max_chunk: usize, @@ -583,7 +475,7 @@ mod tests { let plaintext = b"compressible-rio-v2-block-".repeat(4096); let mut encoder = S2BlockEncoder::new(); let compressed = encode_block(&plaintext, &mut encoder); - let decoded = decode(&compressed).expect("decode payload"); + let decoded = minlz::decode(&compressed).expect("decode payload"); assert_eq!(decoded, plaintext); } @@ -604,6 +496,97 @@ mod tests { 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::::new()), 0, CompressionAlgorithm::default()); + assert_eq!(zero.block_size, DEFAULT_BLOCK_SIZE); + + let oversized = CompressReader::with_block_size( + Cursor::new(Vec::::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; diff --git a/crates/rio/Cargo.toml b/crates/rio/Cargo.toml index 46723ff2b..189d0f328 100644 --- a/crates/rio/Cargo.toml +++ b/crates/rio/Cargo.toml @@ -79,6 +79,7 @@ rustfs-tls-runtime.workspace = true rustfs-utils = { workspace = true, features = ["io", "hash", "compress"] } serde_json = { workspace = true, features = ["raw_value"] } md-5 = { workspace = true } +minlz.workspace = true tracing.workspace = true thiserror.workspace = true base64-simd.workspace = true diff --git a/crates/rio/src/lib.rs b/crates/rio/src/lib.rs index 254ae4918..899352910 100644 --- a/crates/rio/src/lib.rs +++ b/crates/rio/src/lib.rs @@ -100,6 +100,9 @@ mod limit_reader; pub use limit_reader::LimitReader; +mod s2_decoder; +pub use s2_decoder::{MAX_S2_DECOMPRESSED_BLOCK_SIZE, S2Decoder}; + mod etag_reader; pub use etag_reader::EtagReader; diff --git a/crates/rio/src/s2_decoder.rs b/crates/rio/src/s2_decoder.rs new file mode 100644 index 000000000..acfc0ddf8 --- /dev/null +++ b/crates/rio/src/s2_decoder.rs @@ -0,0 +1,858 @@ +// 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. + +//! Bounded asynchronous decoder for the S2/Snappy framed stream format. + +use minlz::{MAX_DECODE_DST_SIZE, crc::crc, decode, decode_into, decode_len}; +use pin_project_lite::pin_project; +use std::cmp::min; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, ReadBuf}; + +const S2_MAGIC_BODY: &[u8] = b"S2sTwO"; +const SNAPPY_MAGIC_BODY: &[u8] = b"sNaPpY"; +const CHUNK_TYPE_COMPRESSED_DATA: u8 = 0x00; +const CHUNK_TYPE_UNCOMPRESSED_DATA: u8 = 0x01; +const CHUNK_TYPE_PADDING: u8 = 0xfe; +const CHUNK_TYPE_STREAM_IDENTIFIER: u8 = 0xff; +const CHECKSUM_SIZE: usize = 4; +const CHUNK_HEADER_LEN: usize = 4; +const MAX_READY_READS_PER_POLL: usize = 64; +const MAX_CHUNKS_PER_POLL: usize = 64; +const MAX_INPUT_BYTES_PER_POLL: usize = 256 * 1024; +const MAX_SNAPPY_DECOMPRESSED_BLOCK_SIZE: usize = 64 << 10; +const MAX_FRAMED_CHUNK_SIZE: usize = (1 << 24) - 1; +const MAX_LEGACY_S2_DECOMPRESSED_BLOCK_SIZE: usize = MAX_DECODE_DST_SIZE; + +// This is checksum size + klauspost/s2 MaxEncodedLen(4 MiB). +// MaxEncodedLen adds a four-byte varint and four-byte literal header. The Go +// reader keeps this encoded-input cap even for Snappy frames, then applies the +// tighter 64 KiB limit to the decoded size. +const MAX_S2_COMPRESSED_CHUNK_SIZE: usize = CHECKSUM_SIZE + MAX_S2_DECOMPRESSED_BLOCK_SIZE + 4 + 4; + +/// S2 writers, including minio-go's Snowball writer, cap a decoded block at +/// 4 MiB. Enforcing the framing limit before allocation prevents a tiny block +/// length varint from requesting the block codec's much larger generic limit. +pub const MAX_S2_DECOMPRESSED_BLOCK_SIZE: usize = 4 << 20; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum FrameMode { + Uninitialized, + S2, + Snappy, + MidstreamS2, +} + +impl FrameMode { + fn max_compressed_chunk_size(self, s2_limit: usize) -> usize { + match self { + Self::Snappy => MAX_S2_COMPRESSED_CHUNK_SIZE, + Self::Uninitialized | Self::S2 | Self::MidstreamS2 => s2_limit, + } + } + + fn max_decompressed_block_size(self, s2_limit: usize) -> usize { + match self { + Self::Snappy => MAX_SNAPPY_DECOMPRESSED_BLOCK_SIZE, + Self::Uninitialized | Self::S2 | Self::MidstreamS2 => s2_limit, + } + } +} + +pin_project! { + /// Decode an S2 or Snappy framed stream without blocking the async reader. + #[derive(Debug)] + pub struct S2Decoder { + #[pin] + inner: R, + output: Vec, + output_pos: usize, + finished: bool, + poisoned: bool, + header_buf: [u8; CHUNK_HEADER_LEN], + header_read: usize, + chunk_type: u8, + chunk_buf: Vec, + chunk_len: usize, + chunk_read: usize, + reading_chunk: bool, + skipping_chunk: bool, + frame_mode: FrameMode, + max_s2_compressed_chunk_size: usize, + max_s2_decompressed_block_size: usize, + } +} + +impl S2Decoder { + pub fn new(inner: R) -> Self { + Self::with_limits( + inner, + FrameMode::Uninitialized, + MAX_S2_COMPRESSED_CHUNK_SIZE, + MAX_S2_DECOMPRESSED_BLOCK_SIZE, + ) + } + + /// Create a decoder positioned at a trusted S2 data-chunk boundary. + /// + /// Indexed range reads start after the stream identifier. General stream + /// consumers should use [`S2Decoder::new`] so a missing identifier remains + /// an error. + pub fn new_at_chunk_boundary(inner: R) -> Self { + Self::with_limits( + inner, + FrameMode::MidstreamS2, + MAX_S2_COMPRESSED_CHUNK_SIZE, + MAX_S2_DECOMPRESSED_BLOCK_SIZE, + ) + } + + /// Create a bounded decoder for rio-v2 data written before its block-size + /// API was capped at 4 MiB. + /// + /// This compatibility mode accepts decoded S2 blocks up to the block + /// decoder's 256 MiB safety cap and an encoded chunk up to the format's + /// 24-bit framing limit. That preserves streams written through rio-v2's + /// former uncapped block-size API without restoring unbounded allocation. + /// New streams and general S2 consumers should use the stricter + /// constructors above. + pub fn new_at_legacy_chunk_boundary(inner: R) -> Self { + Self::with_limits( + inner, + FrameMode::MidstreamS2, + MAX_FRAMED_CHUNK_SIZE, + MAX_LEGACY_S2_DECOMPRESSED_BLOCK_SIZE, + ) + } + + fn with_limits( + inner: R, + frame_mode: FrameMode, + max_s2_compressed_chunk_size: usize, + max_s2_decompressed_block_size: usize, + ) -> Self { + Self { + inner, + output: Vec::new(), + output_pos: 0, + finished: false, + poisoned: false, + header_buf: [0u8; CHUNK_HEADER_LEN], + header_read: 0, + chunk_type: 0, + chunk_buf: Vec::new(), + chunk_len: 0, + chunk_read: 0, + reading_chunk: false, + skipping_chunk: false, + frame_mode, + max_s2_compressed_chunk_size, + max_s2_decompressed_block_size, + } + } + + pub fn get_ref(&self) -> &R { + &self.inner + } + + pub fn get_mut(&mut self) -> &mut R { + &mut self.inner + } + + pub fn into_inner(self) -> R { + self.inner + } +} + +impl AsyncRead for S2Decoder +where + R: AsyncRead, +{ + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + let mut this = self.project(); + let mut ready_reads = 0; + let mut completed_chunks = 0; + let mut input_bytes = 0; + + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + if *this.poisoned { + return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "S2 decoder previously failed"))); + } + + if *this.output_pos < this.output.len() { + let to_copy = min(buf.remaining(), this.output.len() - *this.output_pos); + buf.put_slice(&this.output[*this.output_pos..*this.output_pos + to_copy]); + *this.output_pos += to_copy; + return Poll::Ready(Ok(())); + } + + if *this.finished { + return Poll::Ready(Ok(())); + } + + loop { + if ready_reads >= MAX_READY_READS_PER_POLL + || completed_chunks >= MAX_CHUNKS_PER_POLL + || input_bytes >= MAX_INPUT_BYTES_PER_POLL + { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + + if !*this.reading_chunk { + while *this.header_read < CHUNK_HEADER_LEN { + let remaining_poll_bytes = MAX_INPUT_BYTES_PER_POLL - input_bytes; + if ready_reads >= MAX_READY_READS_PER_POLL || remaining_poll_bytes == 0 { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + let read_end = (*this.header_read + remaining_poll_bytes).min(CHUNK_HEADER_LEN); + let mut read_buf = ReadBuf::new(&mut this.header_buf[*this.header_read..read_end]); + match this.inner.as_mut().poll_read(cx, &mut read_buf) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(())) => { + ready_reads += 1; + let read = read_buf.filled().len(); + if read == 0 { + if *this.header_read == 0 { + *this.finished = true; + return Poll::Ready(Ok(())); + } + *this.poisoned = true; + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "unexpected EOF while reading S2 chunk header", + ))); + } + input_bytes += read; + *this.header_read += read; + } + Poll::Ready(Err(err)) => { + *this.poisoned = true; + return Poll::Ready(Err(err)); + } + } + } + if ready_reads >= MAX_READY_READS_PER_POLL || input_bytes >= MAX_INPUT_BYTES_PER_POLL { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + + *this.chunk_type = this.header_buf[0]; + *this.chunk_len = usize::from(this.header_buf[1]) + | (usize::from(this.header_buf[2]) << 8) + | (usize::from(this.header_buf[3]) << 16); + *this.header_read = 0; + + if *this.frame_mode == FrameMode::Uninitialized && *this.chunk_type != CHUNK_TYPE_STREAM_IDENTIFIER { + *this.poisoned = true; + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidData, + "S2 stream identifier must be the first chunk", + ))); + } + + let skippable = matches!(*this.chunk_type, CHUNK_TYPE_PADDING | 0x80..=0xfd); + let invalid_length = match *this.chunk_type { + CHUNK_TYPE_STREAM_IDENTIFIER => *this.chunk_len != S2_MAGIC_BODY.len(), + CHUNK_TYPE_COMPRESSED_DATA => { + *this.chunk_len < CHECKSUM_SIZE + || *this.chunk_len > this.frame_mode.max_compressed_chunk_size(*this.max_s2_compressed_chunk_size) + } + CHUNK_TYPE_UNCOMPRESSED_DATA => { + *this.chunk_len < CHECKSUM_SIZE + || *this.chunk_len - CHECKSUM_SIZE + > this + .frame_mode + .max_decompressed_block_size(*this.max_s2_decompressed_block_size) + } + _ if skippable => false, + _ => { + *this.poisoned = true; + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("unknown S2 chunk type: 0x{:02x}", *this.chunk_type), + ))); + } + }; + if invalid_length { + *this.poisoned = true; + return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "invalid S2 chunk length"))); + } + + if !skippable && this.chunk_buf.len() < *this.chunk_len { + this.chunk_buf.resize(*this.chunk_len, 0); + } + *this.chunk_read = 0; + *this.reading_chunk = true; + *this.skipping_chunk = skippable; + } + + while *this.chunk_read < *this.chunk_len { + let remaining_poll_bytes = MAX_INPUT_BYTES_PER_POLL - input_bytes; + if ready_reads >= MAX_READY_READS_PER_POLL || remaining_poll_bytes == 0 { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + let mut discard = [0u8; 8192]; + let mut read_buf = if *this.skipping_chunk { + let remaining = *this.chunk_len - *this.chunk_read; + let discard_len = remaining.min(discard.len()).min(remaining_poll_bytes); + ReadBuf::new(&mut discard[..discard_len]) + } else { + let read_end = (*this.chunk_read + remaining_poll_bytes).min(*this.chunk_len); + ReadBuf::new(&mut this.chunk_buf[*this.chunk_read..read_end]) + }; + match this.inner.as_mut().poll_read(cx, &mut read_buf) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Ok(())) => { + ready_reads += 1; + let read = read_buf.filled().len(); + if read == 0 { + *this.poisoned = true; + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "unexpected EOF while reading S2 chunk body", + ))); + } + input_bytes += read; + *this.chunk_read += read; + } + Poll::Ready(Err(err)) => { + *this.poisoned = true; + return Poll::Ready(Err(err)); + } + } + } + + completed_chunks += 1; + *this.reading_chunk = false; + if *this.skipping_chunk { + *this.skipping_chunk = false; + continue; + } + + let chunk = &this.chunk_buf[..*this.chunk_len]; + match *this.chunk_type { + CHUNK_TYPE_STREAM_IDENTIFIER => { + *this.frame_mode = if chunk == S2_MAGIC_BODY { + FrameMode::S2 + } else if chunk == SNAPPY_MAGIC_BODY { + FrameMode::Snappy + } else { + *this.poisoned = true; + return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "invalid S2 stream identifier"))); + }; + continue; + } + CHUNK_TYPE_COMPRESSED_DATA | CHUNK_TYPE_UNCOMPRESSED_DATA => { + if *this.frame_mode == FrameMode::Uninitialized { + *this.poisoned = true; + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidData, + "S2 data chunk before stream identifier", + ))); + } + if let Err(err) = decode_chunk_into( + this.output, + chunk, + *this.chunk_type == CHUNK_TYPE_COMPRESSED_DATA, + this.frame_mode + .max_decompressed_block_size(*this.max_s2_decompressed_block_size), + ) { + this.output.clear(); + *this.output_pos = 0; + *this.poisoned = true; + return Poll::Ready(Err(err)); + } + } + _ => unreachable!("chunk type validated before reading its payload"), + } + + if this.output.is_empty() { + *this.output_pos = 0; + continue; + } + + *this.output_pos = 0; + let to_copy = min(buf.remaining(), this.output.len()); + buf.put_slice(&this.output[..to_copy]); + *this.output_pos += to_copy; + return Poll::Ready(Ok(())); + } + } +} + +delegate_reader_capabilities_generic_no_index!(S2Decoder, inner); + +fn decode_chunk_into(output: &mut Vec, chunk: &[u8], compressed: bool, max_decompressed_block_size: usize) -> io::Result<()> { + let expected_crc = u32::from_le_bytes( + chunk[..CHECKSUM_SIZE] + .try_into() + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "S2 chunk smaller than checksum header"))?, + ); + let payload = &chunk[CHECKSUM_SIZE..]; + if compressed { + let (decoded_len, _) = decode_len(payload) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, format!("S2 length decode error: {err}")))?; + if decoded_len > max_decompressed_block_size { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("S2 decompressed block size exceeds limit: size={decoded_len}, limit={max_decompressed_block_size}"), + )); + } + if output.len() >= decoded_len { + let written = decode_into(&mut output[..decoded_len], payload) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, format!("S2 decode error: {err}")))?; + output.truncate(written); + } else { + *output = + decode(payload).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, format!("S2 decode error: {err}")))?; + } + } else { + output.clear(); + output.extend_from_slice(payload); + } + + let actual_crc = crc(output); + if actual_crc != expected_crc { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("S2 CRC mismatch: expected={expected_crc:08x} actual={actual_crc:08x}"), + )); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Cursor; + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + use std::task::{Wake, Waker}; + use tokio::io::AsyncReadExt; + + #[derive(Default)] + struct WakeCounter(AtomicUsize); + + impl Wake for WakeCounter { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + fn append_chunk_header(stream: &mut Vec, chunk_type: u8, payload_len: usize) { + assert!(payload_len <= 0x00ff_ffff, "test chunk payload must fit the framing length"); + stream.push(chunk_type); + stream.push(u8::try_from(payload_len & 0xff).expect("low length byte must fit u8")); + stream.push(u8::try_from((payload_len >> 8) & 0xff).expect("middle length byte must fit u8")); + stream.push(u8::try_from((payload_len >> 16) & 0xff).expect("high length byte must fit u8")); + } + + fn append_chunk(stream: &mut Vec, chunk_type: u8, payload: &[u8]) { + append_chunk_header(stream, chunk_type, payload.len()); + stream.extend_from_slice(payload); + } + + fn append_uncompressed_chunk(stream: &mut Vec, payload: &[u8]) { + let mut chunk = crc(payload).to_le_bytes().to_vec(); + chunk.extend_from_slice(payload); + append_chunk(stream, CHUNK_TYPE_UNCOMPRESSED_DATA, &chunk); + } + + fn append_compressed_chunk(stream: &mut Vec, payload: &[u8]) { + let mut chunk = crc(payload).to_le_bytes().to_vec(); + chunk.extend_from_slice(&minlz::encode(payload)); + append_chunk(stream, CHUNK_TYPE_COMPRESSED_DATA, &chunk); + } + + fn append_uvarint(output: &mut Vec, mut value: usize) { + loop { + let mut byte = u8::try_from(value & 0x7f).expect("varint byte must fit u8"); + value >>= 7; + if value != 0 { + byte |= 0x80; + } + output.push(byte); + if value == 0 { + return; + } + } + } + + const GO_S2_GOLDEN_PLAINTEXT: &[u8] = b"This is a test file with some repeated content to compress.\nThe quick brown fox jumps over the lazy dog.\nThe quick brown fox jumps over the lazy dog.\nThe quick brown fox jumps over the lazy dog.\nLorem ipsum dolor sit amet, consectetur adipiscing elit.\nLorem ipsum dolor sit amet, consectetur adipiscing elit.\nLorem ipsum dolor sit amet, consectetur adipiscing elit.\nBinary compatibility testing with S2 compression format.\nBinary compatibility testing with S2 compression format.\nBinary compatibility testing with S2 compression format.\nPerformance benchmarking and optimization verification.\nPerformance benchmarking and optimization verification.\nPerformance benchmarking and optimization verification.\n"; + + // Fixed interoperability fixture generated with the same dependency and + // option as minio-go PutObjectsSnowball: + // minio-go commit 0e78d3f18efe14e352e20d3a262b99df97b516b8 + // github.com/klauspost/compress v1.19.2 + // s2.NewWriter(dst, s2.WriterBetterCompression()) + // Generator body: create that writer on os.Stdout, io.Copy from os.Stdin, + // then Close; invoke as `go run generator.go < input | xxd -p -c 100000`. + // The exact input is GO_S2_GOLDEN_PLAINTEXT, so this test never creates its + // expected stream through minlz. + const GO_S2_GOLDEN_HEX: &str = "ff06000053327354774f00120100f8de94fbc105f06654686973206973206120746573742066696c65207769746820736f6d6520726570656174656420636f6e74656e7420746f20636f6d70726573732e0a54686520717569636b2062726f776e20666f78206a756d7073206f76657220746865206c617a7920646f67112d15004c684c6f72656d20697073756d20646f6c6f722073697420616d65742c01b85c73656374657475722061646970697363696e6720656c697411391500641442696e61727925432061746962696c697479257901952577045332356120696f6e20666f726d61113915006508506572057d306e63652062656e63686d61726b01a730616e64206f7074696d697a617401a41876657269666963050d1138150062"; + + #[tokio::test] + async fn decodes_fixed_go_s2_fixture() { + let fixture = hex_simd::decode_to_vec(GO_S2_GOLDEN_HEX).expect("golden fixture must be valid hex"); + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + decoder.read_to_end(&mut output).await.expect("Go S2 fixture must decode"); + + assert_eq!(output, GO_S2_GOLDEN_PLAINTEXT); + } + + #[tokio::test] + async fn indexed_chunk_boundary_mode_decodes_a_headerless_tail() { + let fixture = hex_simd::decode_to_vec(GO_S2_GOLDEN_HEX).expect("golden fixture must be valid hex"); + let header_len = CHUNK_HEADER_LEN + S2_MAGIC_BODY.len(); + let mut decoder = S2Decoder::new_at_chunk_boundary(Cursor::new(fixture[header_len..].to_vec())); + let mut output = Vec::new(); + decoder + .read_to_end(&mut output) + .await + .expect("trusted indexed tail should decode without a stream identifier"); + + assert_eq!(output, GO_S2_GOLDEN_PLAINTEXT); + } + + #[tokio::test] + async fn strict_mode_rejects_a_headerless_data_chunk() { + let fixture = hex_simd::decode_to_vec(GO_S2_GOLDEN_HEX).expect("golden fixture must be valid hex"); + let header_len = CHUNK_HEADER_LEN + S2_MAGIC_BODY.len(); + let mut decoder = S2Decoder::new(Cursor::new(fixture[header_len..].to_vec())); + let mut output = Vec::new(); + let err = decoder + .read_to_end(&mut output) + .await + .expect_err("untrusted stream must include a stream identifier"); + + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("must be the first chunk")); + } + + #[tokio::test] + async fn strict_mode_rejects_skippable_chunk_before_identifier() { + let mut fixture = Vec::new(); + append_chunk(&mut fixture, CHUNK_TYPE_PADDING, &[]); + append_chunk(&mut fixture, CHUNK_TYPE_STREAM_IDENTIFIER, S2_MAGIC_BODY); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + let err = decoder + .read_to_end(&mut output) + .await + .expect_err("strict streams must begin with an identifier even when the first chunk is skippable"); + + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("must be the first chunk")); + assert_eq!( + decoder.get_ref().position(), + u64::try_from(CHUNK_HEADER_LEN).expect("chunk header length must fit u64") + ); + } + + #[tokio::test] + async fn snappy_identifier_enforces_the_64_kib_decoded_limit() { + let payload = vec![0u8; MAX_SNAPPY_DECOMPRESSED_BLOCK_SIZE + 1]; + let mut fixture = Vec::new(); + append_chunk(&mut fixture, CHUNK_TYPE_STREAM_IDENTIFIER, SNAPPY_MAGIC_BODY); + append_compressed_chunk(&mut fixture, &payload); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + let err = decoder + .read_to_end(&mut output) + .await + .expect_err("Snappy frames must reject decoded blocks larger than 64 KiB"); + + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("limit=65536")); + } + + #[tokio::test] + async fn snappy_frame_accepts_large_noncanonical_encoding_with_small_output() { + let payload = vec![b'x'; 40_000]; + let mut encoded = Vec::with_capacity(payload.len() * 2 + 3); + append_uvarint(&mut encoded, payload.len()); + for byte in &payload { + encoded.push(0); // One-byte literal tag. + encoded.push(*byte); + } + assert!(encoded.len() > MAX_SNAPPY_DECOMPRESSED_BLOCK_SIZE); + + let mut chunk = crc(&payload).to_le_bytes().to_vec(); + chunk.extend_from_slice(&encoded); + let mut fixture = Vec::new(); + append_chunk(&mut fixture, CHUNK_TYPE_STREAM_IDENTIFIER, SNAPPY_MAGIC_BODY); + append_chunk(&mut fixture, CHUNK_TYPE_COMPRESSED_DATA, &chunk); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + decoder + .read_to_end(&mut output) + .await + .expect("Snappy encoded input uses the S2 reader cap while decoded output stays below 64 KiB"); + + assert_eq!(output, payload); + } + + #[tokio::test] + async fn repeated_identifier_switches_the_frame_limit() { + let payload = vec![b'x'; MAX_SNAPPY_DECOMPRESSED_BLOCK_SIZE + 1]; + let mut fixture = Vec::new(); + append_chunk(&mut fixture, CHUNK_TYPE_STREAM_IDENTIFIER, SNAPPY_MAGIC_BODY); + append_chunk(&mut fixture, CHUNK_TYPE_STREAM_IDENTIFIER, S2_MAGIC_BODY); + append_uncompressed_chunk(&mut fixture, &payload); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + decoder + .read_to_end(&mut output) + .await + .expect("a later S2 identifier must restore the S2 block limit"); + + assert_eq!(output, payload); + } + + #[tokio::test] + async fn rejects_oversized_compressed_chunk_before_allocation() { + let declared_len = MAX_S2_COMPRESSED_CHUNK_SIZE + 1; + let mut fixture = Vec::new(); + append_chunk(&mut fixture, CHUNK_TYPE_STREAM_IDENTIFIER, S2_MAGIC_BODY); + append_chunk_header(&mut fixture, CHUNK_TYPE_COMPRESSED_DATA, declared_len); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + let err = decoder + .read_to_end(&mut output) + .await + .expect_err("oversized compressed chunks must fail from their header"); + + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert_eq!(err.to_string(), "invalid S2 chunk length"); + assert_eq!( + decoder.get_ref().position(), + u64::try_from(CHUNK_HEADER_LEN * 2 + S2_MAGIC_BODY.len()).expect("fixture prefix length must fit u64") + ); + assert_eq!(decoder.chunk_buf.len(), S2_MAGIC_BODY.len()); + } + + #[tokio::test] + async fn accepts_empty_uncompressed_and_compressed_chunks() { + let mut fixture = Vec::new(); + append_chunk(&mut fixture, CHUNK_TYPE_STREAM_IDENTIFIER, S2_MAGIC_BODY); + append_uncompressed_chunk(&mut fixture, &[]); + append_compressed_chunk(&mut fixture, &[]); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + decoder + .read_to_end(&mut output) + .await + .expect("empty data chunks are valid framed-stream no-ops"); + + assert!(output.is_empty()); + } + + #[test] + fn compressed_blocks_reuse_the_decoded_output_allocation() { + let first_payload = vec![b'a'; 32 * 1024]; + let second_payload = vec![b'b'; first_payload.len()]; + let mut first_chunk = crc(&first_payload).to_le_bytes().to_vec(); + first_chunk.extend_from_slice(&minlz::encode(&first_payload)); + let mut second_chunk = crc(&second_payload).to_le_bytes().to_vec(); + second_chunk.extend_from_slice(&minlz::encode(&second_payload)); + + let mut output = Vec::new(); + decode_chunk_into(&mut output, &first_chunk, true, MAX_S2_DECOMPRESSED_BLOCK_SIZE) + .expect("first compressed block should decode"); + let allocation = output.as_ptr(); + decode_chunk_into(&mut output, &second_chunk, true, MAX_S2_DECOMPRESSED_BLOCK_SIZE) + .expect("same-sized compressed block should decode into the existing allocation"); + + assert_eq!(output, second_payload); + assert_eq!(output.as_ptr(), allocation); + } + + #[tokio::test] + async fn rejects_decompressed_block_length_above_framing_limit() { + let mut fixture = b"\xff\x06\x00\x00S2sTwO".to_vec(); + fixture.extend_from_slice(&[CHUNK_TYPE_COMPRESSED_DATA, 8, 0, 0]); + fixture.extend_from_slice(&[0, 0, 0, 0]); + fixture.extend_from_slice(&[0x81, 0x80, 0x80, 0x02]); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + let err = decoder + .read_to_end(&mut output) + .await + .expect_err("oversized decoded block declaration must fail before allocation"); + + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("exceeds limit")); + } + + #[tokio::test] + async fn legacy_chunk_boundary_mode_rejects_blocks_above_decoder_safety_limit() { + let mut chunk = vec![0u8; CHECKSUM_SIZE]; + append_uvarint(&mut chunk, MAX_LEGACY_S2_DECOMPRESSED_BLOCK_SIZE + 1); + let mut fixture = Vec::new(); + append_chunk(&mut fixture, CHUNK_TYPE_COMPRESSED_DATA, &chunk); + + let mut decoder = S2Decoder::new_at_legacy_chunk_boundary(Cursor::new(fixture)); + let mut output = Vec::new(); + let err = decoder + .read_to_end(&mut output) + .await + .expect_err("legacy compatibility must remain bounded before allocation"); + + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains(&format!("limit={MAX_DECODE_DST_SIZE}"))); + assert!(output.is_empty()); + } + + #[tokio::test] + async fn rejects_crc_mismatch() { + let mut fixture = hex_simd::decode_to_vec(GO_S2_GOLDEN_HEX).expect("golden fixture must be valid hex"); + fixture[14] ^= 0xff; + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = [0u8; 128]; + let err = decoder.read(&mut output).await.expect_err("CRC mismatch must fail"); + + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert!(err.to_string().contains("CRC mismatch")); + assert!(decoder.output.is_empty(), "unverified decoded bytes must be discarded"); + + let err = decoder + .read(&mut output) + .await + .expect_err("a poisoned decoder must remain fail-closed on later reads"); + assert_eq!(err.kind(), io::ErrorKind::InvalidData); + assert_eq!(err.to_string(), "S2 decoder previously failed"); + } + + #[tokio::test] + async fn rejects_truncated_chunk() { + let mut fixture = hex_simd::decode_to_vec(GO_S2_GOLDEN_HEX).expect("golden fixture must be valid hex"); + fixture.truncate(fixture.len() - 1); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + let err = decoder + .read_to_end(&mut output) + .await + .expect_err("truncated S2 chunk must fail"); + + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + } + + #[tokio::test] + async fn skips_large_extension_chunks_without_buffering_them() { + const EXTENSION_SIZE: usize = 64 * 1024; + + let mut fixture = b"\xff\x06\x00\x00S2sTwO".to_vec(); + append_chunk_header(&mut fixture, CHUNK_TYPE_PADDING, EXTENSION_SIZE); + fixture.resize(fixture.len() + EXTENSION_SIZE, 0); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + let mut output = Vec::new(); + decoder + .read_to_end(&mut output) + .await + .expect("skippable extension must be consumed"); + + assert!(output.is_empty()); + assert_eq!(decoder.chunk_buf.len(), S2_MAGIC_BODY.len()); + } + + #[test] + fn yields_after_bounded_number_of_non_data_chunks() { + let mut fixture = b"\xff\x06\x00\x00S2sTwO".to_vec(); + for _ in 0..MAX_CHUNKS_PER_POLL { + fixture.extend_from_slice(&[CHUNK_TYPE_PADDING, 0, 0, 0]); + } + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + 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 decoder).poll_read(&mut cx, &mut read_buf).is_pending()); + assert!(read_buf.filled().is_empty()); + assert_eq!(wake_counter.0.load(Ordering::Relaxed), 1); + } + + #[test] + fn yields_while_streaming_a_large_skippable_chunk() { + const EXTENSION_SIZE: usize = MAX_INPUT_BYTES_PER_POLL * 2; + + let mut fixture = b"\xff\x06\x00\x00S2sTwO".to_vec(); + append_chunk_header(&mut fixture, CHUNK_TYPE_PADDING, EXTENSION_SIZE); + fixture.resize(fixture.len() + EXTENSION_SIZE, 0); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + 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 decoder).poll_read(&mut cx, &mut read_buf).is_pending()); + assert!(read_buf.filled().is_empty()); + assert_eq!( + decoder.get_ref().position(), + u64::try_from(MAX_INPUT_BYTES_PER_POLL).expect("poll byte budget must fit u64") + ); + assert_eq!(wake_counter.0.load(Ordering::Relaxed), 1); + } + + #[test] + fn header_reads_respect_the_remaining_poll_byte_budget() { + const FIRST_EXTENSION_SIZE: usize = MAX_INPUT_BYTES_PER_POLL - 15; + + let mut fixture = b"\xff\x06\x00\x00S2sTwO".to_vec(); + append_chunk_header(&mut fixture, CHUNK_TYPE_PADDING, FIRST_EXTENSION_SIZE); + fixture.resize(fixture.len() + FIRST_EXTENSION_SIZE, 0); + fixture.extend_from_slice(&[CHUNK_TYPE_PADDING, 0, 0, 0]); + + let mut decoder = S2Decoder::new(Cursor::new(fixture)); + 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 decoder).poll_read(&mut cx, &mut read_buf).is_pending()); + assert!(read_buf.filled().is_empty()); + assert_eq!( + decoder.get_ref().position(), + u64::try_from(MAX_INPUT_BYTES_PER_POLL).expect("poll byte budget must fit u64") + ); + assert_eq!(wake_counter.0.load(Ordering::Relaxed), 1); + } +} diff --git a/crates/zip/Cargo.toml b/crates/zip/Cargo.toml index 01ec5d84e..5558ffa44 100644 --- a/crates/zip/Cargo.toml +++ b/crates/zip/Cargo.toml @@ -30,9 +30,9 @@ doctest = false [features] default = [] -hotpath = ["hotpath/hotpath", "hotpath/tokio"] -hotpath-alloc = ["hotpath", "hotpath/hotpath-alloc"] -hotpath-cpu = ["hotpath", "hotpath/hotpath-cpu"] +hotpath = ["hotpath/hotpath", "hotpath/tokio", "rustfs-rio/hotpath"] +hotpath-alloc = ["hotpath", "hotpath/hotpath-alloc", "rustfs-rio/hotpath-alloc"] +hotpath-cpu = ["hotpath", "hotpath/hotpath-cpu", "rustfs-rio/hotpath-cpu"] [dependencies] hotpath.workspace = true @@ -40,10 +40,12 @@ async-compression = { workspace = true, features = [ "tokio", "bzip2", "gzip", + "lz4", "zlib", "zstd", "xz", ] } +rustfs-rio.workspace = true tokio = { workspace = true, features = ["io-util", "macros", "rt"] } thiserror = { workspace = true } diff --git a/crates/zip/src/lib.rs b/crates/zip/src/lib.rs index 1bcc78506..cebfc9309 100644 --- a/crates/zip/src/lib.rs +++ b/crates/zip/src/lib.rs @@ -12,9 +12,35 @@ // See the License for the specific language governing permissions and // limitations under the License. -use async_compression::tokio::bufread::{BzDecoder, GzipDecoder, XzDecoder, ZlibDecoder, ZstdDecoder}; +use async_compression::{ + tokio::bufread::{BzDecoder, GzipDecoder, Lz4Decoder, XzDecoder, ZlibDecoder, ZstdDecoder}, + zstd::DParameter, +}; +use rustfs_rio::S2Decoder; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; use thiserror::Error; -use tokio::io::{AsyncRead, BufReader}; +use tokio::io::{AsyncRead, AsyncReadExt, BufReader, ReadBuf}; + +const MAGIC_SNIFF_LEN: usize = 6; +const TAR_HEADER_LEN: usize = 512; +const TAR_CHECKSUM_START: usize = 148; +const TAR_CHECKSUM_END: usize = 156; +const SHARED_SKIPPABLE_FRAME_HEADER_LEN: usize = 8; +const SNIFF_DISCARD_BUFFER_LEN: usize = 8 * 1024; +const SNIFF_YIELD_AFTER_BYTES: usize = 64 * 1024; +const SNIFF_YIELD_AFTER_FRAMES: usize = 64; +const SNIFF_YIELD_AFTER_READY_READS: usize = 64; + +// XZ declares its LZMA2 dictionary size before producing decoded bytes, so the +// decoded-size guard cannot prevent that allocation. This accepts every +// standard xz preset (the largest needs about 65 MiB to decode) while keeping +// one Snowball decoder well below the multi-gigabyte archive budget. +const XZ_DECODER_MEMORY_LIMIT_BYTES: u64 = 128_u64 * 1024 * 1024; +// Match MinIO's Snowball decoder boundary so a frame cannot reserve an +// oversized history window before the decoded-size guard observes any bytes. +const ZSTD_DECODER_MAX_WINDOW_LOG: u32 = 24; pub type Result = std::result::Result; @@ -25,6 +51,8 @@ pub enum ZipError { format: CompressionFormat, operation: &'static str, }, + #[error("failed to inspect archive compression magic")] + InspectStream(#[source] io::Error), } #[derive(Debug, PartialEq, Eq, Clone, Copy)] @@ -35,10 +63,172 @@ pub enum CompressionFormat { Xz, Zlib, Zstd, + Lz4, + S2, Tar, Unknown, } +/// Reader returned by [`CompressionFormat::sniff`]. It replays the bounded +/// prefix that identified the codec before continuing with the original +/// stream. Codec-neutral leading skippable frames are consumed by `sniff` and +/// intentionally omitted from the replayed stream. +#[derive(Debug)] +pub struct SniffedReader { + inner: R, + prefix: Vec, + prefix_pos: usize, +} + +impl AsyncRead for SniffedReader +where + R: AsyncRead + Unpin, +{ + fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + if self.prefix_pos < self.prefix.len() && buf.remaining() > 0 { + let to_copy = buf.remaining().min(self.prefix.len() - self.prefix_pos); + buf.put_slice(&self.prefix[self.prefix_pos..self.prefix_pos + to_copy]); + self.prefix_pos += to_copy; + return Poll::Ready(Ok(())); + } + + Pin::new(&mut self.inner).poll_read(cx, buf) + } +} + +fn is_shared_skippable_magic(prefix: &[u8]) -> bool { + prefix.len() >= 4 && (0x50..=0x5f).contains(&prefix[0]) && prefix[1..4] == [0x2a, 0x4d, 0x18] +} + +#[derive(Debug, Default)] +struct SniffYieldBudget { + ready_reads: usize, + input_bytes: usize, + frames: usize, +} + +impl SniffYieldBudget { + async fn record_read(&mut self, read: usize) { + self.ready_reads = self.ready_reads.saturating_add(1); + self.input_bytes = self.input_bytes.saturating_add(read); + if self.ready_reads >= SNIFF_YIELD_AFTER_READY_READS || self.input_bytes >= SNIFF_YIELD_AFTER_BYTES { + self.yield_now().await; + } + } + + async fn record_frame(&mut self) { + self.frames = self.frames.saturating_add(1); + if self.frames >= SNIFF_YIELD_AFTER_FRAMES { + self.yield_now().await; + } + } + + async fn yield_now(&mut self) { + tokio::task::yield_now().await; + self.ready_reads = 0; + self.input_bytes = 0; + self.frames = 0; + } +} + +async fn inspect_fill_to(input: &mut R, output: &mut Vec, target_len: usize, budget: &mut SniffYieldBudget) -> Result +where + R: AsyncRead + Unpin, +{ + let mut scratch = [0u8; TAR_HEADER_LEN]; + while output.len() < target_len { + let remaining = target_len - output.len(); + let to_read = remaining.min(scratch.len()); + let read = input.read(&mut scratch[..to_read]).await.map_err(ZipError::InspectStream)?; + if read == 0 { + return Ok(false); + } + output.extend_from_slice(&scratch[..read]); + budget.record_read(read).await; + } + Ok(true) +} + +async fn inspect_discard_fully( + input: &mut R, + mut remaining: u64, + budget: &mut SniffYieldBudget, + truncated_message: &'static str, +) -> Result<()> +where + R: AsyncRead + Unpin, +{ + let mut discard = [0u8; SNIFF_DISCARD_BUFFER_LEN]; + while remaining > 0 { + let to_read = usize::try_from(remaining).map_or(discard.len(), |remaining| remaining.min(discard.len())); + let read = input.read(&mut discard[..to_read]).await.map_err(ZipError::InspectStream)?; + if read == 0 { + return Err(ZipError::InspectStream(io::Error::new(io::ErrorKind::UnexpectedEof, truncated_message))); + } + let read_u64 = + u64::try_from(read).map_err(|_| ZipError::InspectStream(io::Error::other("skippable frame read size overflowed")))?; + remaining -= read_u64; + budget.record_read(read).await; + } + Ok(()) +} + +fn tar_octal_field(field: &[u8]) -> Option { + let mut value = 0u64; + let mut saw_digit = false; + let mut terminated = false; + + for byte in field { + match *byte { + b'0'..=b'7' if !terminated => { + saw_digit = true; + value = value.checked_mul(8)?.checked_add(u64::from(*byte - b'0'))?; + } + b' ' if !saw_digit && !terminated => {} + b' ' | 0 if saw_digit || terminated => terminated = true, + 0 => terminated = true, + _ => return None, + } + } + + saw_digit.then_some(value) +} + +fn tar_numeric_field_has_shape(field: &[u8]) -> bool { + if field.first().is_some_and(|byte| byte & 0x80 != 0) { + return true; + } + tar_octal_field(field).is_some() || field.iter().all(|byte| matches!(*byte, b' ' | 0)) +} + +fn is_tar_header(prefix: &[u8]) -> bool { + let Some(header) = prefix.get(..TAR_HEADER_LEN) else { + return false; + }; + if header.iter().all(|byte| *byte == 0) { + return true; + } + + let has_path = header[..100].iter().any(|byte| *byte != 0); + if !has_path + || ![(100, 108), (108, 116), (116, 124), (124, 136), (136, 148)] + .iter() + .all(|(start, end)| tar_numeric_field_has_shape(&header[*start..*end])) + { + return false; + } + + let Some(expected) = tar_octal_field(&header[TAR_CHECKSUM_START..TAR_CHECKSUM_END]) else { + return false; + }; + let actual = header[..TAR_CHECKSUM_START] + .iter() + .chain(std::iter::repeat_n(&b' ', TAR_CHECKSUM_END - TAR_CHECKSUM_START)) + .chain(&header[TAR_CHECKSUM_END..]) + .fold(0u64, |sum, byte| sum + u64::from(*byte)); + actual == expected +} + /// Archive guardrails. The values are carried here so every archive caller /// shares one default policy; enforcement belongs to the caller, which maps a /// breach onto its own protocol error. @@ -84,6 +274,8 @@ impl CompressionFormat { "xz" | "txz" => CompressionFormat::Xz, "zlib" | "zz" => CompressionFormat::Zlib, "zst" | "zstd" | "tzst" => CompressionFormat::Zstd, + "lz4" | "tlz4" => CompressionFormat::Lz4, + "s2" | "snappy" => CompressionFormat::S2, "tar" => CompressionFormat::Tar, "zip" => CompressionFormat::Zip, _ => CompressionFormat::Unknown, @@ -98,11 +290,147 @@ impl CompressionFormat { CompressionFormat::Xz => "xz", CompressionFormat::Zlib => "zlib", CompressionFormat::Zstd => "zst", + CompressionFormat::Lz4 => "lz4", + CompressionFormat::S2 => "s2", CompressionFormat::Tar => "tar", CompressionFormat::Unknown => "", } } + /// Detect a stream codec from an unambiguous bounded prefix. Ordinary + /// unknown bytes are a raw TAR stream, matching MinIO Snowball behavior; + /// the skippable-frame magic shared by Zstd and LZ4 remains + /// [`CompressionFormat::Unknown`] until [`Self::sniff`] sees a later frame. + /// Object names and suffixes do not participate in detection. + /// + /// Zlib deliberately has no magic match here: its two-byte header can also + /// be the start of a valid TAR member name. Callers preserving historical + /// zlib-by-extension behavior must apply that compatibility fallback only + /// after this method returns [`CompressionFormat::Tar`]. + /// ZIP signatures are also left as TAR because stream ZIP decoding is not + /// supported and those bytes are valid at the start of a TAR member name. + pub fn from_magic(prefix: &[u8]) -> Self { + if prefix.starts_with(&[0x1f, 0x8b, 0x08]) { + return Self::Gzip; + } + if prefix.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) { + return Self::Zstd; + } + if is_shared_skippable_magic(prefix) { + return Self::Unknown; + } + if prefix.starts_with(&[0x04, 0x22, 0x4d, 0x18]) { + return Self::Lz4; + } + if prefix.starts_with(&[0xff, 0x06, 0x00, 0x00]) { + return Self::S2; + } + if prefix.starts_with(b"BZh") { + return Self::Bzip2; + } + if prefix.starts_with(&[0xfd, 0x37, 0x7a, 0x58, 0x5a, 0x00]) { + return Self::Xz; + } + Self::Tar + } + + /// Identify the codec and return a reader that replays the deciding prefix. + /// A complete, checksum-valid TAR header wins over a codec-like member-name + /// prefix. Leading skippable frames shared by the Zstd and LZ4 frame formats + /// are discarded with bounded memory until a non-skippable frame identifies + /// the decoder. The underlying reader still observes every discarded byte, + /// so callers can keep transport length and checksum accounting below this + /// boundary. + pub async fn sniff(mut input: R) -> Result<(Self, SniffedReader)> + where + R: AsyncRead + Unpin, + { + let mut prefix = Vec::with_capacity(TAR_HEADER_LEN); + let mut budget = SniffYieldBudget::default(); + let _ = inspect_fill_to(&mut input, &mut prefix, MAGIC_SNIFF_LEN, &mut budget).await?; + let deciding_prefix_len = prefix.len().min(MAGIC_SNIFF_LEN); + let initial_format = Self::from_magic(&prefix[..deciding_prefix_len]); + let needs_tar_disambiguation = + matches!(initial_format, Self::Bzip2 | Self::Lz4) || is_shared_skippable_magic(&prefix[..deciding_prefix_len]); + if !needs_tar_disambiguation { + return Ok(( + initial_format, + SniffedReader { + inner: input, + prefix, + prefix_pos: 0, + }, + )); + } + + let _ = inspect_fill_to(&mut input, &mut prefix, TAR_HEADER_LEN, &mut budget).await?; + if is_tar_header(&prefix) { + return Ok(( + Self::Tar, + SniffedReader { + inner: input, + prefix, + prefix_pos: 0, + }, + )); + } + + let mut skipped_shared_frame = false; + + loop { + if prefix.len() < MAGIC_SNIFF_LEN { + let _ = inspect_fill_to(&mut input, &mut prefix, MAGIC_SNIFF_LEN, &mut budget).await?; + } + let deciding_prefix_len = prefix.len().min(MAGIC_SNIFF_LEN); + if !is_shared_skippable_magic(&prefix[..deciding_prefix_len]) { + let mut format = Self::from_magic(&prefix[..deciding_prefix_len]); + if skipped_shared_frame && !matches!(format, Self::Zstd | Self::Lz4) { + format = Self::Unknown; + } + return Ok(( + format, + SniffedReader { + inner: input, + prefix, + prefix_pos: 0, + }, + )); + } + + skipped_shared_frame = true; + let complete_header = + inspect_fill_to(&mut input, &mut prefix, SHARED_SKIPPABLE_FRAME_HEADER_LEN, &mut budget).await?; + if !complete_header { + return Err(ZipError::InspectStream(io::Error::new( + io::ErrorKind::UnexpectedEof, + "truncated shared Zstd/LZ4 skippable frame header", + ))); + } + let payload_len = + usize::try_from(u32::from_le_bytes([prefix[4], prefix[5], prefix[6], prefix[7]])).map_err(|_| { + ZipError::InspectStream(io::Error::new( + io::ErrorKind::InvalidData, + "shared Zstd/LZ4 skippable frame length does not fit usize", + )) + })?; + let buffered_payload_len = (prefix.len() - SHARED_SKIPPABLE_FRAME_HEADER_LEN).min(payload_len); + let buffered_frame_len = SHARED_SKIPPABLE_FRAME_HEADER_LEN + .checked_add(buffered_payload_len) + .ok_or_else(|| ZipError::InspectStream(io::Error::other("skippable frame buffered length overflowed")))?; + prefix.drain(..buffered_frame_len); + let remaining_payload_len = payload_len - buffered_payload_len; + inspect_discard_fully( + &mut input, + u64::try_from(remaining_payload_len) + .map_err(|_| ZipError::InspectStream(io::Error::other("skippable frame remaining length exceeds u64")))?, + &mut budget, + "truncated shared Zstd/LZ4 skippable frame payload", + ) + .await?; + budget.record_frame().await; + } + } + pub fn get_decoder(&self, input: R) -> Result> where R: AsyncRead + Send + Unpin + 'static, @@ -126,15 +454,21 @@ impl CompressionFormat { Box::new(decoder) } CompressionFormat::Xz => { - let mut decoder = XzDecoder::new(reader); + let mut decoder = XzDecoder::with_mem_limit(reader, XZ_DECODER_MEMORY_LIMIT_BYTES); decoder.multiple_members(true); Box::new(decoder) } CompressionFormat::Zstd => { - let mut decoder = ZstdDecoder::new(reader); + let mut decoder = ZstdDecoder::with_params(reader, &[DParameter::window_log_max(ZSTD_DECODER_MAX_WINDOW_LOG)]); decoder.multiple_members(true); Box::new(decoder) } + CompressionFormat::Lz4 => { + let mut decoder = Lz4Decoder::new(reader); + decoder.multiple_members(true); + Box::new(decoder) + } + CompressionFormat::S2 => Box::new(S2Decoder::new(reader)), CompressionFormat::Tar => Box::new(reader), CompressionFormat::Zip => { return Err(ZipError::UnsupportedFormat { @@ -157,18 +491,306 @@ impl CompressionFormat { #[cfg(test)] mod tests { use super::*; - use async_compression::tokio::write::GzipEncoder; + use async_compression::{ + Level, + tokio::write::{BzEncoder, GzipEncoder, Lz4Encoder, XzEncoder, ZstdEncoder}, + zstd::CParameter, + }; + use std::future::Future; use std::mem::size_of; + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + use std::task::{Wake, Waker}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; + #[derive(Default)] + struct WakeCounter(AtomicUsize); + + impl Wake for WakeCounter { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + async fn encode_xz(payload: &[u8]) -> Vec { + let mut encoder = XzEncoder::with_quality(Vec::new(), Level::Fastest); + encoder.write_all(payload).await.expect("XZ encode should succeed"); + encoder.shutdown().await.expect("XZ encoder shutdown should succeed"); + encoder.into_inner() + } + + async fn encode_lz4(payload: &[u8]) -> Vec { + let mut encoder = Lz4Encoder::new(Vec::new()); + encoder.write_all(payload).await.expect("LZ4 encode should succeed"); + encoder.shutdown().await.expect("LZ4 encoder shutdown should succeed"); + encoder.into_inner() + } + + async fn encode_zstd_with_window(payload: &[u8], window_log: u32) -> Vec { + let mut encoder = ZstdEncoder::with_quality_and_params( + Vec::new(), + Level::Default, + &[CParameter::window_log(window_log), CParameter::content_size_flag(false)], + ); + encoder.write_all(payload).await.expect("Zstd encode should succeed"); + encoder.shutdown().await.expect("Zstd encoder shutdown should succeed"); + encoder.into_inner() + } + + fn shared_skippable_frame(variant: u8, payload: &[u8]) -> Vec { + assert!(variant <= 0x0f); + let payload_len = u32::try_from(payload.len()).expect("test skippable payload should fit u32"); + let mut frame = Vec::with_capacity(SHARED_SKIPPABLE_FRAME_HEADER_LEN + payload.len()); + frame.extend_from_slice(&[0x50 + variant, 0x2a, 0x4d, 0x18]); + frame.extend_from_slice(&payload_len.to_le_bytes()); + frame.extend_from_slice(payload); + frame + } + + fn raw_tar_with_name(name: &[u8]) -> Vec { + assert!(!name.is_empty() && name.len() <= 100, "test TAR name must fit the legacy name field"); + + let mut header = [0u8; TAR_HEADER_LEN]; + header[..name.len()].copy_from_slice(name); + for (start, end) in [(100, 108), (108, 116), (116, 124), (124, 136), (136, 148)] { + header[start..end].fill(b'0'); + header[end - 1] = 0; + } + header[156] = b'0'; + header[257..263].copy_from_slice(b"ustar\0"); + header[263..265].copy_from_slice(b"00"); + header[TAR_CHECKSUM_START..TAR_CHECKSUM_END].fill(b' '); + let checksum = header.iter().fold(0u64, |sum, byte| sum + u64::from(*byte)); + let checksum_field = format!("{checksum:06o}\0 "); + assert_eq!(checksum_field.len(), TAR_CHECKSUM_END - TAR_CHECKSUM_START); + header[TAR_CHECKSUM_START..TAR_CHECKSUM_END].copy_from_slice(checksum_field.as_bytes()); + + let mut archive = header.to_vec(); + archive.resize(TAR_HEADER_LEN * 3, 0); + archive + } + + struct FragmentedReader { + bytes: Vec, + position: usize, + return_pending: bool, + } + + impl FragmentedReader { + fn new(bytes: Vec) -> Self { + Self { + bytes, + position: 0, + return_pending: true, + } + } + } + + impl AsyncRead for FragmentedReader { + fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, output: &mut ReadBuf<'_>) -> Poll> { + if self.return_pending { + self.return_pending = false; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + if self.position >= self.bytes.len() || output.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + let byte = self.bytes[self.position]; + self.position += 1; + self.return_pending = true; + output.put_slice(&[byte]); + Poll::Ready(Ok(())) + } + } + + struct AlwaysReadyOneByteReader { + bytes: Vec, + position: usize, + ready_reads: Arc, + bytes_read: Arc, + } + + impl AsyncRead for AlwaysReadyOneByteReader { + fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, output: &mut ReadBuf<'_>) -> Poll> { + self.ready_reads.fetch_add(1, Ordering::Relaxed); + if self.position >= self.bytes.len() || output.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + let byte = self.bytes[self.position]; + self.position += 1; + self.bytes_read.fetch_add(1, Ordering::Relaxed); + output.put_slice(&[byte]); + Poll::Ready(Ok(())) + } + } + + #[derive(Debug)] + struct ErrorAfterBytes { + bytes: Vec, + position: usize, + } + + impl AsyncRead for ErrorAfterBytes { + fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, output: &mut ReadBuf<'_>) -> Poll> { + if self.position >= self.bytes.len() { + return Poll::Ready(Err(io::Error::new(io::ErrorKind::ConnectionReset, "sentinel inspect failure"))); + } + let available = self.bytes.len() - self.position; + let to_copy = available.min(output.remaining()); + output.put_slice(&self.bytes[self.position..self.position + to_copy]); + self.position += to_copy; + Poll::Ready(Ok(())) + } + } + + async fn assert_shared_skippable_round_trip(format: CompressionFormat, encoded: Vec, frames: &[Vec]) { + let mut stream = Vec::new(); + for frame in frames { + stream.extend_from_slice(frame); + } + stream.extend_from_slice(&encoded); + + let (detected, mut sniffed) = CompressionFormat::sniff(std::io::Cursor::new(stream.clone())) + .await + .expect("shared skippable prefix should be inspected"); + assert_eq!(detected, format); + let mut replayed_prefix = [0u8; MAGIC_SNIFF_LEN]; + sniffed + .read_exact(&mut replayed_prefix) + .await + .expect("deciding codec prefix should replay completely"); + assert_eq!(replayed_prefix, encoded[..MAGIC_SNIFF_LEN]); + + let (detected, sniffed) = CompressionFormat::sniff(std::io::Cursor::new(stream)) + .await + .expect("shared skippable prefix should be inspected for decoding"); + let mut decoder = detected.get_decoder(sniffed).expect("detected decoder should be created"); + let mut decoded = Vec::new(); + decoder.read_to_end(&mut decoded).await.expect("framed stream should decode"); + assert_eq!(decoded, b"payload"); + } + + async fn assert_ready_sniff_yields(input: Vec, encoded_prefix: &[u8]) { + let wake_counter = Arc::new(WakeCounter::default()); + let waker = Waker::from(wake_counter.clone()); + let mut sniff = Box::pin(CompressionFormat::sniff(std::io::Cursor::new(input))); + { + let mut cx = Context::from_waker(&waker); + assert!(sniff.as_mut().poll(&mut cx).is_pending(), "bounded inspection must yield"); + } + tokio::task::yield_now().await; + assert_eq!(wake_counter.0.load(Ordering::Relaxed), 1, "yield must arrange another poll"); + + let (format, mut sniffed) = sniff.await.expect("inspection should resume after yielding"); + assert_eq!(format, CompressionFormat::Lz4); + let mut replayed_prefix = [0u8; MAGIC_SNIFF_LEN]; + sniffed + .read_exact(&mut replayed_prefix) + .await + .expect("deciding codec prefix should replay after yielding"); + assert_eq!(replayed_prefix, encoded_prefix); + } + + fn xz_crc32(bytes: &[u8]) -> u32 { + let mut crc = u32::MAX; + for byte in bytes { + crc ^= u32::from(*byte); + for _ in 0..8 { + let mask = 0_u32.wrapping_sub(crc & 1); + crc = (crc >> 1) ^ (0xedb8_8320 & mask); + } + } + !crc + } + + async fn xz_with_dictionary_property(payload: &[u8], dictionary_property: u8) -> Vec { + const XZ_STREAM_HEADER_LEN: usize = 12; + + let mut encoded = encode_xz(payload).await; + assert_eq!(&encoded[..MAGIC_SNIFF_LEN], b"\xfd7zXZ\0"); + + let block_header_start = XZ_STREAM_HEADER_LEN; + let block_header_len = (usize::from(encoded[block_header_start]) + 1) * 4; + let block_header_crc_start = block_header_start + block_header_len - 4; + assert_eq!(block_header_len, 12, "default XZ fixture should use one compact LZMA2 filter header"); + assert_eq!( + &encoded[block_header_start + 1..block_header_start + 4], + &[0x00, 0x21, 0x01], + "default XZ fixture should contain one LZMA2 filter with one property byte" + ); + + encoded[block_header_start + 4] = dictionary_property; + let block_header_crc = xz_crc32(&encoded[block_header_start..block_header_crc_start]).to_le_bytes(); + encoded[block_header_crc_start..block_header_crc_start + 4].copy_from_slice(&block_header_crc); + encoded + } + #[test] fn test_compression_format_from_extension() { assert_eq!(CompressionFormat::from_extension("gz"), CompressionFormat::Gzip); assert_eq!(CompressionFormat::from_extension("ZIP"), CompressionFormat::Zip); assert_eq!(CompressionFormat::from_extension("tzst"), CompressionFormat::Zstd); + assert_eq!(CompressionFormat::from_extension("s2"), CompressionFormat::S2); assert_eq!(CompressionFormat::from_extension("txt"), CompressionFormat::Unknown); } + #[test] + fn test_compression_format_from_magic() { + let cases: &[(&[u8], CompressionFormat)] = &[ + (&[0x1f, 0x8b, 0x08, 0x00], CompressionFormat::Gzip), + (b"BZh9", CompressionFormat::Bzip2), + (&[0x28, 0xb5, 0x2f, 0xfd], CompressionFormat::Zstd), + (&[0x50, 0x2a, 0x4d, 0x18], CompressionFormat::Unknown), + (&[0x04, 0x22, 0x4d, 0x18], CompressionFormat::Lz4), + (&[0xff, 0x06, 0x00, 0x00], CompressionFormat::S2), + (&[0xfd, 0x37, 0x7a, 0x58, 0x5a, 0x00], CompressionFormat::Xz), + (&[0x78, 0x9c], CompressionFormat::Tar), + (&[0x78, 0x5e, b'o', b'b'], CompressionFormat::Tar), + (b"PK\x03\x04", CompressionFormat::Tar), + (b"plain tar bytes", CompressionFormat::Tar), + ]; + + for (magic, expected) in cases { + assert_eq!(CompressionFormat::from_magic(magic), *expected, "magic={magic:02x?}"); + } + } + + #[tokio::test] + async fn test_sniff_prefers_checksum_valid_tar_over_codec_like_member_names() { + let cases: &[(&[u8], CompressionFormat)] = &[ + (b"BZh9-report.txt", CompressionFormat::Bzip2), + (b"\x04\x22\x4d\x18-report.txt", CompressionFormat::Lz4), + ]; + + for (name, prefix_format) in cases { + assert_eq!(CompressionFormat::from_magic(name), *prefix_format); + let archive = raw_tar_with_name(name); + let (format, mut sniffed) = CompressionFormat::sniff(std::io::Cursor::new(archive.clone())) + .await + .expect("checksum-valid TAR header should be inspectable"); + assert_eq!(format, CompressionFormat::Tar, "member name prefix={name:02x?}"); + + let mut replayed = Vec::new(); + sniffed + .read_to_end(&mut replayed) + .await + .expect("TAR lookahead should replay without loss"); + assert_eq!(replayed, archive); + } + + let mut bad_checksum = raw_tar_with_name(b"BZh9-bad-checksum.txt"); + bad_checksum[99] = b'x'; + let (format, _) = CompressionFormat::sniff(std::io::Cursor::new(bad_checksum)) + .await + .expect("malformed TAR lookahead should still permit codec detection"); + assert_eq!(format, CompressionFormat::Bzip2, "TAR priority must require a valid checksum"); + } + #[test] fn test_compression_format_size_is_small() { assert!(size_of::() <= 8); @@ -181,15 +803,34 @@ mod tests { encoder.write_all(b"payload").await.expect("gzip encode should succeed"); encoder.shutdown().await.expect("gzip encoder shutdown should succeed"); - let mut decoder = CompressionFormat::Gzip - .get_decoder(std::io::Cursor::new(encoder.into_inner())) - .expect("gzip decoder should be created"); + let (format, sniffed) = CompressionFormat::sniff(std::io::Cursor::new(encoder.into_inner())) + .await + .expect("gzip magic should be inspected"); + assert_eq!(format, CompressionFormat::Gzip); + let mut decoder = format.get_decoder(sniffed).expect("gzip decoder should be created"); let mut decoded = Vec::new(); decoder.read_to_end(&mut decoded).await.expect("gzip decode should succeed"); assert_eq!(decoded, b"payload"); } + #[tokio::test] + async fn test_sniff_still_recognizes_a_real_bzip2_stream() { + let mut encoder = BzEncoder::new(Vec::new()); + encoder.write_all(b"payload").await.expect("Bzip2 encode should succeed"); + encoder.shutdown().await.expect("Bzip2 encoder shutdown should succeed"); + + let (format, sniffed) = CompressionFormat::sniff(std::io::Cursor::new(encoder.into_inner())) + .await + .expect("Bzip2 magic should be inspected after TAR lookahead"); + assert_eq!(format, CompressionFormat::Bzip2); + let mut decoder = format.get_decoder(sniffed).expect("Bzip2 decoder should be created"); + let mut decoded = Vec::new(); + decoder.read_to_end(&mut decoded).await.expect("Bzip2 decode should succeed"); + + assert_eq!(decoded, b"payload"); + } + #[tokio::test] async fn test_get_decoder_consumes_concatenated_gzip_members() { async fn gzip_member(payload: &[u8]) -> Vec { @@ -214,6 +855,303 @@ mod tests { assert_eq!(decoded, b"first-second"); } + #[tokio::test] + async fn test_get_decoder_round_trips_lz4_stream() { + let mut encoder = Lz4Encoder::new(Vec::new()); + encoder.write_all(b"payload").await.expect("LZ4 encode should succeed"); + encoder.shutdown().await.expect("LZ4 encoder shutdown should succeed"); + + let (format, sniffed) = CompressionFormat::sniff(std::io::Cursor::new(encoder.into_inner())) + .await + .expect("LZ4 magic should be inspected"); + assert_eq!(format, CompressionFormat::Lz4); + let mut decoder = format.get_decoder(sniffed).expect("LZ4 decoder should be created"); + let mut decoded = Vec::new(); + decoder.read_to_end(&mut decoded).await.expect("LZ4 decode should succeed"); + + assert_eq!(decoded, b"payload"); + } + + #[tokio::test] + async fn test_sniff_skips_shared_frames_before_lz4_and_zstd() { + let lz4 = encode_lz4(b"payload").await; + assert_shared_skippable_round_trip( + CompressionFormat::Lz4, + lz4, + &[shared_skippable_frame(0, b""), shared_skippable_frame(15, b"lz4 metadata")], + ) + .await; + + let zstd = encode_zstd_with_window(b"payload", ZSTD_DECODER_MAX_WINDOW_LOG).await; + let large_metadata = vec![0x5a; SNIFF_YIELD_AFTER_BYTES + 1]; + assert_shared_skippable_round_trip( + CompressionFormat::Zstd, + zstd, + &[shared_skippable_frame(3, &large_metadata), shared_skippable_frame(4, b"")], + ) + .await; + } + + #[tokio::test] + async fn test_sniff_rejects_truncated_shared_skippable_frames() { + let truncated_header = [0x50, 0x2a, 0x4d, 0x18, 0x04, 0x00]; + let err = CompressionFormat::sniff(std::io::Cursor::new(truncated_header)) + .await + .expect_err("truncated shared frame header must fail"); + assert!(matches!( + err, + ZipError::InspectStream(ref source) if source.kind() == io::ErrorKind::UnexpectedEof + )); + + let mut truncated_payload = shared_skippable_frame(0, b"abcd"); + truncated_payload.truncate(truncated_payload.len() - 2); + let err = CompressionFormat::sniff(std::io::Cursor::new(truncated_payload)) + .await + .expect_err("truncated shared frame payload must fail"); + assert!(matches!( + err, + ZipError::InspectStream(ref source) if source.kind() == io::ErrorKind::UnexpectedEof + )); + } + + #[tokio::test] + async fn test_sniff_shared_frames_handles_fragmented_pending_reads() { + let encoded = encode_lz4(b"payload").await; + let mut input = Vec::new(); + for variant in 0..=SNIFF_YIELD_AFTER_FRAMES { + input.extend_from_slice(&shared_skippable_frame(u8::try_from(variant % 16).expect("variant should fit u8"), b"")); + } + input.extend_from_slice(&encoded); + + let (format, sniffed) = CompressionFormat::sniff(FragmentedReader::new(input)) + .await + .expect("fragmented shared frames should be inspected"); + assert_eq!(format, CompressionFormat::Lz4); + let mut decoder = format.get_decoder(sniffed).expect("LZ4 decoder should be created"); + let mut decoded = Vec::new(); + decoder + .read_to_end(&mut decoded) + .await + .expect("fragmented stream should decode"); + assert_eq!(decoded, b"payload"); + } + + #[tokio::test] + async fn test_sniff_shared_frames_yields_at_frame_budget_with_ready_reader() { + let encoded = encode_lz4(b"payload").await; + let mut input = Vec::new(); + for variant in 0..SNIFF_YIELD_AFTER_FRAMES { + let variant = u8::try_from(variant % 16).expect("variant should fit u8"); + input.extend_from_slice(&shared_skippable_frame(variant, b"")); + } + input.extend_from_slice(&encoded); + + assert_ready_sniff_yields(input, &encoded[..MAGIC_SNIFF_LEN]).await; + } + + #[tokio::test] + async fn test_sniff_shared_frame_yields_at_byte_budget_with_ready_reader() { + let encoded = encode_lz4(b"payload").await; + let metadata = vec![0x5a; SNIFF_YIELD_AFTER_BYTES]; + let mut input = shared_skippable_frame(0, &metadata); + input.extend_from_slice(&encoded); + + assert_ready_sniff_yields(input, &encoded[..MAGIC_SNIFF_LEN]).await; + } + + #[tokio::test] + async fn test_sniff_shared_frame_bounds_always_ready_one_byte_reads_per_poll() { + let encoded = encode_lz4(b"payload").await; + let metadata = vec![0x5a; TAR_HEADER_LEN * 4]; + let mut input = shared_skippable_frame(0, &metadata); + input.extend_from_slice(&encoded); + + let ready_reads = Arc::new(AtomicUsize::new(0)); + let bytes_read = Arc::new(AtomicUsize::new(0)); + let reader = AlwaysReadyOneByteReader { + bytes: input, + position: 0, + ready_reads: ready_reads.clone(), + bytes_read: bytes_read.clone(), + }; + let wake_counter = Arc::new(WakeCounter::default()); + let waker = Waker::from(wake_counter.clone()); + let mut sniff = Box::pin(CompressionFormat::sniff(reader)); + + for poll_number in 0..10 { + let before = ready_reads.load(Ordering::Relaxed); + let mut cx = Context::from_waker(&waker); + assert!( + sniff.as_mut().poll(&mut cx).is_pending(), + "poll {poll_number} must yield at the ready-read budget" + ); + let reads_this_poll = ready_reads.load(Ordering::Relaxed) - before; + assert_eq!( + reads_this_poll, SNIFF_YIELD_AFTER_READY_READS, + "poll {poll_number} exceeded the ready-read budget" + ); + } + assert!( + bytes_read.load(Ordering::Relaxed) > TAR_HEADER_LEN, + "manual polls must progress beyond TAR lookahead into the shared-frame payload" + ); + assert_eq!( + wake_counter.0.load(Ordering::Relaxed), + 0, + "Tokio defers the wake until the scheduler runs" + ); + tokio::task::yield_now().await; + assert_eq!( + wake_counter.0.load(Ordering::Relaxed), + 1, + "repeated voluntary yields with the same waker must coalesce into one scheduled repoll" + ); + + let (format, sniffed) = sniff.await.expect("bounded inspection should eventually complete"); + assert_eq!(format, CompressionFormat::Lz4); + let mut decoder = format.get_decoder(sniffed).expect("LZ4 decoder should be created"); + let mut decoded = Vec::new(); + decoder + .read_to_end(&mut decoded) + .await + .expect("LZ4 stream should survive bounded inspection"); + assert_eq!(decoded, b"payload"); + } + + #[tokio::test] + async fn test_sniff_shared_frame_preserves_underlying_read_error() { + let mut bytes = shared_skippable_frame(0, b"abcd"); + bytes.truncate(bytes.len() - 2); + let err = CompressionFormat::sniff(ErrorAfterBytes { bytes, position: 0 }) + .await + .expect_err("underlying read failure must escape inspection"); + + assert!(matches!( + err, + ZipError::InspectStream(ref source) + if source.kind() == io::ErrorKind::ConnectionReset && source.to_string() == "sentinel inspect failure" + )); + } + + #[tokio::test] + async fn test_sniff_shared_frame_without_lz4_or_zstd_successor_is_unknown() { + let mut input = shared_skippable_frame(0, b"metadata"); + input.extend_from_slice(b"raw tar bytes"); + + let (format, mut sniffed) = CompressionFormat::sniff(std::io::Cursor::new(input)) + .await + .expect("complete shared frame should be inspectable"); + assert_eq!(format, CompressionFormat::Unknown); + let mut replayed = Vec::new(); + sniffed + .read_to_end(&mut replayed) + .await + .expect("successor bytes should replay"); + assert_eq!(replayed, b"raw tar bytes"); + } + + #[tokio::test] + async fn test_get_decoder_round_trips_xz_stream_with_memory_limit() { + let encoded = encode_xz(b"payload").await; + let (format, sniffed) = CompressionFormat::sniff(std::io::Cursor::new(encoded)) + .await + .expect("XZ magic should be inspected"); + assert_eq!(format, CompressionFormat::Xz); + + let mut decoder = format.get_decoder(sniffed).expect("XZ decoder should be created"); + let mut decoded = Vec::new(); + decoder + .read_to_end(&mut decoded) + .await + .expect("XZ decode should succeed within the memory limit"); + + assert_eq!(decoded, b"payload"); + } + + #[tokio::test] + async fn test_get_decoder_rejects_xz_dictionary_over_memory_limit() { + const LZMA2_MAX_DICTIONARY_PROPERTY: u8 = 40; + + let encoded = xz_with_dictionary_property(b"payload", LZMA2_MAX_DICTIONARY_PROPERTY).await; + let mut decoder = CompressionFormat::Xz + .get_decoder(std::io::Cursor::new(encoded)) + .expect("XZ decoder should be created before inspecting the stream header"); + let mut decoded = Vec::new(); + let err = decoder + .read_to_end(&mut decoded) + .await + .expect_err("hostile XZ dictionary request must exceed the decoder memory limit"); + + assert!(decoded.is_empty(), "decoder must reject the hostile dictionary before producing output"); + assert!( + err.to_string().contains("memory limit"), + "hostile XZ dictionary should fail at the decoder memory boundary: {err}" + ); + } + + #[tokio::test] + async fn test_get_decoder_accepts_xz_64_mib_dictionary() { + const LZMA2_64_MIB_DICTIONARY_PROPERTY: u8 = 28; + + let encoded = xz_with_dictionary_property(b"payload", LZMA2_64_MIB_DICTIONARY_PROPERTY).await; + let mut decoder = CompressionFormat::Xz + .get_decoder(std::io::Cursor::new(encoded)) + .expect("XZ decoder should be created before inspecting the stream header"); + let mut decoded = Vec::new(); + decoder + .read_to_end(&mut decoded) + .await + .expect("the memory limit must retain preset-9-compatible dictionary headroom"); + + assert_eq!(decoded, b"payload"); + } + + #[tokio::test] + async fn test_get_decoder_accepts_zstd_window_at_limit() { + let encoded = encode_zstd_with_window(b"payload", ZSTD_DECODER_MAX_WINDOW_LOG).await; + let mut decoder = CompressionFormat::Zstd + .get_decoder(std::io::Cursor::new(encoded)) + .expect("Zstd decoder should be created before inspecting the frame header"); + let mut decoded = Vec::new(); + decoder + .read_to_end(&mut decoded) + .await + .expect("16 MiB Zstd windows must remain compatible"); + + assert_eq!(decoded, b"payload"); + } + + #[tokio::test] + async fn test_get_decoder_rejects_zstd_window_over_limit() { + let encoded = encode_zstd_with_window(b"payload", ZSTD_DECODER_MAX_WINDOW_LOG + 1).await; + let mut decoder = CompressionFormat::Zstd + .get_decoder(std::io::Cursor::new(encoded)) + .expect("Zstd decoder should be created before inspecting the frame header"); + let mut decoded = Vec::new(); + decoder + .read_to_end(&mut decoded) + .await + .expect_err("Zstd windows larger than 16 MiB must be rejected"); + + assert!(decoded.is_empty(), "decoder must reject the oversized window before producing output"); + } + + #[tokio::test] + async fn test_sniff_replays_raw_tar_prefix() { + let input = b"ustar-prefix-and-payload".to_vec(); + let (format, mut sniffed) = CompressionFormat::sniff(std::io::Cursor::new(input.clone())) + .await + .expect("raw TAR prefix should be inspected"); + let mut replayed = Vec::new(); + sniffed + .read_to_end(&mut replayed) + .await + .expect("sniffed prefix should replay"); + + assert_eq!(format, CompressionFormat::Tar); + assert_eq!(replayed, input); + } + #[tokio::test] async fn test_get_decoder_rejects_zip_and_unknown_formats() { let zip_err = CompressionFormat::Zip diff --git a/rustfs/src/app/object/extract.rs b/rustfs/src/app/object/extract.rs index 478489223..e55231b0d 100644 --- a/rustfs/src/app/object/extract.rs +++ b/rustfs/src/app/object/extract.rs @@ -63,6 +63,22 @@ struct ExtractArchiveUploadState { body_complete: bool, } +fn resolve_extract_archive_format(key: &str, detected: CompressionFormat) -> CompressionFormat { + // Zlib has no unambiguous magic. Preserve the legacy zlib/zz suffix + // contract without letting other misleading suffixes override content + // detection. + if detected == CompressionFormat::Tar + && Path::new(key) + .extension() + .and_then(|extension| extension.to_str()) + .is_some_and(|extension| CompressionFormat::from_extension(extension) == CompressionFormat::Zlib) + { + CompressionFormat::Zlib + } else { + detected + } +} + impl ExtractArchiveEtagReader { fn new(inner: R, expected_length: u64, state: Arc>) -> Self { Self { @@ -1008,12 +1024,6 @@ impl DefaultObjectUsecase { let body = tokio::io::BufReader::with_capacity(buffer_size, StreamReader::new(body.map(|f| f.map_err(s3s_body_error_to_io)))); - let Some(ext) = Path::new(&key).extension().and_then(|s| s.to_str()) else { - return Err(s3_error!(InvalidArgument, "key extension not found")); - }; - - let ext = ext.to_owned(); - let md5hex = if let Some(base64_md5) = content_md5 { let md5 = base64_simd::STANDARD .decode_to_vec(base64_md5.as_bytes()) @@ -1036,16 +1046,18 @@ impl DefaultObjectUsecase { let expected_archive_length = u64::try_from(size).map_err(|_| S3Error::new(S3ErrorCode::UnexpectedContent))?; let archive_upload_state = Arc::new(Mutex::new(ExtractArchiveUploadState::default())); let extract_limits = put_object_extract_limits(); - let decoder = CompressionFormat::from_extension(&ext) - .get_decoder(ExtractArchiveEtagReader::new( - archive_reader, - expected_archive_length, - archive_upload_state.clone(), - )) - .map_err(|e| { - error!(error = ?e, "Archive decoder creation failed"); - s3_error!(InvalidArgument, "get_decoder err") + let tracked_archive = + ExtractArchiveEtagReader::new(archive_reader, expected_archive_length, archive_upload_state.clone()); + let (detected_archive_format, sniffed_archive) = + CompressionFormat::sniff(tracked_archive).await.map_err(|err| match err { + ZipError::InspectStream(source) => map_extract_archive_error(source), + _ => s3_error!(InvalidArgument, "Failed to detect archive compression"), })?; + let archive_format = resolve_extract_archive_format(&key, detected_archive_format); + let decoder = archive_format.get_decoder(sniffed_archive).map_err(|e| { + error!(error = ?e, "Archive decoder creation failed"); + s3_error!(InvalidArgument, "get_decoder err") + })?; let decoder = ExtractDecodedLimitReader::new(decoder, extract_limits.max_decoded_size); let mut ar = build_put_object_extract_archive(decoder, extract_limits); @@ -1488,6 +1500,92 @@ mod tests { use tokio::io::AsyncReadExt; use tokio_tar::{Builder, EntryType, Header}; + #[test] + fn archive_format_uses_only_the_ambiguous_zlib_extension_as_a_fallback() { + assert_eq!( + resolve_extract_archive_format("archive.zlib", CompressionFormat::Tar), + CompressionFormat::Zlib + ); + assert_eq!( + resolve_extract_archive_format("archive.zz", CompressionFormat::Tar), + CompressionFormat::Zlib + ); + assert_eq!( + resolve_extract_archive_format("raw-but-named.tar.gz", CompressionFormat::Tar), + CompressionFormat::Tar + ); + assert_eq!( + resolve_extract_archive_format("gzip-but-named.zlib", CompressionFormat::Gzip), + CompressionFormat::Gzip + ); + } + + #[tokio::test] + async fn raw_tar_member_names_starting_with_codec_magic_are_not_misdetected() { + let cases = [ + ("PK\u{3}\u{4}-member.txt", b"PK\x03\x04".as_slice()), + ("BZh9-report.txt", b"BZh9".as_slice()), + ("\u{4}\"M\u{18}-report.txt", b"\x04\x22\x4d\x18".as_slice()), + ]; + + for (path, expected_prefix) in cases { + let mut builder = Builder::new(Vec::new()); + let mut header = Header::new_gnu(); + header.set_size(0); + header.set_cksum(); + builder + .append_data(&mut header, path, &b""[..]) + .await + .expect("raw TAR fixture should accept the codec-like member name"); + let bytes = builder.into_inner().await.expect("raw TAR fixture should finalize"); + assert_eq!(&bytes[..expected_prefix.len()], expected_prefix); + + let (format, sniffed) = CompressionFormat::sniff(std::io::Cursor::new(bytes)) + .await + .expect("raw TAR prefix should be inspected"); + assert_eq!(format, CompressionFormat::Tar, "member path={path:?}"); + let decoder = format.get_decoder(sniffed).expect("raw TAR decoder should be created"); + let mut archive = Archive::new(decoder); + let mut entries = archive.entries().expect("raw TAR entry stream should be created"); + let entry = entries + .next() + .await + .expect("raw TAR should contain its first member") + .expect("raw TAR member should parse"); + + assert_eq!(entry.path_bytes().expect("raw TAR member path should parse").as_ref(), path.as_bytes()); + } + } + + #[tokio::test] + async fn archive_etag_reader_validates_sha256_before_completion() { + let payload = b"archive-with-wrong-sha256".to_vec(); + let expected_length = i64::try_from(payload.len()).expect("fixture length must fit i64"); + let hash_reader = HashReader::from_stream( + std::io::Cursor::new(payload), + expected_length, + expected_length, + None, + Some("00".repeat(32)), + false, + ) + .expect("hash reader should be created"); + let state = Arc::new(Mutex::new(ExtractArchiveUploadState::default())); + let mut reader = ExtractArchiveEtagReader::new( + hash_reader, + u64::try_from(expected_length).expect("fixture length must fit u64"), + state.clone(), + ); + let mut output = Vec::new(); + let err = reader + .read_to_end(&mut output) + .await + .expect_err("SHA-256 must be checked before upload completion"); + + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); + assert!(!state.lock().expect("archive state lock must remain healthy").body_complete); + } + fn pax_record(key: &str, value: &[u8]) -> Vec { let body_len = 1 + key.len() + 1 + value.len() + 1; let mut len = body_len + 1; diff --git a/rustfs/src/app/object/mod.rs b/rustfs/src/app/object/mod.rs index d28777a99..8f75eee9e 100644 --- a/rustfs/src/app/object/mod.rs +++ b/rustfs/src/app/object/mod.rs @@ -161,7 +161,7 @@ use rustfs_utils::http::{ }; use rustfs_utils::path::{encode_dir_object, is_dir_object, path_join_buf}; use rustfs_utils::retry::{DEFAULT_RETRY_CAP, DEFAULT_RETRY_UNIT, MAX_JITTER, RetryTimer}; -use rustfs_zip::{ArchiveLimits, CompressionFormat}; +use rustfs_zip::{ArchiveLimits, CompressionFormat, ZipError}; use s3s::StdError; use s3s::dto::{ CacheControl, Checksum, ChecksumAlgorithm, ChecksumType, ContentDisposition, ContentEncoding, ContentLanguage, ContentType,