libsoliton 0.1.3

//! Server-side storage encryption (§11).
//!
//! Messages are batched, compressed (zstd), then encrypted (XChaCha20-Poly1305).
//!
//! Blob format: [version: 1] [flags: 1] [nonce: 24] [ciphertext + tag]
//!
//! ## AAD Variants
//!
//! Two AAD layouts are supported, selected by the encrypt/decrypt function used:
//!
//! **Community Storage (§11.4.1):** `encrypt_blob` / `decrypt_blob`
//! `"lo-storage-v1" || version || flags || len(channel_id) || channel_id || len(segment_id) || segment_id`
//!
//! **DM Queue (§11.4.2):** `encrypt_dm_queue_blob` / `decrypt_dm_queue_blob`
//! `"lo-dm-queue-v1" || version || flags || len(recipient_fp) || recipient_fp || len(batch_id) || batch_id`
//!
//! where `len()` is a 2-byte big-endian u16 length prefix.

use crate::constants;
use crate::error::{Error, Result};
use crate::primitives::{aead, random};
use std::collections::HashMap;
use std::io::Read;
use subtle::ConstantTimeEq;
use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};

/// Storage blob flags.
const FLAG_COMPRESSED: u8 = 0x01;

/// Zstd compression level. ruzstd currently implements Fastest (~zstd level 1)
/// and Uncompressed only; Default/Better/Best are unimplemented as of 0.8.2.
/// Fastest provides adequate compression for message batches at minimal cost.
const COMPRESSION_LEVEL: ruzstd::encoding::CompressionLevel =
    ruzstd::encoding::CompressionLevel::Fastest;

/// Minimum blob size: version(1) + flags(1) + nonce(AEAD_NONCE_SIZE) + tag(AEAD_TAG_SIZE) = 42.
const MIN_BLOB_LEN: usize =
    1 + 1 + crate::constants::AEAD_NONCE_SIZE + crate::constants::AEAD_TAG_SIZE;

/// Maximum decompressed size.
///
/// Prevents OOM from maliciously crafted zstd payloads with extreme
/// compression ratios ("zip bomb" attacks). Legitimate storage blobs
/// should never approach this limit.
///
/// WASM targets use a lower limit (16 MiB) because wasm32 linear memory
/// defaults to much less than 256 MiB in most runtimes.
///
/// Uses `target_arch = "wasm32"` (not `target_os = "unknown"`) to apply the
/// memory limit to both wasm32-unknown-unknown and wasm32-wasi targets.
#[cfg(target_arch = "wasm32")]
const MAX_DECOMPRESSED_SIZE: u64 = 16 * 1024 * 1024;
#[cfg(not(target_arch = "wasm32"))]
const MAX_DECOMPRESSED_SIZE: u64 = 256 * 1024 * 1024;
// usize-typed equivalents — both 16 MiB and 256 MiB fit in usize on all
// supported targets (minimum 32-bit: usize::MAX = 4 GiB).
#[cfg(target_arch = "wasm32")]
const MAX_DECOMPRESSED_SIZE_USIZE: usize = 16 * 1024 * 1024;
#[cfg(not(target_arch = "wasm32"))]
const MAX_DECOMPRESSED_SIZE_USIZE: usize = 256 * 1024 * 1024;

/// A storage key with its version identifier.
///
/// Fields are private to enforce the version ≠ 0 invariant at construction.
/// Use [`StorageKey::new`] to create instances.
#[derive(Clone, Zeroize, ZeroizeOnDrop)]
pub struct StorageKey {
    /// Key version (1-255, maps to the version byte in the blob).
    #[zeroize(skip)]
    version: u8,
    /// The 256-bit encryption key.
    key: [u8; 32],
}

impl StorageKey {
    /// Create a new storage key with validated version.
    ///
    /// Rejects version 0 (`UnsupportedVersion`) — spec requires 1-255.
    /// Version 0 is reserved as "uninitialized / absent" in the blob wire format.
    /// Rejects all-zero key (`InvalidData`) — zero key provides no confidentiality.
    pub fn new(version: u8, mut key: [u8; 32]) -> Result<Self> {
        if version == 0 {
            // [u8; 32] is Copy — the caller's copy remains on the stack.
            // Zeroize our parameter copy on the rejection path.
            key.zeroize();
            return Err(Error::UnsupportedVersion);
        }
        // All-zero XChaCha20-Poly1305 key provides zero confidentiality.
        // Defense-in-depth: a programming error passing uninitialized memory
        // would silently produce an insecure keyring.
        // Constant-time: the key is secret material.
        if bool::from(key.ct_eq(&[0u8; 32])) {
            // Already all-zero — zeroize is a no-op but consistent.
            key.zeroize();
            return Err(Error::InvalidData);
        }
        let result = Self { version, key };
        // [u8; 32] is Copy — field-init shorthand copied the bytes into the
        // struct; the callee's parameter slot still holds secret key material.
        key.zeroize();
        Ok(result)
    }

    /// Key version (1-255).
    pub fn version(&self) -> u8 {
        self.version
    }

    /// The 256-bit encryption key.
    pub fn key(&self) -> &[u8; 32] {
        &self.key
    }
}

/// A keyring holding multiple storage key versions for rotation (§11.6).
///
/// # Thread Safety
///
/// `StorageKeyRing` auto-derives `Send + Sync` but is not designed for
/// concurrent access. Mutating operations (`add_key`, `remove_key`) require
/// `&mut self`. The CAPI layer adds a runtime reentrancy guard for FFI callers.
pub struct StorageKeyRing {
    /// The active version used for new writes.
    active_version: u8,
    /// All available keys indexed by version.
    keys: HashMap<u8, StorageKey>,
}

