rustfs/crates/rio/src/encrypt_reader.rs

// 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 crate::compress_index::{Index, TryGetIndex};
use aes_gcm::aead::Aead;
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
use pin_project_lite::pin_project;
use rustfs_utils::{put_uvarint, put_uvarint_len};
use std::io::Error;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
use tracing::debug;

const ENCRYPTION_BLOCK_SIZE: usize = 8 * 1024;

pin_project! {
    /// A reader wrapper that encrypts data on the fly using AES-256-GCM.
    /// This is a demonstration. For production, use a secure and audited crypto library.
    pub struct EncryptReader<R> {
        #[pin]
        pub inner: R,
        cipher: Aes256Gcm,
        base_nonce: [u8; 12], // 96-bit base nonce for GCM
        buffer: Vec<u8>,
        buffer_pos: usize,
        read_buffer: Vec<u8>,
        block_index: usize,
        finished: bool,
    }
}

impl<R> EncryptReader<R>
where
    R: AsyncRead + Unpin + Send + Sync,
{
    pub fn new(inner: R, key: [u8; 32], nonce: [u8; 12]) -> Self {
        Self {
            inner,
            cipher: Aes256Gcm::new_from_slice(&key).expect("key"),
            base_nonce: nonce,
            buffer: Vec::new(),
            buffer_pos: 0,
            read_buffer: vec![0u8; ENCRYPTION_BLOCK_SIZE],
            block_index: 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))
    }
}

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<std::io::Result<()>> {
        let mut this = self.project();
        // Serve from buffer if any
        if *this.buffer_pos < this.buffer.len() {
            let to_copy = std::cmp::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(()));
        }
        if *this.finished {
            return Poll::Ready(Ok(()));
        }
        // Read a fixed block size from inner.
        let mut temp_buf = ReadBuf::new(&mut this.read_buffer[..]);
        match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(Ok(())) => {
                let n = temp_buf.filled().len();
                if n == 0 {
                    // EOF, write end header
                    let mut header = [0u8; 8];
                    header[0] = 0xFF; // type: end
                    *this.buffer = header.to_vec();
                    *this.buffer_pos = 0;
                    *this.finished = true;
                    let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
                    buf.put_slice(&this.buffer[..to_copy]);
                    *this.buffer_pos += to_copy;
                    Poll::Ready(Ok(()))
                } else {
                    // Encrypt the chunk
                    let block_nonce = derive_block_nonce(this.base_nonce, *this.block_index);
                    let nonce = Nonce::try_from(block_nonce.as_slice()).map_err(|_| Error::other("invalid nonce length"))?;
                    let plaintext = &this.read_buffer[..n];
                    let plaintext_len = plaintext.len();
                    let crc = {
                        let mut hasher = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
                        hasher.update(plaintext);
                        hasher.finalize() as u32
                    };
                    let ciphertext = this
                        .cipher
                        .encrypt(&nonce, plaintext)
                        .map_err(|e| Error::other(format!("encrypt error: {e}")))?;
                    let int_len = put_uvarint_len(plaintext_len as u64);
                    let clen = int_len + ciphertext.len() + 4;
                    // Header: 8 bytes
                    // 0: type (0 = encrypted, 0xFF = end)
                    // 1-3: length (little endian u24, ciphertext length)
                    // 4-7: CRC32 of plaintext (little endian u32)
                    let mut header = [0u8; 8];
                    header[0] = 0x00; // 0 = encrypted
                    header[1] = (clen & 0xFF) as u8;
                    header[2] = ((clen >> 8) & 0xFF) as u8;
                    header[3] = ((clen >> 16) & 0xFF) as u8;
                    header[4] = (crc & 0xFF) as u8;
                    header[5] = ((crc >> 8) & 0xFF) as u8;
                    header[6] = ((crc >> 16) & 0xFF) as u8;
                    header[7] = ((crc >> 24) & 0xFF) as u8;
                    debug!(
                        "encrypt block header typ=0 len={} header={:?} plaintext_len={} ciphertext_len={}",
                        clen,
                        header,
                        plaintext_len,
                        ciphertext.len()
                    );
                    let mut out = Vec::with_capacity(8 + int_len + ciphertext.len());
                    out.extend_from_slice(&header);
                    let mut plaintext_len_buf = [0u8; 10];
                    let encoded_len = put_uvarint(&mut plaintext_len_buf, plaintext_len as u64);
                    out.extend_from_slice(&plaintext_len_buf[..encoded_len]);
                    out.extend_from_slice(&ciphertext);
                    *this.buffer = out;
                    *this.buffer_pos = 0;
                    *this.block_index += 1;
                    let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
                    buf.put_slice(&this.buffer[..to_copy]);
                    *this.buffer_pos += to_copy;
                    Poll::Ready(Ok(()))
                }
            }
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
        }
    }
}

