b-wing 0.1.1

//! # KWing: Hybrid Key Encapsulation Mechanism
//!
//! KWing is a paranoia-grade, triple-layered hybrid KEM that derives a shared
//! secret with **NIST Level 5 (256-bit)** post-quantum security.  The construction
//! follows a **redundancy-first** policy: the derived Output Keying Material (OKM)
//! remains secret as long as *any single one* of the three component algorithms is
//! unbroken — whether by classical cryptanalysis, a quantum adversary, or a
//! lattice-specific breakthrough.
//!
//! ## Security Composition
//!
//! | Layer | Algorithm | Security assumption |
//! |-------|-----------|---------------------|
//! | 1 | **X25519** | Classical ECDH — Curve25519 discrete log |
//! | 2 | **ML-KEM-1024** (Kyber) | Module-lattice MLWE (NIST PQ Level 5) |
//! | 3 | **FrodoKEM-1344-SHAKE** | Unstructured LWE — conservative lattice |
//!
//! The three independent shared secrets are combined via **HKDF-SHA3-512**,
//! cryptographically binding them to the full transcript (ephemeral public key,
//! all three ciphertexts, the recipient's encapsulation key, and a 32-byte salt).
//!
//! ## Design Principles
//!
//! * **Caller-supplied randomness.** All randomness is given as typed seed arrays.
//!   The API is fully deterministic — no hidden global RNG state — and therefore
//!   compatible with `no_std`, WASM, and embedded targets.
//! * **Heap-cached keys.** [`KWing`] pre-computes the composite encapsulation key
//!   once on construction, amortising the cost across many encapsulate/decapsulate
//!   calls.
//! * **Zeroize on drop.** Classical and ML-KEM secret material is wrapped in
//!   [`zeroize::Zeroizing`], guaranteeing erasure from memory. FrodoKEM
//!   material relies on standard memory cleanup as the underlying crate
//!   does not yet support the `zeroize` trait.
//! * **Strict size validation.** Every public entry-point validates buffer sizes
//!   before performing any cryptographic work.
//!
//! ## Usage
//!
//! ```rust,no_run
//! # #[cfg(feature = "kem")] {
//! use b_wing::{KWing, KemError};
//!
//! // --- Key generation (recipient) -----------------------------------
//! // Fill from a CSPRNG in production (e.g. `getrandom::fill`).
//! // Never reuse the same seed for different recipients.
//! let secret_seed = [0u8; 128];
//! let recipient = KWing::from_seed(&secret_seed).unwrap();
//! let encapsulation_key = recipient.get_pub_key(); // share with senders
//!
//! // --- Encapsulation (sender) ----------------------------------------
//! // MUST be freshly generated for every encapsulation — never reuse!
//! let encaps_seed = [1u8; 128];
//! let (ciphertext, shared_secret) = KWing::encapsulate(&encaps_seed, encapsulation_key).unwrap();
//! // `shared_secret` is a 64-byte OKM suitable for deriving symmetric keys.
//!
//! // --- Decapsulation (recipient) -------------------------------------
//! let recovered = recipient.decapsulate(&ciphertext).unwrap();
//! assert_eq!(shared_secret, recovered);
//! # }
//! ```

use hkdf::Hkdf;
use sha3::Sha3_512;
use zeroize::Zeroizing;

// RNG for FrodoKEM Determinism
use rand_chacha::ChaCha20Rng;
use rand_core::SeedableRng;

// Main imports
use frodo_kem::{
    Algorithm, Ciphertext as FrodoCiphertext, DecryptionKey as FrodoDecryptionKey,
    EncryptionKey as FrodoEncryptionKey,
};
use ml_kem::{
    Ciphertext, MlKem1024, Seed as MlKemSeed,
    array::Array,
    kem::{Decapsulate, KeyExport},
};
use x25519_dalek::{PublicKey, StaticSecret};

// ======================================================================
// Constants & Error Types
// ======================================================================

/// A fixed 64-byte domain-separation tag embedded in the HKDF `info` field.
///
/// Binds every derived OKM to the KWing construction, preventing cross-protocol
/// confusion attacks.  The value is a randomly-generated constant chosen at
/// library design time and MUST NOT change across versions (doing so would
/// silently break all existing key material).
const K_WING_OKM_CONTEXT: &'static [u8; 64] = &[
    23, 18, 198, 136, 205, 78, 247, 102, 135, 178, 234, 65, 223, 184, 208, 126, 20, 210, 94, 166,
    168, 92, 94, 241, 48, 209, 96, 164, 56, 106, 245, 205, 94, 113, 223, 88, 245, 94, 152, 82, 1,
    243, 111, 55, 252, 234, 237, 104, 244, 74, 251, 49, 208, 140, 49, 164, 217, 58, 35, 189, 66, 7,
    225, 167,
];