impl StorageKeyRing {
    /// Create a new keyring with one initial key.
    pub fn new(key: StorageKey) -> Result<Self> {
        let version = key.version();
        Ok(Self {
            active_version: version,
            keys: HashMap::from([(version, key)]),
        })
    }

    /// Add a key and optionally set it as the active version.
    ///
    /// If a key with the same version already exists and it is not the active
    /// version, its material is replaced and `Ok(true)` is returned.
    ///
    /// Replacing the active version requires `make_active: true` — without it,
    /// the call returns `InvalidData`. This prevents accidental key material
    /// replacement that would make blobs encrypted under the old material
    /// undecryptable.
    ///
    /// # Key Rotation
    ///
    /// After migrating existing blobs to a new key version, callers **must**
    /// call [`remove_key`](Self::remove_key) on old versions. Retaining old
    /// keys allows an attacker with storage write access to replay blobs
    /// encrypted under a compromised old key — the keyring will happily
    /// decrypt them because the old version is still present.
    ///
    /// Returns `Ok(false)` if no key with this version existed.
    pub fn add_key(&mut self, key: StorageKey, make_active: bool) -> Result<bool> {
        let version = key.version();
        // Replacing the active key's material without explicitly re-activating
        // would silently invalidate all blobs encrypted under the old material.
        if version == self.active_version && !make_active {
            return Err(Error::InvalidData);
        }
        if make_active {
            self.active_version = version;
        }
        let replaced = self.keys.insert(version, key).is_some();
        Ok(replaced)
    }

    /// Get the active key for new writes.
    pub fn active_key(&self) -> Option<&StorageKey> {
        self.keys.get(&self.active_version)
    }

    /// Look up a key by version (for decryption of old blobs).
    pub fn get_key(&self, version: u8) -> Option<&StorageKey> {
        self.keys.get(&version)
    }

    /// Remove a key version (after migration).
    ///
    /// Returns an error if `version` is the active key — call `add_key` with
    /// `make_active: true` to set a new active key before removing the old one.
    pub fn remove_key(&mut self, version: u8) -> Result<bool> {
        if version == 0 {
            return Err(Error::UnsupportedVersion);
        }
        // InvalidData: removing the active key would leave the keyring in an
        // unusable state (no key for new writes). Callers must set a new active
        // key via add_key(make_active: true) before removing the old one.
        if version == self.active_version {
            return Err(Error::InvalidData);
        }
        Ok(self.keys.remove(&version).is_some())
    }
}

impl Drop for StorageKeyRing {
    fn drop(&mut self) {
        // Defense-in-depth: StorageKey derives ZeroizeOnDrop, so HashMap's drop
        // will zeroize each element. This explicit loop is intentionally
        // redundant as a second layer of key zeroization.
        for key in self.keys.values_mut() {
            key.key.zeroize();
        }
    }
}

/// Encrypt data for storage (§11.2 write pipeline).
///
/// Pipeline: compress (if enabled) → XChaCha20-Poly1305 encrypt → prepend header.
/// The nonce is 24 random bytes from the OS CSPRNG. The AAD binds
/// `channel_id` and `segment_id` to the ciphertext, preventing cross-channel
/// and cross-segment swaps.
///
/// # Security
///
/// **Compression oracle (CRIME/BREACH):** When `compress = true`, ciphertext
/// length reveals the compression ratio. If an attacker can inject content
/// into the same blob alongside secret plaintext and observe blob sizes, they
/// can perform adaptive chosen-plaintext byte extraction. Use `compress = false`
/// for blobs where an attacker might control part of the plaintext.
///
/// **Memory residue:** When `compress = true`, the compressor's internal
/// allocations may retain plaintext in freed heap memory. Callers with
/// sensitive plaintext should pass `compress = false`.
///
/// **Input size:** Plaintext is capped at `MAX_DECOMPRESSED_SIZE` (256 MiB,
/// 16 MiB on WASM) to match the decrypt-side limit. Larger plaintext would
/// produce blobs that are permanently undecryptable. Mixed-platform
/// deployments (native + WASM) should enforce the lower 16 MiB limit at the
/// application layer so that native-encrypted blobs remain decryptable by
/// WASM clients.
///
/// **AAD exact-match:** `channel_id` and `segment_id` are used as-is with no
/// normalization. Unicode NFC vs NFD, trailing whitespace, or case differences
/// produce different AAD values, causing decryption failure. Callers must
/// ensure these strings are byte-identical at encrypt and decrypt time.
#[must_use = "encrypted blob must be stored; discarding it loses the data"]
pub fn encrypt_blob(
    key: &StorageKey,
    plaintext: &[u8],
    channel_id: &str,
    segment_id: &str,
    compress: bool,
) -> Result<Vec<u8>> {
    // Reject plaintext exceeding the decrypt-side decompression cap.
    // Without this, a blob encrypted with oversized plaintext would be
    // permanently undecryptable (decrypt_blob rejects decompressed output
    // exceeding MAX_DECOMPRESSED_SIZE).
    if plaintext.len() as u64 > MAX_DECOMPRESSED_SIZE {
        return Err(Error::InvalidData);
    }

    // Compression before encryption reduces ciphertext size but may leak
    // plaintext length information via compression ratio (CRIME-style).
    let (compressed, flags) = if compress {
        let c = ruzstd::encoding::compress_to_vec(plaintext, COMPRESSION_LEVEL);
        (Some(Zeroizing::new(c)), FLAG_COMPRESSED)
    } else {
        (None, 0u8)
    };
    // Borrow plaintext directly when not compressing — avoids a heap copy.
    let data: &[u8] = match &compressed {
        Some(c) => c,
        None => plaintext,
    };

    // Each blob gets a unique random nonce — reuse would be catastrophic
    // (XChaCha20 is a stream cipher; nonce reuse enables full plaintext XOR).
    // Birthday bound: random 24-byte nonces have ~2^-96 collision probability
    // per pair — nonce collision is effectively impossible for any realistic
    // number of encryptions per key version.
    let mut nonce = [0u8; 24];
    random::random_bytes(&mut nonce);

    // AAD binds version + flags + channel_id + segment_id to the ciphertext,
    // preventing cross-channel/cross-segment swap attacks and flag tampering (§11.4).
    let aad = build_storage_aad(key.version(), flags, channel_id, segment_id)?;

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

    // Wire format: [version][flags][nonce][ciphertext+tag]
    let mut blob = Vec::with_capacity(1 + 1 + 24 + ciphertext.len());
    blob.push(key.version());
    blob.push(flags);
    blob.extend_from_slice(&nonce);
    blob.extend_from_slice(&ciphertext);

    Ok(blob)
}