delegate_reader_capabilities_generic_no_index!(EncryptReader<R>, inner);

impl<R> TryGetIndex for EncryptReader<R>
where
    R: TryGetIndex,
{
    fn try_get_index(&self) -> Option<&Index> {
        self.inner.try_get_index()
    }
}

pin_project! {
    /// A reader wrapper that decrypts data on the fly using AES-256-GCM.
    /// This is a demonstration. For production, use a secure and audited crypto library.
    pub struct DecryptReader<R> {
        #[pin]
        pub inner: R,
        cipher: Aes256Gcm,
        base_nonce: [u8; 12], // Base nonce recorded in object metadata
        current_nonce_base: [u8; 12], // Active base nonce for the current encrypted segment
        multipart_mode: bool,
        multipart_parts: Vec<usize>,
        current_part_index: usize,
        current_part: usize,
        block_index: usize,
        buffer: Vec<u8>,
        buffer_pos: usize,
        finished: bool,
        // For block framing
        header_buf: [u8; 8],
        header_read: usize,
        header_done: bool,
        ciphertext_buf: Vec<u8>,
        ciphertext_read: usize,
        ciphertext_len: usize,
    }
}

impl<R> DecryptReader<R>
where
    R: AsyncRead + Unpin + Send + Sync,
{
    pub fn new(inner: R, key: [u8; 32], nonce: [u8; 12]) -> Self {
        Self {
            inner,
            cipher: Aes256Gcm::new_from_slice(&key).expect("key"),
            base_nonce: nonce,
            current_nonce_base: nonce,
            multipart_mode: false,
            multipart_parts: Vec::new(),
            current_part_index: 0,
            current_part: 0,
            block_index: 0,
            buffer: Vec::new(),
            buffer_pos: 0,
            finished: false,
            header_buf: [0u8; 8],
            header_read: 0,
            header_done: false,
            ciphertext_buf: Vec::new(),
            ciphertext_read: 0,
            ciphertext_len: 0,
        }
    }

    pub fn new_multipart(inner: R, key: [u8; 32], base_nonce: [u8; 12], multipart_parts: Vec<usize>) -> Self {
        let first_part = multipart_parts.first().copied().unwrap_or(1);
        let initial_nonce = derive_part_nonce(&base_nonce, first_part);

        debug!("decrypt_reader: initialized multipart mode");

        Self {
            inner,
            cipher: Aes256Gcm::new_from_slice(&key).expect("key"),
            base_nonce,
            current_nonce_base: initial_nonce,
            multipart_mode: true,
            multipart_parts,
            current_part_index: 0,
            current_part: first_part,
            block_index: 0,
            buffer: Vec::new(),
            buffer_pos: 0,
            finished: false,
            header_buf: [0u8; 8],
            header_read: 0,
            header_done: false,
            ciphertext_buf: Vec::new(),
            ciphertext_read: 0,
            ciphertext_len: 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<std::io::Result<()>> {
        let mut this = self.project();

        loop {
            // Serve buffered plaintext first
            if *this.buffer_pos < this.buffer.len() {
                let to_copy = std::cmp::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(()));
            }

            if *this.finished {
                return Poll::Ready(Ok(()));
            }

            if *this.ciphertext_len == 0 {
                // Read header (8 bytes) only when there is no in-flight payload.
                while !*this.header_done && *this.header_read < 8 {
                    let mut temp = [0u8; 8];
                    let mut temp_buf = ReadBuf::new(&mut temp[0..8 - *this.header_read]);
                    match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
                        Poll::Pending => return Poll::Pending,
                        Poll::Ready(Ok(())) => {
                            let n = temp_buf.filled().len();
                            if n == 0 {
                                if *this.header_read == 0 {
                                    *this.finished = true;
                                    return Poll::Ready(Ok(()));
                                }
                                return Poll::Ready(Err(Error::new(
                                    std::io::ErrorKind::UnexpectedEof,
                                    "unexpected EOF while reading encrypted block header",
                                )));
                            }
                            this.header_buf[*this.header_read..*this.header_read + n].copy_from_slice(&temp_buf.filled()[..n]);
                            *this.header_read += n;
                        }
                        Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
                    }
                }

                if !*this.header_done && *this.header_read == 8 {
                    *this.header_done = true;
                }

                if !*this.header_done {
                    return Poll::Pending;
                }

                let typ = this.header_buf[0];
                let 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;
                *this.header_done = false;

                if typ == 0xFF {
                    if *this.multipart_mode {
                        let next_part = if *this.current_part_index + 1 < this.multipart_parts.len() {
                            *this.current_part_index += 1;
                            this.multipart_parts[*this.current_part_index]
                        } else {
                            *this.current_part + 1
                        };
                        debug!(
                            next_part = next_part,
                            "decrypt_reader: reached segment terminator, advancing to next part"
                        );
                        *this.current_part = next_part;
                        *this.current_nonce_base = derive_part_nonce(this.base_nonce, *this.current_part);
                        *this.block_index = 0;
                        *this.ciphertext_read = 0;
                        *this.ciphertext_len = 0;
                        continue;
                    }

                    *this.finished = true;
                    *this.block_index = 0;
                    *this.ciphertext_read = 0;
                    *this.ciphertext_len = 0;
                    continue;
                }

                tracing::debug!(typ = typ, len = len, "decrypt block header");

                if len == 0 {
                    tracing::warn!("encountered zero-length encrypted block, treating as end of stream");
                    *this.finished = true;
                    *this.ciphertext_read = 0;
                    *this.ciphertext_len = 0;
                    continue;
                }

                let Some(payload_len) = len.checked_sub(4) else {
                    tracing::error!("invalid encrypted block length: typ={} len={} header={:?}", typ, len, this.header_buf);
                    return Poll::Ready(Err(Error::other("Invalid encrypted block length")));
                };

                if this.ciphertext_buf.len() < payload_len {
                    this.ciphertext_buf.resize(payload_len, 0);
                }
                *this.ciphertext_len = payload_len;
                *this.ciphertext_read = 0;
            }

            while *this.ciphertext_read < *this.ciphertext_len {
                let mut temp_buf = ReadBuf::new(&mut this.ciphertext_buf[*this.ciphertext_read..*this.ciphertext_len]);
                match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
                    Poll::Pending => return Poll::Pending,
                    Poll::Ready(Ok(())) => {
                        let n = temp_buf.filled().len();
                        if n == 0 {
                            return Poll::Ready(Err(Error::new(
                                std::io::ErrorKind::UnexpectedEof,
                                "unexpected EOF while reading encrypted block payload",
                            )));
                        }
                        *this.ciphertext_read += n;
                    }
                    Poll::Ready(Err(e)) => {
                        *this.ciphertext_read = 0;
                        *this.ciphertext_len = 0;
                        return Poll::Ready(Err(e));
                    }
                }
            }

            if *this.ciphertext_read < *this.ciphertext_len {
                return Poll::Pending;
            }

            let ciphertext_buf = &this.ciphertext_buf[..*this.ciphertext_len];
            // `ciphertext_buf`'s length derives from the untrusted 24-bit header length field, so
            // it can be shorter than 16 bytes. `uvarint` is safe on any slice length, so pass the
            // whole slice instead of a fixed `[0..16]` index that panics on corrupted/truncated
            // blocks shorter than 16 bytes.
            // `uvarint_len <= 0` means the length varint was empty/unterminated (0) or overflowed
            // (negative — as usize it would index far past the buffer). The `> len` bound is
            // belt-and-suspenders (a positive return is always <= buf.len()).
            let (plaintext_len, uvarint_len) = rustfs_utils::uvarint(ciphertext_buf);
            if uvarint_len <= 0 || uvarint_len as usize > ciphertext_buf.len() {
                *this.ciphertext_read = 0;
                *this.ciphertext_len = 0;
                return Poll::Ready(Err(Error::new(std::io::ErrorKind::InvalidData, "Invalid encrypted block length prefix")));
            }
            let ciphertext = &ciphertext_buf[uvarint_len as usize..];
            let block_nonce = derive_block_nonce(this.current_nonce_base, *this.block_index);
            let nonce = Nonce::try_from(block_nonce.as_slice()).map_err(|_| Error::other("invalid nonce length"))?;
            let legacy_part_nonce = if *this.multipart_mode {
                derive_legacy_part_nonce(this.base_nonce, *this.current_part)
            } else {
                *this.base_nonce
            };
            let legacy_block_nonce = derive_block_nonce(&legacy_part_nonce, *this.block_index);
            let plaintext = match this.cipher.decrypt(&nonce, ciphertext) {
                Ok(plaintext) => plaintext,
                Err(primary_err) => {
                    let legacy_nonce =
                        Nonce::try_from(legacy_block_nonce.as_slice()).map_err(|_| Error::other("invalid nonce length"))?;

                    match this.cipher.decrypt(&legacy_nonce, ciphertext) {
                        Ok(plaintext) => plaintext,
                        Err(_) => {
                            // Accept previously written streams that reused the part nonce
                            // for every block inside a segment.
                            let legacy_part_nonce = Nonce::try_from(legacy_part_nonce.as_slice())
                                .map_err(|_| Error::other("invalid nonce length"))?;
                            this.cipher
                                .decrypt(&legacy_part_nonce, ciphertext)
                                .map_err(|_| Error::other(format!("decrypt error: {primary_err}")))?
                        }
                    }
                }
            };

            debug!(
                part = *this.current_part,
                plaintext_len = plaintext.len(),
                "decrypt_reader: decrypted chunk"
            );

            if plaintext.len() != plaintext_len as usize {
                *this.ciphertext_read = 0;
                *this.ciphertext_len = 0;
                return Poll::Ready(Err(Error::other("Plaintext length mismatch")));
            }

            let expected_crc = (this.header_buf[4] as u32)
                | ((this.header_buf[5] as u32) << 8)
                | ((this.header_buf[6] as u32) << 16)
                | ((this.header_buf[7] as u32) << 24);
            let actual_crc = {
                let mut hasher = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
                hasher.update(&plaintext);
                hasher.finalize() as u32
            };
            if actual_crc != expected_crc {
                *this.ciphertext_read = 0;
                *this.ciphertext_len = 0;
                return Poll::Ready(Err(Error::other("CRC32 mismatch")));
            }

            *this.buffer = plaintext;
            *this.buffer_pos = 0;
            *this.block_index += 1;
            *this.ciphertext_read = 0;
            *this.ciphertext_len = 0;

            let to_copy = std::cmp::min(buf.remaining(), this.buffer.len());
            buf.put_slice(&this.buffer[..to_copy]);
            *this.buffer_pos += to_copy;
            return Poll::Ready(Ok(()));
        }
    }
}

