key-vault 1.0.0

//! [`StandardFragmenter`] — the baseline Layer 3 implementation.
//!
//! `StandardFragmenter` splits the raw key into variable-size chunks whose
//! lengths are sampled uniformly from `[frag_min, frag_max]`, applies a
//! random Fisher-Yates permutation, allocates each chunk in its own
//! [`LockedBytes`] buffer (so chunks are at independent heap addresses),
//! and stores the reconstruction order in a separately-locked layout
//! buffer.
//!
//! Two consecutive calls to `fragment` on the same key produce
//! [`Fragments`] with different chunk counts, different chunk sizes, and
//! different orderings. The randomness is sourced from
//! [`getrandom`](https://docs.rs/getrandom), the OS CSPRNG.

use alloc::borrow::Cow;
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::fmt;

use super::util::{fisher_yates, sample_range, zero_buffer, zero_buffer_owned};
use super::{FragmentStrategy, Fragments};
use crate::Result;
use crate::decoy::DecoyStrategy;
use crate::error::Error;
use crate::fetcher::RawKey;
use crate::memory::LockedBytes;

/// Sentinel layout value marking a chunk as a decoy. Real chunks store
/// their original offset, which is always strictly less than the original
/// key length. Decoy chunks store `u32::MAX`, which the defragment path
/// recognizes and skips. This caps the supported key length at
/// `u32::MAX - 1` bytes (~4 GiB) — well above any realistic key.
const DECOY_OFFSET: u32 = u32::MAX;

/// Default minimum chunk size — small enough to avoid leaking the
/// fragmentation boundary, large enough to keep the chunk count
/// reasonable.
const DEFAULT_MIN_CHUNK: usize = 1;

/// Default maximum chunk size. Eight bytes is large enough to amortize
/// per-chunk overhead and small enough that a 32-byte symmetric key still
/// produces several chunks.
const DEFAULT_MAX_CHUNK: usize = 8;

/// Variable-chunk + shuffle fragmenter. Default Layer 3 implementation.
///
/// Construct with [`StandardFragmenter::new`] for the default chunk-size
/// range, or [`StandardFragmenter::with_chunk_range`] to customize.
///
/// # Examples
///
/// Typical use is through [`KeyVaultBuilder`](crate::KeyVaultBuilder), which
/// owns a `StandardFragmenter` internally. `RawKey` deliberately does not
/// expose its bytes to outside callers, so we verify the round-trip by
/// length:
///
/// ```
/// use key_vault::{KeyVaultBuilder, RawKey};
///
/// let vault = KeyVaultBuilder::new()
///     .normalize_with_blake3(false)
///     .with_chunk_range(2, 4)
///     .build();
///
/// let original_len = b"some key material".len();
/// let raw = RawKey::new(b"some key material".to_vec());
/// let frags = vault.fragment(&raw).unwrap();
/// let recovered = vault.defragment(&frags).unwrap();
/// assert_eq!(recovered.len(), original_len);
/// ```
#[derive(Clone)]
pub struct StandardFragmenter {
    min_chunk: usize,
    max_chunk: usize,
    /// Optional Layer 4 decoy strategy. When set, `fragment` emits
    /// additional decoy chunks alongside the real ones; `defragment`
    /// recognizes and skips them via the [`DECOY_OFFSET`] sentinel in the
    /// layout buffer.
    decoy: Option<Arc<dyn DecoyStrategy>>,
}

impl StandardFragmenter {
    /// Construct a fragmenter with the default chunk-size range
    /// (`min = 1`, `max = 8`) and no decoy strategy.
    #[must_use]
    pub fn new() -> Self {
        Self {
            min_chunk: DEFAULT_MIN_CHUNK,
            max_chunk: DEFAULT_MAX_CHUNK,
            decoy: None,
        }
    }

    /// Construct a fragmenter with a custom chunk-size range. `min` must be
    /// at least 1 and `max` must be at least `min`; both are clamped at
    /// construction.
    ///
    /// Larger maxima reduce the chunk count (lower memory overhead, less
    /// scatter). Smaller maxima increase the chunk count (more scatter,
    /// higher memory overhead). The default of 1-8 strikes a balance
    /// validated by the round-trip + multi-fragmentation tests.
    #[must_use]
    pub fn with_chunk_range(min: usize, max: usize) -> Self {
        let min = min.max(1);
        let max = max.max(min);
        Self {
            min_chunk: min,
            max_chunk: max,
            decoy: None,
        }
    }

