securedrop_protocol_minimal/primitives/
mlkem.rs

1use crate::primitives::provider::kem::{PrivateKey, PublicKey};
2use crate::primitives::provider::mlkem::{KEY_GENERATION_SEED_SIZE, mlkem768};
3use rand_core::{CryptoRng, RngCore};
4
5// From NIST ML-KEM spec:
6// https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.203.pdf
7// See Table 2 which shows k = 3 for ML-KEM-768 and
8// Algorithm 19 which defines the size of the encap and decap keys in terms of k
9pub const MLKEM768_PUBLIC_KEY_LEN: usize = 1184;
10pub const MLKEM768_PRIVATE_KEY_LEN: usize = 2400;
11
12pub(crate) const LEN_MLKEM_SHAREDSECRET_ENCAPS: usize = 1088;
13
14/// MLKEM-768 public key.
15#[derive(Debug, Clone)]
16pub(crate) struct MLKEM768PublicKey([u8; MLKEM768_PUBLIC_KEY_LEN]);
17
18/// MLKEM-768 private key.
19#[derive(Debug, Clone)]
20pub(crate) struct MLKEM768PrivateKey([u8; MLKEM768_PRIVATE_KEY_LEN]);
21
22impl MLKEM768PublicKey {
23    pub(crate) fn as_bytes(&self) -> &[u8; MLKEM768_PUBLIC_KEY_LEN] {
24        &self.0
25    }
26
27    pub(crate) fn from_bytes(bytes: [u8; MLKEM768_PUBLIC_KEY_LEN]) -> Self {
28        Self(bytes)
29    }
30}
31
32impl MLKEM768PrivateKey {
33    pub(crate) fn as_bytes(&self) -> &[u8; MLKEM768_PRIVATE_KEY_LEN] {
34        &self.0
35    }
36
37    pub(crate) fn from_bytes(bytes: [u8; MLKEM768_PRIVATE_KEY_LEN]) -> Self {
38        Self(bytes)
39    }
40}
41
42// TODO: this uses generate_key_pair from mlkem768 library, the other uses keygen derand
43pub(crate) fn from_bytes(
44    seed: [u8; KEY_GENERATION_SEED_SIZE],
45) -> Result<(MLKEM768PrivateKey, MLKEM768PublicKey), anyhow::Error> {
46    let (sk, pk) = mlkem768::generate_key_pair(seed).into_parts();
47    let mlkem_encaps = MLKEM768PublicKey::from_bytes(*pk.as_slice());
48    let mlkem_decaps = MLKEM768PrivateKey::from_bytes(*sk.as_slice());
49
50    Ok((mlkem_decaps, mlkem_encaps))
51}
52
53/// Generate MLKEM-768 keypair from external randomness
54/// FOR TEST PURPOSES ONLY
55pub(crate) fn deterministic_keygen(
56    randomness: [u8; KEY_GENERATION_SEED_SIZE],
57) -> Result<(MLKEM768PrivateKey, MLKEM768PublicKey), anyhow::Error> {
58    use crate::primitives::provider::kem::{Algorithm, key_gen_derand};
59
60    // Generate MLKEM-768 keypair using libcrux_kem with deterministic randomness
61    let (sk, pk) = key_gen_derand(Algorithm::MlKem768, &randomness)
62        .map_err(|e| anyhow::anyhow!("MLKEM-768 deterministic key generation failed: {:?}", e))?;
63
64    typed(sk, pk)
65}
66
67#[cfg_attr(hax, hax_lib::fstar::verification_status(lax))]
68/// Helper, convert libcrux type to our key types
69fn typed(
70    sk: PrivateKey,
71    pk: PublicKey,
72) -> Result<(MLKEM768PrivateKey, MLKEM768PublicKey), anyhow::Error> {
73    // Convert to our types
74    let private_key_bytes = sk.encode();
75    let public_key_bytes = pk.encode();
76
77    // Validate key sizes (MLKEM-768 should have consistent sizes)
78    if private_key_bytes.len() != MLKEM768_PRIVATE_KEY_LEN
79        || public_key_bytes.len() != MLKEM768_PUBLIC_KEY_LEN
80    {
81        // 1. This branch is skipped if the conditions are met dynamically.
82        //
83        // 2. Implication: this branch is *unreachable* based on the
84        //    postconditions of `encode()` assumed (in the stub) or subsequently
85        //    proven (by the crate itself).
86        //
87        // 3. Implication: this conditional and the error-handling below are
88        //    unnecessary if the postconditions of `encode()` are proven (by the
89        //    crate itself), and this function can return `(private_key,
90        //    public_key)` directly, without wrapping it in a `Result`.
91        return Err(anyhow::anyhow!(
92            "Unexpected MLKEM-768 key sizes: private={}, public={}",
93            private_key_bytes.len(),
94            public_key_bytes.len()
95        ));
96    }
97
98    let private_key = MLKEM768PrivateKey::from_bytes(
99        private_key_bytes
100            .as_slice()
101            .try_into()
102            .map_err(|_| anyhow::anyhow!("Failed to convert private key bytes"))?,
103    );
104    let public_key = MLKEM768PublicKey::from_bytes(
105        public_key_bytes
106            .as_slice()
107            .try_into()
108            .map_err(|_| anyhow::anyhow!("Failed to convert public key bytes"))?,
109    );
110
111    Ok((private_key, public_key))
112}
113
114/// Generate a new MLKEM-768 keypair using libcrux_kem
115pub(crate) fn generate_mlkem768_keypair<R: RngCore + CryptoRng>(
116    rng: &mut R,
117) -> Result<(MLKEM768PrivateKey, MLKEM768PublicKey), anyhow::Error> {
118    use crate::primitives::provider::kem::{Algorithm, key_gen};
119
120    // Generate MLKEM-768 keypair using libcrux_kem
121    let (sk, pk) = key_gen(Algorithm::MlKem768, rng)
122        .map_err(|e| anyhow::anyhow!("MLKEM-768 key generation failed: {:?}", e))?;
123
124    typed(sk, pk)
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use rand_chacha::ChaCha20Rng;
131    use rand_core::SeedableRng;
132
133    #[test]
134    fn test_mlkem768_key_generation() {
135        let mut rng = ChaCha20Rng::seed_from_u64(42);
136
137        let (private_key, public_key) =
138            generate_mlkem768_keypair(&mut rng).expect("Should generate MLKEM-768 keypair");
139
140        // Verify key sizes
141        assert_eq!(private_key.as_bytes().len(), MLKEM768_PRIVATE_KEY_LEN);
142        assert_eq!(public_key.as_bytes().len(), MLKEM768_PUBLIC_KEY_LEN);
143
144        // Verify keys are different (they have different sizes anyway)
145        assert_ne!(private_key.as_bytes().len(), public_key.as_bytes().len());
146    }
147
148    #[test]
149    fn test_mlkem768_key_serialization() {
150        let mut rng = ChaCha20Rng::seed_from_u64(42);
151
152        let (private_key, public_key) =
153            generate_mlkem768_keypair(&mut rng).expect("Should generate MLKEM-768 keypair");
154
155        // Test round-trip serialization
156        let private_bytes = *private_key.as_bytes();
157        let public_bytes = *public_key.as_bytes();
158
159        let reconstructed_private = MLKEM768PrivateKey::from_bytes(private_bytes);
160        let reconstructed_public = MLKEM768PublicKey::from_bytes(public_bytes);
161
162        assert_eq!(private_key.as_bytes(), reconstructed_private.as_bytes());
163        assert_eq!(public_key.as_bytes(), reconstructed_public.as_bytes());
164    }
165}
securedrop_protocol_minimal/primitives/mlkem.rs

securedrop_protocol_minimal/primitives/
mlkem.rs