delegate_reader_capabilities_generic_no_index!(DecryptReader<R>, inner);

impl<R> TryGetIndex for DecryptReader<R>
where
    R: TryGetIndex,
{
    fn try_get_index(&self) -> Option<&Index> {
        self.inner.try_get_index()
    }
}

fn derive_block_nonce(base: &[u8; 12], block_index: usize) -> [u8; 12] {
    derive_nonce_offset(base, 8, block_index)
}

pub fn multipart_part_nonce(base_nonce: [u8; 12], part_number: usize) -> [u8; 12] {
    derive_part_nonce(&base_nonce, part_number)
}

fn derive_part_nonce(base: &[u8; 12], part_number: usize) -> [u8; 12] {
    derive_nonce_offset(base, 4, part_number)
}

fn derive_legacy_part_nonce(base: &[u8; 12], part_number: usize) -> [u8; 12] {
    derive_nonce_offset(base, 8, part_number)
}

fn derive_nonce_offset(base: &[u8; 12], start: usize, offset: usize) -> [u8; 12] {
    let mut nonce = *base;
    let mut suffix = [0u8; 4];
    suffix.copy_from_slice(&nonce[start..start + 4]);
    let current = u32::from_be_bytes(suffix);
    let next = current.wrapping_add(offset as u32);
    nonce[start..start + 4].copy_from_slice(&next.to_be_bytes());
    nonce
}