/// Decrypt a storage blob (§11.2 read pipeline).
///
/// Pipeline: parse header → decrypt → decompress (if flagged).
///
/// The returned `Zeroizing<Vec<u8>>` automatically zeroizes plaintext on drop.
///
/// # Security
///
/// XChaCha20-Poly1305 authentication verifies both ciphertext and AAD integrity
/// before any plaintext is returned. Decompression is bounded to
/// `MAX_DECOMPRESSED_SIZE` (256 MiB) to prevent zip-bomb OOM attacks.
///
/// When decompressing, `read_to_end` may allocate and free intermediate
/// buffers that are not covered by `Zeroizing` — decrypted plaintext
/// fragments may persist in freed heap memory until overwritten. The final
/// output buffer is wrapped in `Zeroizing` and zeroized on drop.
#[must_use = "decrypted data contains sensitive plaintext that must be consumed or zeroized"]
pub fn decrypt_blob(
    keyring: &StorageKeyRing,
    blob: &[u8],
    channel_id: &str,
    segment_id: &str,
) -> Result<Zeroizing<Vec<u8>>> {
    // All pre-AEAD checks return AeadFailed to prevent an attacker with
    // storage access from probing blob structure via distinct error codes.
    if blob.len() < MIN_BLOB_LEN {
        return Err(Error::AeadFailed);
    }

    let version = blob[0];
    let flags = blob[1];

    // Reject blobs with unknown flag bits set (bits 1-7 are reserved).
    // Forward-compatibility safety: a v1 reader must not silently ignore
    // unknown flag semantics that a newer writer may have set.
    let known_flags = FLAG_COMPRESSED;
    if flags & !known_flags != 0 {
        return Err(Error::AeadFailed);
    }

    // blob.len() >= MIN_BLOB_LEN (42) guaranteed above; try_into() is structurally
    // infallible. Maps to AeadFailed to keep all decrypt error paths consistent.
    let nonce: &[u8; 24] = blob[2..26].try_into().map_err(|_| Error::AeadFailed)?;
    let ciphertext = &blob[26..];

    // Return AeadFailed (not UnsupportedVersion) for unknown key versions —
    // distinguishing "no key for this version" from "wrong key" would let an
    // attacker with storage access enumerate all key versions in the keyring.
    let key = keyring.get_key(version).ok_or(Error::AeadFailed)?;

    // Map AAD construction errors to AeadFailed — consistent with the
    // design principle that all pre-AEAD checks are indistinguishable.
    // build_storage_aad fails only if channel_id/segment_id exceed u16::MAX;
    // leaking this as InvalidData would reveal that AAD construction failed
    // vs the key being wrong.
    let aad =
        build_storage_aad(version, flags, channel_id, segment_id).map_err(|_| Error::AeadFailed)?;

    // aead_decrypt returns Zeroizing<Vec<u8>>, so plaintext is
    // zeroized on all paths (including early return from decompression).

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

    // Decompress if flagged, with bounded output to prevent zip bombs.
    // Post-AEAD errors are mapped to AeadFailed so that decrypt failures are
    // indistinguishable to external observers — a DecompressionFailed would
    // reveal that AEAD authentication succeeded (1-bit post-auth oracle).
    if flags & FLAG_COMPRESSED != 0 {
        // Defense-in-depth: ruzstd's StreamingDecoder rejects empty input with
        // an I/O error. This cannot arise from encrypt_blob (which always produces
        // a non-empty zstd frame even for empty plaintext), but guards against
        // malformed blobs where AEAD-authenticated content is zero bytes with the
        // compressed flag set.
        if data.is_empty() {
            return Ok(Zeroizing::new(Vec::new()));
        }
        let decoder = ruzstd::decoding::StreamingDecoder::new(data.as_slice())
            .map_err(|_| Error::AeadFailed)?;
        // take(MAX + 1) reads one byte beyond the limit so the subsequent
        // length check can distinguish "exactly at limit" (OK) from "exceeds
        // limit" (reject).
        let mut limited = decoder.take(MAX_DECOMPRESSED_SIZE + 1);
        // Pre-allocate with 4× compressed size to reduce intermediate heap
        // allocations from read_to_end — those intermediates are freed without
        // zeroization (only the final allocation is covered by Zeroizing).
        // saturating_mul: on overflow, saturates to usize::MAX → Vec allocates
        // a smaller initial buffer and read_to_end handles reallocation.
        let hint = data
            .len()
            .saturating_mul(4)
            .min(MAX_DECOMPRESSED_SIZE_USIZE);
        let mut decompressed = Zeroizing::new(Vec::with_capacity(hint));
        limited
            .read_to_end(&mut decompressed)
            .map_err(|_| Error::AeadFailed)?;
        // Cast is lossless: usize::MAX (even on 32-bit) > MAX_DECOMPRESSED_SIZE (256 MiB).
        if decompressed.len() as u64 > MAX_DECOMPRESSED_SIZE {
            return Err(Error::AeadFailed);
        }
        Ok(decompressed)
    } else {
        Ok(data)
    }
}

