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}