Skip to main content

b_wing/
kem.rs

1//! # KWing: Hybrid Key Encapsulation Mechanism
2//!
3//! KWing is a paranoia-grade, triple-layered hybrid KEM that derives a shared
4//! secret with **NIST Level 5 (256-bit)** post-quantum security.  The construction
5//! follows a **redundancy-first** policy: the derived Output Keying Material (OKM)
6//! remains secret as long as *any single one* of the three component algorithms is
7//! unbroken — whether by classical cryptanalysis, a quantum adversary, or a
8//! lattice-specific breakthrough.
9//!
10//! ## Security Composition
11//!
12//! | Layer | Algorithm | Security assumption |
13//! |-------|-----------|---------------------|
14//! | 1 | **X25519** | Classical ECDH — Curve25519 discrete log |
15//! | 2 | **ML-KEM-1024** (Kyber) | Module-lattice MLWE (NIST PQ Level 5) |
16//! | 3 | **FrodoKEM-1344-SHAKE** | Unstructured LWE — conservative lattice |
17//!
18//! The three independent shared secrets are combined via **HKDF-SHA3-512**,
19//! cryptographically binding them to the full transcript (ephemeral public key,
20//! all three ciphertexts, the recipient's encapsulation key, and a 32-byte salt).
21//!
22//! ## Design Principles
23//!
24//! * **Caller-supplied randomness.** All randomness is given as typed seed arrays.
25//!   The API is fully deterministic — no hidden global RNG state — and therefore
26//!   compatible with `no_std`, WASM, and embedded targets.
27//! * **Heap-cached keys.** [`KWing`] pre-computes the composite encapsulation key
28//!   once on construction, amortising the cost across many encapsulate/decapsulate
29//!   calls.
30//! * **Zeroize on drop.** Classical and ML-KEM secret material is wrapped in
31//!   [`zeroize::Zeroizing`], guaranteeing erasure from memory. FrodoKEM
32//!   material relies on standard memory cleanup as the underlying crate
33//!   does not yet support the `zeroize` trait.
34//! * **Strict size validation.** Every public entry-point validates buffer sizes
35//!   before performing any cryptographic work.
36//!
37//! ## Usage
38//!
39//! ```rust,no_run
40//! # #[cfg(feature = "kem")] {
41//! use b_wing::{KWing, KemError};
42//!
43//! // --- Key generation (recipient) -----------------------------------
44//! // Fill from a CSPRNG in production (e.g. `getrandom::fill`).
45//! // Never reuse the same seed for different recipients.
46//! let secret_seed = [0u8; 128];
47//! let recipient = KWing::from_seed(&secret_seed).unwrap();
48//! let encapsulation_key = recipient.get_pub_key(); // share with senders
49//!
50//! // --- Encapsulation (sender) ----------------------------------------
51//! // MUST be freshly generated for every encapsulation — never reuse!
52//! let encaps_seed = [1u8; 128];
53//! let (ciphertext, shared_secret) = KWing::encapsulate(&encaps_seed, encapsulation_key).unwrap();
54//! // `shared_secret` is a 64-byte OKM suitable for deriving symmetric keys.
55//!
56//! // --- Decapsulation (recipient) -------------------------------------
57//! let recovered = recipient.decapsulate(&ciphertext).unwrap();
58//! assert_eq!(shared_secret, recovered);
59//! # }
60//! ```
61
62use hkdf::Hkdf;
63use hybrid_array::Array;
64use sha3::Sha3_512;
65use zeroize::Zeroizing;
66
67// RNG for FrodoKEM Determinism
68use rand_chacha::ChaCha20Rng;
69use rand_core::SeedableRng;
70
71// Main imports
72use frodo_kem::{
73    Algorithm, Ciphertext as FrodoCiphertext, DecryptionKey as FrodoDecryptionKey,
74    EncryptionKey as FrodoEncryptionKey,
75};
76use ml_kem::{
77    B32, Ciphertext as MlKemCiphertext, MlKem1024, Seed as MlKemSeed, SharedKey as MlKemSharedKey,
78    kem::{Decapsulate, KeyExport, TryKeyInit},
79};
80use x25519_dalek::{PublicKey, StaticSecret};
81
82// ======================================================================
83// Constants & Error Types
84// ======================================================================
85
86/// A fixed 64-byte domain-separation tag embedded in the HKDF `info` field.
87///
88/// Binds every derived OKM to the KWing construction, preventing cross-protocol
89/// confusion attacks.  The value is a randomly-generated constant chosen at
90/// library design time and MUST NOT change across versions (doing so would
91/// silently break all existing key material).
92const K_WING_OKM_CONTEXT: &[u8; 64] = &[
93    23, 18, 198, 136, 205, 78, 247, 102, 135, 178, 234, 65, 223, 184, 208, 126, 20, 210, 94, 166,
94    168, 92, 94, 241, 48, 209, 96, 164, 56, 106, 245, 205, 94, 113, 223, 88, 245, 94, 152, 82, 1,
95    243, 111, 55, 252, 234, 237, 104, 244, 74, 251, 49, 208, 140, 49, 164, 217, 58, 35, 189, 66, 7,
96    225, 167,
97];
98
99/// Errors that can occur during KWing encapsulation or decapsulation.
100///
101/// All variants implement [`core::fmt::Display`] and [`core::error::Error`]
102/// (on Rust ≥ 1.81 / when `std` is available) for ergonomic error propagation.
103///
104/// # Example
105///
106/// ```rust,no_run
107/// # #[cfg(feature = "kem")] {
108/// use b_wing::{KWing, KemError};
109/// let bad_pk = vec![0u8; 10]; // wrong size
110/// assert_eq!(
111///     KWing::encapsulate(&[0u8; 128], &bad_pk),
112///     Err(KemError::InvalidFormat),
113/// );
114/// # }
115/// ```
116#[derive(Debug, Copy, Clone, PartialEq, Eq)]
117pub enum Error {
118    /// One of the three underlying encapsulation primitives returned an error.
119    ///
120    /// This is unexpected under normal operation; it indicates a bug in the
121    /// caller-supplied seed or an upstream library fault.
122    EncapsulateError,
123    /// One of the three underlying decapsulation primitives returned an error.
124    ///
125    /// This typically indicates that the ciphertext was produced by a different
126    /// key pair, or that it has been corrupted / tampered with.
127    DecapsulateError,
128    /// The X25519 Diffie-Hellman output was an all-zero (non-contributory) point.
129    ///
130    /// This is a known mathematical degenerate case for Curve25519.  KWing
131    /// rejects it to prevent a class of small-subgroup attacks in which an
132    /// attacker supplies a low-order public key.
133    LowEntropyKey,
134    /// A key, ciphertext, or seed was the wrong length or could not be parsed.
135    ///
136    /// The expected sizes are [`KWing::ENCAPSULATION_KEY_SIZE`] and
137    /// [`KWing::CIPHERTEXT_SIZE`] respectively.
138    InvalidFormat,
139}
140
141impl core::fmt::Display for Error {
142    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
143        match self {
144            Error::EncapsulateError => write!(f, "Encapsulation failed"),
145            Error::DecapsulateError => write!(f, "Decapsulation failed"),
146            Error::LowEntropyKey => write!(f, "Low entropy or non-contributory key detected"),
147            Error::InvalidFormat => write!(f, "Invalid format or size"),
148        }
149    }
150}
151
152// ======================================================================
153// Helper Functions
154// ======================================================================
155
156/// Derives the final 64-byte Output Keying Material (OKM) via HKDF-SHA3-512.
157///
158/// This function implements the **transcript-binding** step of KWing:
159/// all three shared secrets are combined as HKDF input key material (IKM),
160/// while the full protocol transcript is fed as the `info` field to prevent
161/// cross-context key reuse.
162///
163/// # HKDF Construction
164///
165/// ```text
166/// IKM  = dh_ss || ml_kem_ss || frodo_ss   (96 bytes)
167/// salt = 32-byte caller-supplied random salt
168/// info = dh_eph_pub || ml_kem_ct || frodo_ct || ek || K_WING_OKM_CONTEXT
169/// OKM  = HKDF-SHA3-512(IKM, salt, info)[0..64]
170/// ```
171///
172/// # Arguments
173///
174/// * `dh_ss`     — 32-byte X25519 shared secret (zeroized after use).
175/// * `ml_kem_ss` — 32-byte ML-KEM-1024 shared secret (zeroized after use).
176/// * `frodo_ss`  — 32-byte FrodoKEM shared secret (zeroized after use).
177/// * `salt`      — 32-byte encapsulation salt (part of the ciphertext).
178/// * `dh_eph_pub`— 32-byte ephemeral X25519 public key.
179/// * `ml_kem_ct` — 1568-byte ML-KEM ciphertext.
180/// * `frodo_ct`  — FrodoKEM ciphertext byte slice.
181/// * `ek`        — The recipient's composite encapsulation key.
182///
183/// # Returns
184///
185/// The 64-byte OKM, or [`Error::InvalidFormat`] if HKDF expansion fails
186/// (only possible if the output length exceeds the HKDF limit, which cannot
187/// happen with a fixed 64-byte output).
188#[allow(clippy::too_many_arguments)]
189fn derive_key(
190    dh_ss: Zeroizing<[u8; 32]>,
191    ml_kem_ss: Zeroizing<[u8; 32]>,
192    frodo_ss: Zeroizing<[u8; 32]>,
193    salt: &[u8; 32],
194    dh_eph_pub: &[u8; 32],
195    ml_kem_ct: &[u8; 1568],
196    frodo_ct: &[u8],
197    ek: &[u8],
198) -> Result<[u8; 64], Error> {
199    // IKM is now 96 bytes (X25519 + ML-KEM + FrodoKEM)
200    let mut ikm = Zeroizing::new([0u8; 96]);
201    ikm[0..32].copy_from_slice(&*dh_ss);
202    ikm[32..64].copy_from_slice(&*ml_kem_ss);
203    ikm[64..96].copy_from_slice(&*frodo_ss);
204    drop((dh_ss, ml_kem_ss, frodo_ss));
205
206    let hkdf = Hkdf::<Sha3_512>::new(Some(salt), &*ikm);
207    drop(ikm);
208
209    // Build the common transcript prefix once
210    let mut okm_info = Vec::with_capacity(
211        32 + ml_kem_ct.len() + frodo_ct.len() + ek.len() + K_WING_OKM_CONTEXT.len(),
212    );
213    okm_info.extend_from_slice(dh_eph_pub);
214    okm_info.extend_from_slice(ml_kem_ct);
215    okm_info.extend_from_slice(frodo_ct);
216    okm_info.extend_from_slice(ek);
217    okm_info.extend_from_slice(K_WING_OKM_CONTEXT);
218
219    let mut okm = [0u8; 64];
220    hkdf.expand(&okm_info, &mut okm)
221        .map_err(|_| Error::InvalidFormat)?;
222
223    Ok(okm)
224}
225
226// ======================================================================
227// Expanded KWing Key (Stateful / High-Throughput)
228// ======================================================================
229
230/// A stateful, high-throughput KWing key holder for decapsulation.
231///
232/// `KWing` pre-computes and heap-caches all three component secret keys and
233/// the composite encapsulation key at construction time.  Subsequent calls to
234/// [`decapsulate`][KWing::decapsulate] reuse the cached material without any
235/// additional key-derivation overhead.
236///
237/// # Key Sizes
238///
239/// | Constant | Bytes | Layout |
240/// |----------|-------|--------|
241/// | [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE] | 23,120 | `X25519(32) ‖ ML-KEM-1024(1568) ‖ FrodoKEM-1344(21520)` |
242/// | [`CIPHERTEXT_SIZE`][KWing::CIPHERTEXT_SIZE] | 23,328 | `X25519 eph(32) ‖ Salt(32) ‖ ML-KEM CT(1568) ‖ FrodoKEM CT(21696)` |
243///
244/// # Security Note
245///
246/// The encapsulation key (public key) returned by [`get_pub_key`][KWing::get_pub_key]
247/// is safe to distribute freely.  The underlying secret key material stored in
248/// The underlying X25519 and ML-KEM secret material stored in this struct is
249/// wrapped in [`zeroize::Zeroizing`]. FrodoKEM material depends on standard
250/// process memory cleanup.
251///
252/// # Example
253///
254/// ```rust,no_run
255/// # #[cfg(feature = "kem")] {
256/// use b_wing::KWing;
257///
258/// let secret_seed = [0u8; 128]; // use a real CSPRNG in production
259/// let kwing = KWing::from_seed(&secret_seed).unwrap();
260///
261/// // The public encapsulation key can be shared with any sender.
262/// let ek = kwing.get_pub_key();
263/// assert_eq!(ek.len(), KWing::ENCAPSULATION_KEY_SIZE);
264/// # }
265/// ```
266pub struct KWing {
267    dh_secret: Zeroizing<StaticSecret>,
268    ml_kem_dk: ml_kem::DecapsulationKey<MlKem1024>,
269    frodo_sk: FrodoDecryptionKey,
270    composite_pk: Vec<u8>,
271}
272
273impl KWing {
274    /// Byte length of the composite encapsulation (public) key: **23,120 bytes**.
275    ///
276    /// Memory layout:
277    /// ```text
278    /// [ X25519 pub (32) | ML-KEM-1024 ek (1568) | FrodoKEM-1344 pk (21520) ]
279    /// ```
280    pub const ENCAPSULATION_KEY_SIZE: usize = 23120;
281
282    /// Byte length of the composite ciphertext: **23,328 bytes**.
283    ///
284    /// Memory layout:
285    /// ```text
286    /// [ X25519 eph pub (32) | Salt (32) | ML-KEM-1024 ct (1568) | FrodoKEM-1344 ct (21696) ]
287    /// ```
288    pub const CIPHERTEXT_SIZE: usize = 23328;
289
290    /// Expands a 128-byte secret seed into a fully initialized `KWing` key holder.
291    ///
292    /// The seed is partitioned deterministically as follows:
293    ///
294    /// | Bytes | Usage |
295    /// |-------|-------|
296    /// | `[0..32]`   | X25519 static secret |
297    /// | `[32..64]`  | ML-KEM-1024 keygen parameter `d` |
298    /// | `[64..96]`  | ML-KEM-1024 keygen parameter `z` |
299    /// | `[96..128]` | FrodoKEM-1344 keygen seed (fed into ChaCha20) |
300    ///
301    /// # Security Requirements
302    ///
303    /// * `secret_seed` **must** be generated by a cryptographically secure
304    ///   random number generator (CSPRNG) such as `getrandom`.
305    /// * Never reuse the same seed for different recipients or sessions.
306    /// * The seed should be treated with the same care as a private key.
307    ///
308    /// # Errors
309    ///
310    /// Returns [`Error::InvalidFormat`] if an internal slice conversion fails
311    /// (this should be impossible given a correctly-sized input).
312    ///
313    /// # Example
314    ///
315    /// ```rust,no_run
316    /// # #[cfg(feature = "kem")] {
317    /// use b_wing::KWing;
318    ///
319    /// let mut seed = [0u8; 128];
320    /// getrandom::fill(&mut seed).expect("CSPRNG failed");
321    /// let kwing = KWing::from_seed(&seed).expect("key generation failed");
322    /// # }
323    /// ```
324    pub fn from_seed(secret_seed: &[u8; 128]) -> Result<Self, Error> {
325        // 1. X25519
326        let dh_secret = Zeroizing::new(StaticSecret::from(
327            <[u8; 32]>::try_from(&secret_seed[0..32]).map_err(|_| Error::InvalidFormat)?,
328        ));
329        let dh_pub = PublicKey::from(&*dh_secret);
330
331        // 2. ML-KEM-1024
332        let ml_kem_d = Zeroizing::new(
333            <[u8; 32]>::try_from(&secret_seed[32..64]).map_err(|_| Error::InvalidFormat)?,
334        );
335        let ml_kem_z = Zeroizing::new(
336            <[u8; 32]>::try_from(&secret_seed[64..96]).map_err(|_| Error::InvalidFormat)?,
337        );
338        // ML-KEM-1024 Seed = d || z (64 bytes)
339        let mut ml_kem_seed = MlKemSeed::default();
340        ml_kem_seed[..32].copy_from_slice(&*ml_kem_d);
341        ml_kem_seed[32..].copy_from_slice(&*ml_kem_z);
342        let ml_kem_dk = ml_kem::DecapsulationKey::<MlKem1024>::from_seed(ml_kem_seed);
343        let ml_ek = ml_kem_dk.encapsulation_key();
344
345        // 3. FrodoKEM-1344-SHAKE
346        let frodo_seed =
347            <[u8; 32]>::try_from(&secret_seed[96..128]).map_err(|_| Error::InvalidFormat)?;
348        let mut frodo_rng = ChaCha20Rng::from_seed(frodo_seed);
349        let frodo = Algorithm::FrodoKem1344Shake;
350        let (frodo_pk, frodo_sk) = frodo.generate_keypair(&mut frodo_rng);
351
352        // 4. Cache composite PK on Heap
353        let mut composite_pk = Vec::with_capacity(Self::ENCAPSULATION_KEY_SIZE);
354        composite_pk.extend_from_slice(dh_pub.as_bytes());
355        composite_pk.extend_from_slice(ml_ek.to_bytes().as_slice());
356        composite_pk.extend_from_slice(frodo_pk.value());
357
358        Ok(Self {
359            dh_secret,
360            ml_kem_dk,
361            frodo_sk,
362            composite_pk,
363        })
364    }
365
366    /// Returns a reference to the cached composite encapsulation key.
367    ///
368    /// The returned slice is [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE]
369    /// bytes long and is safe to distribute publicly.  Pass it to
370    /// [`encapsulate`][KWing::encapsulate] on the sender's side.
371    #[must_use]
372    pub fn get_pub_key(&self) -> &[u8] {
373        &self.composite_pk
374    }
375
376    /// Encapsulates a fresh shared secret against the recipient's composite public key.
377    ///
378    /// This is the **sender-side** operation.  It runs all three component KEMs
379    /// deterministically from `encaps_seed` and combines their outputs into a
380    /// single composite ciphertext and a 64-byte OKM.
381    ///
382    /// # Seed Layout
383    ///
384    /// | Bytes | Usage |
385    /// |-------|-------|
386    /// | `[0..32]`   | X25519 ephemeral secret |
387    /// | `[32..64]`  | ML-KEM-1024 randomness `m` |
388    /// | `[64..96]`  | FrodoKEM encapsulation randomness (ChaCha20 seed) |
389    /// | `[96..128]` | HKDF salt (transmitted in the ciphertext) |
390    ///
391    /// # Security Requirements
392    ///
393    /// * `encaps_seed` **must** be freshly generated from a CSPRNG for **every**
394    ///   encapsulation.  Reusing the seed against the same recipient leaks the
395    ///   X25519 secret key.
396    ///
397    /// # Errors
398    ///
399    /// | Variant | Cause |
400    /// |---------|-------|
401    /// | [`Error::InvalidFormat`] | `ek` is not exactly [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE] bytes |
402    /// | [`Error::LowEntropyKey`] | X25519 DH output is a low-order (all-zero) point |
403    /// | [`Error::EncapsulateError`] | An underlying KEM primitive failed |
404    ///
405    /// # Example
406    ///
407    /// ```rust,no_run
408    /// # #[cfg(feature = "kem")] {
409    /// use b_wing::KWing;
410    ///
411    /// # let secret_seed = [0u8; 128];
412    /// # let kwing = KWing::from_seed(&secret_seed).unwrap();
413    /// # let ek = kwing.get_pub_key();
414    /// let mut encaps_seed = [0u8; 128];
415    /// getrandom::fill(&mut encaps_seed).expect("CSPRNG failed");
416    ///
417    /// let (ciphertext, shared_secret) = KWing::encapsulate(&encaps_seed, ek).unwrap();
418    /// assert_eq!(ciphertext.len(), KWing::CIPHERTEXT_SIZE);
419    /// assert_eq!(shared_secret.len(), 64);
420    /// # }
421    /// ```
422    pub fn encapsulate(encaps_seed: &[u8; 128], ek: &[u8]) -> Result<(Vec<u8>, [u8; 64]), Error> {
423        if ek.len() != Self::ENCAPSULATION_KEY_SIZE {
424            return Err(Error::InvalidFormat);
425        }
426
427        let frodo = Algorithm::FrodoKem1344Shake;
428
429        // 1. Parse Composite Key — use TryKeyInit::new_from_slice for direct
430        //    byte-slice → EncapsulationKey construction without manual Array wrapping.
431        let dh_pub =
432            PublicKey::from(<[u8; 32]>::try_from(&ek[0..32]).map_err(|_| Error::InvalidFormat)?);
433        let ml_kem_ek = ml_kem::EncapsulationKey::<MlKem1024>::new_from_slice(&ek[32..1600])
434            .map_err(|_| Error::InvalidFormat)?;
435        let frodo_pk =
436            FrodoEncryptionKey::from_bytes(frodo, &ek[1600..]).map_err(|_| Error::InvalidFormat)?;
437
438        // 2. Setup Deterministic RNGs
439        let dh_eph_secret = Zeroizing::new(StaticSecret::from(
440            <[u8; 32]>::try_from(&encaps_seed[0..32]).map_err(|_| Error::InvalidFormat)?,
441        ));
442        let ml_kem_m: B32 = Array(
443            encaps_seed[32..64]
444                .try_into()
445                .map_err(|_| Error::InvalidFormat)?,
446        );
447        let frodo_rng_seed =
448            <[u8; 32]>::try_from(&encaps_seed[64..96]).map_err(|_| Error::InvalidFormat)?;
449        let salt = <[u8; 32]>::try_from(&encaps_seed[96..128]).map_err(|_| Error::InvalidFormat)?;
450
451        let dh_eph_pub = PublicKey::from(&*dh_eph_secret);
452        let dh_eph_pub_bytes = dh_eph_pub.as_bytes();
453
454        // 3. Execute X25519
455        let dh_ss = Zeroizing::new(dh_eph_secret.diffie_hellman(&dh_pub));
456        if !dh_ss.was_contributory() {
457            return Err(Error::LowEntropyKey);
458        }
459
460        // 4. Execute ML-KEM — use native hybrid-array types throughout.
461        let (ml_kem_ct, ml_kem_ss): (MlKemCiphertext<MlKem1024>, MlKemSharedKey) =
462            ml_kem_ek.encapsulate_deterministic(&ml_kem_m);
463        let ml_kem_ss: Zeroizing<[u8; 32]> = Zeroizing::new(ml_kem_ss.into());
464
465        // 5. Execute FrodoKEM
466        let mut frodo_rng = ChaCha20Rng::from_seed(frodo_rng_seed);
467        let (frodo_ct, frodo_ss) = frodo
468            .encapsulate_with_rng(&frodo_pk, &mut frodo_rng)
469            .map_err(|_| Error::EncapsulateError)?;
470        let frodo_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
471            frodo_ss
472                .value()
473                .try_into()
474                .map_err(|_| Error::EncapsulateError)?,
475        );
476        let frodo_ct_bytes = frodo_ct.value();
477
478        // 6. HKDF Derivation — use .as_slice() for zero-copy access to the
479        //    ciphertext bytes from the hybrid-array backed ML-KEM Ciphertext.
480        let ml_kem_ct_slice: &[u8] = ml_kem_ct.as_slice();
481        let ml_kem_ct_arr: [u8; 1568] = ml_kem_ct_slice
482            .try_into()
483            .map_err(|_| Error::EncapsulateError)?;
484        let okm = derive_key(
485            Zeroizing::new(dh_ss.to_bytes()),
486            ml_kem_ss,
487            frodo_ss,
488            &salt,
489            dh_eph_pub_bytes,
490            &ml_kem_ct_arr,
491            frodo_ct_bytes,
492            ek,
493        )?;
494
495        // 7. Assemble Ciphertext
496        let mut ciphertext = Vec::with_capacity(Self::CIPHERTEXT_SIZE);
497        ciphertext.extend_from_slice(dh_eph_pub_bytes);
498        ciphertext.extend_from_slice(&salt);
499        ciphertext.extend_from_slice(&ml_kem_ct_arr);
500        ciphertext.extend_from_slice(frodo_ct_bytes);
501
502        Ok((ciphertext, okm))
503    }
504
505    /// Decapsulates a composite ciphertext to recover the 64-byte Output Keying Material.
506    ///
507    /// This is the **recipient-side** operation.  It parses the composite
508    /// ciphertext, runs all three component decapsulations using the cached
509    /// secret keys, and recomputes the HKDF transcript to produce the OKM.
510    ///
511    /// The OKM is cryptographically bound to the ciphertext and to this specific
512    /// `KWing` instance, so it will not match any other recipient or ciphertext.
513    ///
514    /// # Errors
515    ///
516    /// | Variant | Cause |
517    /// |---------|-------|
518    /// | [`Error::InvalidFormat`] | `ct` is not exactly [`CIPHERTEXT_SIZE`][KWing::CIPHERTEXT_SIZE] bytes |
519    /// | [`Error::LowEntropyKey`] | X25519 DH output is a low-order (all-zero) point |
520    /// | [`Error::DecapsulateError`] | An underlying KEM primitive failed |
521    ///
522    /// # Example
523    ///
524    /// ```rust,no_run
525    /// # #[cfg(feature = "kem")] {
526    /// use b_wing::KWing;
527    ///
528    /// # let secret_seed = [0u8; 128];
529    /// # let encaps_seed = [1u8; 128];
530    /// # let kwing = KWing::from_seed(&secret_seed).unwrap();
531    /// # let ek = kwing.get_pub_key().to_vec();
532    /// # let (ct, _) = KWing::encapsulate(&encaps_seed, &ek).unwrap();
533    /// let okm = kwing.decapsulate(&ct).unwrap();
534    /// assert_eq!(okm.len(), 64);
535    /// // Derive a 32-byte AES-256 key and 32-byte MAC key from the OKM:
536    /// let aes_key = &okm[..32];
537    /// let mac_key = &okm[32..];
538    /// # }
539    /// ```
540    pub fn decapsulate(&self, ct: &[u8]) -> Result<[u8; 64], Error> {
541        if ct.len() != Self::CIPHERTEXT_SIZE {
542            return Err(Error::InvalidFormat);
543        }
544
545        let frodo = Algorithm::FrodoKem1344Shake;
546
547        // 1. Parse Ciphertext
548        let dh_eph_pub =
549            PublicKey::from(<[u8; 32]>::try_from(&ct[0..32]).map_err(|_| Error::InvalidFormat)?);
550        let salt: [u8; 32] = ct[32..64].try_into().map_err(|_| Error::InvalidFormat)?;
551        let ml_kem_ct_slice = &ct[64..1632];
552        let frodo_ct_bytes = &ct[1632..];
553        let frodo_ct = FrodoCiphertext::from_bytes(frodo, frodo_ct_bytes)
554            .map_err(|_| Error::InvalidFormat)?;
555
556        // 2. Execute X25519
557        let dh_ss = Zeroizing::new(self.dh_secret.diffie_hellman(&dh_eph_pub));
558        if !dh_ss.was_contributory() {
559            return Err(Error::LowEntropyKey);
560        }
561
562        // 3. Execute ML-KEM — use decapsulate_slice() for direct byte-slice
563        //    ciphertext handling without manual Array construction.
564        let ml_kem_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
565            self.ml_kem_dk
566                .decapsulate_slice(ml_kem_ct_slice)
567                .map_err(|_| Error::DecapsulateError)?
568                .into(),
569        );
570
571        // 4. Execute FrodoKEM
572        let frodo_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
573            frodo
574                .decapsulate(&self.frodo_sk, &frodo_ct)
575                .map_err(|_| Error::DecapsulateError)?
576                .0
577                .value()
578                .try_into()
579                .map_err(|_| Error::DecapsulateError)?,
580        );
581
582        // 5. HKDF Derivation & Proof
583        let ml_kem_ct_arr: [u8; 1568] = ml_kem_ct_slice
584            .try_into()
585            .map_err(|_| Error::DecapsulateError)?;
586        let okm = derive_key(
587            Zeroizing::new(dh_ss.to_bytes()),
588            ml_kem_ss,
589            frodo_ss,
590            &salt,
591            dh_eph_pub.as_bytes(),
592            &ml_kem_ct_arr,
593            frodo_ct_bytes,
594            self.get_pub_key(),
595        )?;
596
597        Ok(okm)
598    }
599}
600
601#[cfg(test)]
602mod tests {
603    use super::*;
604    use std::sync::LazyLock;
605
606    // ======================================================================
607    // Test Constants & Helpers
608    // ======================================================================
609
610    static SECRET_SEED: [u8; 128] = [0x42; 128];
611    static ENCAPS_SEED: [u8; 128] = [0x84; 128];
612
613    static K_WING: LazyLock<KWing> = LazyLock::new(|| KWing::from_seed(&SECRET_SEED).unwrap());
614
615    static ENCAPS_RESULT: LazyLock<(Vec<u8>, [u8; 64])> =
616        LazyLock::new(|| KWing::encapsulate(&ENCAPS_SEED, K_WING.get_pub_key()).unwrap());
617
618    // ======================================================================
619    // Happy Path & Determinism
620    // ======================================================================
621
622    #[test]
623    fn test_happy_path_round_trip() {
624        // 1. Get cached Public Key
625        let pk = K_WING.get_pub_key();
626        assert_eq!(pk.len(), KWing::ENCAPSULATION_KEY_SIZE);
627
628        // 2. Get cached Encapsulation
629        let (ct, okm_encapsulated) = &*ENCAPS_RESULT;
630        assert_eq!(ct.len(), KWing::CIPHERTEXT_SIZE);
631
632        // 3. Decapsulate
633        let okm_decapsulated = K_WING
634            .decapsulate(ct)
635            .expect("Decapsulation should succeed");
636
637        // 4. Assert Output Keying Material Matches
638        assert_eq!(
639            okm_encapsulated, &okm_decapsulated,
640            "Decapsulated OKM must exactly match the Encapsulated OKM"
641        );
642    }
643
644    #[test]
645    fn test_strict_determinism() {
646        // Deterministic Key Generation
647        let binding = KWing::from_seed(&SECRET_SEED).unwrap();
648        let pk2 = binding.get_pub_key();
649        assert_eq!(
650            K_WING.get_pub_key(),
651            pk2,
652            "Public keys must be identical for the same seed"
653        );
654
655        // Deterministic Encapsulation
656        let (ct2, okm2) = KWing::encapsulate(&ENCAPS_SEED, pk2).unwrap();
657        assert_eq!(
658            ENCAPS_RESULT.0, ct2,
659            "Ciphertexts must be identical for the same seeds"
660        );
661        assert_eq!(
662            ENCAPS_RESULT.1, okm2,
663            "OKMs must be identical for the same seeds"
664        );
665    }
666
667    // ======================================================================
668    // Formatting & Boundary Rejections
669    // ======================================================================
670
671    #[test]
672    fn test_invalid_public_key_length() {
673        let bad_pk = vec![0u8; KWing::ENCAPSULATION_KEY_SIZE - 1]; // 1 byte too short
674
675        let result = KWing::encapsulate(&ENCAPS_SEED, &bad_pk);
676        assert_eq!(
677            result,
678            Err(Error::InvalidFormat),
679            "Encapsulate must reject invalid public key lengths immediately"
680        );
681    }
682
683    #[test]
684    fn test_invalid_ciphertext_length() {
685        let bad_ct = vec![0u8; KWing::CIPHERTEXT_SIZE + 5]; // 5 bytes too long
686
687        let result = K_WING.decapsulate(&bad_ct);
688        assert_eq!(
689            result,
690            Err(Error::InvalidFormat),
691            "Decapsulate must reject invalid ciphertext lengths immediately"
692        );
693    }
694
695    // ======================================================================
696    // Cryptographic Tampering & Mathematical Edge Cases
697    // ======================================================================
698
699    #[test]
700    fn test_low_entropy_key_encapsulate_rejection() {
701        let mut pk = K_WING.get_pub_key().to_vec();
702
703        // Force the X25519 public key part to all zeros.
704        pk[0..32].fill(0);
705
706        let result = KWing::encapsulate(&ENCAPS_SEED, &pk);
707        assert_eq!(
708            result,
709            Err(Error::LowEntropyKey),
710            "Encapsulate must reject mathematical weak points (all-zero DH shared secret)"
711        );
712    }
713
714    #[test]
715    fn test_low_entropy_key_decapsulate_rejection() {
716        let mut ct = ENCAPS_RESULT.0.clone();
717
718        // Force the Ephemeral X25519 public key part in the CT to all zeros.
719        ct[0..32].fill(0);
720
721        let result = K_WING.decapsulate(&ct);
722        assert_eq!(
723            result,
724            Err(Error::LowEntropyKey),
725            "Decapsulate must reject mathematical weak points injected by an attacker"
726        );
727    }
728}