/// Build community storage AAD (§11.4.1).
///
/// Wire format (all lengths are 2-byte big-endian u16):
/// `"lo-storage-v1" || version || flags || len(channel_id) || channel_id || len(segment_id) || segment_id`
fn build_storage_aad(
    version: u8,
    flags: u8,
    channel_id: &str,
    segment_id: &str,
) -> Result<Vec<u8>> {
    let ch = channel_id.as_bytes();
    let seg = segment_id.as_bytes();

    // AAD format uses 2-byte big-endian u16 length prefixes — inputs
    // exceeding u16::MAX cannot be encoded.
    if ch.len() > u16::MAX as usize || seg.len() > u16::MAX as usize {
        return Err(Error::InvalidData);
    }

    let mut aad =
        Vec::with_capacity(constants::STORAGE_AAD.len() + 1 + 1 + 2 + ch.len() + 2 + seg.len());
    aad.extend_from_slice(constants::STORAGE_AAD);
    aad.push(version);
    aad.push(flags);
    // The preceding guard ensures both lengths ≤ u16::MAX; these never fail.
    let ch_len = u16::try_from(ch.len()).expect("ch.len() ≤ u16::MAX validated above");
    let seg_len = u16::try_from(seg.len()).expect("seg.len() ≤ u16::MAX validated above");
    aad.extend_from_slice(&ch_len.to_be_bytes());
    aad.extend_from_slice(ch);
    aad.extend_from_slice(&seg_len.to_be_bytes());
    aad.extend_from_slice(seg);
    Ok(aad)
}

/// Build DM queue AAD (§11.4.2).
///
/// Wire format (all lengths are 2-byte big-endian u16):
/// `"lo-dm-queue-v1" || version || flags || len(recipient_fp) || recipient_fp || len(batch_id) || batch_id`
///
/// `recipient_fp` is a fixed 32-byte identity fingerprint. The 2-byte length
/// prefix is included for format consistency with §11.4.1 and to keep the
/// AAD parser uniform across both variants.
fn build_dm_queue_aad(
    version: u8,
    flags: u8,
    recipient_fp: &[u8; 32],
    batch_id: &str,
) -> Result<Vec<u8>> {
    let bid = batch_id.as_bytes();

    // AAD format uses 2-byte big-endian u16 length prefixes — inputs
    // exceeding u16::MAX cannot be encoded.
    if bid.len() > u16::MAX as usize {
        return Err(Error::InvalidData);
    }

    let mut aad =
        Vec::with_capacity(constants::DM_QUEUE_AAD.len() + 1 + 1 + 2 + 32 + 2 + bid.len());
    aad.extend_from_slice(constants::DM_QUEUE_AAD);
    aad.push(version);
    aad.push(flags);
    // recipient_fp is fixed-size (32), but encode with length prefix for
    // wire format consistency with the community storage AAD layout.
    aad.extend_from_slice(&32u16.to_be_bytes());
    aad.extend_from_slice(recipient_fp);
    // The preceding guard ensures bid.len() ≤ u16::MAX; this never fails.
    let bid_len = u16::try_from(bid.len()).expect("bid.len() ≤ u16::MAX validated above");
    aad.extend_from_slice(&bid_len.to_be_bytes());
    aad.extend_from_slice(bid);
    Ok(aad)
}

/// Encrypt data for DM queue storage (§11.4.2 write pipeline).
///
/// Identical to [`encrypt_blob`] except the AAD binds `recipient_fp` and
/// `batch_id` instead of `channel_id` and `segment_id`. This prevents
/// cross-recipient and cross-batch swaps in the DM relay queue.
///
/// # Security
///
/// All security notes from [`encrypt_blob`] apply. Additionally:
///
/// **Recipient binding:** The recipient's identity fingerprint is bound into
/// the AAD. A relay server cannot serve a blob encrypted for Alice to Bob —
/// Bob's decrypt call will supply his own fingerprint, producing a different
/// AAD and causing AEAD failure.
#[must_use = "encrypted blob must be stored; discarding it loses the data"]
pub fn encrypt_dm_queue_blob(
    key: &StorageKey,
    plaintext: &[u8],
    recipient_fp: &[u8; 32],
    batch_id: &str,
    compress: bool,
) -> Result<Vec<u8>> {
    if plaintext.len() as u64 > MAX_DECOMPRESSED_SIZE {
        return Err(Error::InvalidData);
    }

    let (compressed, flags) = if compress {
        let c = ruzstd::encoding::compress_to_vec(plaintext, COMPRESSION_LEVEL);
        (Some(Zeroizing::new(c)), FLAG_COMPRESSED)
    } else {
        (None, 0u8)
    };
    let data: &[u8] = match &compressed {
        Some(c) => c,
        None => plaintext,
    };

    let mut nonce = [0u8; 24];
    random::random_bytes(&mut nonce);

    let aad = build_dm_queue_aad(key.version(), flags, recipient_fp, batch_id)?;

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

    // Wire format: [version][flags][nonce][ciphertext+tag]
    let mut blob = Vec::with_capacity(1 + 1 + 24 + ciphertext.len());
    blob.push(key.version());
    blob.push(flags);
    blob.extend_from_slice(&nonce);
    blob.extend_from_slice(&ciphertext);

    Ok(blob)
}

