Skip to main content

dryoc/classic/
crypto_kem_mlkem768.rs

1//! # ML-KEM-768 key encapsulation
2//!
3//! Implements libsodium's `crypto_kem_mlkem768_*` functions: ML-KEM-768 from
4//! FIPS 203, a lattice-based key encapsulation mechanism (KEM) believed to
5//! resist attacks by quantum computers.
6//!
7//! A KEM lets a sender create a fresh shared secret for the holder of a
8//! public key. [`crypto_kem_mlkem768_enc`] returns the shared secret and a
9//! ciphertext; the recipient recovers the same secret with
10//! [`crypto_kem_mlkem768_dec`] and its secret key. Feed the shared secret to
11//! a key-derivation function before using it as an encryption key. A KEM
12//! does not authenticate the sender.
13//!
14//! Prefer [`crate::classic::crypto_kem`], which uses X-Wing: ML-KEM-768
15//! combined with X25519, so it stays secure if either one is broken. Use
16//! ML-KEM-768 directly when a protocol requires it.
17//!
18//! Encapsulation rejects a public key that is not a valid FIPS 203
19//! encapsulation key. Decapsulation always succeeds: a ciphertext that was
20//! not created for the key yields an unrelated pseudorandom secret
21//! ("implicit rejection"), so the caller learns nothing from the result.
22//!
23//! ```
24//! use dryoc::classic::crypto_kem_mlkem768::*;
25//!
26//! let (public_key, secret_key) = crypto_kem_mlkem768_keypair();
27//!
28//! let mut ciphertext = [0u8; dryoc::constants::CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES];
29//! let mut sender_secret = SharedSecret::default();
30//! crypto_kem_mlkem768_enc(&mut ciphertext, &mut sender_secret, &public_key)
31//!     .expect("encapsulation failed");
32//!
33//! let mut recipient_secret = SharedSecret::default();
34//! crypto_kem_mlkem768_dec(&mut recipient_secret, &ciphertext, &secret_key);
35//! assert_eq!(sender_secret, recipient_secret);
36//! ```
37
38use zeroize::Zeroizing;
39
40use crate::constants::{
41    CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES, CRYPTO_KEM_MLKEM768_ENCSEEDBYTES,
42    CRYPTO_KEM_MLKEM768_PUBLICKEYBYTES, CRYPTO_KEM_MLKEM768_SECRETKEYBYTES,
43    CRYPTO_KEM_MLKEM768_SEEDBYTES, CRYPTO_KEM_MLKEM768_SHAREDSECRETBYTES,
44};
45use crate::error::Error;
46use crate::mlkem::{self, Arith};
47use crate::rng::copy_randombytes;
48
49/// ML-KEM-768 public (encapsulation) key.
50pub type PublicKey = [u8; CRYPTO_KEM_MLKEM768_PUBLICKEYBYTES];
51/// ML-KEM-768 secret (decapsulation) key, in FIPS 203's expanded form.
52pub type SecretKey = [u8; CRYPTO_KEM_MLKEM768_SECRETKEYBYTES];
53/// ML-KEM-768 ciphertext.
54pub type Ciphertext = [u8; CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES];
55/// Shared secret produced by encapsulation and decapsulation.
56pub type SharedSecret = [u8; CRYPTO_KEM_MLKEM768_SHAREDSECRETBYTES];
57/// Key-generation seed: FIPS 203's `d || z`.
58pub type Seed = [u8; CRYPTO_KEM_MLKEM768_SEEDBYTES];
59/// Encapsulation seed: FIPS 203's message `m`.
60pub type EncSeed = [u8; CRYPTO_KEM_MLKEM768_ENCSEEDBYTES];
61
62/// In-place variant of [`crypto_kem_mlkem768_seed_keypair`].
63pub fn crypto_kem_mlkem768_seed_keypair_inplace(
64    public_key: &mut PublicKey,
65    secret_key: &mut SecretKey,
66    seed: &Seed,
67) {
68    mlkem::keypair(Arith::detect(), public_key, secret_key, seed);
69}
70
71/// Deterministically derives a key pair from `seed`.
72///
73/// Compatible with libsodium's `crypto_kem_mlkem768_seed_keypair`.
74#[must_use]
75pub fn crypto_kem_mlkem768_seed_keypair(seed: &Seed) -> (PublicKey, SecretKey) {
76    let mut public_key = [0u8; CRYPTO_KEM_MLKEM768_PUBLICKEYBYTES];
77    let mut secret_key = [0u8; CRYPTO_KEM_MLKEM768_SECRETKEYBYTES];
78    crypto_kem_mlkem768_seed_keypair_inplace(&mut public_key, &mut secret_key, seed);
79    (public_key, secret_key)
80}
81
82/// In-place variant of [`crypto_kem_mlkem768_keypair`].
83pub fn crypto_kem_mlkem768_keypair_inplace(public_key: &mut PublicKey, secret_key: &mut SecretKey) {
84    let mut seed = Zeroizing::new([0u8; CRYPTO_KEM_MLKEM768_SEEDBYTES]);
85    copy_randombytes(seed.as_mut_slice());
86    crypto_kem_mlkem768_seed_keypair_inplace(public_key, secret_key, &seed);
87}
88
89/// Returns a randomly generated key pair.
90///
91/// Compatible with libsodium's `crypto_kem_mlkem768_keypair`.
92#[must_use]
93pub fn crypto_kem_mlkem768_keypair() -> (PublicKey, SecretKey) {
94    let mut public_key = [0u8; CRYPTO_KEM_MLKEM768_PUBLICKEYBYTES];
95    let mut secret_key = [0u8; CRYPTO_KEM_MLKEM768_SECRETKEYBYTES];
96    crypto_kem_mlkem768_keypair_inplace(&mut public_key, &mut secret_key);
97    (public_key, secret_key)
98}
99
100/// Creates a random shared secret for `public_key`, writing it to
101/// `shared_secret` and its encapsulation to `ciphertext`.
102///
103/// Compatible with libsodium's `crypto_kem_mlkem768_enc`.
104///
105/// # Errors
106///
107/// Returns [`Error::InvalidKey`] if `public_key` is not a valid ML-KEM-768
108/// encapsulation key (a coefficient is not reduced modulo `q`).
109pub fn crypto_kem_mlkem768_enc(
110    ciphertext: &mut Ciphertext,
111    shared_secret: &mut SharedSecret,
112    public_key: &PublicKey,
113) -> Result<(), Error> {
114    let mut seed = Zeroizing::new([0u8; CRYPTO_KEM_MLKEM768_ENCSEEDBYTES]);
115    copy_randombytes(seed.as_mut_slice());
116    crypto_kem_mlkem768_enc_deterministic(ciphertext, shared_secret, public_key, &seed)
117}
118
119/// Deterministic variant of [`crypto_kem_mlkem768_enc`] with the
120/// encapsulation randomness taken from `seed`. For known-answer tests; a
121/// repeated seed repeats the shared secret.
122///
123/// Compatible with libsodium's `crypto_kem_mlkem768_enc_deterministic`.
124///
125/// # Errors
126///
127/// Returns [`Error::InvalidKey`] if `public_key` is not a valid ML-KEM-768
128/// encapsulation key.
129pub fn crypto_kem_mlkem768_enc_deterministic(
130    ciphertext: &mut Ciphertext,
131    shared_secret: &mut SharedSecret,
132    public_key: &PublicKey,
133    seed: &EncSeed,
134) -> Result<(), Error> {
135    mlkem::encapsulate(Arith::detect(), ciphertext, shared_secret, public_key, seed)
136}
137
138/// Recovers the shared secret encapsulated in `ciphertext` with
139/// `secret_key`, writing it to `shared_secret`. A ciphertext not created for
140/// this key yields an unrelated pseudorandom secret instead of an error.
141///
142/// Compatible with libsodium's `crypto_kem_mlkem768_dec`.
143pub fn crypto_kem_mlkem768_dec(
144    shared_secret: &mut SharedSecret,
145    ciphertext: &Ciphertext,
146    secret_key: &SecretKey,
147) {
148    mlkem::decapsulate(Arith::detect(), shared_secret, ciphertext, secret_key);
149}
150
151/// Cross-checks against libsodium 1.0.22's ML-KEM-768.
152#[cfg(all(test, dryoc_native_tests))]
153pub(crate) mod native_tests {
154    use super::*;
155    use crate::native_test_util as sodium;
156    use crate::test_prelude::*;
157    use crate::utils::test_util::XorShift64;
158
159    /// Seeds for key generation or encapsulation: all-zero, all-ones, four
160    /// fixed pseudorandom ones and two fresh random ones.
161    pub(crate) fn seeds<const N: usize>() -> Vec<[u8; N]> {
162        let mut rng = XorShift64::new(0x6d6c_6b65_6d37_3638);
163        let mut seeds = vec![[0u8; N], [0xff; N]];
164        for _ in 0..4 {
165            seeds.push(core::array::from_fn(|_| rng.next_u64() as u8));
166        }
167        for _ in 0..2 {
168            let mut seed = [0u8; N];
169            copy_randombytes(&mut seed);
170            seeds.push(seed);
171        }
172        seeds
173    }
174
175    /// Copies of `key`, whose first 1152 bytes are an ML-KEM-768 encoded
176    /// polynomial vector, with its first or its last 12-bit coefficient set
177    /// to 4095: not reduced modulo q, so FIPS 203's modulus check rejects it.
178    pub(crate) fn with_unreduced_coefficient<const N: usize>(key: &[u8; N]) -> [[u8; N]; 2] {
179        let (mut first, mut last) = (*key, *key);
180        first[0] = 0xff;
181        first[1] |= 0x0f;
182        last[1150] |= 0xf0;
183        last[1151] = 0xff;
184        [first, last]
185    }
186
187    /// Copies of `ciphertext` with one bit flipped at the start, middle and
188    /// end.
189    pub(crate) fn tampered<const N: usize>(ciphertext: &[u8; N]) -> [[u8; N]; 3] {
190        [0, N / 2, N - 1].map(|index| {
191            let mut copy = *ciphertext;
192            copy[index] ^= 0x01;
193            copy
194        })
195    }
196
197    /// For every key seed and encapsulation seed, both libraries derive the
198    /// same key pair, encapsulate to the same ciphertext and shared secret,
199    /// decapsulate it to that secret, and agree on the implicit-rejection
200    /// secret of each tampered ciphertext.
201    #[test]
202    fn test_mlkem768_matches_libsodium() {
203        for seed in seeds::<CRYPTO_KEM_MLKEM768_SEEDBYTES>() {
204            let (public_key, secret_key) = crypto_kem_mlkem768_seed_keypair(&seed);
205            let (so_public_key, so_secret_key) = sodium::crypto_kem_mlkem768_seed_keypair(&seed);
206            assert_eq!(public_key, so_public_key, "seed {seed:02x?}");
207            assert_eq!(secret_key, so_secret_key, "seed {seed:02x?}");
208
209            for enc_seed in seeds::<CRYPTO_KEM_MLKEM768_ENCSEEDBYTES>() {
210                let mut ciphertext = [0u8; CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES];
211                let mut sent = [0u8; CRYPTO_KEM_MLKEM768_SHAREDSECRETBYTES];
212                crypto_kem_mlkem768_enc_deterministic(
213                    &mut ciphertext,
214                    &mut sent,
215                    &public_key,
216                    &enc_seed,
217                )
218                .expect("enc");
219                let (so_ciphertext, so_sent) =
220                    sodium::crypto_kem_mlkem768_enc_deterministic(&public_key, &enc_seed)
221                        .expect("libsodium enc");
222                assert_eq!(ciphertext, so_ciphertext, "enc seed {enc_seed:02x?}");
223                assert_eq!(sent, so_sent, "enc seed {enc_seed:02x?}");
224
225                for ciphertext in core::iter::once(ciphertext).chain(tampered(&ciphertext)) {
226                    let mut received = [0u8; CRYPTO_KEM_MLKEM768_SHAREDSECRETBYTES];
227                    crypto_kem_mlkem768_dec(&mut received, &ciphertext, &secret_key);
228                    let so_received = sodium::crypto_kem_mlkem768_dec(&ciphertext, &secret_key)
229                        .expect("libsodium dec");
230                    assert_eq!(received, so_received, "enc seed {enc_seed:02x?}");
231                    assert_eq!(received == sent, ciphertext == so_ciphertext);
232                }
233            }
234        }
235    }
236
237    /// Both libraries refuse to encapsulate to a key with an unreduced
238    /// coefficient, deterministically or not.
239    #[test]
240    fn test_mlkem768_unreduced_keys_rejected_like_libsodium() {
241        for seed in seeds::<CRYPTO_KEM_MLKEM768_SEEDBYTES>() {
242            let (public_key, _) = crypto_kem_mlkem768_seed_keypair(&seed);
243            for invalid in with_unreduced_coefficient(&public_key) {
244                let mut ciphertext = [0u8; CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES];
245                let mut shared_secret = [0u8; CRYPTO_KEM_MLKEM768_SHAREDSECRETBYTES];
246                assert!(
247                    crypto_kem_mlkem768_enc_deterministic(
248                        &mut ciphertext,
249                        &mut shared_secret,
250                        &invalid,
251                        &[9u8; CRYPTO_KEM_MLKEM768_ENCSEEDBYTES],
252                    )
253                    .is_err()
254                );
255                assert!(
256                    crypto_kem_mlkem768_enc(&mut ciphertext, &mut shared_secret, &invalid).is_err()
257                );
258                assert!(
259                    sodium::crypto_kem_mlkem768_enc_deterministic(
260                        &invalid,
261                        &[9u8; CRYPTO_KEM_MLKEM768_ENCSEEDBYTES]
262                    )
263                    .is_err()
264                );
265                assert!(sodium::crypto_kem_mlkem768_enc(&invalid).is_err());
266            }
267        }
268    }
269
270    /// Randomized encapsulations made by either library decapsulate to the
271    /// same shared secret in the other.
272    #[test]
273    fn test_mlkem768_randomized_interop_with_libsodium() {
274        for _ in 0..8 {
275            let (public_key, secret_key) = crypto_kem_mlkem768_keypair();
276
277            let mut ciphertext = [0u8; CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES];
278            let mut sent = [0u8; CRYPTO_KEM_MLKEM768_SHAREDSECRETBYTES];
279            crypto_kem_mlkem768_enc(&mut ciphertext, &mut sent, &public_key).expect("enc");
280            assert_eq!(
281                sodium::crypto_kem_mlkem768_dec(&ciphertext, &secret_key).expect("libsodium dec"),
282                sent
283            );
284
285            let (so_ciphertext, so_sent) =
286                sodium::crypto_kem_mlkem768_enc(&public_key).expect("libsodium enc");
287            let mut received = [0u8; CRYPTO_KEM_MLKEM768_SHAREDSECRETBYTES];
288            crypto_kem_mlkem768_dec(&mut received, &so_ciphertext, &secret_key);
289            assert_eq!(received, so_sent);
290        }
291    }
292}
293
294#[cfg(all(test, feature = "nightly"))]
295mod benches {
296    extern crate test;
297
298    use super::*;
299
300    #[bench]
301    fn mlkem768_keypair_bench(b: &mut test::Bencher) {
302        let seed = [7u8; CRYPTO_KEM_MLKEM768_SEEDBYTES];
303        b.iter(|| crypto_kem_mlkem768_seed_keypair(test::black_box(&seed)));
304    }
305
306    #[bench]
307    fn mlkem768_enc_bench(b: &mut test::Bencher) {
308        let (public_key, _) =
309            crypto_kem_mlkem768_seed_keypair(&[7u8; CRYPTO_KEM_MLKEM768_SEEDBYTES]);
310        let (mut ciphertext, mut shared_secret) =
311            ([0u8; CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES], [0u8; 32]);
312        b.iter(|| {
313            crypto_kem_mlkem768_enc_deterministic(
314                &mut ciphertext,
315                &mut shared_secret,
316                test::black_box(&public_key),
317                test::black_box(&[9u8; 32]),
318            )
319            .expect("enc")
320        });
321    }
322
323    /// libsodium's `crypto_kem_mlkem768_seed_keypair` on the same seed as
324    /// `mlkem768_keypair_bench`.
325    #[cfg(dryoc_native_tests)]
326    #[bench]
327    fn libsodium_mlkem768_keypair_bench(b: &mut test::Bencher) {
328        crate::native_test_util::init();
329        let seed = [7u8; CRYPTO_KEM_MLKEM768_SEEDBYTES];
330        let mut public_key = [0u8; CRYPTO_KEM_MLKEM768_PUBLICKEYBYTES];
331        let mut secret_key = [0u8; CRYPTO_KEM_MLKEM768_SECRETKEYBYTES];
332        b.iter(|| {
333            // SAFETY: the key buffers and `seed` are arrays of libsodium's
334            // sizes.
335            let rc = unsafe {
336                libsodium_sys::crypto_kem_mlkem768_seed_keypair(
337                    public_key.as_mut_ptr(),
338                    secret_key.as_mut_ptr(),
339                    test::black_box(seed.as_ptr()),
340                )
341            };
342            assert_eq!(rc, 0);
343            test::black_box((&public_key, &secret_key));
344        });
345    }
346
347    /// libsodium's `crypto_kem_mlkem768_enc_deterministic` on the same key
348    /// and seed as `mlkem768_enc_bench`.
349    #[cfg(dryoc_native_tests)]
350    #[bench]
351    fn libsodium_mlkem768_enc_bench(b: &mut test::Bencher) {
352        crate::native_test_util::init();
353        let (public_key, _) =
354            crypto_kem_mlkem768_seed_keypair(&[7u8; CRYPTO_KEM_MLKEM768_SEEDBYTES]);
355        let (mut ciphertext, mut shared_secret) =
356            ([0u8; CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES], [0u8; 32]);
357        let seed = [9u8; 32];
358        b.iter(|| {
359            // SAFETY: the output buffers, `public_key` and `seed` are arrays
360            // of libsodium's sizes.
361            let rc = unsafe {
362                libsodium_sys::crypto_kem_mlkem768_enc_deterministic(
363                    ciphertext.as_mut_ptr(),
364                    shared_secret.as_mut_ptr(),
365                    test::black_box(public_key.as_ptr()),
366                    test::black_box(seed.as_ptr()),
367                )
368            };
369            assert_eq!(rc, 0);
370            test::black_box((&ciphertext, &shared_secret));
371        });
372    }
373
374    #[bench]
375    fn mlkem768_dec_bench(b: &mut test::Bencher) {
376        let (public_key, secret_key) =
377            crypto_kem_mlkem768_seed_keypair(&[7u8; CRYPTO_KEM_MLKEM768_SEEDBYTES]);
378        let mut ciphertext = [0u8; CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES];
379        let mut shared_secret = [0u8; 32];
380        crypto_kem_mlkem768_enc(&mut ciphertext, &mut shared_secret, &public_key).expect("enc");
381        b.iter(|| {
382            crypto_kem_mlkem768_dec(
383                &mut shared_secret,
384                test::black_box(&ciphertext),
385                test::black_box(&secret_key),
386            )
387        });
388    }
389
390    /// libsodium's `crypto_kem_mlkem768_dec` with the same key and
391    /// ciphertext setup as `mlkem768_dec_bench`.
392    #[cfg(dryoc_native_tests)]
393    #[bench]
394    fn libsodium_mlkem768_dec_bench(b: &mut test::Bencher) {
395        crate::native_test_util::init();
396        let (public_key, secret_key) =
397            crypto_kem_mlkem768_seed_keypair(&[7u8; CRYPTO_KEM_MLKEM768_SEEDBYTES]);
398        let mut ciphertext = [0u8; CRYPTO_KEM_MLKEM768_CIPHERTEXTBYTES];
399        let mut shared_secret = [0u8; 32];
400        crypto_kem_mlkem768_enc(&mut ciphertext, &mut shared_secret, &public_key).expect("enc");
401        b.iter(|| {
402            // SAFETY: `shared_secret`, `ciphertext` and `secret_key` are
403            // arrays of libsodium's sizes.
404            let rc = unsafe {
405                libsodium_sys::crypto_kem_mlkem768_dec(
406                    shared_secret.as_mut_ptr(),
407                    test::black_box(ciphertext.as_ptr()),
408                    test::black_box(secret_key.as_ptr()),
409                )
410            };
411            assert_eq!(rc, 0);
412            test::black_box(&shared_secret);
413        });
414    }
415}