libsoliton 0.1.3

//! Streaming/chunked AEAD for large payloads (§15).
//!
//! Chunked AEAD with counter-derived nonces over XChaCha20-Poly1305. Each
//! chunk is independently encrypted with a unique nonce derived from a random
//! base nonce and the chunk index, enabling random-access decryption without
//! processing preceding chunks.
//!
//! Inspired by the STREAM construction (Hoang, Reyhanitabar, Rogaway, Vizár,
//! CRYPTO 2015) but uses counter-based nonce derivation for random access
//! rather than ciphertext chaining.
//!
//! ## Wire Format
//!
//! ```text
//! Stream = Header || Chunk₀ || Chunk₁ || ... || ChunkN
//!
//! Header (26 bytes):
//!   version     (1)  — stream format version (0x01)
//!   flags       (1)  — bit 0: compression (0 = none, 1 = zstd), bits 1-7: reserved
//!   base_nonce  (24) — random, unique per stream
//!
//! Chunk:
//!   tag_byte    (1)        — 0x00 = non-final, 0x01 = final
//!   ciphertext  (variable) — AEAD output: encrypted plaintext + 16-byte Poly1305 tag
//! ```
//!
//! ## Security: Truncation
//!
//! Chunked AEAD authenticates each chunk independently but does not
//! authenticate stream length. An attacker who drops the final chunk leaves
//! the decryptor in a non-finalized state with all prior chunks successfully
//! authenticated. Callers MUST check `is_last` on each returned chunk or
//! call `is_finalized()` after the stream ends to detect truncation.
//!
//! ## AAD Construction
//!
//! ```text
//! aad = "lo-stream-v1" || version || flags || base_nonce
//!    || chunk_index (u64 BE) || tag_byte || caller_aad
//! ```

use crate::constants::{
    self, AEAD_NONCE_SIZE, AEAD_TAG_SIZE, STREAM_CHUNK_OVERHEAD, STREAM_CHUNK_SIZE,
    STREAM_HEADER_SIZE, STREAM_VERSION, STREAM_ZSTD_OVERHEAD,
};
use crate::error::{Error, Result};
use crate::primitives::{aead, random};
use std::io::Read;
use zeroize::Zeroizing;

/// Compression flag in the flags byte (bit 0).
const FLAG_COMPRESSED: u8 = 0x01;

/// Tag byte value for non-final chunks.
const TAG_NON_FINAL: u8 = 0x00;

/// Tag byte value for the final chunk.
const TAG_FINAL: u8 = 0x01;

// ── Nonce derivation ─────────────────────────────────────────────────────

/// Build the per-chunk nonce by XORing a mask into the base nonce.
///
/// Mask layout (24 bytes):
///   bytes 0-7:  chunk_index (u64 big-endian)
///   byte 8:     tag_byte (0x00 non-final, 0x01 final)
///   bytes 9-23: zero padding
///
/// The mask is injective over `(index, tag_byte)` pairs, so XORing with a
/// fixed base nonce produces a unique nonce per chunk. See §Nonce Derivation
/// in STREAMING_AEAD.md for the injectivity proof.
fn build_chunk_nonce(
    base_nonce: &[u8; AEAD_NONCE_SIZE],
    index: u64,
    tag_byte: u8,
) -> [u8; AEAD_NONCE_SIZE] {
    let mut mask = [0u8; AEAD_NONCE_SIZE];
    mask[..8].copy_from_slice(&index.to_be_bytes());
    mask[8] = tag_byte;
    // Remaining bytes are zero — XOR with zero is identity.

    let mut nonce = *base_nonce;
    for i in 0..AEAD_NONCE_SIZE {
        nonce[i] ^= mask[i];
    }
    nonce
}

// ── AAD construction ─────────────────────────────────────────────────────

/// Build the per-chunk AAD.
///
/// Layout: domain_label (12) || version (1) || flags (1) || base_nonce (24)
///      || chunk_index (8, u64 BE) || tag_byte (1) || caller_aad (variable)
///
/// Total fixed size: 47 bytes + caller_aad.len().
/// caller_aad has no length prefix because it is the terminal field.
fn build_chunk_aad(
    version: u8,
    flags: u8,
    base_nonce: &[u8; AEAD_NONCE_SIZE],
    index: u64,
    tag_byte: u8,
    caller_aad: &[u8],
) -> Vec<u8> {
    let fixed_len = constants::STREAM_AAD.len() + 1 + 1 + AEAD_NONCE_SIZE + 8 + 1;
    let mut aad = Vec::with_capacity(fixed_len + caller_aad.len());
    aad.extend_from_slice(constants::STREAM_AAD);
    aad.push(version);
    aad.push(flags);
    aad.extend_from_slice(base_nonce);
    aad.extend_from_slice(&index.to_be_bytes());
    aad.push(tag_byte);
    aad.extend_from_slice(caller_aad);
    aad
}

// ── Compression helpers ──────────────────────────────────────────────────

/// Zstd compression level. Matches storage blob convention (Fastest / ~level 1).
const COMPRESSION_LEVEL: ruzstd::encoding::CompressionLevel =
    ruzstd::encoding::CompressionLevel::Fastest;

/// Compress plaintext with zstd. Returns `Internal` if compressed output
/// exceeds `plaintext.len() + STREAM_ZSTD_OVERHEAD` (indicates a zstd
/// implementation bug — the 256-byte bound provides ~5× safety margin).
fn compress_chunk(plaintext: &[u8]) -> Result<Zeroizing<Vec<u8>>> {
    let compressed = ruzstd::encoding::compress_to_vec(plaintext, COMPRESSION_LEVEL);
    // Guard against zstd expansion exceeding the documented overhead bound.
    // This cannot happen with a conforming zstd implementation for 1 MiB inputs.
    if compressed.len() > plaintext.len() + STREAM_ZSTD_OVERHEAD {
        return Err(Error::Internal);
    }
    // Wrap in Zeroizing — compressed plaintext is trivially decompressible to the
    // original plaintext and must not linger on the heap after use. Mirrors the
    // decompress_chunk return type.
    Ok(Zeroizing::new(compressed))
}

/// Decompress a chunk payload, bounded to `STREAM_CHUNK_SIZE`.
/// Returns `AeadFailed` on decompression failure or size overflow
/// (post-auth error collapse).
fn decompress_chunk(data: &[u8]) -> Result<Zeroizing<Vec<u8>>> {
    let decoder = ruzstd::decoding::StreamingDecoder::new(data).map_err(|_| Error::AeadFailed)?;
    // Read one byte beyond STREAM_CHUNK_SIZE to distinguish "exactly at limit"
    // from "exceeds limit".
    let mut limited = decoder.take(STREAM_CHUNK_SIZE as u64 + 1);
    let hint = data.len().saturating_mul(4).min(STREAM_CHUNK_SIZE);
    let mut decompressed = Zeroizing::new(Vec::with_capacity(hint));
    limited
        .read_to_end(&mut decompressed)
        .map_err(|_| Error::AeadFailed)?;
    if decompressed.len() > STREAM_CHUNK_SIZE {
        return Err(Error::AeadFailed);
    }
    Ok(decompressed)
}

// ── StreamEncryptor ──────────────────────────────────────────────────────

/// Streaming encryptor state.
///
/// Encrypts plaintext in fixed-size chunks (1 MiB) with per-chunk
/// XChaCha20-Poly1305. The key is zeroized on drop; other fields are
/// non-secret metadata.
pub struct StreamEncryptor {
    /// Encryption key — zeroized on drop.
    key: Zeroizing<[u8; 32]>,
    /// Random 24-byte base nonce, unique per stream.
    base_nonce: [u8; AEAD_NONCE_SIZE],
    /// Stream format version (always STREAM_VERSION).
    version: u8,
    /// Flags byte (bit 0 = compression).
    flags: u8,
    /// Next chunk index to encrypt.
    next_index: u64,
    /// Caller-supplied AAD, constant across all chunks.
    caller_aad: Vec<u8>,
    /// Whether compression is enabled.
    compress: bool,
    /// Whether a final chunk has been encrypted.
    finalized: bool,
}

