feat(rio): rio_v2 is compatible with minio for storing data. (#3115)

* Set up a compatibility layer for replacing old Rio components with new ones.

* fix(rio). compress range

* feat(rio). Add the experimental feature rio_v2 to support minio data at the binary level.

* feat(rio_v2): add sse-c test

* test compression component

* simple fix

* fix minlz encode

* fix metadata

* fix kms key cache error

* Update launch.json

* ci: set nix crate download user agent

* fix: gate obs pyroscope backend

* ignore minio test

* fix encrypt check

* fix

* fix

* fix

* Update object_usecase.rs

* Update ci.yml

* fix

* ci add rio-v2 test

* fix

* ci fix

* fix

* Reconstructed into a more reasonable compatibility mode

* fix

* fix

---------

Signed-off-by: houseme <housemecn@gmail.com>
Signed-off-by: 唐小鸭 <tangtang1251@qq.com>
Co-authored-by: houseme <housemecn@gmail.com>
Co-authored-by: cxymds <Cxymds@qq.com>
Co-authored-by: 安正超 <anzhengchao@gmail.com>
This commit is contained in:
唐小鸭
2026-06-08 19:59:14 +08:00
committed by GitHub
parent 9504dff595
commit f7724d223b
47 changed files with 8742 additions and 682 deletions
+755
View File
@@ -0,0 +1,755 @@
// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use minlz::{Encoder as MinlzEncoder, crc::crc, decode};
use pin_project_lite::pin_project;
use rand::RngExt;
use rustfs_rio::{EtagResolvable, HashReaderDetector, HashReaderMut, Index, TryGetIndex};
use rustfs_utils::CompressionAlgorithm;
use std::cmp::min;
use std::fmt;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
const MAGIC_CHUNK: &[u8] = b"\xff\x06\x00\x00S2sTwO";
const 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;
pin_project! {
#[derive(Debug)]
pub struct CompressReader<R> {
#[pin]
inner: R,
buffer: Vec<u8>,
pos: usize,
done: bool,
block_size: usize,
index: Index,
written: usize,
uncompressed_written: usize,
temp_buffer: Vec<u8>,
read_buffer: Vec<u8>,
wrote_stream_header: bool,
padding_multiple: Option<usize>,
block_encoder: S2BlockEncoder,
}
}
struct S2BlockEncoder {
inner: MinlzEncoder,
}
impl S2BlockEncoder {
fn new() -> Self {
Self {
inner: MinlzEncoder::new(),
}
}
fn encode(&mut self, uncompressed: &[u8]) -> Vec<u8> {
self.inner.encode(uncompressed)
}
}
impl fmt::Debug for S2BlockEncoder {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("S2BlockEncoder").finish_non_exhaustive()
}
}
impl<R> CompressReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
pub fn new(inner: R, _compression_algorithm: CompressionAlgorithm) -> Self {
Self::with_block_size(inner, DEFAULT_BLOCK_SIZE, CompressionAlgorithm::default())
}
pub fn with_block_size(inner: R, block_size: usize, _compression_algorithm: CompressionAlgorithm) -> Self {
Self {
inner,
buffer: Vec::new(),
pos: 0,
done: false,
block_size,
index: Index::new(),
written: 0,
uncompressed_written: 0,
temp_buffer: Vec::with_capacity(block_size),
read_buffer: vec![0u8; block_size],
wrote_stream_header: false,
padding_multiple: None,
block_encoder: S2BlockEncoder::new(),
}
}
pub fn with_encrypted_padding(inner: R, _compression_algorithm: CompressionAlgorithm) -> Self {
let mut reader = Self::new(inner, CompressionAlgorithm::default());
reader.padding_multiple = Some(ENCRYPTED_PADDING_MULTIPLE);
reader
}
}
impl<R> TryGetIndex for CompressReader<R> {
fn try_get_index(&self) -> Option<&Index> {
(self.uncompressed_written > MIN_INDEX_SIZE).then_some(&self.index)
}
}
impl<R> AsyncRead for CompressReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
let mut this = self.project();
if *this.pos < this.buffer.len() {
let to_copy = min(buf.remaining(), this.buffer.len() - *this.pos);
buf.put_slice(&this.buffer[*this.pos..*this.pos + to_copy]);
*this.pos += to_copy;
if *this.pos == this.buffer.len() {
this.buffer.clear();
*this.pos = 0;
}
return Poll::Ready(Ok(()));
}
if *this.done {
return Poll::Ready(Ok(()));
}
while this.temp_buffer.len() < *this.block_size {
let remaining = *this.block_size - this.temp_buffer.len();
let mut read_buf = ReadBuf::new(&mut this.read_buffer[..remaining]);
match this.inner.as_mut().poll_read(cx, &mut read_buf) {
Poll::Pending => {
return Poll::Pending;
}
Poll::Ready(Ok(())) => {
let n = read_buf.filled().len();
if n == 0 {
break;
}
this.temp_buffer.extend_from_slice(read_buf.filled());
}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
}
}
if this.temp_buffer.is_empty() {
if let Some(padding_multiple) = *this.padding_multiple
&& let Some(padding_chunk) = build_padding_chunk(*this.written, padding_multiple)?
{
*this.written += padding_chunk.len();
this.index.total_compressed = *this.written as i64;
*this.buffer = padding_chunk;
*this.pos = 0;
*this.done = true;
let to_copy = min(buf.remaining(), this.buffer.len());
buf.put_slice(&this.buffer[..to_copy]);
*this.pos += to_copy;
if *this.pos == this.buffer.len() {
this.buffer.clear();
*this.pos = 0;
}
return Poll::Ready(Ok(()));
}
*this.done = true;
return Poll::Ready(Ok(()));
}
let mut out = Vec::new();
if !*this.wrote_stream_header {
out.extend_from_slice(MAGIC_CHUNK);
*this.written += MAGIC_CHUNK.len();
*this.wrote_stream_header = true;
}
if let Err(err) = this.index.add(*this.written as i64, *this.uncompressed_written as i64) {
return Poll::Ready(Err(err));
}
let block = build_s2_chunk(this.temp_buffer.as_slice(), this.block_encoder)?;
*this.uncompressed_written += this.temp_buffer.len();
*this.written += block.len();
this.index.total_uncompressed = *this.uncompressed_written as i64;
this.index.total_compressed = *this.written as i64;
out.extend_from_slice(&block);
this.temp_buffer.clear();
*this.buffer = out;
*this.pos = 0;
let to_copy = min(buf.remaining(), this.buffer.len());
buf.put_slice(&this.buffer[..to_copy]);
*this.pos += to_copy;
if *this.pos == this.buffer.len() {
this.buffer.clear();
*this.pos = 0;
}
Poll::Ready(Ok(()))
}
}
impl<R> EtagResolvable for CompressReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for CompressReader<R>
where
R: HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
pin_project! {
#[derive(Debug)]
pub struct DecompressReader<R> {
#[pin]
inner: R,
buffer: Vec<u8>,
buffer_pos: usize,
finished: bool,
header_buf: [u8; CHUNK_HEADER_LEN],
header_read: usize,
chunk_type: u8,
chunk_buf: Vec<u8>,
chunk_len: usize,
chunk_read: usize,
reading_chunk: bool,
stream_initialized: bool,
}
}
impl<R> DecompressReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
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,
}
}
}
impl<R> AsyncRead for DecompressReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
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(()));
}
}
}
impl<R> EtagResolvable for DecompressReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for DecompressReader<R>
where
R: HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
fn build_s2_chunk(uncompressed: &[u8], encoder: &mut S2BlockEncoder) -> io::Result<Vec<u8>> {
let compressed = encode_block(uncompressed, encoder);
let checksum = crc(uncompressed);
let dst_limit = uncompressed.len().saturating_sub(uncompressed.len() / 32).saturating_sub(5);
let (chunk_type, payload) = if compressed.len() <= dst_limit {
(CHUNK_TYPE_COMPRESSED_DATA, compressed)
} else {
(CHUNK_TYPE_UNCOMPRESSED_DATA, uncompressed.to_vec())
};
let chunk_len = payload.len() + CHECKSUM_SIZE;
if chunk_len > MAX_CHUNK_SIZE {
return Err(io::Error::new(io::ErrorKind::InvalidData, "S2 chunk exceeds 24-bit framing limit"));
}
let mut out = Vec::with_capacity(CHUNK_HEADER_LEN + chunk_len);
out.push(chunk_type);
out.push((chunk_len & 0xff) as u8);
out.push(((chunk_len >> 8) & 0xff) as u8);
out.push(((chunk_len >> 16) & 0xff) as u8);
out.extend_from_slice(&checksum.to_le_bytes());
out.extend_from_slice(&payload);
Ok(out)
}
fn encode_block(uncompressed: &[u8], encoder: &mut S2BlockEncoder) -> Vec<u8> {
encoder.encode(uncompressed)
}
fn build_padding_chunk(current_size: usize, padding_multiple: usize) -> io::Result<Option<Vec<u8>>> {
if padding_multiple == 0 || current_size.is_multiple_of(padding_multiple) {
return Ok(None);
}
let padding_len = (padding_multiple - ((current_size + CHUNK_HEADER_LEN) % padding_multiple)) % padding_multiple;
if padding_len > MAX_CHUNK_SIZE {
return Err(io::Error::new(io::ErrorKind::InvalidData, "S2 padding exceeds 24-bit framing limit"));
}
let mut out = Vec::with_capacity(CHUNK_HEADER_LEN + padding_len);
out.push(CHUNK_TYPE_PADDING);
out.push((padding_len & 0xff) as u8);
out.push(((padding_len >> 8) & 0xff) as u8);
out.push(((padding_len >> 16) & 0xff) as u8);
if padding_len > 0 {
let mut padding = vec![0u8; padding_len];
rand::rng().fill(padding.as_mut_slice());
out.extend_from_slice(&padding);
}
Ok(Some(out))
}
fn decode_chunk(chunk: &[u8], compressed: bool) -> io::Result<Vec<u8>> {
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 tokio::io::AsyncReadExt;
struct PendingAfterBytes<R> {
inner: R,
max_chunk: usize,
pending_next: bool,
}
impl<R> PendingAfterBytes<R> {
fn new(inner: R, max_chunk: usize) -> Self {
Self {
inner,
max_chunk,
pending_next: false,
}
}
}
impl<R: AsyncRead + Unpin> AsyncRead for PendingAfterBytes<R> {
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
if self.pending_next {
self.pending_next = false;
cx.waker().wake_by_ref();
return Poll::Pending;
}
let allowed = self.max_chunk.min(buf.remaining());
if allowed == 0 {
return Poll::Ready(Ok(()));
}
let mut scratch = vec![0u8; allowed];
let mut limited = ReadBuf::new(&mut scratch);
match Pin::new(&mut self.inner).poll_read(cx, &mut limited) {
Poll::Pending => Poll::Pending,
Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
Poll::Ready(Ok(())) => {
let filled = limited.filled();
if !filled.is_empty() {
buf.put_slice(filled);
self.pending_next = true;
}
Poll::Ready(Ok(()))
}
}
}
}
fn s2_chunk_types(stream: &[u8]) -> Vec<u8> {
let mut chunk_types = Vec::new();
let mut offset = 0usize;
while offset + CHUNK_HEADER_LEN <= stream.len() {
let chunk_type = stream[offset];
let chunk_len =
(stream[offset + 1] as usize) | ((stream[offset + 2] as usize) << 8) | ((stream[offset + 3] as usize) << 16);
chunk_types.push(chunk_type);
offset += CHUNK_HEADER_LEN + chunk_len;
}
chunk_types
}
#[test]
fn minlz_encoded_payload_decodes_with_minlz() {
let plaintext = b"compressible-rio-v2-block-".repeat(4096);
let mut encoder = S2BlockEncoder::new();
let compressed = encode_block(&plaintext, &mut encoder);
let decoded = decode(&compressed).expect("decode payload");
assert_eq!(decoded, plaintext);
}
#[tokio::test]
async fn s2_compress_reader_roundtrip() {
let plaintext = b"hello-rio-v2-s2-".repeat(32_768);
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
let mut compressed = Vec::new();
reader.read_to_end(&mut compressed).await.expect("read compressed data");
assert!(compressed.starts_with(MAGIC_CHUNK));
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
let mut actual = Vec::new();
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
assert_eq!(actual, plaintext);
}
#[tokio::test]
async fn s2_compress_reader_roundtrip_near_erasure_boundary() {
let size = 4 * 1024 * 1024 - 97;
let plaintext = pseudo_random_bytes(size);
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
let mut compressed = Vec::new();
reader.read_to_end(&mut compressed).await.expect("read compressed data");
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
let mut actual = Vec::new();
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
assert_eq!(actual, plaintext);
}
fn pseudo_random_bytes(size: usize) -> Vec<u8> {
(0..size)
.scan(0x9e37_79b9_7f4a_7c15u64, |state, _| {
*state ^= *state << 7;
*state ^= *state >> 9;
*state = state.wrapping_mul(0xbf58_476d_1ce4_e5b9);
Some((*state >> 32) as u8)
})
.collect()
}
#[tokio::test]
async fn s2_compress_reader_roundtrip_large_random() {
let plaintext = pseudo_random_bytes(8 * 1024 * 1024 + 123);
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
let mut compressed = Vec::new();
reader.read_to_end(&mut compressed).await.expect("read compressed data");
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
let mut actual = Vec::new();
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
assert_eq!(actual, plaintext);
}
#[tokio::test]
async fn s2_compress_reader_roundtrip_with_pending_source() {
let plaintext = pseudo_random_bytes(2 * 1024 * 1024 + 17);
let pending_reader = PendingAfterBytes::new(Cursor::new(plaintext.clone()), 257);
let mut reader = CompressReader::new(pending_reader, CompressionAlgorithm::default());
let mut compressed = Vec::new();
reader.read_to_end(&mut compressed).await.expect("read compressed data");
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
let mut actual = Vec::new();
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
assert_eq!(actual, plaintext);
}
#[tokio::test]
async fn s2_decompress_reader_returns_bytes_on_first_read() {
let plaintext = b"abcdefghijklmnopqrstuvwxyz".to_vec();
let mut compressed = Vec::new();
CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default())
.read_to_end(&mut compressed)
.await
.expect("compress plaintext");
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
let mut buf = [0u8; 64];
let n = decompressor.read(&mut buf).await.expect("read first decompressed chunk");
assert!(n > 0);
assert_eq!(&buf[..n], plaintext.as_slice());
}
#[tokio::test]
async fn s2_decompress_reader_resumes_chunk_body_after_pending() {
let plaintext = pseudo_random_bytes(1024 * 1024 + 123);
let mut compressed = Vec::new();
CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default())
.read_to_end(&mut compressed)
.await
.expect("compress plaintext");
let pending_reader = PendingAfterBytes::new(Cursor::new(compressed), 257);
let mut decompressor = DecompressReader::new(pending_reader, CompressionAlgorithm::default());
let mut actual = Vec::new();
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
assert_eq!(actual, plaintext);
}
#[tokio::test]
async fn s2_decompress_reader_handles_concatenated_streams_after_pending() {
let first = pseudo_random_bytes(16 * 1024 * 1024);
let second = pseudo_random_bytes(12 * 1024 * 1024 + 123);
let mut first_compressed = Vec::new();
CompressReader::new(Cursor::new(first.clone()), CompressionAlgorithm::default())
.read_to_end(&mut first_compressed)
.await
.expect("compress first stream");
let mut second_compressed = Vec::new();
CompressReader::new(Cursor::new(second.clone()), CompressionAlgorithm::default())
.read_to_end(&mut second_compressed)
.await
.expect("compress second stream");
first_compressed.extend_from_slice(&second_compressed);
let pending_reader = PendingAfterBytes::new(Cursor::new(first_compressed), 4096);
let mut decompressor = DecompressReader::new(pending_reader, CompressionAlgorithm::default());
let mut actual = Vec::new();
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
let mut expected = first;
expected.extend_from_slice(&second);
assert_eq!(actual, expected);
}
#[tokio::test]
async fn s2_compress_reader_with_encrypted_padding_emits_padding_frame() {
let plaintext = b"encrypted-padding-check-".repeat(8192);
let mut reader = CompressReader::with_encrypted_padding(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
let mut compressed = Vec::new();
reader.read_to_end(&mut compressed).await.expect("read compressed data");
assert_eq!(compressed.len() % ENCRYPTED_PADDING_MULTIPLE, 0);
assert!(s2_chunk_types(&compressed).contains(&CHUNK_TYPE_PADDING));
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
let mut actual = Vec::new();
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
assert_eq!(actual, plaintext);
}
#[tokio::test]
async fn s2_compress_reader_skips_index_for_small_streams() {
let plaintext = b"index-threshold-check-".repeat(16_384);
let mut reader = CompressReader::new(Cursor::new(plaintext.clone()), CompressionAlgorithm::default());
let mut compressed = Vec::new();
reader.read_to_end(&mut compressed).await.expect("read compressed data");
assert!(reader.try_get_index().is_none());
let mut decompressor = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
let mut actual = Vec::new();
decompressor.read_to_end(&mut actual).await.expect("read decompressed data");
assert_eq!(actual, plaintext);
}
}
+759
View File
@@ -0,0 +1,759 @@
// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use aes_gcm::aead::Aead;
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
use hmac::{Hmac, Mac};
use pin_project_lite::pin_project;
use rand::RngExt;
use rustfs_rio::{EtagResolvable, HashReaderDetector, HashReaderMut, Index, TryGetIndex, multipart_part_nonce};
use sha2::Sha256;
use std::cmp::min;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
const DARE_VERSION_20: u8 = 0x20;
const DARE_CIPHER_AES_256_GCM: u8 = 0x00;
const DARE_HEADER_SIZE: usize = 16;
const DARE_TAG_SIZE: usize = 16;
const DARE_PAYLOAD_SIZE: usize = 64 * 1024;
type HmacSha256 = Hmac<Sha256>;
#[derive(Clone, Copy)]
enum MultipartKeySource {
LegacyNonce { base_nonce: [u8; 12] },
ObjectKey { object_key: [u8; 32] },
}
pin_project! {
pub struct EncryptReader<R> {
#[pin]
inner: R,
cipher: Aes256Gcm,
base_nonce: [u8; 12],
sequence_number: u32,
temp_buffer: Vec<u8>,
read_buffer: Vec<u8>,
output_buffer: Vec<u8>,
output_pos: usize,
finished: bool,
}
}
impl<R> EncryptReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
pub fn new_with_object_key(inner: R, object_key: [u8; 32]) -> Self {
Self::new(inner, object_key, random_stream_nonce())
}
pub fn new(inner: R, key: [u8; 32], nonce: [u8; 12]) -> Self {
Self::new_with_sequence(inner, key, nonce, 0)
}
pub fn new_with_sequence(inner: R, key: [u8; 32], nonce: [u8; 12], sequence_number: u32) -> Self {
Self {
inner,
cipher: Aes256Gcm::new_from_slice(&key).expect("valid AES-256-GCM key"),
base_nonce: nonce,
sequence_number,
temp_buffer: Vec::with_capacity(DARE_PAYLOAD_SIZE + 1),
read_buffer: vec![0u8; DARE_PAYLOAD_SIZE + 1],
output_buffer: Vec::new(),
output_pos: 0,
finished: false,
}
}
pub fn new_multipart(inner: R, key: [u8; 32], base_nonce: [u8; 12], part_number: usize) -> Self {
Self::new(inner, key, multipart_part_nonce(base_nonce, part_number))
}
pub fn new_multipart_with_object_key(inner: R, object_key: [u8; 32], part_number: u32) -> Self {
Self::new(inner, derive_part_key(object_key, part_number), random_stream_nonce())
}
pub fn new_multipart_with_sequence(
inner: R,
key: [u8; 32],
base_nonce: [u8; 12],
part_number: usize,
sequence_number: u32,
) -> Self {
Self::new_with_sequence(inner, key, multipart_part_nonce(base_nonce, part_number), sequence_number)
}
}
impl<R> AsyncRead for EncryptReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
let mut this = self.project();
if *this.output_pos < this.output_buffer.len() {
let to_copy = min(buf.remaining(), this.output_buffer.len() - *this.output_pos);
buf.put_slice(&this.output_buffer[*this.output_pos..*this.output_pos + to_copy]);
*this.output_pos += to_copy;
if *this.output_pos == this.output_buffer.len() {
this.output_buffer.clear();
*this.output_pos = 0;
}
return Poll::Ready(Ok(()));
}
if *this.finished {
return Poll::Ready(Ok(()));
}
loop {
if this.temp_buffer.len() > DARE_PAYLOAD_SIZE {
let package = build_dare_package(
this.cipher,
*this.sequence_number,
*this.base_nonce,
&this.temp_buffer[..DARE_PAYLOAD_SIZE],
false,
)?;
let carry = this.temp_buffer.split_off(DARE_PAYLOAD_SIZE);
*this.temp_buffer = carry;
*this.sequence_number = this.sequence_number.wrapping_add(1);
*this.output_buffer = package;
break;
}
let remaining = DARE_PAYLOAD_SIZE + 1 - this.temp_buffer.len();
let mut read_buf = ReadBuf::new(&mut this.read_buffer[..remaining]);
match this.inner.as_mut().poll_read(cx, &mut read_buf) {
Poll::Pending => {
if this.temp_buffer.len() > DARE_PAYLOAD_SIZE {
continue;
}
return Poll::Pending;
}
Poll::Ready(Ok(())) => {
let n = read_buf.filled().len();
if n == 0 {
if this.temp_buffer.is_empty() {
*this.finished = true;
return Poll::Ready(Ok(()));
}
let package =
build_dare_package(this.cipher, *this.sequence_number, *this.base_nonce, this.temp_buffer, true)?;
this.temp_buffer.clear();
*this.sequence_number = this.sequence_number.wrapping_add(1);
*this.output_buffer = package;
*this.finished = true;
break;
}
this.temp_buffer.extend_from_slice(read_buf.filled());
}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
}
}
let to_copy = min(buf.remaining(), this.output_buffer.len());
buf.put_slice(&this.output_buffer[..to_copy]);
*this.output_pos += to_copy;
if *this.output_pos == this.output_buffer.len() {
this.output_buffer.clear();
*this.output_pos = 0;
}
Poll::Ready(Ok(()))
}
}
impl<R> EtagResolvable for EncryptReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for EncryptReader<R>
where
R: HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
impl<R> TryGetIndex for EncryptReader<R>
where
R: TryGetIndex,
{
fn try_get_index(&self) -> Option<&Index> {
self.inner.try_get_index()
}
}
pin_project! {
pub struct DecryptReader<R> {
#[pin]
inner: R,
cipher: Aes256Gcm,
expected_base_nonce: Option<[u8; 12]>,
sequence_number: u32,
multipart_parts: Vec<usize>,
current_part_index: usize,
multipart_mode: bool,
multipart_key_source: Option<MultipartKeySource>,
header_buf: [u8; DARE_HEADER_SIZE],
header_read: usize,
ciphertext_buf: Vec<u8>,
ciphertext_len: usize,
ciphertext_read: usize,
ref_nonce: Option<[u8; 12]>,
finalized: bool,
finished: bool,
output_buffer: Vec<u8>,
output_pos: usize,
}
}
impl<R> DecryptReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
pub fn new_with_object_key(inner: R, object_key: [u8; 32]) -> Self {
Self::new_with_object_key_and_sequence(inner, object_key, 0)
}
pub fn new(inner: R, key: [u8; 32], nonce: [u8; 12]) -> Self {
Self::new_with_sequence(inner, key, nonce, 0)
}
pub fn new_with_object_key_and_sequence(inner: R, object_key: [u8; 32], sequence_number: u32) -> Self {
Self {
inner,
cipher: Aes256Gcm::new_from_slice(&object_key).expect("valid AES-256-GCM key"),
expected_base_nonce: None,
sequence_number,
multipart_parts: Vec::new(),
current_part_index: 0,
multipart_mode: false,
multipart_key_source: None,
header_buf: [0u8; DARE_HEADER_SIZE],
header_read: 0,
ciphertext_buf: Vec::new(),
ciphertext_len: 0,
ciphertext_read: 0,
ref_nonce: None,
finalized: false,
finished: false,
output_buffer: Vec::new(),
output_pos: 0,
}
}
pub fn new_with_sequence(inner: R, key: [u8; 32], nonce: [u8; 12], sequence_number: u32) -> Self {
Self {
inner,
cipher: Aes256Gcm::new_from_slice(&key).expect("valid AES-256-GCM key"),
expected_base_nonce: Some(nonce),
sequence_number,
multipart_parts: Vec::new(),
current_part_index: 0,
multipart_mode: false,
multipart_key_source: None,
header_buf: [0u8; DARE_HEADER_SIZE],
header_read: 0,
ciphertext_buf: Vec::new(),
ciphertext_len: 0,
ciphertext_read: 0,
ref_nonce: None,
finalized: false,
finished: false,
output_buffer: Vec::new(),
output_pos: 0,
}
}
pub fn new_multipart(inner: R, key: [u8; 32], base_nonce: [u8; 12], multipart_parts: Vec<usize>) -> Self {
Self::new_multipart_with_sequence(inner, key, base_nonce, multipart_parts, 0)
}
pub fn new_multipart_with_object_key(inner: R, object_key: [u8; 32], multipart_parts: Vec<usize>) -> Self {
Self::new_multipart_with_object_key_and_sequence(inner, object_key, multipart_parts, 0)
}
pub fn new_multipart_with_sequence(
inner: R,
key: [u8; 32],
base_nonce: [u8; 12],
multipart_parts: Vec<usize>,
sequence_number: u32,
) -> Self {
let first_part = multipart_parts.first().copied().unwrap_or(1);
Self {
inner,
cipher: Aes256Gcm::new_from_slice(&key).expect("valid AES-256-GCM key"),
expected_base_nonce: Some(multipart_part_nonce(base_nonce, first_part)),
sequence_number,
multipart_parts,
current_part_index: 0,
multipart_mode: true,
multipart_key_source: Some(MultipartKeySource::LegacyNonce { base_nonce }),
header_buf: [0u8; DARE_HEADER_SIZE],
header_read: 0,
ciphertext_buf: Vec::new(),
ciphertext_len: 0,
ciphertext_read: 0,
ref_nonce: None,
finalized: false,
finished: false,
output_buffer: Vec::new(),
output_pos: 0,
}
}
pub fn new_multipart_with_object_key_and_sequence(
inner: R,
object_key: [u8; 32],
multipart_parts: Vec<usize>,
sequence_number: u32,
) -> Self {
let first_part = multipart_parts.first().copied().unwrap_or(1);
let first_key = derive_part_key(object_key, first_part as u32);
Self {
inner,
cipher: Aes256Gcm::new_from_slice(&first_key).expect("valid AES-256-GCM key"),
expected_base_nonce: None,
sequence_number,
multipart_parts,
current_part_index: 0,
multipart_mode: true,
multipart_key_source: Some(MultipartKeySource::ObjectKey { object_key }),
header_buf: [0u8; DARE_HEADER_SIZE],
header_read: 0,
ciphertext_buf: Vec::new(),
ciphertext_len: 0,
ciphertext_read: 0,
ref_nonce: None,
finalized: false,
finished: false,
output_buffer: Vec::new(),
output_pos: 0,
}
}
}
impl<R> AsyncRead for DecryptReader<R>
where
R: AsyncRead + Unpin + Send + Sync,
{
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
let mut this = self.project();
if *this.output_pos < this.output_buffer.len() {
let to_copy = min(buf.remaining(), this.output_buffer.len() - *this.output_pos);
buf.put_slice(&this.output_buffer[*this.output_pos..*this.output_pos + to_copy]);
*this.output_pos += to_copy;
if *this.output_pos == this.output_buffer.len() {
this.output_buffer.clear();
*this.output_pos = 0;
}
return Poll::Ready(Ok(()));
}
loop {
if *this.finished {
return Poll::Ready(Ok(()));
}
if *this.finalized {
if !this.multipart_parts.is_empty() && *this.current_part_index + 1 < this.multipart_parts.len() {
let next_part = this.multipart_parts[*this.current_part_index + 1];
*this.current_part_index += 1;
match this.multipart_key_source.as_ref().copied() {
Some(MultipartKeySource::LegacyNonce { base_nonce }) => {
*this.expected_base_nonce = Some(multipart_part_nonce(base_nonce, next_part));
}
Some(MultipartKeySource::ObjectKey { object_key }) => {
let part_key = derive_part_key(object_key, next_part as u32);
*this.cipher = Aes256Gcm::new_from_slice(&part_key).expect("valid AES-256-GCM key");
*this.expected_base_nonce = None;
}
None => {}
}
*this.sequence_number = 0;
*this.ref_nonce = None;
*this.finalized = false;
continue;
}
*this.finished = true;
return Poll::Ready(Ok(()));
}
while *this.header_read < DARE_HEADER_SIZE {
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 DARE header",
)));
}
*this.header_read += n;
}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
}
}
let header = this.header_buf;
let payload_len = usize::from(u16::from_le_bytes([header[2], header[3]])) + 1;
let package_len = payload_len + DARE_TAG_SIZE;
if payload_len == 0 || payload_len > DARE_PAYLOAD_SIZE {
return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "invalid DARE payload size")));
}
if !is_final_header(*header) && payload_len != DARE_PAYLOAD_SIZE {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::InvalidData,
"non-final DARE package must carry a full 64KiB payload",
)));
}
if this.ciphertext_buf.len() < package_len {
this.ciphertext_buf.resize(package_len, 0);
}
*this.ciphertext_len = package_len;
*this.ciphertext_read = 0;
while *this.ciphertext_read < *this.ciphertext_len {
let mut read_buf = ReadBuf::new(&mut this.ciphertext_buf[*this.ciphertext_read..*this.ciphertext_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 DARE ciphertext",
)));
}
*this.ciphertext_read += n;
}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
}
}
match open_dare_package(
this.cipher,
*this.sequence_number,
*this.expected_base_nonce,
*header,
&this.ciphertext_buf[..*this.ciphertext_len],
this.ref_nonce,
) {
Ok((plaintext, ref_nonce, finalized)) => {
*this.ref_nonce = Some(ref_nonce);
*this.finalized = finalized;
*this.header_read = 0;
*this.ciphertext_len = 0;
*this.ciphertext_read = 0;
*this.sequence_number = this.sequence_number.wrapping_add(1);
*this.output_buffer = plaintext;
let to_copy = min(buf.remaining(), this.output_buffer.len());
buf.put_slice(&this.output_buffer[..to_copy]);
*this.output_pos += to_copy;
if *this.output_pos == this.output_buffer.len() {
this.output_buffer.clear();
*this.output_pos = 0;
}
return Poll::Ready(Ok(()));
}
Err(err) => return Poll::Ready(Err(err)),
}
}
}
}
impl<R> EtagResolvable for DecryptReader<R>
where
R: EtagResolvable,
{
fn try_resolve_etag(&mut self) -> Option<String> {
self.inner.try_resolve_etag()
}
}
impl<R> HashReaderDetector for DecryptReader<R>
where
R: HashReaderDetector,
{
fn is_hash_reader(&self) -> bool {
self.inner.is_hash_reader()
}
fn as_hash_reader_mut(&mut self) -> Option<&mut dyn HashReaderMut> {
self.inner.as_hash_reader_mut()
}
}
impl<R> TryGetIndex for DecryptReader<R>
where
R: TryGetIndex,
{
fn try_get_index(&self) -> Option<&Index> {
self.inner.try_get_index()
}
}
fn build_dare_package(
cipher: &Aes256Gcm,
sequence_number: u32,
base_nonce: [u8; 12],
plaintext: &[u8],
final_package: bool,
) -> io::Result<Vec<u8>> {
let mut header = [0u8; DARE_HEADER_SIZE];
header[0] = DARE_VERSION_20;
header[1] = DARE_CIPHER_AES_256_GCM;
header[2..4].copy_from_slice(
&u16::try_from(plaintext.len() - 1)
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "DARE payload too large for Version 2.0 framing"))?
.to_le_bytes(),
);
header[4..16].copy_from_slice(&base_nonce);
if final_package {
header[4] |= 0x80;
} else {
header[4] &= 0x7F;
}
let mut package_nonce = header[4..16].try_into().expect("nonce slice");
xor_sequence_into_nonce(&mut package_nonce, sequence_number);
let nonce = Nonce::try_from(package_nonce.as_slice())
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid DARE nonce length"))?;
let ciphertext = cipher
.encrypt(
&nonce,
aes_gcm::aead::Payload {
msg: plaintext,
aad: &header[..4],
},
)
.map_err(|err| io::Error::other(format!("failed to encrypt DARE package: {err}")))?;
let mut package = Vec::with_capacity(DARE_HEADER_SIZE + ciphertext.len());
package.extend_from_slice(&header);
package.extend_from_slice(&ciphertext);
Ok(package)
}
fn open_dare_package(
cipher: &Aes256Gcm,
sequence_number: u32,
expected_base_nonce: Option<[u8; 12]>,
header: [u8; DARE_HEADER_SIZE],
ciphertext: &[u8],
ref_nonce: &mut Option<[u8; 12]>,
) -> io::Result<(Vec<u8>, [u8; 12], bool)> {
if header[0] != DARE_VERSION_20 {
return Err(io::Error::new(io::ErrorKind::InvalidData, "unsupported DARE version"));
}
if header[1] != DARE_CIPHER_AES_256_GCM {
return Err(io::Error::new(io::ErrorKind::InvalidData, "unsupported DARE cipher suite"));
}
let header_nonce: [u8; 12] = header[4..16].try_into().expect("nonce slice");
if let Some(expected_base_nonce) = expected_base_nonce {
let masked_expected = apply_final_flag(expected_base_nonce, is_final_header(header));
if header_nonce != masked_expected {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"DARE package nonce does not match the configured stream nonce",
));
}
}
let current_ref = ref_nonce.get_or_insert(header_nonce);
let expected_ref = apply_final_flag(*current_ref, is_final_header(header));
if header_nonce != expected_ref {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"DARE package nonce does not match the stream reference nonce",
));
}
let mut package_nonce = header_nonce;
xor_sequence_into_nonce(&mut package_nonce, sequence_number);
let nonce = Nonce::try_from(package_nonce.as_slice())
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "invalid DARE nonce length"))?;
let plaintext = cipher
.decrypt(
&nonce,
aes_gcm::aead::Payload {
msg: ciphertext,
aad: &header[..4],
},
)
.map_err(|err| io::Error::new(io::ErrorKind::InvalidData, format!("DARE authentication failed: {err}")))?;
Ok((plaintext, *current_ref, is_final_header(header)))
}
fn xor_sequence_into_nonce(nonce: &mut [u8; 12], sequence_number: u32) {
let last = u32::from_le_bytes([nonce[8], nonce[9], nonce[10], nonce[11]]) ^ sequence_number;
nonce[8..12].copy_from_slice(&last.to_le_bytes());
}
fn apply_final_flag(mut nonce: [u8; 12], final_package: bool) -> [u8; 12] {
if final_package {
nonce[0] |= 0x80;
} else {
nonce[0] &= 0x7F;
}
nonce
}
fn is_final_header(header: [u8; DARE_HEADER_SIZE]) -> bool {
header[4] & 0x80 != 0
}
pub fn derive_part_key(object_key: [u8; 32], part_number: u32) -> [u8; 32] {
let mut mac = HmacSha256::new_from_slice(&object_key).expect("HMAC-SHA256 accepts 32-byte object keys");
mac.update(&part_number.to_le_bytes());
let mut part_key = [0u8; 32];
part_key.copy_from_slice(mac.finalize().into_bytes().as_slice());
part_key
}
fn random_stream_nonce() -> [u8; 12] {
let mut nonce = [0u8; 12];
rand::rng().fill(&mut nonce);
nonce
}
#[cfg(test)]
mod tests {
use super::*;
use hex::encode as hex_encode;
use std::io::Cursor;
use tokio::io::AsyncReadExt;
const DARE_PACKAGE_SIZE: usize = DARE_HEADER_SIZE + DARE_PAYLOAD_SIZE + DARE_TAG_SIZE;
#[tokio::test]
async fn decrypt_reader_can_start_from_non_zero_sequence_number() {
let plaintext = vec![0xAB; DARE_PAYLOAD_SIZE * 2 + 19];
let key_bytes = [0x44u8; 32];
let base_nonce = [0x66u8; 12];
let mut encrypted = Vec::new();
EncryptReader::new(Cursor::new(plaintext.clone()), key_bytes, base_nonce)
.read_to_end(&mut encrypted)
.await
.expect("encrypt plaintext");
let tail = encrypted[DARE_PACKAGE_SIZE..].to_vec();
let mut decrypted = Vec::new();
DecryptReader::new_with_sequence(Cursor::new(tail), key_bytes, base_nonce, 1)
.read_to_end(&mut decrypted)
.await
.expect("decrypt tail packages with non-zero sequence");
assert_eq!(decrypted, plaintext[DARE_PAYLOAD_SIZE..]);
}
#[tokio::test]
async fn singlepart_object_key_roundtrip_uses_header_nonce() {
let object_key = [0x51u8; 32];
let plaintext = vec![0xAC; DARE_PAYLOAD_SIZE + 17];
let mut encrypted = Vec::new();
EncryptReader::new_with_object_key(Cursor::new(plaintext.clone()), object_key)
.read_to_end(&mut encrypted)
.await
.expect("encrypt object-key singlepart stream");
let mut decrypted = Vec::new();
DecryptReader::new_with_object_key(Cursor::new(encrypted), object_key)
.read_to_end(&mut decrypted)
.await
.expect("decrypt object-key singlepart stream");
assert_eq!(decrypted, plaintext);
}
#[test]
fn derive_part_key_matches_minio_test_vectors() {
assert_eq!(
hex_encode(derive_part_key([0u8; 32], 0)),
"aa7855e13839dd767cd5da7c1ff5036540c9264b7a803029315e55375287b4af"
);
assert_eq!(
hex_encode(derive_part_key([0u8; 32], 1)),
"a3e7181c6eed030fd52f79537c56c4d07da92e56d374ff1dd2043350785b37d8"
);
assert_eq!(
hex_encode(derive_part_key([0u8; 32], 10_000)),
"f86e65c396ed52d204ee44bd1a0bbd86eb8b01b7354e67a3b3ae0e34dd5bd115"
);
}
#[tokio::test]
async fn multipart_object_key_roundtrip_resets_sequence_per_part() {
let object_key = [0x19u8; 32];
let part_one_plaintext = vec![0xA1; DARE_PAYLOAD_SIZE + 31];
let part_two_plaintext = vec![0xB2; DARE_PAYLOAD_SIZE * 2 + 7];
let mut encrypted_one = Vec::new();
EncryptReader::new_multipart_with_object_key(Cursor::new(part_one_plaintext.clone()), object_key, 1)
.read_to_end(&mut encrypted_one)
.await
.expect("encrypt multipart part one with object key");
let mut encrypted_two = Vec::new();
EncryptReader::new_multipart_with_object_key(Cursor::new(part_two_plaintext.clone()), object_key, 2)
.read_to_end(&mut encrypted_two)
.await
.expect("encrypt multipart part two with object key");
let mut encrypted = encrypted_one.clone();
encrypted.extend_from_slice(&encrypted_two);
let mut decrypted = Vec::new();
DecryptReader::new_multipart_with_object_key(Cursor::new(encrypted), object_key, vec![1, 2])
.read_to_end(&mut decrypted)
.await
.expect("decrypt multipart object-key stream");
let mut expected = part_one_plaintext;
expected.extend_from_slice(&part_two_plaintext);
assert_eq!(decrypted, expected);
assert_ne!(encrypted_one, encrypted_two[..encrypted_one.len()]);
}
}
+91
View File
@@ -0,0 +1,91 @@
// 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.
//! rio_v2 selectively replaces legacy rio components while keeping the
//! remaining API surface stable for feature-gated integration.
mod compress_reader;
mod encrypt_reader;
mod s2_index;
pub use compress_reader::{CompressReader, DecompressReader};
pub use encrypt_reader::{DecryptReader, EncryptReader, derive_part_key};
pub use rustfs_rio::DEFAULT_ENCRYPTION_BLOCK_SIZE;
pub use rustfs_rio::DynReader;
pub use rustfs_rio::EtagReader;
pub use rustfs_rio::EtagResolvable;
pub use rustfs_rio::HardLimitReader;
pub use rustfs_rio::HashReader;
pub use rustfs_rio::HashReaderDetector;
pub use rustfs_rio::HashReaderMut;
pub use rustfs_rio::Index;
pub use rustfs_rio::LimitReader;
pub use rustfs_rio::ReadStream;
pub use rustfs_rio::Reader;
pub use rustfs_rio::ReaderCapabilities;
pub use rustfs_rio::TryGetIndex;
pub use rustfs_rio::WarpReader;
pub use rustfs_rio::boxed_reader;
pub use rustfs_rio::read_checksums;
pub use rustfs_rio::resolve_etag_generic;
pub use rustfs_rio::wrap_reader;
pub use s2_index::{decode_minio_index_bytes, minio_index_storage_bytes};
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
use tokio::io::AsyncReadExt;
const DARE_VERSION_20: u8 = 0x20;
const DARE_HEADER_SIZE: usize = 16;
const DARE_TAG_SIZE: usize = 16;
const DARE_PAYLOAD_SIZE: usize = 64 * 1024;
const DARE_PACKAGE_SIZE: usize = DARE_HEADER_SIZE + DARE_PAYLOAD_SIZE + DARE_TAG_SIZE;
#[tokio::test]
async fn encrypt_reader_emits_dare_v2_packages() {
let plaintext = vec![0x5Au8; DARE_PAYLOAD_SIZE + 17];
let key_bytes = [0x11u8; 32];
let base_nonce = [0x22u8; 12];
let mut encrypted = Vec::new();
EncryptReader::new(Cursor::new(plaintext), key_bytes, base_nonce)
.read_to_end(&mut encrypted)
.await
.expect("encrypt plaintext into rio_v2 stream");
assert!(
encrypted.len() > DARE_PACKAGE_SIZE,
"expected at least one full DARE package and one final package"
);
assert_eq!(encrypted[0], DARE_VERSION_20, "rio_v2 encrypted streams must start with a DARE V2 header");
assert_eq!(
&encrypted[4..16],
&base_nonce,
"rio_v2 should preserve the configured nonce in the first DARE header"
);
let second_header_offset = DARE_PACKAGE_SIZE;
assert_eq!(
encrypted[second_header_offset], DARE_VERSION_20,
"rio_v2 should emit subsequent DARE V2 package headers at 64KiB boundaries"
);
assert_ne!(
encrypted[second_header_offset + 4] & 0x80,
0,
"the final DARE package must set the final flag"
);
}
}
+434
View File
@@ -0,0 +1,434 @@
// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use bytes::Bytes;
use rustfs_rio::Index;
use serde::Deserialize;
use std::io;
const S2_INDEX_HEADER: &[u8] = b"s2idx\x00";
const S2_INDEX_TRAILER: &[u8] = b"\x00xdi2s";
const CHUNK_TYPE_INDEX: u8 = 0x99;
const SKIPPABLE_FRAME_HEADER: usize = 4;
const MAX_INDEX_ENTRIES: usize = 1 << 16;
#[derive(Debug, Deserialize)]
struct LegacyIndexJson {
total_uncompressed: i64,
total_compressed: i64,
offsets: Vec<LegacyIndexOffset>,
est_block_uncompressed: i64,
}
#[derive(Debug, Deserialize)]
struct LegacyIndexOffset {
compressed: i64,
uncompressed: i64,
}
#[derive(Debug, Clone)]
struct S2IndexInfo {
compressed_offset: i64,
uncompressed_offset: i64,
}
#[derive(Debug, Clone)]
struct S2Index {
total_uncompressed: i64,
total_compressed: i64,
est_block_uncompressed: i64,
info: Vec<S2IndexInfo>,
}
pub fn minio_index_storage_bytes(index: &Index) -> Bytes {
let decoded = legacy_index_to_s2_index(index).unwrap_or_else(|_| S2Index {
total_uncompressed: index.total_uncompressed,
total_compressed: index.total_compressed,
est_block_uncompressed: 0,
info: Vec::new(),
});
let encoded = decoded.into_full_bytes();
remove_index_headers(encoded.as_ref())
.map(Bytes::copy_from_slice)
.unwrap_or(encoded)
}
pub fn decode_minio_index_bytes(bytes: &Bytes) -> Option<Index> {
let decoded = S2Index::load(bytes.as_ref())
.or_else(|_| S2Index::load(&restore_index_headers(bytes.as_ref())))
.ok()?;
let mut index = Index::new();
for info in decoded.info {
index.add(info.compressed_offset, info.uncompressed_offset).ok()?;
}
index.total_uncompressed = decoded.total_uncompressed;
index.total_compressed = decoded.total_compressed;
Some(index)
}
fn legacy_index_to_s2_index(index: &Index) -> io::Result<S2Index> {
let json = index
.to_json()
.map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
let decoded: LegacyIndexJson =
serde_json::from_slice(&json).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
Ok(S2Index {
total_uncompressed: decoded.total_uncompressed,
total_compressed: decoded.total_compressed,
est_block_uncompressed: decoded.est_block_uncompressed,
info: decoded
.offsets
.into_iter()
.map(|offset| S2IndexInfo {
compressed_offset: offset.compressed,
uncompressed_offset: offset.uncompressed,
})
.collect(),
})
}
impl S2Index {
fn into_full_bytes(self) -> Bytes {
let mut out = Vec::new();
out.extend_from_slice(&[CHUNK_TYPE_INDEX, 0, 0, 0]);
out.extend_from_slice(S2_INDEX_HEADER);
write_varint(&mut out, self.total_uncompressed);
write_varint(&mut out, self.total_compressed);
write_varint(&mut out, self.est_block_uncompressed);
write_varint(&mut out, self.info.len() as i64);
let has_uncompressed = self.has_explicit_uncompressed_offsets();
out.push(u8::from(has_uncompressed));
if has_uncompressed {
for (idx, info) in self.info.iter().enumerate() {
let mut offset = info.uncompressed_offset;
if idx > 0 {
let prev = &self.info[idx - 1];
offset -= prev.uncompressed_offset + self.est_block_uncompressed;
}
write_varint(&mut out, offset);
}
}
let mut compressed_predict = self.est_block_uncompressed / 2;
for (idx, info) in self.info.iter().enumerate() {
let mut offset = info.compressed_offset;
if idx > 0 {
let prev = &self.info[idx - 1];
offset -= prev.compressed_offset + compressed_predict;
compressed_predict += offset / 2;
}
write_varint(&mut out, offset);
}
let mut total_size = [0u8; 4];
total_size.copy_from_slice(&((out.len() + 4 + S2_INDEX_TRAILER.len()) as u32).to_le_bytes());
out.extend_from_slice(&total_size);
out.extend_from_slice(S2_INDEX_TRAILER);
let chunk_len = out.len() - SKIPPABLE_FRAME_HEADER;
out[1] = chunk_len as u8;
out[2] = (chunk_len >> 8) as u8;
out[3] = (chunk_len >> 16) as u8;
Bytes::from(out)
}
fn has_explicit_uncompressed_offsets(&self) -> bool {
for (idx, info) in self.info.iter().enumerate() {
if idx == 0 {
if info.uncompressed_offset != 0 {
return true;
}
continue;
}
if info.uncompressed_offset != self.info[idx - 1].uncompressed_offset + self.est_block_uncompressed {
return true;
}
}
false
}
fn load(mut bytes: &[u8]) -> io::Result<Self> {
if bytes.len() <= SKIPPABLE_FRAME_HEADER + S2_INDEX_HEADER.len() + S2_INDEX_TRAILER.len() {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buffer too small"));
}
if bytes[0] != CHUNK_TYPE_INDEX {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid index chunk type"));
}
let chunk_len = (bytes[1] as usize) | ((bytes[2] as usize) << 8) | ((bytes[3] as usize) << 16);
bytes = &bytes[SKIPPABLE_FRAME_HEADER..];
if bytes.len() < chunk_len {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buffer too small"));
}
bytes = &bytes[..chunk_len];
if !bytes.starts_with(S2_INDEX_HEADER) {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid index header"));
}
bytes = &bytes[S2_INDEX_HEADER.len()..];
let (total_uncompressed, used) = read_varint(bytes)?;
if total_uncompressed < 0 {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid uncompressed size"));
}
bytes = &bytes[used..];
let (total_compressed, used) = read_varint(bytes)?;
bytes = &bytes[used..];
let (est_block_uncompressed, used) = read_varint(bytes)?;
if est_block_uncompressed < 0 {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid block size"));
}
bytes = &bytes[used..];
let (entries, used) = read_varint(bytes)?;
if entries < 0 || entries > MAX_INDEX_ENTRIES as i64 {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid number of entries"));
}
bytes = &bytes[used..];
if bytes.is_empty() {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buffer too small"));
}
let has_uncompressed = bytes[0];
if has_uncompressed & 1 != has_uncompressed {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid uncompressed flag"));
}
bytes = &bytes[1..];
let mut info = vec![
S2IndexInfo {
compressed_offset: 0,
uncompressed_offset: 0,
};
entries as usize
];
for idx in 0..info.len() {
let mut uncompressed_offset = 0_i64;
if has_uncompressed != 0 {
let (value, used) = read_varint(bytes)?;
uncompressed_offset = value;
bytes = &bytes[used..];
}
if idx > 0 {
let prev = info[idx - 1].uncompressed_offset;
uncompressed_offset += prev + est_block_uncompressed;
if uncompressed_offset <= prev {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid uncompressed offset"));
}
}
if uncompressed_offset < 0 {
return Err(io::Error::new(io::ErrorKind::InvalidData, "negative uncompressed offset"));
}
info[idx].uncompressed_offset = uncompressed_offset;
}
let mut compressed_predict = est_block_uncompressed / 2;
for idx in 0..info.len() {
let (mut compressed_offset, used) = read_varint(bytes)?;
bytes = &bytes[used..];
if idx > 0 {
let next_compressed_predict = compressed_predict + compressed_offset / 2;
let prev = info[idx - 1].compressed_offset;
compressed_offset += prev + compressed_predict;
if compressed_offset <= prev {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid compressed offset"));
}
compressed_predict = next_compressed_predict;
}
if compressed_offset < 0 {
return Err(io::Error::new(io::ErrorKind::InvalidData, "negative compressed offset"));
}
info[idx].compressed_offset = compressed_offset;
}
if bytes.len() < 4 + S2_INDEX_TRAILER.len() {
return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buffer too small"));
}
bytes = &bytes[4..];
if !bytes.starts_with(S2_INDEX_TRAILER) {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid index trailer"));
}
Ok(Self {
total_uncompressed,
total_compressed,
est_block_uncompressed,
info,
})
}
}
fn remove_index_headers(bytes: &[u8]) -> Option<&[u8]> {
let save = SKIPPABLE_FRAME_HEADER + S2_INDEX_HEADER.len() + S2_INDEX_TRAILER.len() + 4;
if bytes.len() <= save || bytes[0] != CHUNK_TYPE_INDEX {
return None;
}
let chunk_len = (bytes[1] as usize) | ((bytes[2] as usize) << 8) | ((bytes[3] as usize) << 16);
let bytes = &bytes[SKIPPABLE_FRAME_HEADER..];
if bytes.len() < chunk_len {
return None;
}
let bytes = &bytes[..chunk_len];
let bytes = bytes.strip_prefix(S2_INDEX_HEADER)?;
let bytes = bytes.strip_suffix(S2_INDEX_TRAILER)?;
if bytes.len() < 4 {
return None;
}
Some(&bytes[..bytes.len() - 4])
}
fn restore_index_headers(input: &[u8]) -> Vec<u8> {
if input.is_empty() {
return Vec::new();
}
let mut bytes = Vec::with_capacity(SKIPPABLE_FRAME_HEADER + S2_INDEX_HEADER.len() + input.len() + 4 + S2_INDEX_TRAILER.len());
bytes.extend_from_slice(&[CHUNK_TYPE_INDEX, 0, 0, 0]);
bytes.extend_from_slice(S2_INDEX_HEADER);
bytes.extend_from_slice(input);
bytes.extend_from_slice(&((bytes.len() + 4 + S2_INDEX_TRAILER.len()) as u32).to_le_bytes());
bytes.extend_from_slice(S2_INDEX_TRAILER);
let chunk_len = bytes.len() - SKIPPABLE_FRAME_HEADER;
bytes[1] = chunk_len as u8;
bytes[2] = (chunk_len >> 8) as u8;
bytes[3] = (chunk_len >> 16) as u8;
bytes
}
fn write_varint(out: &mut Vec<u8>, value: i64) {
let mut unsigned = ((value as u64) << 1) ^ ((value >> 63) as u64);
while unsigned >= 0x80 {
out.push((unsigned as u8) | 0x80);
unsigned >>= 7;
}
out.push(unsigned as u8);
}
fn read_varint(bytes: &[u8]) -> io::Result<(i64, usize)> {
let (unsigned, used) = read_uvarint(bytes)?;
let value = ((unsigned >> 1) as i64) ^ (-((unsigned & 1) as i64));
Ok((value, used))
}
fn read_uvarint(bytes: &[u8]) -> io::Result<(u64, usize)> {
let mut value = 0_u64;
let mut shift = 0_u32;
for (idx, byte) in bytes.iter().copied().enumerate() {
if byte < 0x80 {
if idx > 9 || (idx == 9 && byte > 1) {
return Err(io::Error::new(io::ErrorKind::InvalidData, "varint overflow"));
}
return Ok((value | ((byte as u64) << shift), idx + 1));
}
value |= ((byte & 0x7f) as u64) << shift;
shift += 7;
}
Err(io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected EOF while reading varint"))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn signed_varint_matches_go_binary_varint_examples() {
let cases = [
(0, vec![0x00]),
(-1, vec![0x01]),
(1, vec![0x02]),
(-2, vec![0x03]),
(64, vec![0x80, 0x01]),
(-64, vec![0x7f]),
];
for (value, expected) in cases {
let mut encoded = Vec::new();
write_varint(&mut encoded, value);
assert_eq!(encoded, expected);
assert_eq!(read_varint(&encoded).unwrap(), (value, encoded.len()));
}
}
#[test]
fn minio_storage_bytes_round_trip_through_headerless_form() {
let mut source = Index::new();
source.add(0, 0).unwrap();
source.add(1_048_576, 2_097_152).unwrap();
source.total_uncompressed = 4_194_304;
source.total_compressed = 3_000_000;
let stored = minio_index_storage_bytes(&source);
assert!(!stored.starts_with(&[CHUNK_TYPE_INDEX, 0x2a, 0x4d, 0x18]));
assert_eq!(
S2Index::load(&restore_index_headers(&stored))
.expect("restore full index")
.info
.len(),
2
);
let decoded = decode_minio_index_bytes(&stored).expect("decode headerless MinIO index");
assert_eq!(decoded.total_uncompressed, source.total_uncompressed);
assert_eq!(decoded.total_compressed, source.total_compressed);
assert_eq!(decoded.find(2_097_152).unwrap(), (1_048_576, 2_097_152));
}
#[test]
fn minio_index_allows_unknown_total_compressed_size() {
let index = S2Index {
total_uncompressed: 4_194_304,
total_compressed: -1,
est_block_uncompressed: 0,
info: vec![
S2IndexInfo {
compressed_offset: 0,
uncompressed_offset: 0,
},
S2IndexInfo {
compressed_offset: 1_048_576,
uncompressed_offset: 2_097_152,
},
],
};
let full = index.into_full_bytes();
assert_eq!(full[0], CHUNK_TYPE_INDEX);
let headerless = Bytes::copy_from_slice(remove_index_headers(full.as_ref()).expect("strip index headers"));
let restored = restore_index_headers(&headerless);
assert_eq!(restored[0], CHUNK_TYPE_INDEX);
let decoded = decode_minio_index_bytes(&headerless).expect("decode index with unknown compressed total");
assert_eq!(decoded.total_uncompressed, 4_194_304);
assert_eq!(decoded.total_compressed, -1);
assert_eq!(decoded.find(2_097_152).unwrap(), (1_048_576, 2_097_152));
}
}