    /// Attach a [`DecoyStrategy`] to this fragmenter.
    ///
    /// When a decoy is set, every call to `fragment` produces real chunks
    /// **plus** decoy chunks generated by the strategy. The decoy chunks
    /// are interleaved with the real ones via the same Fisher-Yates shuffle
    /// and are recognized at defragment time by a sentinel value
    /// (`u32::MAX`) in the locked layout buffer.
    ///
    /// The decoy strategy is held in an `Arc<dyn DecoyStrategy>` so the
    /// fragmenter remains `Clone` and the same strategy can be shared
    /// across multiple builders.
    #[must_use]
    pub fn with_decoy<D>(mut self, decoy: D) -> Self
    where
        D: DecoyStrategy + 'static,
    {
        self.decoy = Some(Arc::new(decoy));
        self
    }
}

impl Default for StandardFragmenter {
    /// Same as [`StandardFragmenter::new`] — the default-range
    /// (`min = 1`, `max = 8`) fragmenter with no decoy strategy.
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Debug for StandardFragmenter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("StandardFragmenter")
            .field("min_chunk", &self.min_chunk)
            .field("max_chunk", &self.max_chunk)
            .field("decoy", &self.decoy.as_ref().map(|d| d.describe()))
            .finish()
    }
}

impl FragmentStrategy for StandardFragmenter {
    fn fragment(&self, key: &RawKey) -> Result<Fragments> {
        let bytes = key.as_bytes();
        let total_len = bytes.len();
        if total_len == 0 {
            return Err(Error::Fragment(alloc::string::ToString::to_string(
                "empty key cannot be fragmented",
            )));
        }
        // Defensive cap: real-chunk offsets must be < DECOY_OFFSET so
        // defragment can distinguish them. `DECOY_OFFSET = u32::MAX`, so
        // any key shorter than u32::MAX bytes is safe.
        if total_len >= DECOY_OFFSET as usize {
            return Err(Error::Fragment(alloc::string::ToString::to_string(
                "key too large for fragmentation",
            )));
        }

        // Step 1: choose chunk sizes summing to total_len for the real key.
        let sizes = sample_chunk_sizes(total_len, self.min_chunk, self.max_chunk)?;
        let n_real = sizes.len();

        // Step 2: walk the key in original order, building
        // (offset, slice) pairs.
        let mut real_pairs: Vec<(u32, &[u8])> = Vec::with_capacity(n_real);
        {
            let mut offset = 0usize;
            for &size in &sizes {
                let offset_u32 = u32::try_from(offset).map_err(|_| {
                    Error::Fragment(alloc::string::ToString::to_string(
                        "key too large for fragmentation",
                    ))
                })?;
                real_pairs.push((offset_u32, &bytes[offset..offset + size]));
                offset += size;
            }
        }

        // Step 3 (optional): if a Layer-4 decoy is configured, generate
        // matching decoy chunks. We choose a count equal to the real chunk
        // count by default — a 1:1 mix between real and decoy gives an
        // attacker a 50/50 prior on any specific chunk even before they
        // try to defeat the layout map.
        //
        // Each decoy chunk gets an independent size sampled from the same
        // range so its size distribution matches the real chunks.
        let n_decoy = if self.decoy.is_some() { n_real } else { 0 };
        let mut decoy_chunks: Vec<LockedBytes> = Vec::with_capacity(n_decoy);
        if let Some(ref decoy) = self.decoy {
            for _ in 0..n_decoy {
                let size = sample_range(self.min_chunk, self.max_chunk)?;
                let bytes = decoy.generate(key, size)?;
                decoy_chunks.push(LockedBytes::from_slice(&bytes));
                // Zero the temporary plaintext decoy buffer before drop.
                zero_buffer_owned(bytes);
            }
        }

        // Step 4: combine real (offset, slice) pairs and decoy (DECOY_OFFSET,
        // pre-allocated LockedBytes) entries into one indexable vector,
        // then Fisher-Yates shuffle.
        //
        // We can't mix `&[u8]` slices and owned `LockedBytes` in one Vec, so
        // we shuffle their *indices* instead and walk the result twice.
        let total_chunks = n_real + n_decoy;
        let mut order: Vec<ChunkKind> = Vec::with_capacity(total_chunks);
        for i in 0..n_real {
            order.push(ChunkKind::Real(i));
        }
        for i in 0..n_decoy {
            order.push(ChunkKind::Decoy(i));
        }
        fisher_yates(&mut order)?;

        // Step 5: walk the shuffled order, materializing chunks and the
        // layout buffer in lockstep.
        let mut chunks: Vec<LockedBytes> = Vec::with_capacity(total_chunks);
        let mut layout_bytes: Vec<u8> = Vec::with_capacity(total_chunks * 4);
        // We need to pull decoy chunks out one at a time; convert to Option
        // so we can `.take()` each slot without moving the Vec.
        let mut decoy_slots: Vec<Option<LockedBytes>> =
            decoy_chunks.into_iter().map(Some).collect();
        for kind in &order {
            match *kind {
                ChunkKind::Real(idx) => {
                    let (offset, slice) = real_pairs[idx];
                    chunks.push(LockedBytes::from_slice(slice));
                    layout_bytes.extend_from_slice(&offset.to_le_bytes());
                }
                ChunkKind::Decoy(idx) => {
                    let lb = decoy_slots[idx].take().ok_or(Error::Internal(
                        "decoy slot taken twice during fragmentation",
                    ))?;
                    chunks.push(lb);
                    layout_bytes.extend_from_slice(&DECOY_OFFSET.to_le_bytes());
                }
            }
        }

        let layout = LockedBytes::from_slice(&layout_bytes);
        zero_buffer(&mut layout_bytes);
        drop(layout_bytes);
        drop(real_pairs);
        drop(decoy_slots);
        drop(order);

        Ok(Fragments::from_parts(chunks, layout, total_len))
    }