/// Create a streaming encryptor.
///
/// `key`: 32-byte encryption key (caller-provided, typically random per file).
/// `aad`: optional caller-supplied AAD context (bound into every chunk's AAD).
/// `compress`: if true, each chunk is zstd-compressed before encryption.
///
/// Generates a random 24-byte base nonce from the OS CSPRNG.
///
/// # Panics
///
/// Panics if the OS CSPRNG is unavailable (matching all other randomness-
/// consuming call sites in the library).
#[must_use = "dropping the encryptor loses the CSPRNG-based nonce and the encryption handle"]
pub fn stream_encrypt_init(key: &[u8; 32], aad: &[u8], compress: bool) -> Result<StreamEncryptor> {
    let key_copy = Zeroizing::new(*key);
    let base_nonce: [u8; AEAD_NONCE_SIZE] = random::random_array();
    // [u8; 32] is Copy — the caller's array is unaffected. Our Zeroizing
    // wrapper owns its own copy; no source zeroization needed here because
    // the caller owns the original.
    Ok(StreamEncryptor {
        key: key_copy,
        base_nonce,
        version: STREAM_VERSION,
        flags: if compress { FLAG_COMPRESSED } else { 0 },
        next_index: 0,
        caller_aad: aad.to_vec(),
        compress,
        finalized: false,
    })
}

impl StreamEncryptor {
    /// Return the 26-byte header. Caller writes this before any chunks.
    pub fn header(&self) -> [u8; STREAM_HEADER_SIZE] {
        let mut hdr = [0u8; STREAM_HEADER_SIZE];
        hdr[0] = self.version;
        hdr[1] = self.flags;
        hdr[2..].copy_from_slice(&self.base_nonce);
        hdr
    }

    /// Encrypt one chunk sequentially.
    ///
    /// `is_last` MUST be `true` for exactly the final chunk. Plaintext MUST be
    /// exactly `STREAM_CHUNK_SIZE` bytes for non-final chunks. Final chunk may
    /// be `0..=STREAM_CHUNK_SIZE` bytes.
    ///
    /// Returns the chunk wire bytes: `tag_byte (1) || aead_output (plaintext_len + 16)`.
    ///
    /// # Error Recovery
    ///
    /// On error, internal state is unchanged — the chunk index is not advanced
    /// and `finalized` is not set. State advances only on successful return.
    #[must_use = "dropping the ciphertext loses encrypted data and desynchronizes the stream (chunk index was already advanced)"]
    pub fn encrypt_chunk(&mut self, plaintext: &[u8], is_last: bool) -> Result<Vec<u8>> {
        if self.finalized {
            return Err(Error::InvalidData);
        }
        if self.next_index == u64::MAX {
            return Err(Error::ChainExhausted);
        }

        let index = self.next_index;
        let output = self.encrypt_chunk_inner(index, is_last, plaintext)?;

        // State advances only after successful encryption.
        // wrapping_add: next_index was checked != u64::MAX above, so this
        // cannot wrap. Using wrapping_add to satisfy the compiler without
        // a redundant overflow check.
        self.next_index = index.wrapping_add(1);
        if is_last {
            self.finalized = true;
        }

        Ok(output)
    }

    /// Encrypt a specific chunk by index (random access).
    ///
    /// Stateless: does not advance `next_index`, does not set `finalized`, and
    /// does not enforce the post-finalization guard. The caller is responsible
    /// for assigning unique indices and identifying the final chunk.
    ///
    /// Enables parallel chunk encryption: independent calls with distinct
    /// indices may be dispatched concurrently. Because all per-chunk state
    /// (nonce, AAD) is derived from the index alone, the results are assembled
    /// in index order after all encryptions complete without synchronization.
    /// The CAPI reentrancy guard prevents concurrent calls on the same handle —
    /// parallel encryption requires one handle per thread (or use the Rust API
    /// directly, where `&self` enables concurrent calls from safe code).
    ///
    /// `is_last` MUST be `true` for exactly the final chunk. Plaintext MUST be
    /// exactly `STREAM_CHUNK_SIZE` bytes for non-final chunks. Final chunk may
    /// be `0..=STREAM_CHUNK_SIZE` bytes.
    ///
    /// Returns the chunk wire bytes: `tag_byte (1) || aead_output (plaintext_len + 16)`.
    #[must_use = "dropping the ciphertext loses encrypted data; re-encrypting the same index reuses the nonce"]
    pub fn encrypt_chunk_at(&self, index: u64, is_last: bool, plaintext: &[u8]) -> Result<Vec<u8>> {
        self.encrypt_chunk_inner(index, is_last, plaintext)
    }

    /// Return whether the encryptor has been finalized.
    pub fn is_finalized(&self) -> bool {
        self.finalized
    }

    /// Raw pointer to the internal key bytes (test-utils only).
    #[cfg(all(feature = "test-utils", debug_assertions))]
    #[deprecated(note = "test-utils only — do not call in production code")]
    pub fn key_ptr(&self) -> *const u8 {
        self.key.as_ptr()
    }

    /// Core encryption logic shared by sequential and random-access paths.
    ///
    /// Validates plaintext size, optionally compresses, derives the per-chunk
    /// nonce and AAD from `index` and `is_last`, encrypts with
    /// XChaCha20-Poly1305, and returns the wire bytes: `tag_byte || ciphertext`.
    ///
    /// Does not read or mutate `next_index` or `finalized` — those invariants
    /// are enforced by `encrypt_chunk` for the sequential path; `encrypt_chunk_at`
    /// leaves them to the application layer.
    fn encrypt_chunk_inner(&self, index: u64, is_last: bool, plaintext: &[u8]) -> Result<Vec<u8>> {
        // Size validation: non-final chunks must be exactly STREAM_CHUNK_SIZE;
        // the final chunk may be 0..=STREAM_CHUNK_SIZE.
        if is_last {
            if plaintext.len() > STREAM_CHUNK_SIZE {
                return Err(Error::InvalidData);
            }
        } else if plaintext.len() != STREAM_CHUNK_SIZE {
            return Err(Error::InvalidData);
        }

        let tag_byte = if is_last { TAG_FINAL } else { TAG_NON_FINAL };

        // Compress if enabled and plaintext is non-empty.
        // Empty final chunks bypass compression (degenerate zstd edge case).
        let compressed_buf;
        let to_encrypt: &[u8] = if self.compress && !plaintext.is_empty() {
            compressed_buf = compress_chunk(plaintext)?;
            &compressed_buf
        } else {
            plaintext
        };

        let nonce = build_chunk_nonce(&self.base_nonce, index, tag_byte);
        let aad = build_chunk_aad(
            self.version,
            self.flags,
            &self.base_nonce,
            index,
            tag_byte,
            &self.caller_aad,
        );

        let ciphertext = aead::aead_encrypt(&self.key, &nonce, to_encrypt, &aad)?;

        // Build wire output: tag_byte || ciphertext (includes Poly1305 tag).
        let mut output = Vec::with_capacity(1 + ciphertext.len());
        output.push(tag_byte);
        output.extend_from_slice(&ciphertext);

        Ok(output)
    }
}

impl Drop for StreamEncryptor {
    fn drop(&mut self) {
        // key is Zeroizing<[u8; 32]> — zeroized automatically.
        // caller_aad is non-secret metadata, not zeroized.
    }
}

// ── StreamDecryptor ──────────────────────────────────────────────────────

/// Streaming decryptor state.
///
/// Decrypts chunks either sequentially (`decrypt_chunk`) or randomly
/// (`decrypt_chunk_at`). The key is zeroized on drop.
pub struct StreamDecryptor {
    /// Decryption key — zeroized on drop.
    key: Zeroizing<[u8; 32]>,
    /// Base nonce from the stream header.
    base_nonce: [u8; AEAD_NONCE_SIZE],
    /// Stream format version (validated at init).
    version: u8,
    /// Flags byte (validated at init — reserved bits rejected).
    flags: u8,
    /// Next expected sequential chunk index.
    next_index: u64,
    /// Caller-supplied AAD, constant across all chunks.
    caller_aad: Vec<u8>,
    /// Whether compression is enabled (derived from flags bit 0).
    compressed: bool,
    /// Whether a final chunk has been decrypted sequentially.
    finalized: bool,
}

/// Create a streaming decryptor from a 26-byte header.
///
/// # Errors
///
/// - `UnsupportedVersion` if the header version byte is not `0x01`.
/// - `AeadFailed` if reserved flag bits are set (error-oracle collapse —
///   an attacker with write access to the header could use a distinct error
///   to confirm header delivery without attempting chunk decryption).
pub fn stream_decrypt_init(
    key: &[u8; 32],
    header: &[u8; STREAM_HEADER_SIZE],
    aad: &[u8],
) -> Result<StreamDecryptor> {
    let version = header[0];
    let flags = header[1];

    // Not collapsed to AeadFailed (unlike the flags check below): version is
    // a forward-compatibility signal, and its existence is not secret — a
    // distinct error lets callers distinguish "newer format" from "corrupt
    // data" without leaking information an attacker doesn't already have.
    if version != STREAM_VERSION {
        return Err(Error::UnsupportedVersion);
    }

    // Reject reserved flag bits. Collapsed to AeadFailed to prevent an error
    // oracle distinguishing "unsupported flag" from "authentication failure"
    // on attacker-controlled input.
    if flags & !FLAG_COMPRESSED != 0 {
        return Err(Error::AeadFailed);
    }

    let mut base_nonce = [0u8; AEAD_NONCE_SIZE];
    base_nonce.copy_from_slice(&header[2..]);

    Ok(StreamDecryptor {
        key: Zeroizing::new(*key),
        base_nonce,
        version,
        flags,
        next_index: 0,
        caller_aad: aad.to_vec(),
        compressed: flags & FLAG_COMPRESSED != 0,
        finalized: false,
    })
}