#[cfg(test)]
mod tests {
    use aes_gcm::aead::Aead;
    use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
    use std::io::Cursor;
    use std::pin::Pin;
    use std::task::{Context, Poll};

    use crate::HardLimitReader;

    use super::*;
    use futures::StreamExt;
    use rand::{Rng, RngExt};
    use tokio::io::{AsyncRead, AsyncReadExt, BufReader, ReadBuf};
    use tokio_util::io::ReaderStream;

    struct ChunkedCursor {
        inner: Cursor<Vec<u8>>,
        max_chunk: usize,
    }

    impl ChunkedCursor {
        fn new(data: Vec<u8>, max_chunk: usize) -> Self {
            Self {
                inner: Cursor::new(data),
                max_chunk,
            }
        }
    }

    impl AsyncRead for ChunkedCursor {
        fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
            if self.max_chunk == 0 || buf.remaining() == 0 {
                return Poll::Ready(Ok(()));
            }

            let remaining = self.inner.get_ref().len() as u64 - self.inner.position();
            if remaining == 0 {
                return Poll::Ready(Ok(()));
            }

            let to_read = remaining.min(self.max_chunk as u64).min(buf.remaining() as u64) as usize;
            let start = self.inner.position() as usize;
            let end = start + to_read;
            buf.put_slice(&self.inner.get_ref()[start..end]);
            self.inner.set_position(end as u64);
            Poll::Ready(Ok(()))
        }
    }

    struct PendingChunkedCursor {
        inner: Cursor<Vec<u8>>,
        max_chunk: usize,
        should_pending: bool,
    }

    impl PendingChunkedCursor {
        fn new(data: Vec<u8>, max_chunk: usize) -> Self {
            Self {
                inner: Cursor::new(data),
                max_chunk,
                should_pending: true,
            }
        }
    }

    impl AsyncRead for PendingChunkedCursor {
        fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
            if self.should_pending {
                self.should_pending = false;
                cx.waker().wake_by_ref();
                return Poll::Pending;
            }

            if self.max_chunk == 0 || buf.remaining() == 0 {
                return Poll::Ready(Ok(()));
            }

            let remaining = self.inner.get_ref().len() as u64 - self.inner.position();
            if remaining == 0 {
                return Poll::Ready(Ok(()));
            }

            let to_read = remaining.min(self.max_chunk as u64).min(buf.remaining() as u64) as usize;
            let start = self.inner.position() as usize;
            let end = start + to_read;
            buf.put_slice(&self.inner.get_ref()[start..end]);
            self.inner.set_position(end as u64);
            self.should_pending = true;
            Poll::Ready(Ok(()))
        }
    }

    fn encrypt_with_legacy_nonce_reuse(data: &[u8], key: [u8; 32], nonce: [u8; 12]) -> Vec<u8> {
        let cipher = Aes256Gcm::new_from_slice(&key).expect("valid key");
        let nonce = Nonce::try_from(nonce.as_slice()).expect("valid nonce");
        let mut encrypted = Vec::new();

        for chunk in data.chunks(ENCRYPTION_BLOCK_SIZE) {
            let crc = {
                let mut hasher = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
                hasher.update(chunk);
                hasher.finalize() as u32
            };
            let ciphertext = cipher.encrypt(&nonce, chunk).expect("legacy encrypt");
            let int_len = put_uvarint_len(chunk.len() as u64);
            let clen = int_len + ciphertext.len() + 4;
            let mut header = [0u8; 8];
            header[1] = (clen & 0xFF) as u8;
            header[2] = ((clen >> 8) & 0xFF) as u8;
            header[3] = ((clen >> 16) & 0xFF) as u8;
            header[4] = (crc & 0xFF) as u8;
            header[5] = ((crc >> 8) & 0xFF) as u8;
            header[6] = ((crc >> 16) & 0xFF) as u8;
            header[7] = ((crc >> 24) & 0xFF) as u8;
            encrypted.extend_from_slice(&header);
            let mut plaintext_len_buf = [0u8; 10];
            let encoded_len = put_uvarint(&mut plaintext_len_buf, chunk.len() as u64);
            encrypted.extend_from_slice(&plaintext_len_buf[..encoded_len]);
            encrypted.extend_from_slice(&ciphertext);
        }

        encrypted.extend_from_slice(&[0xFF, 0, 0, 0, 0, 0, 0, 0]);
        encrypted
    }

    async fn encrypt_part_with_legacy_nonce_layout(
        data: &[u8],
        key: [u8; 32],
        base_nonce: [u8; 12],
        part_number: usize,
    ) -> Vec<u8> {
        let nonce = derive_legacy_part_nonce(&base_nonce, part_number);
        let reader = BufReader::new(Cursor::new(data.to_vec()));
        let mut encrypt_reader = EncryptReader::new(reader, key, nonce);
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");
        encrypted
    }

    fn extract_encrypted_payloads(encrypted: &[u8]) -> Vec<Vec<u8>> {
        let mut payloads = Vec::new();
        let mut pos = 0;

        while pos + 8 <= encrypted.len() {
            let header = &encrypted[pos..pos + 8];
            pos += 8;
            if header[0] == 0xFF {
                break;
            }

            let len = (header[1] as usize) | ((header[2] as usize) << 8) | ((header[3] as usize) << 16);
            let payload_len = len - 4;
            payloads.push(encrypted[pos..pos + payload_len].to_vec());
            pos += payload_len;
        }

        payloads
    }

    #[tokio::test]
    async fn test_encrypt_decrypt_reader_aes256gcm() {
        let data = b"hello sse encrypt";
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        let reader = BufReader::new(&data[..]);
        let encrypt_reader = EncryptReader::new(reader, key, nonce);

        // Encrypt
        let mut encrypt_reader = encrypt_reader;
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");

        // Decrypt using DecryptReader
        let reader = Cursor::new(encrypted.clone());
        let decrypt_reader = DecryptReader::new(reader, key, nonce);
        let mut decrypt_reader = decrypt_reader;
        let mut decrypted = Vec::new();
        decrypt_reader
            .read_to_end(&mut decrypted)
            .await
            .expect("operation should succeed");

        assert_eq!(&decrypted, data);
    }

    #[tokio::test]
    async fn test_decrypt_reader_only() {
        // Encrypt some data first
        let data = b"test decrypt only";
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        // Encrypt
        let reader = BufReader::new(&data[..]);
        let encrypt_reader = EncryptReader::new(reader, key, nonce);
        let mut encrypt_reader = encrypt_reader;
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");

        // Now test DecryptReader

        let reader = Cursor::new(encrypted.clone());
        let decrypt_reader = DecryptReader::new(reader, key, nonce);
        let mut decrypt_reader = decrypt_reader;
        let mut decrypted = Vec::new();
        decrypt_reader
            .read_to_end(&mut decrypted)
            .await
            .expect("operation should succeed");

        assert_eq!(&decrypted, data);
    }

    #[tokio::test]
    async fn test_encrypt_decrypt_reader_large() {
        use rand::Rng;
        let size = 1024 * 1024;
        let mut data = vec![0u8; size];
        rand::rng().fill(&mut data[..]);
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        let reader = std::io::Cursor::new(data.clone());
        let encrypt_reader = EncryptReader::new(reader, key, nonce);
        let mut encrypt_reader = encrypt_reader;
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");

        let reader = std::io::Cursor::new(encrypted.clone());
        let decrypt_reader = DecryptReader::new(reader, key, nonce);
        let mut decrypt_reader = decrypt_reader;
        let mut decrypted = Vec::new();
        decrypt_reader
            .read_to_end(&mut decrypted)
            .await
            .expect("operation should succeed");

        assert_eq!(&decrypted, &data);
    }

    #[tokio::test]
    async fn test_decrypt_reader_large_with_small_chunks() {
        let size = 1024 * 1024;
        let mut data = vec![0u8; size];
        rand::rng().fill(&mut data[..]);
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        let reader = Cursor::new(data.clone());
        let mut encrypt_reader = EncryptReader::new(reader, key, nonce);
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");

        let reader = ChunkedCursor::new(encrypted, 3);
        let mut decrypt_reader = DecryptReader::new(reader, key, nonce);
        let mut decrypted = Vec::new();
        decrypt_reader
            .read_to_end(&mut decrypted)
            .await
            .expect("operation should succeed");

        assert_eq!(decrypted, data);
    }

    #[tokio::test]
    async fn test_decrypt_reader_large_with_pending_chunks() {
        let size = 1024 * 1024;
        let mut data = vec![0u8; size];
        rand::rng().fill(&mut data[..]);
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        let reader = Cursor::new(data.clone());
        let mut encrypt_reader = EncryptReader::new(reader, key, nonce);
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");

        let reader = PendingChunkedCursor::new(encrypted, 3);
        let mut decrypt_reader = DecryptReader::new(reader, key, nonce);
        let mut decrypted = Vec::new();
        decrypt_reader
            .read_to_end(&mut decrypted)
            .await
            .expect("operation should succeed");

        assert_eq!(decrypted, data);
    }

    #[tokio::test]
    async fn test_decrypt_reader_large_through_reader_stream() {
        let size = 1024 * 1024;
        let mut data = vec![0u8; size];
        rand::rng().fill(&mut data[..]);
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        let reader = Cursor::new(data.clone());
        let mut encrypt_reader = EncryptReader::new(reader, key, nonce);
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");

        let reader = ChunkedCursor::new(encrypted, 8192);
        let decrypt_reader = DecryptReader::new(reader, key, nonce);
        let mut stream = ReaderStream::with_capacity(Box::new(decrypt_reader), 262_144);

        let mut decrypted = Vec::new();
        while let Some(chunk) = stream.next().await {
            let bytes = chunk.expect("operation should succeed");
            decrypted.extend_from_slice(&bytes);
        }

        assert_eq!(decrypted, data);
    }

    #[tokio::test]
    async fn test_decrypt_reader_large_through_hard_limit_reader_stream() {
        let size = 1024 * 1024;
        let mut data = vec![0u8; size];
        rand::rng().fill(&mut data[..]);
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        let reader = Cursor::new(data.clone());
        let mut encrypt_reader = EncryptReader::new(reader, key, nonce);
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");

        let reader = ChunkedCursor::new(encrypted, 8192);
        let decrypt_reader = DecryptReader::new(reader, key, nonce);
        let limit_reader = HardLimitReader::new(decrypt_reader, size as i64);
        let mut stream = ReaderStream::with_capacity(Box::new(limit_reader), 262_144);

        let mut decrypted = Vec::new();
        while let Some(chunk) = stream.next().await {
            let bytes = chunk.expect("operation should succeed");
            decrypted.extend_from_slice(&bytes);
        }

        assert_eq!(decrypted, data);
    }

    #[tokio::test]
    async fn test_decrypt_reader_multipart_segments() {
        let mut key = [0u8; 32];
        let mut base_nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut base_nonce);

        let part_one = vec![0xA5; 512 * 1024];
        let part_two = vec![0x5A; 256 * 1024];

        async fn encrypt_part(data: &[u8], key: [u8; 32], base_nonce: [u8; 12], part_number: usize) -> Vec<u8> {
            let nonce = derive_part_nonce(&base_nonce, part_number);
            let reader = BufReader::new(Cursor::new(data.to_vec()));
            let mut encrypt_reader = EncryptReader::new(reader, key, nonce);
            let mut encrypted = Vec::new();
            encrypt_reader
                .read_to_end(&mut encrypted)
                .await
                .expect("operation should succeed");
            encrypted
        }

        let encrypted_one = encrypt_part(&part_one, key, base_nonce, 1).await;
        let encrypted_two = encrypt_part(&part_two, key, base_nonce, 2).await;

        let mut combined = Vec::with_capacity(encrypted_one.len() + encrypted_two.len());
        combined.extend_from_slice(&encrypted_one);
        combined.extend_from_slice(&encrypted_two);

        let reader = BufReader::new(Cursor::new(combined));
        let mut decrypt_reader = DecryptReader::new_multipart(reader, key, base_nonce, vec![1, 2]);
        let mut decrypted = Vec::new();
        decrypt_reader
            .read_to_end(&mut decrypted)
            .await
            .expect("operation should succeed");

        let mut expected = Vec::with_capacity(part_one.len() + part_two.len());
        expected.extend_from_slice(&part_one);
        expected.extend_from_slice(&part_two);

        assert_eq!(decrypted, expected);
    }

    #[tokio::test]
    async fn test_encrypt_reader_uses_distinct_nonces_per_block() {
        let data = vec![0xAB; ENCRYPTION_BLOCK_SIZE * 2];
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        let reader = Cursor::new(data);
        let mut encrypt_reader = EncryptReader::new(reader, key, nonce);
        let mut encrypted = Vec::new();
        encrypt_reader
            .read_to_end(&mut encrypted)
            .await
            .expect("operation should succeed");

        let payloads = extract_encrypted_payloads(&encrypted);
        assert!(payloads.len() >= 2);
        assert_ne!(payloads[0], payloads[1]);
    }

    #[test]
    fn test_part_and_block_nonces_do_not_collide_across_parts() {
        let base_nonce = [0u8; 12];
        let part_one_block_one = derive_block_nonce(&derive_part_nonce(&base_nonce, 1), 1);
        let part_two_block_zero = derive_block_nonce(&derive_part_nonce(&base_nonce, 2), 0);

        assert_ne!(part_one_block_one, part_two_block_zero);
    }

    #[tokio::test]
    async fn test_decrypt_reader_accepts_legacy_single_nonce_streams() {
        let mut data = vec![0u8; ENCRYPTION_BLOCK_SIZE * 3 + 17];
        rand::rng().fill(&mut data[..]);
        let mut key = [0u8; 32];
        let mut nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut nonce);

        let encrypted = encrypt_with_legacy_nonce_reuse(&data, key, nonce);
        let reader = Cursor::new(encrypted);
        let mut decrypt_reader = DecryptReader::new(reader, key, nonce);
        let mut decrypted = Vec::new();
        decrypt_reader
            .read_to_end(&mut decrypted)
            .await
            .expect("operation should succeed");

        assert_eq!(decrypted, data);
    }

    #[tokio::test]
    async fn test_decrypt_reader_accepts_legacy_multipart_nonce_layout() {
        let mut key = [0u8; 32];
        let mut base_nonce = [0u8; 12];
        rand::rng().fill_bytes(&mut key);
        rand::rng().fill_bytes(&mut base_nonce);

        let part_one = vec![0x11; ENCRYPTION_BLOCK_SIZE + 97];
        let part_two = vec![0x22; ENCRYPTION_BLOCK_SIZE + 33];

        let encrypted_one = encrypt_part_with_legacy_nonce_layout(&part_one, key, base_nonce, 1).await;
        let encrypted_two = encrypt_part_with_legacy_nonce_layout(&part_two, key, base_nonce, 2).await;

        let mut combined = Vec::with_capacity(encrypted_one.len() + encrypted_two.len());
        combined.extend_from_slice(&encrypted_one);
        combined.extend_from_slice(&encrypted_two);

        let reader = BufReader::new(Cursor::new(combined));
        let mut decrypt_reader = DecryptReader::new_multipart(reader, key, base_nonce, vec![1, 2]);
        let mut decrypted = Vec::new();
        decrypt_reader
            .read_to_end(&mut decrypted)
            .await
            .expect("operation should succeed");

        let mut expected = Vec::with_capacity(part_one.len() + part_two.len());
        expected.extend_from_slice(&part_one);
        expected.extend_from_slice(&part_two);

        assert_eq!(decrypted, expected);
    }

    // Regression: a corrupted block header whose length yields a payload shorter than 16 bytes
    // must not panic. Header (8 bytes): [typ, len_lo, len_mid, len_hi, crc0..crc3]; payload is
    // `len - 4` bytes. Pre-fix, poll_read sliced `ciphertext_buf[0..16]` unconditionally,
    // panicking with "range end index 16 out of range for slice of length N" when N < 16.
    #[tokio::test]
    async fn test_decrypt_reader_short_block_no_panic() {
        let key = [0u8; 32];
        let nonce = [0u8; 12];

        // len = 8 -> payload_len = 4 (< 16). Provide exactly 4 payload bytes.
        let len: usize = 8;
        let mut input = vec![
            0x00u8, // typ (regular block)
            (len & 0xFF) as u8,
            ((len >> 8) & 0xFF) as u8,
            ((len >> 16) & 0xFF) as u8,
        ];
        input.extend_from_slice(&[0u8; 4]); // crc (unused before the panic site)
        input.extend_from_slice(&[0x01u8, 0x02, 0x03, 0x04]); // 4-byte payload

        let mut decrypt_reader = DecryptReader::new(Cursor::new(input), key, nonce);
        let mut out = Vec::new();
        let res = decrypt_reader.read_to_end(&mut out).await;
        assert!(res.is_err(), "corrupted short encrypted block must return an error, not panic");
    }
}