    fn defragment(&self, fragments: &Fragments) -> Result<RawKey> {
        let mut out = alloc::vec![0u8; fragments.total_len()];
        self.defragment_into(fragments, &mut out)?;
        Ok(RawKey::new(out))
    }

    fn defragment_into(&self, fragments: &Fragments, out: &mut [u8]) -> Result<()> {
        let n_chunks = fragments.chunk_count();
        let layout = fragments.layout().as_bytes();
        let total_len = fragments.total_len();
        if layout.len() != n_chunks * 4 {
            return Err(Error::Defragment(alloc::string::ToString::to_string(
                "layout buffer length does not match chunk count",
            )));
        }
        if out.len() != total_len {
            return Err(Error::Defragment(alloc::string::ToString::to_string(
                "scratch buffer size does not match fragments.total_len()",
            )));
        }

        // Single-pass: write each real chunk directly into `out` at the
        // offset recorded in the layout buffer. Decoy chunks
        // (`offset == DECOY_OFFSET`) are skipped. No intermediate
        // (offset, chunk) Vec is allocated — this is the change that
        // removes one heap allocation per `with_key` call.
        let mut written = 0usize;
        for (i, chunk) in fragments.chunks().iter().enumerate() {
            let raw: [u8; 4] = layout[i * 4..i * 4 + 4].try_into().map_err(|_| {
                Error::Defragment(alloc::string::ToString::to_string(
                    "layout buffer slice did not size to u32",
                ))
            })?;
            let offset = u32::from_le_bytes(raw);
            if offset == DECOY_OFFSET {
                continue;
            }
            let chunk_bytes = chunk.as_bytes();
            let start = offset as usize;
            let end = start.checked_add(chunk_bytes.len()).ok_or_else(|| {
                Error::Defragment(alloc::string::ToString::to_string(
                    "chunk offset overflowed when added to chunk length",
                ))
            })?;
            if end > total_len {
                return Err(Error::Defragment(alloc::string::ToString::to_string(
                    "chunk would write past end of output buffer",
                )));
            }
            out[start..end].copy_from_slice(chunk_bytes);
            written = written.saturating_add(chunk_bytes.len());
        }

        if written != total_len {
            return Err(Error::Defragment(alloc::string::ToString::to_string(
                "reassembled length does not match recorded total",
            )));
        }

        Ok(())
    }

    fn describe(&self) -> Cow<'_, str> {
        Cow::Borrowed("standard")
    }
}

/// Tag used during Fisher-Yates of the combined real + decoy ordering.
/// We shuffle this tag list rather than the chunks themselves so we can
/// keep the borrowed real-key slices and the pre-allocated decoy
/// `LockedBytes` in their original storage until we walk the shuffled
/// order to materialize the final `chunks` Vec.
#[derive(Clone, Copy)]
enum ChunkKind {
    Real(usize),
    Decoy(usize),
}

/// Sample chunk sizes summing exactly to `total`, with each size in
/// `[min, max]` except possibly the last (which absorbs any short
/// remainder).
fn sample_chunk_sizes(total: usize, min: usize, max: usize) -> Result<Vec<usize>> {
    if min == 0 || max < min {
        return Err(Error::Fragment(alloc::string::ToString::to_string(
            "invalid chunk-size range",
        )));
    }
    let mut sizes: Vec<usize> = Vec::new();
    let mut remaining = total;
    while remaining > 0 {
        if remaining <= max {
            sizes.push(remaining);
            remaining = 0;
        } else {
            let pick = sample_range(min, max)?;
            let pick = pick.min(remaining.saturating_sub(min));
            let pick = pick.max(min).min(max).min(remaining);
            sizes.push(pick);
            remaining -= pick;
        }
    }
    Ok(sizes)
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::cast_possible_truncation,
    clippy::cast_sign_loss
)]
mod tests {
    use super::*;