impl StreamDecryptor {
    /// Decrypt the next sequential chunk.
    ///
    /// Returns `(plaintext, is_last)`. Plaintext is wrapped in `Zeroizing`.
    ///
    /// # Security
    ///
    /// Chunked AEAD does not authenticate stream length. Callers MUST verify
    /// that `is_last == true` was received for exactly one chunk, or check
    /// `is_finalized()` after the stream ends, to detect truncation attacks.
    ///
    /// # Error Recovery
    ///
    /// On error, internal state is unchanged.
    pub fn decrypt_chunk(&mut self, chunk: &[u8]) -> Result<(Zeroizing<Vec<u8>>, bool)> {
        if self.finalized {
            return Err(Error::InvalidData);
        }
        if self.next_index == u64::MAX {
            return Err(Error::ChainExhausted);
        }

        let index = self.next_index;
        let (plaintext, is_last) = self.decrypt_chunk_inner(index, chunk)?;

        // State advances only after successful decryption.
        // wrapping_add: next_index was checked != u64::MAX above, so this
        // cannot wrap. Using wrapping_add to satisfy the compiler without
        // a redundant overflow check.
        self.next_index = index.wrapping_add(1);
        if is_last {
            self.finalized = true;
        }

        Ok((plaintext, is_last))
    }

    /// Decrypt a specific chunk by index (random access).
    ///
    /// Does NOT advance the sequential counter. Does NOT set the `finalized`
    /// flag, regardless of `tag_byte`. Can be called after sequential
    /// finalization.
    pub fn decrypt_chunk_at(&self, index: u64, chunk: &[u8]) -> Result<(Zeroizing<Vec<u8>>, bool)> {
        self.decrypt_chunk_inner(index, chunk)
    }

    /// Return whether the decryptor has been finalized (a chunk with
    /// `tag_byte=0x01` was successfully decrypted via `decrypt_chunk`).
    pub fn is_finalized(&self) -> bool {
        self.finalized
    }

    /// Return the next expected sequential chunk index.
    pub fn expected_index(&self) -> u64 {
        self.next_index
    }

    /// Raw pointer to the internal key bytes (test-utils only).
    #[cfg(all(feature = "test-utils", debug_assertions))]
    #[deprecated(note = "test-utils only — do not call in production code")]
    pub fn key_ptr(&self) -> *const u8 {
        self.key.as_ptr()
    }

    /// Core decrypt logic shared by sequential and random-access paths.
    fn decrypt_chunk_inner(&self, index: u64, chunk: &[u8]) -> Result<(Zeroizing<Vec<u8>>, bool)> {
        // Minimum chunk: tag_byte (1) + AEAD tag (16) = 17 bytes.
        // Collapsed to AeadFailed — chunk bytes are attacker-controlled.
        if chunk.len() < STREAM_CHUNK_OVERHEAD {
            return Err(Error::AeadFailed);
        }

        let tag_byte = chunk[0];
        let ciphertext = &chunk[1..];
        let is_last = tag_byte == TAG_FINAL;

        // Pre-AEAD framing check for uncompressed non-final chunks.
        // Without compression, ciphertext size is deterministic:
        // STREAM_CHUNK_SIZE + AEAD_TAG_SIZE bytes. This check does not reveal
        // whether AEAD would have succeeded — the plaintext size is derivable
        // from the ciphertext size without any secret.
        if !self.compressed && !is_last && ciphertext.len() != STREAM_CHUNK_SIZE + AEAD_TAG_SIZE {
            return Err(Error::InvalidData);
        }

        let nonce = build_chunk_nonce(&self.base_nonce, index, tag_byte);
        let aad = build_chunk_aad(
            self.version,
            self.flags,
            &self.base_nonce,
            index,
            tag_byte,
            &self.caller_aad,
        );

        let decrypted = aead::aead_decrypt(&self.key, &nonce, ciphertext, &aad)?;

        // Post-AEAD processing: decompress if needed.
        if self.compressed && !decrypted.is_empty() {
            let decompressed = decompress_chunk(&decrypted)?;

            // Post-auth size enforcement for compressed non-final chunks.
            // Size is only known after decrypt + decompress, so this is
            // collapsed to AeadFailed to prevent a 1-bit oracle.
            if !is_last && decompressed.len() != STREAM_CHUNK_SIZE {
                return Err(Error::AeadFailed);
            }

            Ok((decompressed, is_last))
        } else {
            Ok((decrypted, is_last))
        }
    }
}

impl Drop for StreamDecryptor {
    fn drop(&mut self) {
        // key is Zeroizing<[u8; 32]> — zeroized automatically.
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::Error;
    use std::collections::HashSet;

    // ── Helper ───────────────────────────────────────────────────────────

    fn test_key() -> [u8; 32] {
        [0x42u8; 32]
    }

    fn chunk_plaintext(byte: u8) -> Vec<u8> {
        vec![byte; STREAM_CHUNK_SIZE]
    }

    // ── Round-trip ───────────────────────────────────────────────────────

    #[test]
    fn encrypt_decrypt_uncompressed() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let pt = chunk_plaintext(0xAA);
        let chunk = enc.encrypt_chunk(&pt, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
        assert!(is_last);
        assert_eq!(&*decrypted, &pt);
    }

    #[test]
    fn encrypt_decrypt_compressed() {
        let key = test_key();
        // Compressible data (all same byte).
        let mut enc = stream_encrypt_init(&key, b"", true).unwrap();
        let header = enc.header();
        let pt = chunk_plaintext(0xBB);
        let chunk = enc.encrypt_chunk(&pt, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
        assert!(is_last);
        assert_eq!(&*decrypted, &pt);
    }

    #[test]
    fn encrypt_decrypt_multi_chunk() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let pt0 = chunk_plaintext(0x01);
        let pt1 = chunk_plaintext(0x02);
        let pt2 = vec![0x03u8; 500]; // partial final

        let c0 = enc.encrypt_chunk(&pt0, false).unwrap();
        let c1 = enc.encrypt_chunk(&pt1, false).unwrap();
        let c2 = enc.encrypt_chunk(&pt2, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (d0, last0) = dec.decrypt_chunk(&c0).unwrap();
        assert!(!last0);
        assert_eq!(&*d0, &pt0);

        let (d1, last1) = dec.decrypt_chunk(&c1).unwrap();
        assert!(!last1);
        assert_eq!(&*d1, &pt1);

        let (d2, last2) = dec.decrypt_chunk(&c2).unwrap();
        assert!(last2);
        assert_eq!(&*d2, &pt2);
    }

    #[test]
    fn encrypt_decrypt_multi_chunk_compressed() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", true).unwrap();
        let header = enc.header();

        let pt0 = chunk_plaintext(0x10);
        let pt1 = chunk_plaintext(0x20);
        let pt2 = vec![0x30u8; 1000];

        let c0 = enc.encrypt_chunk(&pt0, false).unwrap();
        let c1 = enc.encrypt_chunk(&pt1, false).unwrap();
        let c2 = enc.encrypt_chunk(&pt2, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (d0, _) = dec.decrypt_chunk(&c0).unwrap();
        assert_eq!(&*d0, &pt0);
        let (d1, _) = dec.decrypt_chunk(&c1).unwrap();
        assert_eq!(&*d1, &pt1);
        let (d2, last) = dec.decrypt_chunk(&c2).unwrap();
        assert!(last);
        assert_eq!(&*d2, &pt2);
    }

    #[test]
    fn encrypt_decrypt_empty_file() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk(b"", true).unwrap();
        // tag_byte (1) + AEAD tag (16) = 17 bytes
        assert_eq!(chunk.len(), STREAM_CHUNK_OVERHEAD);

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
        assert!(is_last);
        assert!(decrypted.is_empty());
    }

    #[test]
    fn encrypt_decrypt_partial_final() {
        let key = test_key();
        for size in [1, STREAM_CHUNK_SIZE / 2, STREAM_CHUNK_SIZE] {
            let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
            let header = enc.header();
            let pt = vec![0xCCu8; size];
            let chunk = enc.encrypt_chunk(&pt, true).unwrap();

            let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
            let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
            assert!(is_last);
            assert_eq!(&*decrypted, &pt);
        }
    }

    #[test]
    fn encrypt_decrypt_with_caller_aad() {
        let key = test_key();
        let aad = b"file-id-12345";
        let mut enc = stream_encrypt_init(&key, aad, false).unwrap();
        let header = enc.header();
        let pt = chunk_plaintext(0xDD);
        let chunk = enc.encrypt_chunk(&pt, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, aad).unwrap();
        let (decrypted, _) = dec.decrypt_chunk(&chunk).unwrap();
        assert_eq!(&*decrypted, &pt);
    }

    // ── Random access ────────────────────────────────────────────────────

    #[test]
    fn decrypt_chunk_at_middle() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let mut chunks = Vec::new();
        let mut plaintexts = Vec::new();
        for i in 0..4u8 {
            let pt = chunk_plaintext(i);
            chunks.push(enc.encrypt_chunk(&pt, false).unwrap());
            plaintexts.push(pt);
        }
        let final_pt = vec![0xFFu8; 100];
        chunks.push(enc.encrypt_chunk(&final_pt, true).unwrap());
        plaintexts.push(final_pt);

        let dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk_at(2, &chunks[2]).unwrap();
        assert!(!is_last);
        assert_eq!(&*decrypted, &plaintexts[2]);
    }