/// Errors that can occur during KWing encapsulation or decapsulation.
///
/// All variants implement [`core::fmt::Display`] and [`core::error::Error`]
/// (on Rust ≥ 1.81 / when `std` is available) for ergonomic error propagation.
///
/// # Example
///
/// ```rust,no_run
/// # #[cfg(feature = "kem")] {
/// use b_wing::{KWing, KemError};
/// let bad_pk = vec![0u8; 10]; // wrong size
/// assert_eq!(
///     KWing::encapsulate(&[0u8; 128], &bad_pk),
///     Err(KemError::InvalidFormat),
/// );
/// # }
/// ```
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Error {
    /// One of the three underlying encapsulation primitives returned an error.
    ///
    /// This is unexpected under normal operation; it indicates a bug in the
    /// caller-supplied seed or an upstream library fault.
    EncapsulateError,
    /// One of the three underlying decapsulation primitives returned an error.
    ///
    /// This typically indicates that the ciphertext was produced by a different
    /// key pair, or that it has been corrupted / tampered with.
    DecapsulateError,
    /// The X25519 Diffie-Hellman output was an all-zero (non-contributory) point.
    ///
    /// This is a known mathematical degenerate case for Curve25519.  KWing
    /// rejects it to prevent a class of small-subgroup attacks in which an
    /// attacker supplies a low-order public key.
    LowEntropyKey,
    /// A key, ciphertext, or seed was the wrong length or could not be parsed.
    ///
    /// The expected sizes are [`KWing::ENCAPSULATION_KEY_SIZE`] and
    /// [`KWing::CIPHERTEXT_SIZE`] respectively.
    InvalidFormat,
}

impl core::fmt::Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Error::EncapsulateError => write!(f, "Encapsulation failed"),
            Error::DecapsulateError => write!(f, "Decapsulation failed"),
            Error::LowEntropyKey => write!(f, "Low entropy or non-contributory key detected"),
            Error::InvalidFormat => write!(f, "Invalid format or size"),
        }
    }
}

// ======================================================================
// Helper Functions
// ======================================================================

/// Derives the final 64-byte Output Keying Material (OKM) via HKDF-SHA3-512.
///
/// This function implements the **transcript-binding** step of KWing:
/// all three shared secrets are combined as HKDF input key material (IKM),
/// while the full protocol transcript is fed as the `info` field to prevent
/// cross-context key reuse.
///
/// # HKDF Construction
///
/// ```text
/// IKM  = dh_ss || ml_kem_ss || frodo_ss   (96 bytes)
/// salt = 32-byte caller-supplied random salt
/// info = dh_eph_pub || ml_kem_ct || frodo_ct || ek || K_WING_OKM_CONTEXT
/// OKM  = HKDF-SHA3-512(IKM, salt, info)[0..64]
/// ```
///
/// # Arguments
///
/// * `dh_ss`     — 32-byte X25519 shared secret (zeroized after use).
/// * `ml_kem_ss` — 32-byte ML-KEM-1024 shared secret (zeroized after use).
/// * `frodo_ss`  — 32-byte FrodoKEM shared secret (zeroized after use).
/// * `salt`      — 32-byte encapsulation salt (part of the ciphertext).
/// * `dh_eph_pub`— 32-byte ephemeral X25519 public key.
/// * `ml_kem_ct` — 1568-byte ML-KEM ciphertext.
/// * `frodo_ct`  — FrodoKEM ciphertext byte slice.
/// * `ek`        — The recipient's composite encapsulation key.
///
/// # Returns
///
/// The 64-byte OKM, or [`Error::InvalidFormat`] if HKDF expansion fails
/// (only possible if the output length exceeds the HKDF limit, which cannot
/// happen with a fixed 64-byte output).
fn derive_key(
    dh_ss: Zeroizing<[u8; 32]>,
    ml_kem_ss: Zeroizing<[u8; 32]>,
    frodo_ss: Zeroizing<[u8; 32]>,
    salt: &[u8; 32],
    dh_eph_pub: &[u8; 32],
    ml_kem_ct: &[u8; 1568],
    frodo_ct: &[u8],
    ek: &[u8],
) -> Result<[u8; 64], Error> {
    // IKM is now 96 bytes (X25519 + ML-KEM + FrodoKEM)
    let mut ikm = Zeroizing::new([0u8; 96]);
    ikm[0..32].copy_from_slice(&*dh_ss);
    ikm[32..64].copy_from_slice(&*ml_kem_ss);
    ikm[64..96].copy_from_slice(&*frodo_ss);
    drop((dh_ss, ml_kem_ss, frodo_ss));

    let hkdf = Hkdf::<Sha3_512>::new(Some(salt), &*ikm);
    drop(ikm);

    // Build the common transcript prefix once
    let mut okm_info = Vec::with_capacity(
        32 + ml_kem_ct.len() + frodo_ct.len() + ek.len() + K_WING_OKM_CONTEXT.len(),
    );
    okm_info.extend_from_slice(dh_eph_pub);
    okm_info.extend_from_slice(ml_kem_ct);
    okm_info.extend_from_slice(frodo_ct);
    okm_info.extend_from_slice(ek);
    okm_info.extend_from_slice(K_WING_OKM_CONTEXT);

    let mut okm = [0u8; 64];
    hkdf.expand(&okm_info, &mut okm)
        .map_err(|_| Error::InvalidFormat)?;

    Ok(okm)
}