/// Decrypt a DM queue storage blob (§11.4.2 read pipeline).
///
/// Identical to [`decrypt_blob`] except the AAD binds `recipient_fp` and
/// `batch_id` instead of `channel_id` and `segment_id`.
///
/// # Security
///
/// All security notes from [`decrypt_blob`] apply. The returned
/// `Zeroizing<Vec<u8>>` automatically zeroizes plaintext on drop.
#[must_use = "decrypted data contains sensitive plaintext that must be consumed or zeroized"]
pub fn decrypt_dm_queue_blob(
    keyring: &StorageKeyRing,
    blob: &[u8],
    recipient_fp: &[u8; 32],
    batch_id: &str,
) -> Result<Zeroizing<Vec<u8>>> {
    // All pre-AEAD checks return AeadFailed to prevent an attacker with
    // storage access from probing blob structure via distinct error codes.
    if blob.len() < MIN_BLOB_LEN {
        return Err(Error::AeadFailed);
    }

    let version = blob[0];
    let flags = blob[1];

    let known_flags = FLAG_COMPRESSED;
    if flags & !known_flags != 0 {
        return Err(Error::AeadFailed);
    }

    // blob.len() >= MIN_BLOB_LEN (42) guaranteed above; try_into() is structurally
    // infallible. Maps to AeadFailed to keep all decrypt error paths consistent.
    let nonce: &[u8; 24] = blob[2..26].try_into().map_err(|_| Error::AeadFailed)?;
    let ciphertext = &blob[26..];

    let key = keyring.get_key(version).ok_or(Error::AeadFailed)?;

    // Map AAD construction errors to AeadFailed — consistent with the
    // design principle that all pre-AEAD checks are indistinguishable.
    let aad = build_dm_queue_aad(version, flags, recipient_fp, batch_id)
        .map_err(|_| Error::AeadFailed)?;

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

    if flags & FLAG_COMPRESSED != 0 {
        if data.is_empty() {
            return Ok(Zeroizing::new(Vec::new()));
        }
        let decoder = ruzstd::decoding::StreamingDecoder::new(data.as_slice())
            .map_err(|_| Error::AeadFailed)?;
        let mut limited = decoder.take(MAX_DECOMPRESSED_SIZE + 1);
        let hint = data
            .len()
            .saturating_mul(4)
            .min(MAX_DECOMPRESSED_SIZE_USIZE);
        let mut decompressed = Zeroizing::new(Vec::with_capacity(hint));
        limited
            .read_to_end(&mut decompressed)
            .map_err(|_| Error::AeadFailed)?;
        if decompressed.len() as u64 > MAX_DECOMPRESSED_SIZE {
            return Err(Error::AeadFailed);
        }
        Ok(decompressed)
    } else {
        Ok(data)
    }
}

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

    fn test_key(version: u8) -> StorageKey {
        StorageKey::new(version, crate::primitives::random::random_array()).unwrap()
    }

    fn test_keyring(version: u8) -> (StorageKey, StorageKeyRing) {
        let key = test_key(version);
        let key_copy = key.clone();
        let ring = StorageKeyRing::new(key).unwrap();
        (key_copy, ring)
    }

    #[test]
    fn encrypt_decrypt_uncompressed() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_blob(key, b"hello world", "chan", "seg", false).unwrap();
        let pt = decrypt_blob(&ring, &blob, "chan", "seg").unwrap();
        assert_eq!(&*pt, b"hello world");
    }

    #[test]
    fn encrypt_decrypt_compressed() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_blob(key, b"hello world", "chan", "seg", true).unwrap();
        let pt = decrypt_blob(&ring, &blob, "chan", "seg").unwrap();
        assert_eq!(&*pt, b"hello world");
    }

    #[test]
    fn wrong_channel_id() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_blob(key, b"data", "chan_a", "seg", false).unwrap();
        assert!(matches!(
            decrypt_blob(&ring, &blob, "chan_b", "seg"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn wrong_segment_id() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_blob(key, b"data", "chan", "seg_1", false).unwrap();
        assert!(matches!(
            decrypt_blob(&ring, &blob, "chan", "seg_2"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn tampered_blob() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let mut blob = encrypt_blob(key, b"data", "chan", "seg", false).unwrap();
        // Flip a byte in the ciphertext region (after 14-byte header).
        blob[14] ^= 0xFF;
        assert!(matches!(
            decrypt_blob(&ring, &blob, "chan", "seg"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn truncated_blob() {
        // Pre-AEAD errors collapsed to AeadFailed to prevent error oracle.
        let (_, ring) = test_keyring(1);
        let short = vec![0u8; MIN_BLOB_LEN - 1];
        assert!(matches!(
            decrypt_blob(&ring, &short, "chan", "seg"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn unknown_flags() {
        // Pre-AEAD errors collapsed to AeadFailed to prevent error oracle.
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let mut blob = encrypt_blob(key, b"data", "chan", "seg", false).unwrap();
        // Set an unknown flag bit.
        blob[1] |= 0x02;
        assert!(matches!(
            decrypt_blob(&ring, &blob, "chan", "seg"),
            Err(Error::AeadFailed)
        ));
        // Combined: FLAG_COMPRESSED (0x01) | unknown (0x02) — full byte returned.
        blob[1] |= 0x01;
        assert!(matches!(
            decrypt_blob(&ring, &blob, "chan", "seg"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn storage_key_version_0_rejected() {
        assert!(matches!(
            StorageKey::new(0, [0u8; 32]),
            Err(Error::UnsupportedVersion)
        ));
    }

    #[test]
    fn storage_key_zero_key_rejected() {
        assert!(matches!(
            StorageKey::new(1, [0u8; 32]),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn flag_tampering_detected_by_aead() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        // Encrypt uncompressed, then flip the compressed flag in the header.
        let mut blob = encrypt_blob(key, b"not zstd data", "chan", "seg", false).unwrap();
        blob[1] |= FLAG_COMPRESSED;
        // Flags are bound into the AAD — flipping a flag causes AEAD failure.
        assert!(matches!(
            decrypt_blob(&ring, &blob, "chan", "seg"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn key_rotation() {
        let (_, mut ring) = test_keyring(1);
        let key1 = ring.active_key().unwrap();
        let blob_v1 = encrypt_blob(key1, b"v1 data", "chan", "seg", false).unwrap();
        // Add v2 as active.
        ring.add_key(test_key(2), true).unwrap();
        // v1 blob still decryptable.
        let pt = decrypt_blob(&ring, &blob_v1, "chan", "seg").unwrap();
        assert_eq!(&*pt, b"v1 data");
        // v2 can encrypt new blobs.
        let key2 = ring.active_key().unwrap();
        assert_eq!(key2.version(), 2);
        let blob_v2 = encrypt_blob(key2, b"v2 data", "chan", "seg", false).unwrap();
        let pt2 = decrypt_blob(&ring, &blob_v2, "chan", "seg").unwrap();
        assert_eq!(&*pt2, b"v2 data");
    }

    #[test]
    fn key_not_found() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let mut blob = encrypt_blob(key, b"data", "chan", "seg", false).unwrap();
        // Change version byte to one not in the keyring.
        // Returns AeadFailed (not UnsupportedVersion) to prevent an attacker
        // with storage access from enumerating key versions in the keyring.
        blob[0] = 99;
        assert!(matches!(
            decrypt_blob(&ring, &blob, "chan", "seg"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn keyring_add_replace_active_without_flag_rejected() {
        let (_, mut ring) = test_keyring(1);
        assert!(matches!(
            ring.add_key(test_key(1), false),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn keyring_add_replace_active_with_flag_succeeds() {
        let (_, mut ring) = test_keyring(1);
        let replaced = ring.add_key(test_key(1), true).unwrap();
        assert!(replaced);
    }

    #[test]
    fn keyring_remove_active_rejected() {
        let (_, mut ring) = test_keyring(1);
        assert!(matches!(ring.remove_key(1), Err(Error::InvalidData)));
    }

    #[test]
    fn keyring_remove_nonexistent() {
        let (_, mut ring) = test_keyring(1);
        let removed = ring.remove_key(99).unwrap();
        assert!(!removed);
    }

    #[test]
    fn keyring_remove_version_0_rejected() {
        let (_, mut ring) = test_keyring(1);
        assert!(matches!(ring.remove_key(0), Err(Error::UnsupportedVersion)));
    }

    #[test]
    fn keyring_active_key_returns_correct_version() {
        let (_, ring) = test_keyring(5);
        let active = ring.active_key().unwrap();
        assert_eq!(active.version(), 5);
    }

    #[test]
    fn keyring_get_key_returns_none_for_absent() {
        let (_, ring) = test_keyring(1);
        assert!(ring.get_key(99).is_none());
    }

    #[test]
    fn empty_plaintext_both_modes() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        for compress in [false, true] {
            let blob = encrypt_blob(key, b"", "chan", "seg", compress).unwrap();
            let pt = decrypt_blob(&ring, &blob, "chan", "seg").unwrap();
            assert!(pt.is_empty());
        }
    }

    // Build a blob from raw primitives, bypassing encrypt_blob, so that the
    // decrypted plaintext can be anything (e.g. garbage zstd for error tests).
    fn make_blob_with_plaintext(
        key: &StorageKey,
        flags: u8,
        plaintext: &[u8],
        channel_id: &str,
        segment_id: &str,
    ) -> Vec<u8> {
        let nonce = [0x42u8; 24];
        let aad = build_storage_aad(key.version(), flags, channel_id, segment_id).unwrap();
        let ct = crate::primitives::aead::aead_encrypt(key.key(), &nonce, plaintext, &aad).unwrap();
        let mut blob = Vec::with_capacity(1 + 1 + 24 + ct.len());
        blob.push(key.version());
        blob.push(flags);
        blob.extend_from_slice(&nonce);
        blob.extend_from_slice(&ct);
        blob
    }

    #[test]
    fn invalid_compressed_data() {
        // A blob whose FLAG_COMPRESSED flag is set but whose decrypted content
        // is not valid zstd. Returns AeadFailed (not DecompressionFailed) to
        // prevent a post-authentication oracle — all decrypt failures are
        // indistinguishable to external observers.
        let (key, ring) = test_keyring(1);
        let garbage = b"this is definitely not valid zstd";
        let blob = make_blob_with_plaintext(&key, FLAG_COMPRESSED, garbage, "chan", "seg");
        assert!(matches!(
            decrypt_blob(&ring, &blob, "chan", "seg"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    #[ignore = "allocates 256 MiB — run explicitly to verify zip-bomb limit"]
    fn decompression_bomb_rejected() {
        // Verify that a compressed blob expanding to > MAX_DECOMPRESSED_SIZE
        // is rejected without exhausting memory. Returns AeadFailed (not
        // DecompressionFailed) to prevent a post-authentication oracle.
        let plaintext = vec![0u8; MAX_DECOMPRESSED_SIZE_USIZE + 1];
        let compressed = ruzstd::encoding::compress_to_vec(plaintext.as_slice(), COMPRESSION_LEVEL);
        let (key, ring) = test_keyring(1);
        let blob = make_blob_with_plaintext(&key, FLAG_COMPRESSED, &compressed, "bomb", "0");
        assert!(matches!(
            decrypt_blob(&ring, &blob, "bomb", "0"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn storage_aad_structure() {
        let aad = build_storage_aad(1, FLAG_COMPRESSED, "my-channel", "seg-42").unwrap();
        let mut expected = Vec::new();
        expected.extend_from_slice(b"lo-storage-v1");
        expected.push(1); // version
        expected.push(FLAG_COMPRESSED); // flags
        expected.extend_from_slice(&10u16.to_be_bytes()); // len("my-channel")
        expected.extend_from_slice(b"my-channel");
        expected.extend_from_slice(&6u16.to_be_bytes()); // len("seg-42")
        expected.extend_from_slice(b"seg-42");
        assert_eq!(aad, expected);
    }

    #[test]
    fn storage_aad_empty_channel_and_segment_are_distinct() {
        // ("", "x") and ("x", "") must produce different AAD — length prefixes
        // prevent ambiguity even with empty strings.
        let aad1 = build_storage_aad(1, 0, "", "x").unwrap();
        let aad2 = build_storage_aad(1, 0, "x", "").unwrap();
        assert_ne!(aad1, aad2, "empty channel vs empty segment must differ");

        // Empty-string AAD: 0x0000 length prefix + no payload bytes.
        let aad = build_storage_aad(1, 0, "", "").unwrap();
        let prefix_len = b"lo-storage-v1".len() + 1 + 1; // label + version + flags
        assert_eq!(aad.len(), prefix_len + 2 + 2);
    }

    #[test]
    fn build_storage_aad_rejects_oversized_ids() {
        let long = "x".repeat(u16::MAX as usize + 1);
        // Oversized channel_id.
        assert!(matches!(
            build_storage_aad(1, 0, &long, "seg"),
            Err(Error::InvalidData)
        ));
        // Oversized segment_id.
        assert!(matches!(
            build_storage_aad(1, 0, "ch", &long),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn storage_encrypt_decrypt_empty_strings() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_blob(key, b"hello", "", "", false).unwrap();
        let decrypted = decrypt_blob(&ring, &blob, "", "").unwrap();
        assert_eq!(&*decrypted, b"hello");
    }

    // --- DM Queue tests ---

    const TEST_FP: [u8; 32] = [0xABu8; 32];

    #[test]
    fn dm_queue_encrypt_decrypt_uncompressed() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_dm_queue_blob(key, b"dm payload", &TEST_FP, "batch-0", false).unwrap();
        let pt = decrypt_dm_queue_blob(&ring, &blob, &TEST_FP, "batch-0").unwrap();
        assert_eq!(&*pt, b"dm payload");
    }

    #[test]
    fn dm_queue_encrypt_decrypt_compressed() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_dm_queue_blob(key, b"dm payload", &TEST_FP, "batch-0", true).unwrap();
        let pt = decrypt_dm_queue_blob(&ring, &blob, &TEST_FP, "batch-0").unwrap();
        assert_eq!(&*pt, b"dm payload");
    }

    #[test]
    fn dm_queue_wrong_recipient_fp() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_dm_queue_blob(key, b"data", &TEST_FP, "batch-0", false).unwrap();
        let wrong_fp = [0xCDu8; 32];
        assert!(matches!(
            decrypt_dm_queue_blob(&ring, &blob, &wrong_fp, "batch-0"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn dm_queue_wrong_batch_id() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob = encrypt_dm_queue_blob(key, b"data", &TEST_FP, "batch-0", false).unwrap();
        assert!(matches!(
            decrypt_dm_queue_blob(&ring, &blob, &TEST_FP, "batch-1"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn dm_queue_tampered_blob() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let mut blob = encrypt_dm_queue_blob(key, b"data", &TEST_FP, "batch-0", false).unwrap();
        blob[14] ^= 0xFF;
        assert!(matches!(
            decrypt_dm_queue_blob(&ring, &blob, &TEST_FP, "batch-0"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn dm_queue_empty_plaintext() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        for compress in [false, true] {
            let blob = encrypt_dm_queue_blob(key, b"", &TEST_FP, "batch-0", compress).unwrap();
            let pt = decrypt_dm_queue_blob(&ring, &blob, &TEST_FP, "batch-0").unwrap();
            assert!(pt.is_empty());
        }
    }

    #[test]
    fn dm_queue_key_rotation() {
        let (_, mut ring) = test_keyring(1);
        let key1 = ring.active_key().unwrap();
        let blob_v1 = encrypt_dm_queue_blob(key1, b"v1 dm", &TEST_FP, "b0", false).unwrap();
        ring.add_key(test_key(2), true).unwrap();
        // v1 blob still decryptable.
        let pt = decrypt_dm_queue_blob(&ring, &blob_v1, &TEST_FP, "b0").unwrap();
        assert_eq!(&*pt, b"v1 dm");
        // v2 can encrypt new blobs.
        let key2 = ring.active_key().unwrap();
        let blob_v2 = encrypt_dm_queue_blob(key2, b"v2 dm", &TEST_FP, "b0", false).unwrap();
        let pt2 = decrypt_dm_queue_blob(&ring, &blob_v2, &TEST_FP, "b0").unwrap();
        assert_eq!(&*pt2, b"v2 dm");
    }

    #[test]
    fn dm_queue_aad_structure() {
        let fp = [0x42u8; 32];
        let aad = build_dm_queue_aad(1, FLAG_COMPRESSED, &fp, "batch-7").unwrap();
        let mut expected = Vec::new();
        expected.extend_from_slice(b"lo-dm-queue-v1");
        expected.push(1); // version
        expected.push(FLAG_COMPRESSED); // flags
        expected.extend_from_slice(&32u16.to_be_bytes()); // len(recipient_fp)
        expected.extend_from_slice(&fp);
        expected.extend_from_slice(&7u16.to_be_bytes()); // len("batch-7")
        expected.extend_from_slice(b"batch-7");
        assert_eq!(aad, expected);
    }

    #[test]
    fn dm_queue_aad_rejects_oversized_batch_id() {
        let long = "x".repeat(u16::MAX as usize + 1);
        let fp = [0x00u8; 32];
        assert!(matches!(
            build_dm_queue_aad(1, 0, &fp, &long),
            Err(Error::InvalidData)
        ));
    }

    #[test]
    fn community_and_dm_queue_blobs_not_interchangeable() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        // Encrypt as community storage, try to decrypt as DM queue.
        let community_blob = encrypt_blob(key, b"data", "chan", "seg", false).unwrap();
        assert!(matches!(
            decrypt_dm_queue_blob(&ring, &community_blob, &TEST_FP, "seg"),
            Err(Error::AeadFailed)
        ));
        // Encrypt as DM queue, try to decrypt as community storage.
        let dm_blob = encrypt_dm_queue_blob(key, b"data", &TEST_FP, "batch-0", false).unwrap();
        assert!(matches!(
            decrypt_blob(&ring, &dm_blob, "chan", "batch-0"),
            Err(Error::AeadFailed)
        ));
    }

    #[test]
    fn encrypt_blob_nonce_freshness() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let blob1 = encrypt_blob(key, b"same input", "chan", "seg", false).unwrap();
        let blob2 = encrypt_blob(key, b"same input", "chan", "seg", false).unwrap();
        // Each call generates a fresh random 24-byte nonce, so identical
        // plaintext + key + AAD must still produce different ciphertexts.
        // A CSPRNG regression returning zeros would cause nonce reuse and
        // identical blobs.
        assert_ne!(
            blob1, blob2,
            "encrypt_blob must produce different blobs due to random nonce"
        );
    }

    #[test]
    fn encrypt_dm_queue_blob_nonce_freshness() {
        let (_, ring) = test_keyring(1);
        let key = ring.active_key().unwrap();
        let fp = [0xABu8; 32];
        let blob1 = encrypt_dm_queue_blob(key, b"same input", &fp, "batch", false).unwrap();
        let blob2 = encrypt_dm_queue_blob(key, b"same input", &fp, "batch", false).unwrap();
        assert_ne!(
            blob1, blob2,
            "encrypt_dm_queue_blob must produce different blobs due to random nonce"
        );
    }

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

        proptest! {
            #[test]
            fn encrypt_decrypt_roundtrip(
                plaintext in proptest::collection::vec(any::<u8>(), 0..4096),
                compress in any::<bool>(),
                channel in "[a-z]{1,16}",
                segment in "[a-z]{1,16}",
                // Version byte is part of the blob AAD — varying it ensures the
                // version is correctly round-tripped, not hardcoded.
                version in 1u8..=255u8,
            ) {
                let (_, ring) = test_keyring(version);
                let key = ring.active_key().unwrap();

                let blob = encrypt_blob(key, &plaintext, &channel, &segment, compress).unwrap();
                let decrypted = decrypt_blob(&ring, &blob, &channel, &segment).unwrap();
                prop_assert_eq!(&*decrypted, &plaintext);
            }

            #[test]
            fn dm_queue_encrypt_decrypt_roundtrip(
                plaintext in proptest::collection::vec(any::<u8>(), 0..4096),
                compress in any::<bool>(),
                batch_id in "[a-z0-9]{1,16}",
                version in 1u8..=255u8,
            ) {
                let fp = [0xABu8; 32];
                let (_, ring) = test_keyring(version);
                let key = ring.active_key().unwrap();

                let blob = encrypt_dm_queue_blob(key, &plaintext, &fp, &batch_id, compress).unwrap();
                let decrypted = decrypt_dm_queue_blob(&ring, &blob, &fp, &batch_id).unwrap();
                prop_assert_eq!(&*decrypted, &plaintext);
            }
        }
    }
}