    #[test]
    fn decrypt_chunk_at_final() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let pt0 = chunk_plaintext(0x01);
        let _ = enc.encrypt_chunk(&pt0, false).unwrap();
        let final_pt = vec![0xEEu8; 50];
        let final_chunk = enc.encrypt_chunk(&final_pt, true).unwrap();

        let dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk_at(1, &final_chunk).unwrap();
        assert!(is_last);
        assert_eq!(&*decrypted, &final_pt);
    }

    #[test]
    fn decrypt_chunk_at_wrong_index() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let pt0 = chunk_plaintext(0x01);
        let pt1 = chunk_plaintext(0x02);
        let pt2 = vec![0x03u8; 10];
        let _ = enc.encrypt_chunk(&pt0, false).unwrap();
        let _ = enc.encrypt_chunk(&pt1, false).unwrap();
        let c2 = enc.encrypt_chunk(&pt2, true).unwrap();

        let dec = stream_decrypt_init(&key, &header, b"").unwrap();
        // Chunk 2 at index 1 — AAD mismatch.
        assert!(matches!(
            dec.decrypt_chunk_at(1, &c2),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn decrypt_chunk_at_does_not_advance_counter() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let pt0 = chunk_plaintext(0x01);
        let c0 = enc.encrypt_chunk(&pt0, false).unwrap();
        let pt1 = vec![0x02u8; 100];
        let c1 = enc.encrypt_chunk(&pt1, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        // Random-access chunk 1 first.
        let _ = dec.decrypt_chunk_at(1, &c1).unwrap();
        assert_eq!(dec.expected_index(), 0);

        // Sequential decrypt still starts at 0.
        let (d0, _) = dec.decrypt_chunk(&c0).unwrap();
        assert_eq!(&*d0, &pt0);
        assert_eq!(dec.expected_index(), 1);
    }

    #[test]
    fn decrypt_chunk_at_does_not_finalize() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let pt = vec![0xAAu8; 10];
        let chunk = enc.encrypt_chunk(&pt, true).unwrap();

        let dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (_, is_last) = dec.decrypt_chunk_at(0, &chunk).unwrap();
        assert!(is_last);
        assert!(!dec.is_finalized());
    }

    #[test]
    fn decrypt_chunk_at_after_sequential_finalization() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let pt0 = chunk_plaintext(0x01);
        let c0 = enc.encrypt_chunk(&pt0, false).unwrap();
        let pt1 = vec![0x02u8; 100];
        let c1 = enc.encrypt_chunk(&pt1, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let _ = dec.decrypt_chunk(&c0).unwrap();
        let _ = dec.decrypt_chunk(&c1).unwrap();
        assert!(dec.is_finalized());

        // Random access still works after finalization.
        let (d0, _) = dec.decrypt_chunk_at(0, &c0).unwrap();
        assert_eq!(&*d0, &pt0);
    }

    // ── Tamper detection ─────────────────────────────────────────────────

    #[test]
    fn tampered_ciphertext() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let pt = chunk_plaintext(0xAA);
        let mut chunk = enc.encrypt_chunk(&pt, true).unwrap();
        // Flip a byte in the ciphertext region (after tag_byte).
        chunk[5] ^= 0xFF;

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        assert!(matches!(dec.decrypt_chunk(&chunk), Err(Error::AeadFailed)));
    }

    #[test]
    fn tampered_tag_byte() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let pt = chunk_plaintext(0xAA);
        let mut chunk = enc.encrypt_chunk(&pt, true).unwrap();
        // Flip tag_byte from 0x01 (final) to 0x00 (non-final).
        chunk[0] ^= 0x01;

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        // Nonce AND AAD both change — AEAD authentication fails.
        assert!(matches!(
            dec.decrypt_chunk(&chunk),
            Err(Error::AeadFailed | Error::InvalidData)
        ));
    }

    #[test]
    fn wrong_key() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let pt = chunk_plaintext(0xAA);
        let chunk = enc.encrypt_chunk(&pt, true).unwrap();

        let wrong_key = [0x99u8; 32];
        let mut dec = stream_decrypt_init(&wrong_key, &header, b"").unwrap();
        assert!(matches!(dec.decrypt_chunk(&chunk), Err(Error::AeadFailed)));
    }

    #[test]
    fn wrong_caller_aad() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"file-123", false).unwrap();
        let header = enc.header();
        let pt = chunk_plaintext(0xAA);
        let chunk = enc.encrypt_chunk(&pt, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"file-456").unwrap();
        assert!(matches!(dec.decrypt_chunk(&chunk), Err(Error::AeadFailed)));
    }

    // ── Ordering attacks ─────────────────────────────────────────────────

    #[test]
    fn reordered_chunks() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let _c0 = enc.encrypt_chunk(&chunk_plaintext(0x01), false).unwrap();
        let c1 = enc.encrypt_chunk(&chunk_plaintext(0x02), false).unwrap();
        let _c2 = enc.encrypt_chunk(&[0x03u8; 100], true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        // Feed chunk 1 at index 0 — AAD mismatch.
        assert!(matches!(dec.decrypt_chunk(&c1), Err(Error::AeadFailed)));
    }

    #[test]
    fn duplicated_chunk() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let c0 = enc.encrypt_chunk(&chunk_plaintext(0x01), false).unwrap();
        let _c1 = enc.encrypt_chunk(&[0x02u8; 100], true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let _ = dec.decrypt_chunk(&c0).unwrap();
        // Feed chunk 0 again at index 1.
        assert!(matches!(dec.decrypt_chunk(&c0), Err(Error::AeadFailed)));
    }

    #[test]
    fn truncation_detected() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let c0 = enc.encrypt_chunk(&chunk_plaintext(0x01), false).unwrap();
        let c1 = enc.encrypt_chunk(&chunk_plaintext(0x02), false).unwrap();
        let _c2 = enc.encrypt_chunk(&[0x03u8; 100], true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let _ = dec.decrypt_chunk(&c0).unwrap();
        let _ = dec.decrypt_chunk(&c1).unwrap();
        // Stop — never received the final chunk.
        assert!(!dec.is_finalized());
    }

    #[test]
    fn cross_stream_splice() {
        let key_a = [0x11u8; 32];
        let key_b = [0x22u8; 32];
        let pt = chunk_plaintext(0xAA);

        let mut enc_a = stream_encrypt_init(&key_a, b"", false).unwrap();
        let _header_a = enc_a.header();
        let chunk_a = enc_a.encrypt_chunk(&pt, true).unwrap();

        let enc_b = stream_encrypt_init(&key_b, b"", false).unwrap();
        let header_b = enc_b.header();

        // Feed stream A's chunk to stream B's decryptor.
        let mut dec_b = stream_decrypt_init(&key_b, &header_b, b"").unwrap();
        assert!(matches!(
            dec_b.decrypt_chunk(&chunk_a),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn cross_stream_splice_same_key() {
        // Same key, different streams — isolation relies on the random base
        // nonce bound into the AAD. Verifies that two streams sharing a key
        // cannot cross-decrypt.
        let key = test_key();
        let pt = chunk_plaintext(0xCC);

        let mut enc_1 = stream_encrypt_init(&key, b"", false).unwrap();
        let _header_1 = enc_1.header();
        let chunk_1 = enc_1.encrypt_chunk(&pt, true).unwrap();

        let enc_2 = stream_encrypt_init(&key, b"", false).unwrap();
        let header_2 = enc_2.header();

        let mut dec_2 = stream_decrypt_init(&key, &header_2, b"").unwrap();
        assert!(matches!(
            dec_2.decrypt_chunk(&chunk_1),
            Err(Error::AeadFailed)
        ));
    }

    // ── Compression edge cases ───────────────────────────────────────────

    #[test]
    fn compression_bypassed_for_empty_final() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", true).unwrap();
        let chunk = enc.encrypt_chunk(b"", true).unwrap();
        // tag_byte (1) + AEAD tag (16) = 17, regardless of compress flag.
        assert_eq!(chunk.len(), STREAM_CHUNK_OVERHEAD);

        let header = enc.header();
        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
        assert!(is_last);
        assert!(decrypted.is_empty());
    }

    #[test]
    #[allow(clippy::cast_possible_truncation)]
    fn incompressible_data_round_trips() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", true).unwrap();
        let header = enc.header();
        // Random-ish data that won't compress well.
        let pt: Vec<u8> = (0..STREAM_CHUNK_SIZE).map(|i| (i % 256) as u8).collect();
        let chunk = enc.encrypt_chunk(&pt, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, _) = dec.decrypt_chunk(&chunk).unwrap();
        assert_eq!(&*decrypted, &pt);
    }

    #[test]
    fn compression_flag_flip_detected() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", true).unwrap();
        let mut header = enc.header();
        let pt = chunk_plaintext(0xCC);
        let chunk = enc.encrypt_chunk(&pt, true).unwrap();

        // Flip compression flag in header.
        header[1] ^= FLAG_COMPRESSED;

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        // AAD mismatch — flags are bound into AAD.
        assert!(matches!(dec.decrypt_chunk(&chunk), Err(Error::AeadFailed)));
    }

    // ── Error paths ──────────────────────────────────────────────────────

    #[test]
    fn non_final_chunk_short() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let short = vec![0u8; STREAM_CHUNK_SIZE - 1];
        assert!(matches!(
            enc.encrypt_chunk(&short, false),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn non_final_chunk_long() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let long = vec![0u8; STREAM_CHUNK_SIZE + 1];
        assert!(matches!(
            enc.encrypt_chunk(&long, false),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn final_chunk_over_chunk_size() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let big = vec![0u8; STREAM_CHUNK_SIZE + 1];
        assert!(matches!(
            enc.encrypt_chunk(&big, true),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn encrypt_after_finalization() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        enc.encrypt_chunk(b"done", true).unwrap();
        assert!(matches!(
            enc.encrypt_chunk(b"more", true),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn decrypt_after_finalization() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk(b"final", true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let _ = dec.decrypt_chunk(&chunk).unwrap();
        assert!(dec.is_finalized());
        assert!(matches!(dec.decrypt_chunk(&chunk), Err(Error::InvalidData)));
    }

    #[test]
    fn encrypt_chain_exhausted() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        enc.next_index = u64::MAX;
        assert!(matches!(
            enc.encrypt_chunk(b"data", true),
            Err(Error::ChainExhausted)
        ));
        // Random-access API has no sequential guard — still usable.
        assert!(enc.encrypt_chunk_at(0, true, b"data").is_ok());
    }

    #[test]
    fn decrypt_chain_exhausted() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk(b"data", true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        dec.next_index = u64::MAX;
        assert!(matches!(
            dec.decrypt_chunk(&chunk),
            Err(Error::ChainExhausted)
        ));
        // Random-access API has no sequential guard — still usable at index 0.
        let (pt, is_last) = dec.decrypt_chunk_at(0, &chunk).unwrap();
        assert_eq!(&*pt, b"data");
        assert!(is_last);
    }

    #[test]
    fn chunk_too_short() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let _ = enc.encrypt_chunk(b"x", true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let short = vec![0u8; 16]; // < 17 minimum
        assert!(matches!(dec.decrypt_chunk(&short), Err(Error::AeadFailed)));
    }

    #[test]
    fn uncompressed_non_final_wrong_ciphertext_size() {
        let key = test_key();
        let enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        // Construct a fake chunk with tag_byte=0x00 (non-final) but wrong
        // ciphertext size.
        let mut fake_chunk = vec![TAG_NON_FINAL];
        fake_chunk.extend_from_slice(&[0u8; 100]); // wrong size

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        assert!(matches!(
            dec.decrypt_chunk(&fake_chunk),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn decrypt_chunk_at_uncompressed_non_final_wrong_size() {
        let key = test_key();
        let enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let mut fake_chunk = vec![TAG_NON_FINAL];
        fake_chunk.extend_from_slice(&[0u8; 100]); // wrong size

        let dec = stream_decrypt_init(&key, &header, b"").unwrap();
        assert!(matches!(
            dec.decrypt_chunk_at(0, &fake_chunk),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn compressed_non_final_wrong_decompressed_size() {
        // Exercises the defense-in-depth path at decrypt_chunk_inner line 501:
        // a compressed non-final chunk that passes AEAD but decompresses to a
        // size != STREAM_CHUNK_SIZE returns AeadFailed (oracle prevention).
        //
        // Craft a valid AEAD ciphertext for a compressed non-final chunk whose
        // compressed payload decompresses to only 100 bytes (not STREAM_CHUNK_SIZE).
        use crate::primitives::aead;

        let key = test_key();
        let enc = stream_encrypt_init(&key, b"", true).unwrap();
        let header = enc.header();
        let base_nonce: [u8; AEAD_NONCE_SIZE] = header[2..].try_into().unwrap();

        // Compress a 100-byte plaintext — decompressed size will be 100, not STREAM_CHUNK_SIZE.
        let short_pt = vec![0xAA; 100];
        let compressed = compress_chunk(&short_pt).unwrap();

        let index = 0u64;
        let tag_byte = TAG_NON_FINAL;
        let nonce = build_chunk_nonce(&base_nonce, index, tag_byte);
        let aad = build_chunk_aad(
            STREAM_VERSION,
            FLAG_COMPRESSED,
            &base_nonce,
            index,
            tag_byte,
            b"",
        );
        let ciphertext = aead::aead_encrypt(&key, &nonce, &compressed, &aad).unwrap();

        let mut fake_chunk = vec![tag_byte];
        fake_chunk.extend_from_slice(&ciphertext);

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        assert!(matches!(
            dec.decrypt_chunk(&fake_chunk),
            Err(Error::AeadFailed)
        ));
    }

    // ── State & accessors ────────────────────────────────────────────────

    #[test]
    fn is_finalized_false_before_last() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let _ = enc.encrypt_chunk(&chunk_plaintext(0x01), false).unwrap();
        assert!(!enc.is_finalized());
    }

    #[test]
    fn is_finalized_true_after_last() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let _ = enc.encrypt_chunk(b"done", true).unwrap();
        assert!(enc.is_finalized());
    }

    #[test]
    fn expected_index_increments() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let c0 = enc.encrypt_chunk(&chunk_plaintext(0x01), false).unwrap();
        let c1 = enc.encrypt_chunk(&chunk_plaintext(0x02), false).unwrap();
        let c2 = enc.encrypt_chunk(&[0x03u8; 100], true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        assert_eq!(dec.expected_index(), 0);
        let _ = dec.decrypt_chunk(&c0).unwrap();
        assert_eq!(dec.expected_index(), 1);
        let _ = dec.decrypt_chunk(&c1).unwrap();
        assert_eq!(dec.expected_index(), 2);
        let _ = dec.decrypt_chunk(&c2).unwrap();
        assert_eq!(dec.expected_index(), 3);
    }

    #[test]
    fn error_does_not_advance_state() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let c0 = enc.encrypt_chunk(&chunk_plaintext(0x01), false).unwrap();
        let _c1 = enc.encrypt_chunk(&[0x02u8; 100], true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();

        // Feed corrupted chunk at index 0.
        let mut bad = c0.clone();
        bad[5] ^= 0xFF;
        assert!(dec.decrypt_chunk(&bad).is_err());
        assert_eq!(dec.expected_index(), 0);

        // Retry with the correct chunk succeeds.
        let (d0, _) = dec.decrypt_chunk(&c0).unwrap();
        assert_eq!(&*d0, &chunk_plaintext(0x01));
        assert_eq!(dec.expected_index(), 1);
    }

    // ── Header validation ────────────────────────────────────────────────

    #[test]
    fn decrypt_init_wrong_version() {
        let key = test_key();
        let mut header = [0u8; STREAM_HEADER_SIZE];
        header[0] = 0x00; // version 0
        assert!(matches!(
            stream_decrypt_init(&key, &header, b""),
            Err(Error::UnsupportedVersion)
        ));
    }

    #[test]
    fn decrypt_init_version_0x02() {
        let key = test_key();
        let mut header = [0u8; STREAM_HEADER_SIZE];
        header[0] = 0x02;
        assert!(matches!(
            stream_decrypt_init(&key, &header, b""),
            Err(Error::UnsupportedVersion)
        ));
    }

    #[test]
    fn decrypt_init_reserved_flags() {
        let key = test_key();
        let mut header = [0u8; STREAM_HEADER_SIZE];
        header[0] = STREAM_VERSION;
        header[1] = 0x02; // reserved bit set
        assert!(matches!(
            stream_decrypt_init(&key, &header, b""),
            Err(Error::AeadFailed)
        ));
    }

    // ── Nonce derivation ─────────────────────────────────────────────────

    #[test]
    fn nonce_uniqueness_same_index_different_tag() {
        let base = [0u8; AEAD_NONCE_SIZE];
        let n0 = build_chunk_nonce(&base, 5, TAG_NON_FINAL);
        let n1 = build_chunk_nonce(&base, 5, TAG_FINAL);
        assert_ne!(n0, n1);
    }

    #[test]
    fn nonce_uniqueness_different_indices() {
        let base = [0u8; AEAD_NONCE_SIZE];
        let n0 = build_chunk_nonce(&base, 0, TAG_NON_FINAL);
        let n1 = build_chunk_nonce(&base, 1, TAG_NON_FINAL);
        assert_ne!(n0, n1);
    }

    #[test]
    fn nonce_injectivity() {
        let base: [u8; AEAD_NONCE_SIZE] = [0x77; AEAD_NONCE_SIZE];
        let mut seen = HashSet::new();
        for idx in 0..1024u64 {
            for &tag in &[TAG_NON_FINAL, TAG_FINAL] {
                let nonce = build_chunk_nonce(&base, idx, tag);
                seen.insert(nonce);
            }
        }
        assert_eq!(seen.len(), 2048);
    }

    // ── Nonce freshness ──────────────────────────────────────────────────

    #[test]
    fn base_nonce_freshness() {
        let key = test_key();
        let enc1 = stream_encrypt_init(&key, b"", false).unwrap();
        let enc2 = stream_encrypt_init(&key, b"", false).unwrap();
        let h1 = enc1.header();
        let h2 = enc2.header();
        // Base nonces (bytes 2..26) should differ.
        assert_ne!(&h1[2..], &h2[2..]);
    }

    // ── KAT (Known-Answer Tests) ────────────────────────────────────────

    #[test]
    fn nonce_derivation_kat() {
        let base_nonce: [u8; AEAD_NONCE_SIZE] = [
            0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E,
            0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18,
        ];
        // Index 3, non-final: mask = [0,0,0,0,0,0,0,3, 0, 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
        let nonce = build_chunk_nonce(&base_nonce, 3, TAG_NON_FINAL);
        let mut expected = base_nonce;
        expected[7] ^= 3; // only byte 7 differs (index=3 in BE u64)
        assert_eq!(nonce, expected);

        // Index 3, final: mask[8] = 0x01
        let nonce_final = build_chunk_nonce(&base_nonce, 3, TAG_FINAL);
        let mut expected_final = base_nonce;
        expected_final[7] ^= 3;
        expected_final[8] ^= 1;
        assert_eq!(nonce_final, expected_final);
    }

    #[test]
    fn aad_construction_kat() {
        let base_nonce = [0xAA; AEAD_NONCE_SIZE];
        let caller_aad = b"test-ctx";
        let aad = build_chunk_aad(0x01, 0x00, &base_nonce, 7, TAG_FINAL, caller_aad);

        // Use literal label to catch constant drift.
        let mut expected = Vec::new();
        expected.extend_from_slice(b"lo-stream-v1");
        expected.push(0x01); // version
        expected.push(0x00); // flags
        expected.extend_from_slice(&[0xAA; 24]); // base_nonce
        expected.extend_from_slice(&7u64.to_be_bytes()); // index
        expected.push(TAG_FINAL); // tag_byte
        expected.extend_from_slice(b"test-ctx"); // caller_aad

        assert_eq!(aad, expected);
    }

    #[test]
    fn nonce_derivation_f9_spec_vectors() {
        // Specification.md Appendix F.9 — all six nonce derivation vectors.
        use hex_literal::hex;

        let base_nonce: [u8; AEAD_NONCE_SIZE] =
            hex!("101112131415161718191a1b1c1d1e1f2021222324252627");

        // (index=0, tag=0x00)
        assert_eq!(
            build_chunk_nonce(&base_nonce, 0, 0x00),
            hex!("101112131415161718191a1b1c1d1e1f2021222324252627")
        );
        // (index=2, tag=0x00)
        assert_eq!(
            build_chunk_nonce(&base_nonce, 2, 0x00),
            hex!("101112131415161518191a1b1c1d1e1f2021222324252627")
        );
        // (index=0, tag=0x01) — single-chunk final
        assert_eq!(
            build_chunk_nonce(&base_nonce, 0, 0x01),
            hex!("101112131415161719191a1b1c1d1e1f2021222324252627")
        );
        // (index=2, tag=0x01)
        assert_eq!(
            build_chunk_nonce(&base_nonce, 2, 0x01),
            hex!("101112131415161519191a1b1c1d1e1f2021222324252627")
        );
        // u64::MAX boundary — detects 32-bit truncation
        // (index=u64::MAX, tag=0x00)
        assert_eq!(
            build_chunk_nonce(&base_nonce, u64::MAX, 0x00),
            hex!("efeeedecebeae9e818191a1b1c1d1e1f2021222324252627")
        );
        // (index=u64::MAX, tag=0x01)
        assert_eq!(
            build_chunk_nonce(&base_nonce, u64::MAX, 0x01),
            hex!("efeeedecebeae9e819191a1b1c1d1e1f2021222324252627")
        );
    }

    #[test]
    fn aad_construction_f10_spec_vectors() {
        // Specification.md Appendix F.10 — all three AAD construction vectors.
        use hex_literal::hex;

        let base_nonce: [u8; AEAD_NONCE_SIZE] =
            hex!("101112131415161718191a1b1c1d1e1f2021222324252627");

        // (index=0, tag=0x00, caller_aad="")
        assert_eq!(
            build_chunk_aad(0x01, 0x00, &base_nonce, 0, 0x00, b""),
            hex!("6c6f2d73747265616d2d76310100101112131415161718191a1b1c1d1e1f2021222324252627000000000000000000").to_vec()
        );
        // (index=0, tag=0x00, caller_aad="file-abc-123")
        assert_eq!(
            build_chunk_aad(0x01, 0x00, &base_nonce, 0, 0x00, b"file-abc-123"),
            hex!("6c6f2d73747265616d2d76310100101112131415161718191a1b1c1d1e1f202122232425262700000000000000000066696c652d6162632d313233").to_vec()
        );
        // (index=2, tag=0x01, caller_aad="file-abc-123")
        assert_eq!(
            build_chunk_aad(0x01, 0x00, &base_nonce, 2, 0x01, b"file-abc-123"),
            hex!("6c6f2d73747265616d2d76310100101112131415161718191a1b1c1d1e1f202122232425262700000000000000020166696c652d6162632d313233").to_vec()
        );
    }

    #[test]
    fn header_construction_kat() {
        let key = test_key();
        let enc = stream_encrypt_init(&key, b"", false).unwrap();
        let hdr = enc.header();
        assert_eq!(hdr[0], STREAM_VERSION);
        assert_eq!(hdr[1], 0x00); // no compression
        assert_eq!(hdr.len(), STREAM_HEADER_SIZE);

        let enc2 = stream_encrypt_init(&key, b"", true).unwrap();
        let hdr2 = enc2.header();
        assert_eq!(hdr2[1], FLAG_COMPRESSED);
    }

    #[test]
    fn encrypt_chunk_kat() {
        // Genuine KAT: expected ciphertext computed by an independent
        // XChaCha20-Poly1305 implementation (PyNaCl / libsodium) with
        // manually constructed nonce and AAD matching the Rust helpers.
        use hex_literal::hex;

        let key: [u8; 32] =
            hex!("0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20");
        let base_nonce: [u8; AEAD_NONCE_SIZE] =
            hex!("aabbccddeeff00112233445566778899aabbccddeeff0011");
        let plaintext = b"soliton-stream-kat";
        let caller_aad = b"kat-ctx";

        // Expected chunk wire bytes: tag_byte (0x01) || aead_output (18 + 16).
        // Produced by PyNaCl crypto_aead_xchacha20poly1305_ietf_encrypt with:
        //   nonce = base_nonce XOR (index_be8 || tag_byte || zeros)
        //   aad   = "lo-stream-v1" || 0x01 || 0x00 || base_nonce || 0u64_be || 0x01 || "kat-ctx"
        let expected_chunk: [u8; 35] =
            hex!("01544029f4e46ef22cab476995f622910ae950ab67efff477ce0b46f25bf4237e9a9b6");

        // Manually compute what encrypt_chunk would produce for chunk 0, final, no compression.
        let tag_byte = TAG_FINAL;
        let index = 0u64;
        let nonce = build_chunk_nonce(&base_nonce, index, tag_byte);
        let aad = build_chunk_aad(
            STREAM_VERSION,
            0x00,
            &base_nonce,
            index,
            tag_byte,
            caller_aad,
        );
        let ciphertext =
            crate::primitives::aead::aead_encrypt(&key, &nonce, plaintext, &aad).unwrap();

        let mut output = Vec::with_capacity(1 + ciphertext.len());
        output.push(tag_byte);
        output.extend_from_slice(&ciphertext);

        assert_eq!(
            output, expected_chunk,
            "encrypt_chunk output does not match independent reference vector"
        );

        // Verify round-trip via decrypt helpers.
        let decrypted =
            crate::primitives::aead::aead_decrypt(&key, &nonce, &ciphertext, &aad).unwrap();
        assert_eq!(&*decrypted, plaintext);
    }

    // ── Small-data tests (MIRI-friendly) ────────────────────────────────
    //
    // These duplicate coverage of the full-chunk tests but use tiny
    // plaintext (≤ 64 bytes), keeping MIRI runtime under a few seconds.
    // They exercise the encrypt→decrypt round-trip, multi-chunk
    // sequencing, random access, compression, and error recovery paths
    // that would otherwise be excluded from MIRI due to 1 MiB AEAD
    // timeout.

    #[test]
    fn miri_round_trip_small_uncompressed() {
        let key = test_key();
        let pt = b"hello streaming aead";
        let mut enc = stream_encrypt_init(&key, b"ctx", false).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk(pt, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"ctx").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
        assert!(is_last);
        assert_eq!(&*decrypted, pt);
    }

    #[test]
    fn miri_round_trip_small_compressed() {
        let key = test_key();
        let pt = b"compress me please";
        let mut enc = stream_encrypt_init(&key, b"", true).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk(pt, true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
        assert!(is_last);
        assert_eq!(&*decrypted, pt);
    }

    #[test]
    fn miri_multi_chunk_small() {
        let key = test_key();
        let pt_final = b"final chunk data";
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        // Only a final chunk (no non-final, since those require exactly CHUNK_SIZE).
        let c0 = enc.encrypt_chunk(pt_final, true).unwrap();
        assert!(enc.is_finalized());

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (d0, is_last) = dec.decrypt_chunk(&c0).unwrap();
        assert!(is_last);
        assert_eq!(&*d0, pt_final);
        assert!(dec.is_finalized());
    }

    #[test]
    fn miri_random_access_small() {
        let key = test_key();
        let pt = b"random access chunk";
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk(pt, true).unwrap();

        let dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk_at(0, &chunk).unwrap();
        assert!(is_last);
        assert_eq!(&*decrypted, pt);
    }

    #[test]
    fn miri_wrong_key_small() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk(b"secret", true).unwrap();

        let wrong_key = [0xBB; 32];
        let mut dec = stream_decrypt_init(&wrong_key, &header, b"").unwrap();
        assert!(matches!(dec.decrypt_chunk(&chunk), Err(Error::AeadFailed)));
    }

    #[test]
    fn miri_tampered_ciphertext_small() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let mut chunk = enc.encrypt_chunk(b"data", true).unwrap();
        // Flip a byte in the ciphertext payload.
        chunk[5] ^= 0xFF;

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        assert!(matches!(dec.decrypt_chunk(&chunk), Err(Error::AeadFailed)));
    }

    #[test]
    fn miri_error_does_not_advance_state_small() {
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();
        let good = enc.encrypt_chunk(b"good", true).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        assert_eq!(dec.expected_index(), 0);

        // Bad chunk should not advance state.
        let bad = vec![0u8; 17];
        let _ = dec.decrypt_chunk(&bad);
        assert_eq!(dec.expected_index(), 0);
        assert!(!dec.is_finalized());

        // Good chunk should still work.
        let (pt, is_last) = dec.decrypt_chunk(&good).unwrap();
        assert!(is_last);
        assert_eq!(&*pt, b"good");
        assert_eq!(dec.expected_index(), 1);
    }

    // ── encrypt_chunk_at ─────────────────────────────────────────────────

    #[test]
    fn encrypt_chunk_at_single_chunk_round_trip() {
        let key = test_key();
        let pt = b"random-access encrypt";
        let enc = stream_encrypt_init(&key, b"ctx", false).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk_at(0, true, pt).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"ctx").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
        assert!(is_last);
        assert_eq!(&*decrypted, pt);
    }

    #[test]
    fn encrypt_chunk_at_multi_chunk_in_order() {
        // Encrypt chunks 0, 1, 2 via encrypt_chunk_at in index order,
        // assemble the stream, decrypt sequentially.
        let key = test_key();
        let enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let pt0 = chunk_plaintext(0xAA);
        let pt1 = chunk_plaintext(0xBB);
        let pt2 = b"final small";

        let c0 = enc.encrypt_chunk_at(0, false, &pt0).unwrap();
        let c1 = enc.encrypt_chunk_at(1, false, &pt1).unwrap();
        let c2 = enc.encrypt_chunk_at(2, true, pt2).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (d0, last0) = dec.decrypt_chunk(&c0).unwrap();
        assert!(!last0);
        assert_eq!(&*d0, &pt0);
        let (d1, last1) = dec.decrypt_chunk(&c1).unwrap();
        assert!(!last1);
        assert_eq!(&*d1, &pt1);
        let (d2, last2) = dec.decrypt_chunk(&c2).unwrap();
        assert!(last2);
        assert_eq!(&*d2, pt2);
    }

    #[test]
    fn encrypt_chunk_at_out_of_order_encrypt() {
        // Encrypt chunks 2, 0, 1 in that order, then assemble in index order
        // and decrypt sequentially. Validates that out-of-order encrypt
        // produces valid ciphertext: nonces are index-derived, not order-dependent.
        let key = test_key();
        let enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let pt0 = chunk_plaintext(0x11);
        let pt1 = chunk_plaintext(0x22);
        let pt2 = b"last";

        // Encrypt out of order.
        let c2 = enc.encrypt_chunk_at(2, true, pt2).unwrap();
        let c0 = enc.encrypt_chunk_at(0, false, &pt0).unwrap();
        let c1 = enc.encrypt_chunk_at(1, false, &pt1).unwrap();

        // Decrypt in index order.
        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (d0, _) = dec.decrypt_chunk(&c0).unwrap();
        assert_eq!(&*d0, &pt0);
        let (d1, _) = dec.decrypt_chunk(&c1).unwrap();
        assert_eq!(&*d1, &pt1);
        let (d2, last) = dec.decrypt_chunk(&c2).unwrap();
        assert!(last);
        assert_eq!(&*d2, pt2);
    }

    #[test]
    fn encrypt_chunk_at_wrong_index_aead_fails() {
        // Ciphertext produced at index 0 cannot be decrypted at index 1:
        // both nonce and AAD embed the chunk index, so any mismatch fails AEAD.
        let key = test_key();
        let enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        let pt = b"index mismatch test";
        let chunk_at_0 = enc.encrypt_chunk_at(0, true, pt).unwrap();

        let dec = stream_decrypt_init(&key, &header, b"").unwrap();
        // Feed chunk 0's ciphertext to decrypt_chunk_at with index 1 — must fail.
        assert!(matches!(
            dec.decrypt_chunk_at(1, &chunk_at_0),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn encrypt_chunk_at_does_not_advance_sequential_counter() {
        // Calling encrypt_chunk_at must not mutate next_index: the sequential
        // path must still produce a valid chunk at index 0 afterward.
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let header = enc.header();

        // Random-access encrypt at index 5 — sequential counter stays at 0.
        let pt5 = b"skipped ahead";
        let _ = enc.encrypt_chunk_at(5, true, pt5).unwrap();

        // Sequential encrypt must still start at index 0.
        let pt0 = b"sequential start";
        let c0 = enc.encrypt_chunk(pt0, true).unwrap();
        assert!(enc.is_finalized());

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (d0, last) = dec.decrypt_chunk(&c0).unwrap();
        assert!(last);
        assert_eq!(&*d0, pt0);
    }

    #[test]
    fn encrypt_chunk_at_does_not_finalize() {
        // encrypt_chunk_at must not set finalized, even when is_last=true.
        let key = test_key();
        let mut enc = stream_encrypt_init(&key, b"", false).unwrap();
        let pt = b"not yet final";
        let _ = enc.encrypt_chunk_at(0, true, pt).unwrap();
        // finalized must still be false — encrypt_chunk_at does not seal the stream.
        assert!(!enc.is_finalized());
        // Sequential encrypt_chunk must still accept calls.
        let _ = enc.encrypt_chunk(pt, true).unwrap();
        assert!(enc.is_finalized());
    }

    #[test]
    fn encrypt_chunk_at_invalid_sizes() {
        let key = test_key();
        let enc = stream_encrypt_init(&key, b"", false).unwrap();

        // Non-final chunk not exactly STREAM_CHUNK_SIZE.
        let short = vec![0u8; STREAM_CHUNK_SIZE - 1];
        assert!(matches!(
            enc.encrypt_chunk_at(0, false, &short),
            Err(Error::InvalidData)
        ));

        // Final chunk larger than STREAM_CHUNK_SIZE.
        let big = vec![0u8; STREAM_CHUNK_SIZE + 1];
        assert!(matches!(
            enc.encrypt_chunk_at(0, true, &big),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn encrypt_chunk_at_with_compression() {
        // Verify that compress=true streams round-trip correctly via encrypt_chunk_at.
        let key = test_key();
        let pt = b"compressible payload for random-access encrypt";
        let enc = stream_encrypt_init(&key, b"", true).unwrap();
        let header = enc.header();
        let chunk = enc.encrypt_chunk_at(0, true, pt).unwrap();

        let mut dec = stream_decrypt_init(&key, &header, b"").unwrap();
        let (decrypted, is_last) = dec.decrypt_chunk(&chunk).unwrap();
        assert!(is_last);
        assert_eq!(&*decrypted, pt);
    }

    // ── Proptest ─────────────────────────────────────────────────────────

    mod proptests {
        use super::*;
        use proptest::prelude::*;

        // These proptests generate up to ~4-5 MiB of plaintext per case
        // (multi-chunk streams) and run 256 cases by default. Expect
        // ~30-60s total. Ctrl+C may be delayed — proptest installs a
        // signal handler to persist its failure database before exiting.
        proptest! {
            #[test]
            fn proptest_stream_round_trip(
                key in prop::array::uniform32(0u8..),
                plaintext in prop::collection::vec(any::<u8>(), 0..STREAM_CHUNK_SIZE * 4),
                compress in any::<bool>(),
                caller_aad in prop::collection::vec(any::<u8>(), 0..64),
            ) {
                let mut enc = stream_encrypt_init(&key, &caller_aad, compress).unwrap();
                let header = enc.header();

                let mut chunks = Vec::new();
                if plaintext.is_empty() {
                    chunks.push(enc.encrypt_chunk(b"", true).unwrap());
                } else {
                    let full_count = plaintext.len() / STREAM_CHUNK_SIZE;
                    for i in 0..full_count {
                        let start = i * STREAM_CHUNK_SIZE;
                        let end = start + STREAM_CHUNK_SIZE;
                        let remainder = &plaintext[end..];
                        if remainder.is_empty() {
                            // This full chunk is also the last.
                            chunks.push(enc.encrypt_chunk(&plaintext[start..end], true).unwrap());
                        } else {
                            chunks.push(enc.encrypt_chunk(&plaintext[start..end], false).unwrap());
                        }
                    }
                    if !enc.is_finalized() {
                        let start = full_count * STREAM_CHUNK_SIZE;
                        chunks.push(enc.encrypt_chunk(&plaintext[start..], true).unwrap());
                    }
                }

                let mut dec = stream_decrypt_init(&key, &header, &caller_aad).unwrap();
                let mut recovered = Vec::new();
                for ct in &chunks {
                    let (pt, _) = dec.decrypt_chunk(ct).unwrap();
                    recovered.extend_from_slice(&pt);
                }
                prop_assert_eq!(recovered, plaintext);
            }

            #[test]
            fn proptest_random_access_any_chunk(
                key in prop::array::uniform32(0u8..),
                compress in any::<bool>(),
                // 1..=5 chunks worth of data to keep runtime reasonable.
                plaintext in prop::collection::vec(any::<u8>(), 1..STREAM_CHUNK_SIZE * 5),
            ) {
                let mut enc = stream_encrypt_init(&key, b"", compress).unwrap();
                let header = enc.header();

                let mut chunks = Vec::new();
                let mut chunk_plaintexts: Vec<Vec<u8>> = Vec::new();
                let full_count = plaintext.len() / STREAM_CHUNK_SIZE;
                for i in 0..full_count {
                    let start = i * STREAM_CHUNK_SIZE;
                    let end = start + STREAM_CHUNK_SIZE;
                    let remainder = &plaintext[end..];
                    if remainder.is_empty() {
                        chunks.push(enc.encrypt_chunk(&plaintext[start..end], true).unwrap());
                    } else {
                        chunks.push(enc.encrypt_chunk(&plaintext[start..end], false).unwrap());
                    }
                    chunk_plaintexts.push(plaintext[start..end].to_vec());
                }
                if !enc.is_finalized() {
                    let start = full_count * STREAM_CHUNK_SIZE;
                    chunks.push(enc.encrypt_chunk(&plaintext[start..], true).unwrap());
                    chunk_plaintexts.push(plaintext[start..].to_vec());
                }

                let dec = stream_decrypt_init(&key, &header, b"").unwrap();
                // Pick the middle chunk for random access.
                let target = chunks.len() / 2;
                let (pt, _) = dec.decrypt_chunk_at(target as u64, &chunks[target]).unwrap();
                prop_assert_eq!(&*pt, &chunk_plaintexts[target]);
            }

            #[test]
            fn proptest_encrypt_chunk_at_round_trip(
                key in prop::array::uniform32(0u8..),
                compress in any::<bool>(),
                plaintext in prop::collection::vec(any::<u8>(), 0..STREAM_CHUNK_SIZE * 4),
                caller_aad in prop::collection::vec(any::<u8>(), 0..64),
            ) {
                // Encrypt every chunk via encrypt_chunk_at, assemble in index
                // order, then decrypt sequentially. Verifies that random-access
                // encrypt produces the same wire format as sequential encrypt.
                let enc = stream_encrypt_init(&key, &caller_aad, compress).unwrap();
                let header = enc.header();

                let chunk_size = STREAM_CHUNK_SIZE;
                let full_count = plaintext.len() / chunk_size;
                // An empty plaintext or an exact multiple both produce one final
                // chunk beyond the full_count full-size chunks, so total is always
                // full_count + 1.
                let total = full_count + 1;

                let mut chunks: Vec<Vec<u8>> = Vec::with_capacity(total);
                for i in 0..full_count {
                    let start = i * chunk_size;
                    let end = start + chunk_size;
                    let is_last = end == plaintext.len();
                    chunks.push(enc.encrypt_chunk_at(i as u64, is_last, &plaintext[start..end]).unwrap());
                    if is_last {
                        // plaintext length is exactly a multiple of chunk_size.
                        let mut dec = stream_decrypt_init(&key, &header, &caller_aad).unwrap();
                        let mut recovered = Vec::new();
                        for ct in &chunks {
                            let (pt, _) = dec.decrypt_chunk(ct).unwrap();
                            recovered.extend_from_slice(&pt);
                        }
                        prop_assert_eq!(recovered, plaintext);
                        return Ok(());
                    }
                }
                // Final (possibly partial) chunk.
                let start = full_count * chunk_size;
                chunks.push(enc.encrypt_chunk_at(full_count as u64, true, &plaintext[start..]).unwrap());

                let mut dec = stream_decrypt_init(&key, &header, &caller_aad).unwrap();
                let mut recovered = Vec::new();
                for ct in &chunks {
                    let (pt, _) = dec.decrypt_chunk(ct).unwrap();
                    recovered.extend_from_slice(&pt);
                }
                prop_assert_eq!(recovered, plaintext);
            }
        }
    }
}