    /// Build a `RawKey` with arbitrary bytes.
    fn key(bytes: &[u8]) -> RawKey {
        RawKey::new(bytes.to_vec())
    }

    #[test]
    fn round_trip_short_key() {
        let frag = StandardFragmenter::new();
        let original = key(&[0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
        let fragments = frag.fragment(&original).unwrap();
        let recovered = frag.defragment(&fragments).unwrap();
        assert_eq!(recovered.len(), 10);
        assert_eq!(recovered.as_bytes(), original.as_bytes());
    }

    #[test]
    fn round_trip_256_bit_key() {
        let frag = StandardFragmenter::new();
        let bytes: Vec<u8> = (0..32).map(|i| (i * 7) as u8).collect();
        let original = key(&bytes);
        let fragments = frag.fragment(&original).unwrap();
        let recovered = frag.defragment(&fragments).unwrap();
        assert_eq!(recovered.as_bytes(), &bytes[..]);
    }

    #[test]
    fn round_trip_for_many_sizes() {
        let frag = StandardFragmenter::new();
        for len in [1usize, 7, 16, 32, 64, 128, 255, 256, 500, 1024, 4096] {
            let bytes: Vec<u8> = (0..len).map(|i| (i & 0xff) as u8).collect();
            let original = key(&bytes);
            let fragments = frag.fragment(&original).expect("fragment");
            let recovered = frag.defragment(&fragments).expect("defragment");
            assert_eq!(
                recovered.as_bytes(),
                &bytes[..],
                "round-trip mismatch for len = {len}"
            );
        }
    }

    #[test]
    fn two_calls_produce_different_layouts() {
        let frag = StandardFragmenter::new();
        let bytes: Vec<u8> = (0..32).map(|i| (i ^ 0x5a) as u8).collect();
        let original = key(&bytes);

        let a = frag.fragment(&original).unwrap();
        let b = frag.fragment(&original).unwrap();

        // At 32 bytes with chunk sizes 1..=8, the probability of two
        // consecutive fragmentations producing identical layout AND
        // identical chunk counts is astronomically small. Treat a match
        // as a strong signal of broken randomness rather than coincidence.
        let same_count = a.chunk_count() == b.chunk_count();
        let same_layout = same_count && a.layout().as_bytes() == b.layout().as_bytes();
        assert!(
            !(same_count && same_layout),
            "two consecutive fragmentations produced the same layout"
        );

        // Both still round-trip cleanly.
        assert_eq!(frag.defragment(&a).unwrap().as_bytes(), &bytes[..]);
        assert_eq!(frag.defragment(&b).unwrap().as_bytes(), &bytes[..]);
    }

    #[test]
    fn chunk_sizes_respect_configured_range() {
        let frag = StandardFragmenter::with_chunk_range(2, 4);
        let bytes: Vec<u8> = (0..32).collect();
        let original = key(&bytes);
        let fragments = frag.fragment(&original).unwrap();

        // Chunks are Fisher-Yates shuffled, so the "remainder" chunk (which
        // may fall below `min` when the total length doesn't divide cleanly)
        // can land at any index. We verify the post-shuffle invariants:
        //   1. Every chunk size is in [1, max].
        //   2. At most one chunk falls below `min` (the remainder).
        //   3. Total bytes sum to the original length.
        let chunks = fragments.chunks();
        let mut below_min = 0;
        let mut total = 0usize;
        for c in chunks {
            assert!(
                c.len() >= 1 && c.len() <= 4,
                "chunk size {} not in [1,4]",
                c.len()
            );
            if c.len() < 2 {
                below_min += 1;
            }
            total += c.len();
        }
        assert!(
            below_min <= 1,
            "more than one chunk below min size: {below_min}"
        );
        assert_eq!(total, 32);

        assert_eq!(frag.defragment(&fragments).unwrap().as_bytes(), &bytes[..]);
    }

    #[test]
    fn empty_key_rejected() {
        let frag = StandardFragmenter::new();
        let empty = key(&[]);
        let err = frag.fragment(&empty).unwrap_err();
        assert!(matches!(err, Error::Fragment(_)));
    }

    #[test]
    fn describe_returns_standard() {
        let frag = StandardFragmenter::new();
        assert_eq!(frag.describe(), "standard");
    }

    #[test]
    fn stress_round_trip_thousand_iterations() {
        let frag = StandardFragmenter::new();
        let bytes: Vec<u8> = (0..32).map(|i| ((i * 13) ^ 0xa5) as u8).collect();
        let original = key(&bytes);
        for _ in 0..1000 {
            let fragments = frag.fragment(&original).expect("fragment");
            let recovered = frag.defragment(&fragments).expect("defragment");
            assert_eq!(recovered.as_bytes(), &bytes[..]);
        }
    }
}