// ======================================================================
// Expanded KWing Key (Stateful / High-Throughput)
// ======================================================================

/// A stateful, high-throughput KWing key holder for decapsulation.
///
/// `KWing` pre-computes and heap-caches all three component secret keys and
/// the composite encapsulation key at construction time.  Subsequent calls to
/// [`decapsulate`][KWing::decapsulate] reuse the cached material without any
/// additional key-derivation overhead.
///
/// # Key Sizes
///
/// | Constant | Bytes | Layout |
/// |----------|-------|--------|
/// | [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE] | 23,120 | `X25519(32) ‖ ML-KEM-1024(1568) ‖ FrodoKEM-1344(21520)` |
/// | [`CIPHERTEXT_SIZE`][KWing::CIPHERTEXT_SIZE] | 23,328 | `X25519 eph(32) ‖ Salt(32) ‖ ML-KEM CT(1568) ‖ FrodoKEM CT(21696)` |
///
/// # Security Note
///
/// The encapsulation key (public key) returned by [`get_pub_key`][KWing::get_pub_key]
/// is safe to distribute freely.  The underlying secret key material stored in
/// The underlying X25519 and ML-KEM secret material stored in this struct is
/// wrapped in [`zeroize::Zeroizing`]. FrodoKEM material depends on standard
/// process memory cleanup.
///
/// # Example
///
/// ```rust,no_run
/// # #[cfg(feature = "kem")] {
/// use b_wing::KWing;
///
/// let secret_seed = [0u8; 128]; // use a real CSPRNG in production
/// let kwing = KWing::from_seed(&secret_seed).unwrap();
///
/// // The public encapsulation key can be shared with any sender.
/// let ek = kwing.get_pub_key();
/// assert_eq!(ek.len(), KWing::ENCAPSULATION_KEY_SIZE);
/// # }
/// ```
pub struct KWing {
    dh_secret: Zeroizing<StaticSecret>,
    ml_kem_dk: ml_kem::DecapsulationKey<MlKem1024>,
    frodo_sk: FrodoDecryptionKey,
    composite_pk: Vec<u8>,
}

impl KWing {
    /// Byte length of the composite encapsulation (public) key: **23,120 bytes**.
    ///
    /// Memory layout:
    /// ```text
    /// [ X25519 pub (32) | ML-KEM-1024 ek (1568) | FrodoKEM-1344 pk (21520) ]
    /// ```
    pub const ENCAPSULATION_KEY_SIZE: usize = 23120;

    /// Byte length of the composite ciphertext: **23,328 bytes**.
    ///
    /// Memory layout:
    /// ```text
    /// [ X25519 eph pub (32) | Salt (32) | ML-KEM-1024 ct (1568) | FrodoKEM-1344 ct (21696) ]
    /// ```
    pub const CIPHERTEXT_SIZE: usize = 23328;

