Skip to main content

ml_kem_roundtrip/
ml_kem_roundtrip.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2026 Cédric Mesnil <cslashm@pm.me>
3
4// ML-KEM Round-Trip Example
5//
6// Demonstrates the full key encapsulation flow for all three ML-KEM parameter
7// sets defined in FIPS 203:
8//
9//   - ML-KEM-512  (NIST Security Category 1, ~128-bit classical)
10//   - ML-KEM-768  (NIST Security Category 3, ~192-bit classical)
11//   - ML-KEM-1024 (NIST Security Category 5, ~256-bit classical)
12//
13// Each round-trip consists of:
14//   1. Key generation  -- produces an encapsulation key (public) and a
15//                         decapsulation key (private).
16//   2. Encapsulation   -- the sender uses the encapsulation key to create a
17//                         shared secret and a ciphertext.
18//   3. Decapsulation   -- the receiver uses the decapsulation key and the
19//                         ciphertext to recover the same shared secret.
20//
21// Run:
22//   cargo run --example roundtrip -p ml_kem
23
24use quantica::ml_kem::*;
25use std::time::Instant;
26
27/// Pretty-print a byte slice as a hex string.
28fn hex(bytes: &[u8]) -> String {
29    bytes.iter().map(|b| format!("{:02x}", b)).collect()
30}
31
32/// Run a full ML-KEM round-trip for a given parameter set and print results.
33fn roundtrip<P: Params>(name: &str, rng: &mut impl CryptoRng) {
34    println!("--- {} ---", name);
35
36    // Step 1: Key generation
37    // The encapsulation key (ek) is public; the decapsulation key (dk) is secret.
38    let t0 = Instant::now();
39    let (ek, dk) = MlKem::<P>::keygen(rng).expect("keygen failed");
40    let keygen_us = t0.elapsed().as_micros();
41    println!(
42        "  keygen:  ek = {} bytes, dk = {} bytes  ({} us)",
43        ek.len(),
44        dk.len(),
45        keygen_us
46    );
47
48    // Step 2: Encapsulation (performed by the sender)
49    // Produces a 32-byte shared secret (ss_sender) and a ciphertext (ct).
50    let t0 = Instant::now();
51    let (ss_sender, ct) = MlKem::<P>::encaps(&ek, rng).expect("encaps failed");
52    let encaps_us = t0.elapsed().as_micros();
53    println!(
54        "  encaps:  ct = {} bytes, ss = {} bytes  ({} us)",
55        ct.len(),
56        ss_sender.len(),
57        encaps_us
58    );
59
60    // Step 3: Decapsulation (performed by the receiver)
61    // Recovers the shared secret from the ciphertext using the private key.
62    // The `rng` parameter supports side-channel countermeasures (NTT shuffling).
63    let t0 = Instant::now();
64    let ss_receiver = MlKem::<P>::decaps(&dk, &ct, rng).expect("decaps failed");
65    let decaps_us = t0.elapsed().as_micros();
66    println!("  decaps:  ({} us)", decaps_us);
67
68    // Verify that both parties derived the same shared secret.
69    if ss_sender == ss_receiver {
70        println!("  result:  MATCH");
71    } else {
72        println!("  result:  MISMATCH -- this should never happen!");
73    }
74
75    // Print the shared secret in hexadecimal.
76    println!("  shared secret: {}", hex(&ss_sender));
77    println!();
78}
79
80fn main() {
81    println!("=== ML-KEM (FIPS 203) Round-Trip Example ===\n");
82
83    let mut rng = OsRng;
84
85    // Run all three parameter sets.
86    roundtrip::<MlKem512>("ML-KEM-512  (Category 1)", &mut rng);
87    roundtrip::<MlKem768>("ML-KEM-768  (Category 3)", &mut rng);
88    roundtrip::<MlKem1024>("ML-KEM-1024 (Category 5)", &mut rng);
89
90    println!("All round-trips completed successfully.");
91}