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 sha3::Sha3_512;
64use zeroize::Zeroizing;
65
66// RNG for FrodoKEM Determinism
67use rand_chacha::ChaCha20Rng;
68use rand_core::SeedableRng;
69
70// Main imports
71use frodo_kem::{
72    Algorithm, Ciphertext as FrodoCiphertext, DecryptionKey as FrodoDecryptionKey,
73    EncryptionKey as FrodoEncryptionKey,
74};
75use ml_kem::{
76    Ciphertext, MlKem1024, Seed as MlKemSeed,
77    array::Array,
78    kem::{Decapsulate, KeyExport},
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: &'static [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).
188fn derive_key(
189    dh_ss: Zeroizing<[u8; 32]>,
190    ml_kem_ss: Zeroizing<[u8; 32]>,
191    frodo_ss: Zeroizing<[u8; 32]>,
192    salt: &[u8; 32],
193    dh_eph_pub: &[u8; 32],
194    ml_kem_ct: &[u8; 1568],
195    frodo_ct: &[u8],
196    ek: &[u8],
197) -> Result<[u8; 64], Error> {
198    // IKM is now 96 bytes (X25519 + ML-KEM + FrodoKEM)
199    let mut ikm = Zeroizing::new([0u8; 96]);
200    ikm[0..32].copy_from_slice(&*dh_ss);
201    ikm[32..64].copy_from_slice(&*ml_kem_ss);
202    ikm[64..96].copy_from_slice(&*frodo_ss);
203    drop((dh_ss, ml_kem_ss, frodo_ss));
204
205    let hkdf = Hkdf::<Sha3_512>::new(Some(salt), &*ikm);
206    drop(ikm);
207
208    // Build the common transcript prefix once
209    let mut okm_info = Vec::with_capacity(
210        32 + ml_kem_ct.len() + frodo_ct.len() + ek.len() + K_WING_OKM_CONTEXT.len(),
211    );
212    okm_info.extend_from_slice(dh_eph_pub);
213    okm_info.extend_from_slice(ml_kem_ct);
214    okm_info.extend_from_slice(frodo_ct);
215    okm_info.extend_from_slice(ek);
216    okm_info.extend_from_slice(K_WING_OKM_CONTEXT);
217
218    let mut okm = [0u8; 64];
219    hkdf.expand(&okm_info, &mut okm)
220        .map_err(|_| Error::InvalidFormat)?;
221
222    Ok(okm)
223}
224
225// ======================================================================
226// Expanded KWing Key (Stateful / High-Throughput)
227// ======================================================================
228
229/// A stateful, high-throughput KWing key holder for decapsulation.
230///
231/// `KWing` pre-computes and heap-caches all three component secret keys and
232/// the composite encapsulation key at construction time.  Subsequent calls to
233/// [`decapsulate`][KWing::decapsulate] reuse the cached material without any
234/// additional key-derivation overhead.
235///
236/// # Key Sizes
237///
238/// | Constant | Bytes | Layout |
239/// |----------|-------|--------|
240/// | [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE] | 23,120 | `X25519(32) ‖ ML-KEM-1024(1568) ‖ FrodoKEM-1344(21520)` |
241/// | [`CIPHERTEXT_SIZE`][KWing::CIPHERTEXT_SIZE] | 23,328 | `X25519 eph(32) ‖ Salt(32) ‖ ML-KEM CT(1568) ‖ FrodoKEM CT(21696)` |
242///
243/// # Security Note
244///
245/// The encapsulation key (public key) returned by [`get_pub_key`][KWing::get_pub_key]
246/// is safe to distribute freely.  The underlying secret key material stored in
247/// The underlying X25519 and ML-KEM secret material stored in this struct is
248/// wrapped in [`zeroize::Zeroizing`]. FrodoKEM material depends on standard
249/// process memory cleanup.
250///
251/// # Example
252///
253/// ```rust,no_run
254/// # #[cfg(feature = "kem")] {
255/// use b_wing::KWing;
256///
257/// let secret_seed = [0u8; 128]; // use a real CSPRNG in production
258/// let kwing = KWing::from_seed(&secret_seed).unwrap();
259///
260/// // The public encapsulation key can be shared with any sender.
261/// let ek = kwing.get_pub_key();
262/// assert_eq!(ek.len(), KWing::ENCAPSULATION_KEY_SIZE);
263/// # }
264/// ```
265pub struct KWing {
266    dh_secret: Zeroizing<StaticSecret>,
267    ml_kem_dk: ml_kem::DecapsulationKey<MlKem1024>,
268    frodo_sk: FrodoDecryptionKey,
269    composite_pk: Vec<u8>,
270}
271
272impl KWing {
273    /// Byte length of the composite encapsulation (public) key: **23,120 bytes**.
274    ///
275    /// Memory layout:
276    /// ```text
277    /// [ X25519 pub (32) | ML-KEM-1024 ek (1568) | FrodoKEM-1344 pk (21520) ]
278    /// ```
279    pub const ENCAPSULATION_KEY_SIZE: usize = 23120;
280
281    /// Byte length of the composite ciphertext: **23,328 bytes**.
282    ///
283    /// Memory layout:
284    /// ```text
285    /// [ X25519 eph pub (32) | Salt (32) | ML-KEM-1024 ct (1568) | FrodoKEM-1344 ct (21696) ]
286    /// ```
287    pub const CIPHERTEXT_SIZE: usize = 23328;
288
289    /// Expands a 128-byte secret seed into a fully initialized `KWing` key holder.
290    ///
291    /// The seed is partitioned deterministically as follows:
292    ///
293    /// | Bytes | Usage |
294    /// |-------|-------|
295    /// | `[0..32]`   | X25519 static secret |
296    /// | `[32..64]`  | ML-KEM-1024 keygen parameter `d` |
297    /// | `[64..96]`  | ML-KEM-1024 keygen parameter `z` |
298    /// | `[96..128]` | FrodoKEM-1344 keygen seed (fed into ChaCha20) |
299    ///
300    /// # Security Requirements
301    ///
302    /// * `secret_seed` **must** be generated by a cryptographically secure
303    ///   random number generator (CSPRNG) such as `getrandom`.
304    /// * Never reuse the same seed for different recipients or sessions.
305    /// * The seed should be treated with the same care as a private key.
306    ///
307    /// # Errors
308    ///
309    /// Returns [`Error::InvalidFormat`] if an internal slice conversion fails
310    /// (this should be impossible given a correctly-sized input).
311    ///
312    /// # Example
313    ///
314    /// ```rust,no_run
315    /// # #[cfg(feature = "kem")] {
316    /// use b_wing::KWing;
317    ///
318    /// let mut seed = [0u8; 128];
319    /// getrandom::fill(&mut seed).expect("CSPRNG failed");
320    /// let kwing = KWing::from_seed(&seed).expect("key generation failed");
321    /// # }
322    /// ```
323    pub fn from_seed(secret_seed: &[u8; 128]) -> Result<Self, Error> {
324        // 1. X25519
325        let dh_secret = Zeroizing::new(StaticSecret::from(
326            <[u8; 32]>::try_from(&secret_seed[0..32]).map_err(|_| Error::InvalidFormat)?,
327        ));
328        let dh_pub = PublicKey::from(&*dh_secret);
329
330        // 2. ML-KEM-1024
331        let ml_kem_d = Zeroizing::new(
332            <[u8; 32]>::try_from(&secret_seed[32..64]).map_err(|_| Error::InvalidFormat)?,
333        );
334        let ml_kem_z = Zeroizing::new(
335            <[u8; 32]>::try_from(&secret_seed[64..96]).map_err(|_| Error::InvalidFormat)?,
336        );
337        // ML-KEM-1024 Seed = d || z (64 bytes)
338        let mut ml_kem_seed = MlKemSeed::default();
339        ml_kem_seed[..32].copy_from_slice(&*ml_kem_d);
340        ml_kem_seed[32..].copy_from_slice(&*ml_kem_z);
341        let ml_kem_dk = ml_kem::DecapsulationKey::<MlKem1024>::from_seed(ml_kem_seed);
342        let ml_ek = ml_kem_dk.encapsulation_key();
343
344        // 3. FrodoKEM-1344-SHAKE
345        let frodo_seed =
346            <[u8; 32]>::try_from(&secret_seed[96..128]).map_err(|_| Error::InvalidFormat)?;
347        let mut frodo_rng = ChaCha20Rng::from_seed(frodo_seed);
348        let frodo = Algorithm::FrodoKem1344Shake;
349        let (frodo_pk, frodo_sk) = frodo.generate_keypair(&mut frodo_rng);
350
351        // 4. Cache composite PK on Heap
352        let mut composite_pk = Vec::with_capacity(Self::ENCAPSULATION_KEY_SIZE);
353        composite_pk.extend_from_slice(dh_pub.as_bytes());
354        composite_pk.extend_from_slice(&ml_ek.to_bytes());
355        composite_pk.extend_from_slice(frodo_pk.value());
356
357        Ok(Self {
358            dh_secret,
359            ml_kem_dk,
360            frodo_sk,
361            composite_pk,
362        })
363    }
364
365    /// Returns a reference to the cached composite encapsulation key.
366    ///
367    /// The returned slice is [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE]
368    /// bytes long and is safe to distribute publicly.  Pass it to
369    /// [`encapsulate`][KWing::encapsulate] on the sender's side.
370    #[must_use]
371    pub fn get_pub_key(&self) -> &[u8] {
372        &self.composite_pk
373    }
374
375    /// Encapsulates a fresh shared secret against the recipient's composite public key.
376    ///
377    /// This is the **sender-side** operation.  It runs all three component KEMs
378    /// deterministically from `encaps_seed` and combines their outputs into a
379    /// single composite ciphertext and a 64-byte OKM.
380    ///
381    /// # Seed Layout
382    ///
383    /// | Bytes | Usage |
384    /// |-------|-------|
385    /// | `[0..32]`   | X25519 ephemeral secret |
386    /// | `[32..64]`  | ML-KEM-1024 randomness `m` |
387    /// | `[64..96]`  | FrodoKEM encapsulation randomness (ChaCha20 seed) |
388    /// | `[96..128]` | HKDF salt (transmitted in the ciphertext) |
389    ///
390    /// # Security Requirements
391    ///
392    /// * `encaps_seed` **must** be freshly generated from a CSPRNG for **every**
393    ///   encapsulation.  Reusing the seed against the same recipient leaks the
394    ///   X25519 secret key.
395    ///
396    /// # Errors
397    ///
398    /// | Variant | Cause |
399    /// |---------|-------|
400    /// | [`Error::InvalidFormat`] | `ek` is not exactly [`ENCAPSULATION_KEY_SIZE`][KWing::ENCAPSULATION_KEY_SIZE] bytes |
401    /// | [`Error::LowEntropyKey`] | X25519 DH output is a low-order (all-zero) point |
402    /// | [`Error::EncapsulateError`] | An underlying KEM primitive failed |
403    ///
404    /// # Example
405    ///
406    /// ```rust,no_run
407    /// # #[cfg(feature = "kem")] {
408    /// use b_wing::KWing;
409    ///
410    /// # let secret_seed = [0u8; 128];
411    /// # let kwing = KWing::from_seed(&secret_seed).unwrap();
412    /// # let ek = kwing.get_pub_key();
413    /// let mut encaps_seed = [0u8; 128];
414    /// getrandom::fill(&mut encaps_seed).expect("CSPRNG failed");
415    ///
416    /// let (ciphertext, shared_secret) = KWing::encapsulate(&encaps_seed, ek).unwrap();
417    /// assert_eq!(ciphertext.len(), KWing::CIPHERTEXT_SIZE);
418    /// assert_eq!(shared_secret.len(), 64);
419    /// # }
420    /// ```
421    pub fn encapsulate(encaps_seed: &[u8; 128], ek: &[u8]) -> Result<(Vec<u8>, [u8; 64]), Error> {
422        if ek.len() != Self::ENCAPSULATION_KEY_SIZE {
423            return Err(Error::InvalidFormat);
424        }
425
426        let frodo = Algorithm::FrodoKem1344Shake;
427
428        // 1. Parse Composite Key
429        let dh_pub =
430            PublicKey::from(<[u8; 32]>::try_from(&ek[0..32]).map_err(|_| Error::InvalidFormat)?);
431        let ml_kem_ek_bytes: ml_kem::kem::Key<ml_kem::EncapsulationKey<MlKem1024>> = Array(
432            ek[32..1600].try_into().map_err(|_| Error::InvalidFormat)?,
433        );
434        let ml_kem_ek = ml_kem::EncapsulationKey::<MlKem1024>::new(&ml_kem_ek_bytes)
435            .map_err(|_| Error::InvalidFormat)?;
436        let frodo_pk =
437            FrodoEncryptionKey::from_bytes(frodo, &ek[1600..]).map_err(|_| Error::InvalidFormat)?;
438
439        // 2. Setup Deterministic RNGs
440        let dh_eph_secret = Zeroizing::new(StaticSecret::from(
441            <[u8; 32]>::try_from(&encaps_seed[0..32]).map_err(|_| Error::InvalidFormat)?,
442        ));
443        let ml_kem_m = Zeroizing::new(
444            <[u8; 32]>::try_from(&encaps_seed[32..64]).map_err(|_| Error::InvalidFormat)?,
445        );
446        let frodo_rng_seed =
447            <[u8; 32]>::try_from(&encaps_seed[64..96]).map_err(|_| Error::InvalidFormat)?;
448        let salt = <[u8; 32]>::try_from(&encaps_seed[96..128]).map_err(|_| Error::InvalidFormat)?;
449
450        let dh_eph_pub = PublicKey::from(&*dh_eph_secret);
451        let dh_eph_pub_bytes = dh_eph_pub.as_bytes();
452
453        // 3. Execute X25519
454        let dh_ss = Zeroizing::new(dh_eph_secret.diffie_hellman(&dh_pub));
455        if !dh_ss.was_contributory() {
456            return Err(Error::LowEntropyKey);
457        }
458
459        // 4. Execute ML-KEM
460        let (ml_kem_ct, ml_kem_ss) = ml_kem_ek
461            .encapsulate_deterministic(&Array(*ml_kem_m));
462        let ml_kem_ss: Zeroizing<[u8; 32]> = Zeroizing::new(ml_kem_ss.into());
463
464        // 5. Execute FrodoKEM
465        let mut frodo_rng = ChaCha20Rng::from_seed(frodo_rng_seed);
466        let (frodo_ct, frodo_ss) = frodo
467            .encapsulate_with_rng(&frodo_pk, &mut frodo_rng)
468            .map_err(|_| Error::EncapsulateError)?;
469        let frodo_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
470            frodo_ss
471                .value()
472                .try_into()
473                .map_err(|_| Error::EncapsulateError)?,
474        );
475        let frodo_ct_arr = frodo_ct.value().to_vec();
476
477        // 6. HKDF Derivation
478        let ml_kem_ct_bytes: [u8; 1568] = ml_kem_ct.0
479            .try_into()
480            .map_err(|_| Error::EncapsulateError)?;
481        let okm = derive_key(
482            Zeroizing::new(dh_ss.to_bytes()),
483            ml_kem_ss,
484            frodo_ss,
485            &salt,
486            dh_eph_pub_bytes,
487            &ml_kem_ct_bytes,
488            &frodo_ct_arr,
489            ek,
490        )?;
491
492        // 7. Assemble Ciphertext
493        let mut ciphertext = Vec::with_capacity(Self::CIPHERTEXT_SIZE);
494        ciphertext.extend_from_slice(dh_eph_pub_bytes);
495        ciphertext.extend_from_slice(&salt);
496        ciphertext.extend_from_slice(&ml_kem_ct_bytes);
497        ciphertext.extend_from_slice(&frodo_ct_arr);
498
499        Ok((ciphertext, okm))
500    }
501
502    /// Decapsulates a composite ciphertext to recover the 64-byte Output Keying Material.
503    ///
504    /// This is the **recipient-side** operation.  It parses the composite
505    /// ciphertext, runs all three component decapsulations using the cached
506    /// secret keys, and recomputes the HKDF transcript to produce the OKM.
507    ///
508    /// The OKM is cryptographically bound to the ciphertext and to this specific
509    /// `KWing` instance, so it will not match any other recipient or ciphertext.
510    ///
511    /// # Errors
512    ///
513    /// | Variant | Cause |
514    /// |---------|-------|
515    /// | [`Error::InvalidFormat`] | `ct` is not exactly [`CIPHERTEXT_SIZE`][KWing::CIPHERTEXT_SIZE] bytes |
516    /// | [`Error::LowEntropyKey`] | X25519 DH output is a low-order (all-zero) point |
517    /// | [`Error::DecapsulateError`] | An underlying KEM primitive failed |
518    ///
519    /// # Example
520    ///
521    /// ```rust,no_run
522    /// # #[cfg(feature = "kem")] {
523    /// use b_wing::KWing;
524    ///
525    /// # let secret_seed = [0u8; 128];
526    /// # let encaps_seed = [1u8; 128];
527    /// # let kwing = KWing::from_seed(&secret_seed).unwrap();
528    /// # let ek = kwing.get_pub_key().to_vec();
529    /// # let (ct, _) = KWing::encapsulate(&encaps_seed, &ek).unwrap();
530    /// let okm = kwing.decapsulate(&ct).unwrap();
531    /// assert_eq!(okm.len(), 64);
532    /// // Derive a 32-byte AES-256 key and 32-byte MAC key from the OKM:
533    /// let aes_key = &okm[..32];
534    /// let mac_key = &okm[32..];
535    /// # }
536    /// ```
537    pub fn decapsulate(&self, ct: &[u8]) -> Result<[u8; 64], Error> {
538        if ct.len() != Self::CIPHERTEXT_SIZE {
539            return Err(Error::InvalidFormat);
540        }
541
542        let frodo = Algorithm::FrodoKem1344Shake;
543
544        // 1. Parse Ciphertext
545        let dh_eph_pub =
546            PublicKey::from(<[u8; 32]>::try_from(&ct[0..32]).map_err(|_| Error::InvalidFormat)?);
547        let salt: [u8; 32] = ct[32..64].try_into().map_err(|_| Error::InvalidFormat)?;
548        let ml_kem_ct: Ciphertext<MlKem1024> = Array(ct[64..1632].try_into().map_err(|_| Error::InvalidFormat)?);
549        let frodo_ct_bytes = &ct[1632..];
550        let frodo_ct = FrodoCiphertext::from_bytes(frodo, &frodo_ct_bytes)
551            .map_err(|_| Error::InvalidFormat)?;
552        // 2. Execute X25519
553        let dh_ss = Zeroizing::new(self.dh_secret.diffie_hellman(&dh_eph_pub));
554        if !dh_ss.was_contributory() {
555            return Err(Error::LowEntropyKey);
556        }
557
558        // 3. Execute ML-KEM
559        let ml_kem_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
560            self.ml_kem_dk
561                .decapsulate(&ml_kem_ct)
562                .into(),
563        );
564
565        // 4. Execute FrodoKEM
566        let frodo_ss: Zeroizing<[u8; 32]> = Zeroizing::new(
567            frodo
568                .decapsulate(&self.frodo_sk, &frodo_ct)
569                .map_err(|_| Error::DecapsulateError)?
570                .0
571                .value()
572                .try_into()
573                .map_err(|_| Error::DecapsulateError)?,
574        );
575
576        // 5. HKDF Derivation & Proof
577        let ml_kem_ct_bytes: [u8; 1568] = ml_kem_ct.0
578            .try_into()
579            .map_err(|_| Error::DecapsulateError)?;
580        let okm = derive_key(
581            Zeroizing::new(dh_ss.to_bytes()),
582            ml_kem_ss,
583            frodo_ss,
584            &salt,
585            dh_eph_pub.as_bytes(),
586            &ml_kem_ct_bytes,
587            frodo_ct_bytes,
588            &self.get_pub_key(),
589        )?;
590
591        Ok(okm)
592    }
593}
594
595#[cfg(test)]
596mod tests {
597    use super::*;
598    use std::sync::LazyLock;
599
600    // ======================================================================
601    // Test Constants & Helpers
602    // ======================================================================
603
604    static SECRET_SEED: [u8; 128] = [0x42; 128];
605    static ENCAPS_SEED: [u8; 128] = [0x84; 128];
606
607    static K_WING: LazyLock<KWing> = LazyLock::new(|| KWing::from_seed(&SECRET_SEED).unwrap());
608
609    static ENCAPS_RESULT: LazyLock<(Vec<u8>, [u8; 64])> =
610        LazyLock::new(|| KWing::encapsulate(&ENCAPS_SEED, K_WING.get_pub_key()).unwrap());
611
612    // ======================================================================
613    // Happy Path & Determinism
614    // ======================================================================
615
616    #[test]
617    fn test_happy_path_round_trip() {
618        // 1. Get cached Public Key
619        let pk = K_WING.get_pub_key();
620        assert_eq!(pk.len(), KWing::ENCAPSULATION_KEY_SIZE);
621
622        // 2. Get cached Encapsulation
623        let (ct, okm_encapsulated) = &*ENCAPS_RESULT;
624        assert_eq!(ct.len(), KWing::CIPHERTEXT_SIZE);
625
626        // 3. Decapsulate
627        let okm_decapsulated = K_WING
628            .decapsulate(ct)
629            .expect("Decapsulation should succeed");
630
631        // 4. Assert Output Keying Material Matches
632        assert_eq!(
633            okm_encapsulated, &okm_decapsulated,
634            "Decapsulated OKM must exactly match the Encapsulated OKM"
635        );
636    }
637
638    #[test]
639    fn test_strict_determinism() {
640        // Deterministic Key Generation
641        let binding = KWing::from_seed(&SECRET_SEED).unwrap();
642        let pk2 = binding.get_pub_key();
643        assert_eq!(
644            K_WING.get_pub_key(),
645            pk2,
646            "Public keys must be identical for the same seed"
647        );
648
649        // Deterministic Encapsulation
650        let (ct2, okm2) = KWing::encapsulate(&ENCAPS_SEED, pk2).unwrap();
651        assert_eq!(
652            ENCAPS_RESULT.0, ct2,
653            "Ciphertexts must be identical for the same seeds"
654        );
655        assert_eq!(
656            ENCAPS_RESULT.1, okm2,
657            "OKMs must be identical for the same seeds"
658        );
659    }
660
661    // ======================================================================
662    // Formatting & Boundary Rejections
663    // ======================================================================
664
665    #[test]
666    fn test_invalid_public_key_length() {
667        let bad_pk = vec![0u8; KWing::ENCAPSULATION_KEY_SIZE - 1]; // 1 byte too short
668
669        let result = KWing::encapsulate(&ENCAPS_SEED, &bad_pk);
670        assert_eq!(
671            result,
672            Err(Error::InvalidFormat),
673            "Encapsulate must reject invalid public key lengths immediately"
674        );
675    }
676
677    #[test]
678    fn test_invalid_ciphertext_length() {
679        let bad_ct = vec![0u8; KWing::CIPHERTEXT_SIZE + 5]; // 5 bytes too long
680
681        let result = K_WING.decapsulate(&bad_ct);
682        assert_eq!(
683            result,
684            Err(Error::InvalidFormat),
685            "Decapsulate must reject invalid ciphertext lengths immediately"
686        );
687    }
688
689    // ======================================================================
690    // Cryptographic Tampering & Mathematical Edge Cases
691    // ======================================================================
692
693    #[test]
694    fn test_low_entropy_key_encapsulate_rejection() {
695        let mut pk = K_WING.get_pub_key().to_vec();
696
697        // Force the X25519 public key part to all zeros.
698        pk[0..32].fill(0);
699
700        let result = KWing::encapsulate(&ENCAPS_SEED, &pk);
701        assert_eq!(
702            result,
703            Err(Error::LowEntropyKey),
704            "Encapsulate must reject mathematical weak points (all-zero DH shared secret)"
705        );
706    }
707
708    #[test]
709    fn test_low_entropy_key_decapsulate_rejection() {
710        let mut ct = ENCAPS_RESULT.0.clone();
711
712        // Force the Ephemeral X25519 public key part in the CT to all zeros.
713        ct[0..32].fill(0);
714
715        let result = K_WING.decapsulate(&ct);
716        assert_eq!(
717            result,
718            Err(Error::LowEntropyKey),
719            "Decapsulate must reject mathematical weak points injected by an attacker"
720        );
721    }
722}