    /// Expands a 128-byte secret seed into a fully initialized `KWing` key holder.
    ///
    /// The seed is partitioned deterministically as follows:
    ///
    /// | Bytes | Usage |
    /// |-------|-------|
    /// | `[0..32]`   | X25519 static secret |
    /// | `[32..64]`  | ML-KEM-1024 keygen parameter `d` |
    /// | `[64..96]`  | ML-KEM-1024 keygen parameter `z` |
    /// | `[96..128]` | FrodoKEM-1344 keygen seed (fed into ChaCha20) |
    ///
    /// # Security Requirements
    ///
    /// * `secret_seed` **must** be generated by a cryptographically secure
    ///   random number generator (CSPRNG) such as `getrandom`.
    /// * Never reuse the same seed for different recipients or sessions.
    /// * The seed should be treated with the same care as a private key.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidFormat`] if an internal slice conversion fails
    /// (this should be impossible given a correctly-sized input).
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "kem")] {
    /// use b_wing::KWing;
    ///
    /// let mut seed = [0u8; 128];
    /// getrandom::fill(&mut seed).expect("CSPRNG failed");
    /// let kwing = KWing::from_seed(&seed).expect("key generation failed");
    /// # }
    /// ```
    pub fn from_seed(secret_seed: &[u8; 128]) -> Result<Self, Error> {
        // 1. X25519
        let dh_secret = Zeroizing::new(StaticSecret::from(
            <[u8; 32]>::try_from(&secret_seed[0..32]).map_err(|_| Error::InvalidFormat)?,
        ));
        let dh_pub = PublicKey::from(&*dh_secret);

        // 2. ML-KEM-1024
        let ml_kem_d = Zeroizing::new(
            <[u8; 32]>::try_from(&secret_seed[32..64]).map_err(|_| Error::InvalidFormat)?,
        );
        let ml_kem_z = Zeroizing::new(
            <[u8; 32]>::try_from(&secret_seed[64..96]).map_err(|_| Error::InvalidFormat)?,
        );
        // ML-KEM-1024 Seed = d || z (64 bytes)
        let mut ml_kem_seed = MlKemSeed::default();
        ml_kem_seed[..32].copy_from_slice(&*ml_kem_d);
        ml_kem_seed[32..].copy_from_slice(&*ml_kem_z);
        let ml_kem_dk = ml_kem::DecapsulationKey::<MlKem1024>::from_seed(ml_kem_seed);
        let ml_ek = ml_kem_dk.encapsulation_key();

        // 3. FrodoKEM-1344-SHAKE
        let frodo_seed =
            <[u8; 32]>::try_from(&secret_seed[96..128]).map_err(|_| Error::InvalidFormat)?;
        let mut frodo_rng = ChaCha20Rng::from_seed(frodo_seed);
        let frodo = Algorithm::FrodoKem1344Shake;
        let (frodo_pk, frodo_sk) = frodo.generate_keypair(&mut frodo_rng);

        // 4. Cache composite PK on Heap
        let mut composite_pk = Vec::with_capacity(Self::ENCAPSULATION_KEY_SIZE);
        composite_pk.extend_from_slice(dh_pub.as_bytes());
        composite_pk.extend_from_slice(&ml_ek.to_bytes());
        composite_pk.extend_from_slice(frodo_pk.value());

        Ok(Self {
            dh_secret,
            ml_kem_dk,
            frodo_sk,
            composite_pk,
        })
    }

    /// Returns a reference to the cached composite encapsulation key.
    ///
    /// The returned slice is [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE]
    /// bytes long and is safe to distribute publicly.  Pass it to
    /// [`encapsulate`][KWing::encapsulate] on the sender's side.
    #[must_use]
    pub fn get_pub_key(&self) -> &[u8] {
        &self.composite_pk
    }

    /// Encapsulates a fresh shared secret against the recipient's composite public key.
    ///
    /// This is the **sender-side** operation.  It runs all three component KEMs
    /// deterministically from `encaps_seed` and combines their outputs into a
    /// single composite ciphertext and a 64-byte OKM.
    ///
    /// # Seed Layout
    ///
    /// | Bytes | Usage |
    /// |-------|-------|
    /// | `[0..32]`   | X25519 ephemeral secret |
    /// | `[32..64]`  | ML-KEM-1024 randomness `m` |
    /// | `[64..96]`  | FrodoKEM encapsulation randomness (ChaCha20 seed) |
    /// | `[96..128]` | HKDF salt (transmitted in the ciphertext) |
    ///
    /// # Security Requirements
    ///
    /// * `encaps_seed` **must** be freshly generated from a CSPRNG for **every**
    ///   encapsulation.  Reusing the seed against the same recipient leaks the
    ///   X25519 secret key.
    ///
    /// # Errors
    ///
    /// | Variant | Cause |
    /// |---------|-------|
    /// | [`Error::InvalidFormat`] | `ek` is not exactly [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE] bytes |
    /// | [`Error::LowEntropyKey`] | X25519 DH output is a low-order (all-zero) point |
    /// | [`Error::EncapsulateError`] | An underlying KEM primitive failed |
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "kem")] {
    /// use b_wing::KWing;
    ///
    /// # let secret_seed = [0u8; 128];
    /// # let kwing = KWing::from_seed(&secret_seed).unwrap();
    /// # let ek = kwing.get_pub_key();
    /// let mut encaps_seed = [0u8; 128];
    /// getrandom::fill(&mut encaps_seed).expect("CSPRNG failed");
    ///
    /// let (ciphertext, shared_secret) = KWing::encapsulate(&encaps_seed, ek).unwrap();
    /// assert_eq!(ciphertext.len(), KWing::CIPHERTEXT_SIZE);
    /// assert_eq!(shared_secret.len(), 64);
    /// # }
    /// ```
    pub fn encapsulate(encaps_seed: &[u8; 128], ek: &[u8]) -> Result<(Vec<u8>, [u8; 64]), Error> {
        if ek.len() != Self::ENCAPSULATION_KEY_SIZE {
            return Err(Error::InvalidFormat);
        }

        let frodo = Algorithm::FrodoKem1344Shake;

        // 1. Parse Composite Key
        let dh_pub =
            PublicKey::from(<[u8; 32]>::try_from(&ek[0..32]).map_err(|_| Error::InvalidFormat)?);
        let ml_kem_ek_bytes: ml_kem::kem::Key<ml_kem::EncapsulationKey<MlKem1024>> = Array(
            ek[32..1600].try_into().map_err(|_| Error::InvalidFormat)?,
        );
        let ml_kem_ek = ml_kem::EncapsulationKey::<MlKem1024>::new(&ml_kem_ek_bytes)
            .map_err(|_| Error::InvalidFormat)?;
        let frodo_pk =
            FrodoEncryptionKey::from_bytes(frodo, &ek[1600..]).map_err(|_| Error::InvalidFormat)?;

        // 2. Setup Deterministic RNGs
        let dh_eph_secret = Zeroizing::new(StaticSecret::from(
            <[u8; 32]>::try_from(&encaps_seed[0..32]).map_err(|_| Error::InvalidFormat)?,
        ));
        let ml_kem_m = Zeroizing::new(
            <[u8; 32]>::try_from(&encaps_seed[32..64]).map_err(|_| Error::InvalidFormat)?,
        );
        let frodo_rng_seed =
            <[u8; 32]>::try_from(&encaps_seed[64..96]).map_err(|_| Error::InvalidFormat)?;
        let salt = <[u8; 32]>::try_from(&encaps_seed[96..128]).map_err(|_| Error::InvalidFormat)?;

        let dh_eph_pub = PublicKey::from(&*dh_eph_secret);
        let dh_eph_pub_bytes = dh_eph_pub.as_bytes();

        // 3. Execute X25519
        let dh_ss = Zeroizing::new(dh_eph_secret.diffie_hellman(&dh_pub));
        if !dh_ss.was_contributory() {
            return Err(Error::LowEntropyKey);
        }

        // 4. Execute ML-KEM
        let (ml_kem_ct, ml_kem_ss) = ml_kem_ek
            .encapsulate_deterministic(&Array(*ml_kem_m));
        let ml_kem_ss: Zeroizing<[u8; 32]> = Zeroizing::new(ml_kem_ss.into());

        // 5. Execute FrodoKEM
        let mut frodo_rng = ChaCha20Rng::from_seed(frodo_rng_seed);
        let (frodo_ct, frodo_ss) = frodo
            .encapsulate_with_rng(&frodo_pk, &mut frodo_rng)
            .map_err(|_| Error::EncapsulateError)?;
        let frodo_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
            frodo_ss
                .value()
                .try_into()
                .map_err(|_| Error::EncapsulateError)?,
        );
        let frodo_ct_arr = frodo_ct.value().to_vec();

        // 6. HKDF Derivation
        let ml_kem_ct_bytes: [u8; 1568] = ml_kem_ct.0
            .try_into()
            .map_err(|_| Error::EncapsulateError)?;
        let okm = derive_key(
            Zeroizing::new(dh_ss.to_bytes()),
            ml_kem_ss,
            frodo_ss,
            &salt,
            dh_eph_pub_bytes,
            &ml_kem_ct_bytes,
            &frodo_ct_arr,
            ek,
        )?;

        // 7. Assemble Ciphertext
        let mut ciphertext = Vec::with_capacity(Self::CIPHERTEXT_SIZE);
        ciphertext.extend_from_slice(dh_eph_pub_bytes);
        ciphertext.extend_from_slice(&salt);
        ciphertext.extend_from_slice(&ml_kem_ct_bytes);
        ciphertext.extend_from_slice(&frodo_ct_arr);

        Ok((ciphertext, okm))
    }

    /// Decapsulates a composite ciphertext to recover the 64-byte Output Keying Material.
    ///
    /// This is the **recipient-side** operation.  It parses the composite
    /// ciphertext, runs all three component decapsulations using the cached
    /// secret keys, and recomputes the HKDF transcript to produce the OKM.
    ///
    /// The OKM is cryptographically bound to the ciphertext and to this specific
    /// `KWing` instance, so it will not match any other recipient or ciphertext.
    ///
    /// # Errors
    ///
    /// | Variant | Cause |
    /// |---------|-------|
    /// | [`Error::InvalidFormat`] | `ct` is not exactly [`CIPHERTEXT_SIZE`][KWing::CIPHERTEXT_SIZE] bytes |
    /// | [`Error::LowEntropyKey`] | X25519 DH output is a low-order (all-zero) point |
    /// | [`Error::DecapsulateError`] | An underlying KEM primitive failed |
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "kem")] {
    /// use b_wing::KWing;
    ///
    /// # let secret_seed = [0u8; 128];
    /// # let encaps_seed = [1u8; 128];
    /// # let kwing = KWing::from_seed(&secret_seed).unwrap();
    /// # let ek = kwing.get_pub_key().to_vec();
    /// # let (ct, _) = KWing::encapsulate(&encaps_seed, &ek).unwrap();
    /// let okm = kwing.decapsulate(&ct).unwrap();
    /// assert_eq!(okm.len(), 64);
    /// // Derive a 32-byte AES-256 key and 32-byte MAC key from the OKM:
    /// let aes_key = &okm[..32];
    /// let mac_key = &okm[32..];
    /// # }
    /// ```
    pub fn decapsulate(&self, ct: &[u8]) -> Result<[u8; 64], Error> {
        if ct.len() != Self::CIPHERTEXT_SIZE {
            return Err(Error::InvalidFormat);
        }

        let frodo = Algorithm::FrodoKem1344Shake;

        // 1. Parse Ciphertext
        let dh_eph_pub =
            PublicKey::from(<[u8; 32]>::try_from(&ct[0..32]).map_err(|_| Error::InvalidFormat)?);
        let salt: [u8; 32] = ct[32..64].try_into().map_err(|_| Error::InvalidFormat)?;
        let ml_kem_ct: Ciphertext<MlKem1024> = Array(ct[64..1632].try_into().map_err(|_| Error::InvalidFormat)?);
        let frodo_ct_bytes = &ct[1632..];
        let frodo_ct = FrodoCiphertext::from_bytes(frodo, &frodo_ct_bytes)
            .map_err(|_| Error::InvalidFormat)?;
        // 2. Execute X25519
        let dh_ss = Zeroizing::new(self.dh_secret.diffie_hellman(&dh_eph_pub));
        if !dh_ss.was_contributory() {
            return Err(Error::LowEntropyKey);
        }

        // 3. Execute ML-KEM
        let ml_kem_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
            self.ml_kem_dk
                .decapsulate(&ml_kem_ct)
                .into(),
        );

        // 4. Execute FrodoKEM
        let frodo_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
            frodo
                .decapsulate(&self.frodo_sk, &frodo_ct)
                .map_err(|_| Error::DecapsulateError)?
                .0
                .value()
                .try_into()
                .map_err(|_| Error::DecapsulateError)?,
        );

        // 5. HKDF Derivation & Proof
        let ml_kem_ct_bytes: [u8; 1568] = ml_kem_ct.0
            .try_into()
            .map_err(|_| Error::DecapsulateError)?;
        let okm = derive_key(
            Zeroizing::new(dh_ss.to_bytes()),
            ml_kem_ss,
            frodo_ss,
            &salt,
            dh_eph_pub.as_bytes(),
            &ml_kem_ct_bytes,
            frodo_ct_bytes,
            &self.get_pub_key(),
        )?;

        Ok(okm)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::LazyLock;

    // ======================================================================
    // Test Constants & Helpers
    // ======================================================================

    static SECRET_SEED: [u8; 128] = [0x42; 128];
    static ENCAPS_SEED: [u8; 128] = [0x84; 128];

    static K_WING: LazyLock<KWing> = LazyLock::new(|| KWing::from_seed(&SECRET_SEED).unwrap());

    static ENCAPS_RESULT: LazyLock<(Vec<u8>, [u8; 64])> =
        LazyLock::new(|| KWing::encapsulate(&ENCAPS_SEED, K_WING.get_pub_key()).unwrap());

    // ======================================================================
    // Happy Path & Determinism
    // ======================================================================

    #[test]
    fn test_happy_path_round_trip() {
        // 1. Get cached Public Key
        let pk = K_WING.get_pub_key();
        assert_eq!(pk.len(), KWing::ENCAPSULATION_KEY_SIZE);

        // 2. Get cached Encapsulation
        let (ct, okm_encapsulated) = &*ENCAPS_RESULT;
        assert_eq!(ct.len(), KWing::CIPHERTEXT_SIZE);

        // 3. Decapsulate
        let okm_decapsulated = K_WING
            .decapsulate(ct)
            .expect("Decapsulation should succeed");

        // 4. Assert Output Keying Material Matches
        assert_eq!(
            okm_encapsulated, &okm_decapsulated,
            "Decapsulated OKM must exactly match the Encapsulated OKM"
        );
    }

    #[test]
    fn test_strict_determinism() {
        // Deterministic Key Generation
        let binding = KWing::from_seed(&SECRET_SEED).unwrap();
        let pk2 = binding.get_pub_key();
        assert_eq!(
            K_WING.get_pub_key(),
            pk2,
            "Public keys must be identical for the same seed"
        );

        // Deterministic Encapsulation
        let (ct2, okm2) = KWing::encapsulate(&ENCAPS_SEED, pk2).unwrap();
        assert_eq!(
            ENCAPS_RESULT.0, ct2,
            "Ciphertexts must be identical for the same seeds"
        );
        assert_eq!(
            ENCAPS_RESULT.1, okm2,
            "OKMs must be identical for the same seeds"
        );
    }

    // ======================================================================
    // Formatting & Boundary Rejections
    // ======================================================================

    #[test]
    fn test_invalid_public_key_length() {
        let bad_pk = vec![0u8; KWing::ENCAPSULATION_KEY_SIZE - 1]; // 1 byte too short

        let result = KWing::encapsulate(&ENCAPS_SEED, &bad_pk);
        assert_eq!(
            result,
            Err(Error::InvalidFormat),
            "Encapsulate must reject invalid public key lengths immediately"
        );
    }

    #[test]
    fn test_invalid_ciphertext_length() {
        let bad_ct = vec![0u8; KWing::CIPHERTEXT_SIZE + 5]; // 5 bytes too long

        let result = K_WING.decapsulate(&bad_ct);
        assert_eq!(
            result,
            Err(Error::InvalidFormat),
            "Decapsulate must reject invalid ciphertext lengths immediately"
        );
    }

    // ======================================================================
    // Cryptographic Tampering & Mathematical Edge Cases
    // ======================================================================

    #[test]
    fn test_low_entropy_key_encapsulate_rejection() {
        let mut pk = K_WING.get_pub_key().to_vec();

        // Force the X25519 public key part to all zeros.
        pk[0..32].fill(0);

        let result = KWing::encapsulate(&ENCAPS_SEED, &pk);
        assert_eq!(
            result,
            Err(Error::LowEntropyKey),
            "Encapsulate must reject mathematical weak points (all-zero DH shared secret)"
        );
    }

    #[test]
    fn test_low_entropy_key_decapsulate_rejection() {
        let mut ct = ENCAPS_RESULT.0.clone();

        // Force the Ephemeral X25519 public key part in the CT to all zeros.
        ct[0..32].fill(0);

        let result = K_WING.decapsulate(&ct);
        assert_eq!(
            result,
            Err(Error::LowEntropyKey),
            "Decapsulate must reject mathematical weak points injected by an attacker"
        );
    }
}