Skip to main content

ml_kem_roundtrip/
ml_kem_roundtrip.rs

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