bouncycastle_mlkem/
mlkem_keys.rs

1use crate::aux_functions::{byte_decode, byte_encode, expandA};
2use crate::matrix::{Matrix, Vector};
3use crate::mlkem::{H, POLY_BYTES, q};
4use crate::mlkem::{MLKEM512_PK_LEN, MLKEM512_SK_LEN, MLKEM512_k};
5use crate::mlkem::{MLKEM768_PK_LEN, MLKEM768_SK_LEN, MLKEM768_k};
6use crate::mlkem::{MLKEM1024_PK_LEN, MLKEM1024_SK_LEN, MLKEM1024_k};
7use crate::{ML_KEM_512_NAME, ML_KEM_768_NAME, ML_KEM_1024_NAME};
8use bouncycastle_core::errors::KEMError;
9use bouncycastle_core::key_material;
10use bouncycastle_core::key_material::{KeyMaterial, KeyMaterialTrait, KeyType};
11use bouncycastle_core::traits::{Hash, KEMPrivateKey, KEMPublicKey, SecurityStrength};
12use bouncycastle_sha3::SHA3_256;
13use bouncycastle_utils::secret::Secret;
14use core::fmt;
15use core::fmt::{Debug, Display, Formatter};
16
17// imports just for docs
18#[allow(unused_imports)]
19use crate::mlkem::MLKEMTrait;
20#[allow(unused_imports)]
21use crate::polynomial::Polynomial;
22
23/* Pub Types */
24
25/// ML-KEM-512 Public Key
26pub type MLKEM512PublicKey = MLKEMPublicKey<MLKEM512_k, MLKEM512_PK_LEN>;
27/// ML-KEM-512 Private Key
28pub type MLKEM512PrivateKey =
29    MLKEMPrivateKey<MLKEM512_k, MLKEM512PublicKey, MLKEM512_SK_LEN, MLKEM512_PK_LEN>;
30/// ML-KEM-768 Public Key
31pub type MLKEM768PublicKey = MLKEMPublicKey<MLKEM768_k, MLKEM768_PK_LEN>;
32/// ML-KEM-768 Private Key
33pub type MLKEM768PrivateKey =
34    MLKEMPrivateKey<MLKEM768_k, MLKEM768PublicKey, MLKEM768_SK_LEN, MLKEM768_PK_LEN>;
35/// ML-KEM-1024 Public Key
36pub type MLKEM1024PublicKey = MLKEMPublicKey<MLKEM1024_k, MLKEM1024_PK_LEN>;
37/// ML-KEM-1024 Private Key
38pub type MLKEM1024PrivateKey =
39    MLKEMPrivateKey<MLKEM1024_k, MLKEM1024PublicKey, MLKEM1024_SK_LEN, MLKEM1024_PK_LEN>;
40
41/* Pre-expanded keys for repeated operations */
42
43/// ML-KEM-512 Public Key with a pre-expanded public matrix A for repeated encaps operations.
44pub type MLKEM512PublicKeyExpanded =
45    MLKEMPublicKeyExpanded<MLKEM512_k, MLKEM512PublicKey, MLKEM512_PK_LEN>;
46/// ML-KEM-512 Private Key with a pre-expanded public matrix A for repeated decaps operations.
47pub type MLKEM512PrivateKeyExpanded = MLKEMPrivateKeyExpanded<
48    MLKEM512_k,
49    MLKEM512PublicKey,
50    MLKEM512PrivateKey,
51    MLKEM512_SK_LEN,
52    MLKEM512_PK_LEN,
53>;
54/// ML-KEM-768 Public Key with a pre-expanded public matrix A for repeated encaps operations.
55pub type MLKEM768PublicKeyExpanded =
56    MLKEMPublicKeyExpanded<MLKEM768_k, MLKEM768PublicKey, MLKEM768_PK_LEN>;
57/// ML-KEM-768 Private Key with a pre-expanded public matrix A for repeated decaps operations.
58pub type MLKEM768PrivateKeyExpanded = MLKEMPrivateKeyExpanded<
59    MLKEM768_k,
60    MLKEM768PublicKey,
61    MLKEM768PrivateKey,
62    MLKEM768_SK_LEN,
63    MLKEM768_PK_LEN,
64>;
65/// ML-KEM-1024 Public Key with a pre-expanded public matrix A for repeated encaps operations.
66pub type MLKEM1024PublicKeyExpanded =
67    MLKEMPublicKeyExpanded<MLKEM1024_k, MLKEM1024PublicKey, MLKEM1024_PK_LEN>;
68/// ML-KEM-1024 Private Key with a pre-expanded public matrix A for repeated decaps operations.
69pub type MLKEM1024PrivateKeyExpanded = MLKEMPrivateKeyExpanded<
70    MLKEM1024_k,
71    MLKEM1024PublicKey,
72    MLKEM1024PrivateKey,
73    MLKEM1024_SK_LEN,
74    MLKEM1024_PK_LEN,
75>;
76
77/// An ML-KEM public key.
78#[derive(Clone)]
79pub struct MLKEMPublicKey<const k: usize, const PK_LEN: usize> {
80    t_hat: Vector<k>,
81    rho: [u8; 32],
82}
83
84/// General trait for all ML-KEM public keys types.
85pub trait MLKEMPublicKeyTrait<const k: usize, const PK_LEN: usize>: KEMPublicKey<PK_LEN> {
86    /// Algorithm 23 pkDecode(𝑝𝑘)
87    /// Reverses the procedure pkEncode.
88    /// Input: Public key 𝑝𝑘 ∈ 𝔹32+32𝑘(bitlen (𝑞−1)−𝑑).
89    /// Output: 𝜌 ∈ 𝔹32, 𝐭1 ∈ 𝑅𝑘 with coefficients in [0, 2bitlen (𝑞−1)−𝑑 − 1].
90    fn pk_decode(pk: &[u8; PK_LEN]) -> Result<Self, KEMError>;
91    /// Get a copy of the expanded public matrix A_hat
92    fn A_hat(&self) -> Matrix<k, k>;
93    /// Get the hash of the public key
94    fn compute_hash(&self) -> [u8; 32];
95}
96
97pub(crate) trait MLKEMPublicKeyInternalTrait<const k: usize, const PK_LEN: usize>:
98    MLKEMPublicKeyTrait<k, PK_LEN>
99{
100    /// Not exposing a constructor publicly because you should have to get an instance either by
101    /// running a keygen, or by decoding an existing key.
102    fn new(t_hat: Vector<k>, rho: [u8; 32]) -> Self;
103
104    /// Get a ref to t1
105    fn t_hat(&self) -> &Vector<k>;
106}
107
108impl<const k: usize, const PK_LEN: usize> MLKEMPublicKeyTrait<k, PK_LEN>
109    for MLKEMPublicKey<k, PK_LEN>
110{
111    fn pk_decode(pk: &[u8; PK_LEN]) -> Result<Self, KEMError> {
112        let (pk_chunks, last_chunk) = pk.as_chunks::<POLY_BYTES>();
113
114        // that should divide evenly the remainder of the array, leaving space for rho at the end
115        debug_assert_eq!(pk_chunks.len(), k);
116        debug_assert_eq!(last_chunk.len(), 32);
117
118        let t_hat = {
119            let mut t_hat = Vector::<k>::new();
120
121            for (t_i, pk_chunk) in t_hat.vec.iter_mut().zip(pk_chunks) {
122                t_i.coeffs.copy_from_slice(&byte_decode::<12, POLY_BYTES>(pk_chunk).coeffs);
123
124                // FIPS 203 says:
125                //      "Specifically, ByteDecode12 converts each 12-bit
126                //      segment of its input into an integer modulo 2^{12} = 4096 and then reduces the result
127                //      modulo 𝑞. This is no longer a one-to-one operation. Indeed, some 12-bit segments could
128                //      correspond to an integer greater than 𝑞 − 1 = 3328 but less than 4096."
129                //  Since this concerns to the case d=12, it should be checked that all coeffs are less than q-1
130                for coeff in t_i.coeffs.iter() {
131                    if *coeff < 0 || *coeff >= q {
132                        return Err(KEMError::DecodingError("Invalid or corrupted key"));
133                    }
134                }
135            }
136
137            t_hat
138        };
139        let rho = last_chunk.try_into().unwrap();
140
141        Ok(Self::new(t_hat, rho))
142    }
143
144    fn A_hat(&self) -> Matrix<k, k> {
145        expandA(&self.rho)
146    }
147
148    fn compute_hash(&self) -> [u8; 32] {
149        let mut out = [0u8; 32];
150        let bytes_written = H::default().hash_out(&self.encode(), &mut out);
151        debug_assert_eq!(bytes_written, 32);
152        out
153    }
154}
155
156impl<const k: usize, const PK_LEN: usize> MLKEMPublicKeyInternalTrait<k, PK_LEN>
157    for MLKEMPublicKey<k, PK_LEN>
158{
159    fn new(t_hat: Vector<k>, rho: [u8; 32]) -> Self {
160        Self { rho, t_hat }
161    }
162
163    fn t_hat(&self) -> &Vector<k> {
164        &self.t_hat
165    }
166}
167
168impl<const k: usize, const PK_LEN: usize> KEMPublicKey<PK_LEN> for MLKEMPublicKey<k, PK_LEN> {
169    /// Encodes the public key as per FIPS 203 Algorithm 13
170    /// 19: ekPKE ← ByteEncode12(𝐭)‖𝜌
171    fn encode(&self) -> [u8; PK_LEN] {
172        let mut pk = [0u8; PK_LEN];
173        self.encode_out(&mut pk);
174
175        pk
176    }
177    /// Encodes the public key as per FIPS 203 Algorithm 13
178    /// 19: ekPKE ← ByteEncode12(𝐭)‖𝜌
179    fn encode_out(&self, out: &mut [u8; PK_LEN]) -> usize {
180        debug_assert_eq!(PK_LEN, 12 * k * 32 + 32);
181        debug_assert_eq!(POLY_BYTES, 12 * 32);
182
183        out.fill(0);
184
185        let (pk_chunks, last_chunk) = out.as_chunks_mut::<POLY_BYTES>();
186
187        // that should divide evenly the remainder of the array, leaving space for rho at the end
188        debug_assert_eq!(pk_chunks.len(), k);
189        debug_assert_eq!(last_chunk.len(), 32);
190
191        for (pk_chunk, t_i) in pk_chunks.into_iter().zip(&self.t_hat.vec) {
192            pk_chunk.copy_from_slice(&byte_encode::<12, POLY_BYTES>(t_i));
193        }
194        last_chunk.copy_from_slice(&self.rho);
195
196        PK_LEN
197    }
198
199    fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
200        if bytes.len() != PK_LEN {
201            return Err(KEMError::DecodingError("Provided key bytes are the incorrect length"));
202        }
203        let bytes_sized: [u8; PK_LEN] = bytes[..PK_LEN].try_into().unwrap();
204        Self::pk_decode(&bytes_sized)
205    }
206}
207
208impl<const k: usize, const PK_LEN: usize> Eq for MLKEMPublicKey<k, PK_LEN> {}
209
210impl<const k: usize, const PK_LEN: usize> PartialEq for MLKEMPublicKey<k, PK_LEN> {
211    fn eq(&self, other: &Self) -> bool {
212        bouncycastle_utils::ct::ct_eq_bytes(&self.encode(), &other.encode())
213    }
214}
215
216impl<const k: usize, const PK_LEN: usize> Debug for MLKEMPublicKey<k, PK_LEN> {
217    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
218        let alg = match k {
219            2 => ML_KEM_512_NAME,
220            3 => ML_KEM_768_NAME,
221            4 => ML_KEM_1024_NAME,
222            _ => panic!("Unsupported key length"),
223        };
224        let hash = SHA3_256::new().hash(&self.encode());
225        write!(f, "MLKEMPublicKey {{ alg: {}, pub_key_hash: {:x?} }}", alg, hash)
226    }
227}
228
229impl<const k: usize, const PK_LEN: usize> Display for MLKEMPublicKey<k, PK_LEN> {
230    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
231        let alg = match k {
232            2 => ML_KEM_512_NAME,
233            3 => ML_KEM_768_NAME,
234            4 => ML_KEM_1024_NAME,
235            _ => panic!("Unsupported key length"),
236        };
237        let hash = SHA3_256::new().hash(&self.encode());
238        write!(f, "MLKEMPublicKey {{ alg: {}, pub_key_hash: {:x?} }}", alg, hash)
239    }
240}
241
242/// A fully expanded ML-KEM public key that includes the intermediate values needed for performing multiple encaps operations
243/// against the same public key, which causes the MLKEMPublicKey struct to take up more memory, but results
244/// in more efficient repeated encaps() operations.
245#[derive(Clone)]
246pub struct MLKEMPublicKeyExpanded<
247    const k: usize,
248    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
249    const PK_LEN: usize,
250> {
251    pub(crate) ek: PK,
252    pub(crate) A_hat: Matrix<k, k>,
253}
254
255impl<const k: usize, PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>, const PK_LEN: usize>
256    MLKEMPublicKeyInternalTrait<k, PK_LEN> for MLKEMPublicKeyExpanded<k, PK, PK_LEN>
257{
258    fn new(t_hat: Vector<k>, rho: [u8; 32]) -> Self {
259        let ek = PK::new(t_hat, rho);
260        let A_hat = ek.A_hat();
261
262        Self { ek, A_hat }
263    }
264
265    fn t_hat(&self) -> &Vector<k> {
266        self.ek.t_hat()
267    }
268}
269
270impl<const k: usize, PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>, const PK_LEN: usize>
271    KEMPublicKey<PK_LEN> for MLKEMPublicKeyExpanded<k, PK, PK_LEN>
272{
273    fn encode(&self) -> [u8; PK_LEN] {
274        let mut pk = [0u8; PK_LEN];
275        self.encode_out(&mut pk);
276
277        pk
278    }
279
280    fn encode_out(&self, out: &mut [u8; PK_LEN]) -> usize {
281        out.fill(0);
282
283        self.ek.encode_out(out)
284    }
285
286    fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
287        if bytes.len() != PK_LEN {
288            return Err(KEMError::DecodingError("Provided key bytes are the incorrect length"));
289        }
290        let bytes_sized: [u8; PK_LEN] = bytes[..PK_LEN].try_into().unwrap();
291        Self::pk_decode(&bytes_sized)
292    }
293}
294
295impl<const k: usize, PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>, const PK_LEN: usize> PartialEq
296    for MLKEMPublicKeyExpanded<k, PK, PK_LEN>
297{
298    fn eq(&self, other: &Self) -> bool {
299        self.encode() == other.encode()
300    }
301}
302
303impl<const k: usize, PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>, const PK_LEN: usize> Eq
304    for MLKEMPublicKeyExpanded<k, PK, PK_LEN>
305{
306}
307
308impl<const k: usize, PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>, const PK_LEN: usize> Debug
309    for MLKEMPublicKeyExpanded<k, PK, PK_LEN>
310{
311    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
312        let alg = match k {
313            2 => ML_KEM_512_NAME,
314            3 => ML_KEM_768_NAME,
315            4 => ML_KEM_1024_NAME,
316            _ => panic!("Unsupported key length"),
317        };
318        let hash = SHA3_256::new().hash(&self.encode());
319        write!(f, "MLKEMPublicKeyExpanded {{ alg: {}, pub_key_hash: {:x?} }}", alg, hash)
320    }
321}
322
323impl<const k: usize, PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>, const PK_LEN: usize> Display
324    for MLKEMPublicKeyExpanded<k, PK, PK_LEN>
325{
326    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
327        let alg = match k {
328            2 => ML_KEM_512_NAME,
329            3 => ML_KEM_768_NAME,
330            4 => ML_KEM_1024_NAME,
331            _ => panic!("Unsupported key length"),
332        };
333        let hash = SHA3_256::new().hash(&self.encode());
334        write!(f, "MLKEMPublicKeyExpanded {{ alg: {}, pub_key_hash: {:x?} }}", alg, hash)
335    }
336}
337
338impl<const k: usize, PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>, const PK_LEN: usize>
339    MLKEMPublicKeyTrait<k, PK_LEN> for MLKEMPublicKeyExpanded<k, PK, PK_LEN>
340{
341    fn pk_decode(pk: &[u8; PK_LEN]) -> Result<Self, KEMError> {
342        let ek = PK::pk_decode(pk)?;
343        let A_hat = ek.A_hat();
344        Ok(Self { ek, A_hat })
345    }
346
347    fn A_hat(&self) -> Matrix<k, k> {
348        self.A_hat.clone()
349    }
350
351    fn compute_hash(&self) -> [u8; 32] {
352        self.ek.compute_hash()
353    }
354}
355
356impl<const k: usize, PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>, const PK_LEN: usize> From<&PK>
357    for MLKEMPublicKeyExpanded<k, PK, PK_LEN>
358{
359    /// Fully expands the intermediate values needed for performing multiple encaps operations
360    /// against the same public key, which causes the MLKEMPublicKey struct to take up
361    fn from(ek: &PK) -> Self {
362        let A_hat = ek.A_hat();
363
364        Self { ek: ek.clone(), A_hat }
365    }
366}
367
368/// An ML-KEM private key.
369///
370/// This will automatically inherit the [`Secret`] protections because [`Polynomial`] wraps the underlying data with [`Secret`].
371#[derive(Clone)]
372pub struct MLKEMPrivateKey<
373    const k: usize,
374    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
375    const SK_LEN: usize,
376    const PK_LEN: usize,
377> {
378    s_hat: Secret<Vector<k>>,
379    ek: PK,
380    pk_hash: [u8; 32],
381    z: Secret<[u8; 32]>,
382    seed_d: Option<Secret<[u8; 32]>>,
383}
384
385impl<
386    const k: usize,
387    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
388    const SK_LEN: usize,
389    const PK_LEN: usize,
390> MLKEMPrivateKey<k, PK, SK_LEN, PK_LEN>
391{
392    /// As described on Algorithm 16 line
393    ///   3: dk ← (dkPKE ‖ ek ‖ H(ek) ‖ 𝑧)
394    fn sk_encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
395        out.fill(0);
396
397        debug_assert_eq!(SK_LEN, /* dk_pke*/ 12*k*32 + /*ek*/PK_LEN + /*H(ek)*/32 + /*z*/32);
398
399        let mut pos = 0usize;
400
401        /* dk_pke */
402        // Alg 13; line 20: dkPKE ← ByteEncode12(𝐬)
403        for i in 0..k {
404            out[i * POLY_BYTES..(i + 1) * POLY_BYTES]
405                .copy_from_slice(&byte_encode::<12, POLY_BYTES>(&self.s_hat[i]));
406        }
407        pos += k * POLY_BYTES;
408
409        /* ek */
410        // Alg 13; line 19: ekPKE ← ByteEncode12(𝐭)‖𝜌
411        debug_assert_eq!(self.ek.encode().len(), PK_LEN);
412        out[pos..pos + PK_LEN].copy_from_slice(&self.ek.encode());
413        pos += PK_LEN;
414
415        /* H(ek) */
416        out[pos..pos + 32].copy_from_slice(&self.pk_hash);
417        pos += 32;
418
419        /* z */
420        out[pos..pos + 32].copy_from_slice(&*self.z);
421
422        debug_assert_eq!(pos + 32, SK_LEN);
423        SK_LEN
424    }
425}
426
427/// General trait for all ML-KEM private keys types.
428pub trait MLKEMPrivateKeyTrait<
429    const k: usize,
430    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
431    const SK_LEN: usize,
432    const PK_LEN: usize,
433>: KEMPrivateKey<SK_LEN>
434{
435    /// Get a ref to the seed, if there is one stored with this private key
436    fn seed(&self) -> Option<KeyMaterial<64>>;
437
438    /// This is a partial implementation of keygen_internal(), and probably not allowed in FIPS mode.
439    fn pk(&self) -> &PK;
440    /// Get a ref to the stored public key hash.
441    fn pk_hash(&self) -> &[u8; 32];
442    /// Decode the private key.
443    fn sk_decode(sk: &[u8; SK_LEN]) -> Result<Self, KEMError>;
444}
445
446pub(crate) trait MLKEMPrivateKeyInternalTrait<
447    const k: usize,
448    PK: MLKEMPublicKeyTrait<k, PK_LEN>,
449    const SK_LEN: usize,
450    const PK_LEN: usize,
451>
452{
453    /// Not exposing a constructor publicly because you should have to get an instance either by
454    /// running a keygen, or by decoding an existing key.
455    fn new(
456        s_hat: Secret<Vector<k>>,
457        ek: PK,
458        h: [u8; 32],
459        z: Secret<[u8; 32]>,
460        seed_d: Option<Secret<[u8; 32]>>,
461    ) -> Self;
462
463    /// Get a ref to s_hat
464    fn s_hat(&self) -> &Vector<k>;
465
466    fn z(&self) -> &Secret<[u8; 32]>;
467}
468
469impl<
470    const k: usize,
471    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
472    const SK_LEN: usize,
473    const PK_LEN: usize,
474> MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN> for MLKEMPrivateKey<k, PK, SK_LEN, PK_LEN>
475{
476    fn seed(&self) -> Option<KeyMaterial<64>> {
477        if self.seed_d.is_none() {
478            None
479        } else {
480            let mut tmp = Secret::<[u8; 64]>::new();
481            tmp[..32].copy_from_slice(&self.seed_d.clone().unwrap().as_ref());
482            tmp[32..].copy_from_slice(&*self.z);
483            let mut seed = KeyMaterial::<64>::from_bytes_as_type(&*tmp, KeyType::Seed).unwrap();
484
485            key_material::do_hazardous_operations(&mut seed, |seed| {
486                seed.set_security_strength(match k {
487                    2 => SecurityStrength::_128bit,
488                    3 => SecurityStrength::_192bit,
489                    4 => SecurityStrength::_256bit,
490                    _ => unreachable!("Invalid mlkem param set"),
491                })
492            })
493            .unwrap();
494
495            Some(seed)
496        }
497    }
498
499    fn pk(&self) -> &PK {
500        &self.ek
501    }
502
503    fn pk_hash(&self) -> &[u8; 32] {
504        &self.pk_hash
505    }
506
507    fn sk_decode(sk: &[u8; SK_LEN]) -> Result<Self, KEMError> {
508        debug_assert_eq!(SK_LEN, /* dk_pke*/ 12*k*32 + /*ek*/PK_LEN + /*H(ek)*/32 + /*z*/32);
509
510        let mut pos = 0usize;
511
512        /* dk_pke */
513        let mut s_hat: Secret<Vector<k>> = Secret::new();
514        // for (s_i, sk_chunk) in s_hat.0.iter_mut().zip(sk_chunks) {
515        for i in 0..k {
516            s_hat[i] = byte_decode::<12, POLY_BYTES>(
517                sk[i * POLY_BYTES..(i + 1) * POLY_BYTES].try_into().unwrap(),
518            );
519
520            // FIPS 203 says:
521            //      "Specifically, ByteDecode12 converts each 12-bit
522            //      segment of its input into an integer modulo 2^{12} = 4096 and then reduces the result
523            //      modulo 𝑞. This is no longer a one-to-one operation. Indeed, some 12-bit segments could
524            //      correspond to an integer greater than 𝑞 − 1 = 3328 but less than 4096."
525            //  Since this concerns to the case d=12, it should be checked that all coeffs are less than q-1
526            for coeff in s_hat[i].coeffs.iter() {
527                if *coeff < 0 || *coeff >= q {
528                    return Err(KEMError::DecodingError("Invalid or corrupted key"));
529                }
530            }
531        }
532        pos += k * POLY_BYTES;
533
534        /* ek */
535        let ek = PK::pk_decode(sk[pos..pos + PK_LEN].try_into().unwrap())?;
536        pos += PK_LEN;
537
538        /* H(ek) */
539        let h_pk: [u8; 32] = sk[pos..pos + 32].try_into().unwrap();
540        pos += 32;
541
542        // This satisfies the "Decapsulation input check #3) in FIPS 203 section 7.3.
543        // It is done here on key load rather than as part of the decapsulation for performance
544        // because if multiple decapsulations are being performed, this check needs to be done only once.
545        if h_pk != ek.compute_hash() {
546            return Err(KEMError::ConsistencyCheckFailed(
547                "Corrupted private key: computed hash of ek != h_ek stored in private key",
548            ));
549        }
550
551        /* z */
552        let mut z = Secret::<[u8; 32]>::new();
553        z.copy_from_slice(sk[pos..pos + 32].try_into().unwrap());
554
555        Ok(Self::new(s_hat, ek, h_pk, z, None))
556    }
557}
558
559impl<
560    const k: usize,
561    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
562    const SK_LEN: usize,
563    const PK_LEN: usize,
564> MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN> for MLKEMPrivateKey<k, PK, SK_LEN, PK_LEN>
565{
566    /// Note to future maintainers: FIPS 203 section 7.3 requires that ek be hashed and compared to pk_hash.
567    fn new(
568        s_hat: Secret<Vector<k>>,
569        ek: PK,
570        pk_hash: [u8; 32],
571        z: Secret<[u8; 32]>,
572        seed_d: Option<Secret<[u8; 32]>>,
573    ) -> Self {
574        Self { s_hat, ek, pk_hash, z, seed_d: seed_d.clone() }
575    }
576
577    fn s_hat(&self) -> &Vector<k> {
578        &self.s_hat
579    }
580
581    fn z(&self) -> &Secret<[u8; 32]> {
582        &self.z
583    }
584}
585
586impl<
587    const k: usize,
588    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
589    const SK_LEN: usize,
590    const PK_LEN: usize,
591> KEMPrivateKey<SK_LEN> for MLKEMPrivateKey<k, PK, SK_LEN, PK_LEN>
592{
593    fn encode(&self) -> [u8; SK_LEN] {
594        let mut out = [0u8; SK_LEN];
595        self.encode_out(&mut out);
596
597        out
598    }
599
600    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
601        out.fill(0);
602
603        self.sk_encode_out(out)
604    }
605
606    fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
607        if bytes.len() != SK_LEN {
608            return Err(KEMError::DecodingError("Provided key bytes are the incorrect length"));
609        }
610        if bytes.len() != SK_LEN {
611            return Err(KEMError::DecodingError("Provided key bytes are the incorrect length"));
612        }
613        let bytes_sized: [u8; SK_LEN] = bytes[..SK_LEN].try_into().unwrap();
614
615        Self::sk_decode(&bytes_sized)
616    }
617}
618
619impl<
620    const k: usize,
621    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
622    const SK_LEN: usize,
623    const PK_LEN: usize,
624> Eq for MLKEMPrivateKey<k, PK, SK_LEN, PK_LEN>
625{
626}
627
628impl<
629    const k: usize,
630    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
631    const SK_LEN: usize,
632    const PK_LEN: usize,
633> PartialEq for MLKEMPrivateKey<k, PK, SK_LEN, PK_LEN>
634{
635    fn eq(&self, other: &Self) -> bool {
636        let self_encoded = self.encode();
637        let other_encoded = other.encode();
638        bouncycastle_utils::ct::ct_eq_bytes(self_encoded.as_ref(), other_encoded.as_ref())
639    }
640}
641
642/// Debug impl mainly to prevent the secret key from being printed in logs.
643impl<
644    const k: usize,
645    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
646    const SK_LEN: usize,
647    const PK_LEN: usize,
648> fmt::Debug for MLKEMPrivateKey<k, PK, SK_LEN, PK_LEN>
649{
650    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
651        let alg = match k {
652            2 => ML_KEM_512_NAME,
653            3 => ML_KEM_768_NAME,
654            4 => ML_KEM_1024_NAME,
655            _ => panic!("Unsupported key length"),
656        };
657        write!(
658            f,
659            "MLKEMPrivateKey {{ alg: {}, pub_key_hash: {:x?}, has_seed: {} }}",
660            alg,
661            self.pk_hash,
662            self.seed_d.is_some(),
663        )
664    }
665}
666
667/// Display impl mainly to prevent the secret key from being printed in logs.
668impl<
669    const k: usize,
670    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
671    const SK_LEN: usize,
672    const PK_LEN: usize,
673> Display for MLKEMPrivateKey<k, PK, SK_LEN, PK_LEN>
674{
675    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
676        let alg = match k {
677            2 => ML_KEM_512_NAME,
678            3 => ML_KEM_768_NAME,
679            4 => ML_KEM_1024_NAME,
680            _ => panic!("Unsupported key length"),
681        };
682        write!(
683            f,
684            "MLKEMPrivateKey {{ alg: {}, pub_key_hash: {:x?}, has_seed: {} }}",
685            alg,
686            self.pk_hash,
687            self.seed_d.is_some(),
688        )
689    }
690}
691
692/// A fully expanded ML-KEM private key that includes the intermediate values needed for performing
693/// multiple decaps operations with the same private key, which causes the private key struct to
694/// take up more memory, but results in more efficient repeated decaps() operations.
695#[derive(Clone)]
696pub struct MLKEMPrivateKeyExpanded<
697    const k: usize,
698    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
699    SK: MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
700        + MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN>,
701    const SK_LEN: usize,
702    const PK_LEN: usize,
703> {
704    _phantom: core::marker::PhantomData<PK>,
705    pub(crate) dk: SK,
706    pub(crate) A_hat: Matrix<k, k>,
707}
708
709impl<
710    const k: usize,
711    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
712    SK: MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
713        + MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN>,
714    const SK_LEN: usize,
715    const PK_LEN: usize,
716> From<&SK> for MLKEMPrivateKeyExpanded<k, PK, SK, SK_LEN, PK_LEN>
717{
718    /// Fully expands the intermediate values needed for performing multiple encaps operations
719    /// against the same public key, which causes the MLKEMPublicKey struct to take up
720    fn from(dk: &SK) -> Self {
721        let A_hat = dk.pk().A_hat();
722
723        Self { _phantom: core::marker::PhantomData, dk: dk.clone(), A_hat }
724    }
725}
726
727impl<
728    const k: usize,
729    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
730    SK: MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
731        + MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN>,
732    const SK_LEN: usize,
733    const PK_LEN: usize,
734> KEMPrivateKey<SK_LEN> for MLKEMPrivateKeyExpanded<k, PK, SK, SK_LEN, PK_LEN>
735{
736    fn encode(&self) -> [u8; SK_LEN] {
737        self.dk.encode()
738    }
739
740    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
741        out.fill(0);
742
743        self.dk.encode_out(out)
744    }
745
746    fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
747        Ok(Self::from(&SK::from_bytes(bytes)?))
748    }
749}
750
751impl<
752    const k: usize,
753    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
754    SK: MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
755        + MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN>,
756    const SK_LEN: usize,
757    const PK_LEN: usize,
758> PartialEq for MLKEMPrivateKeyExpanded<k, PK, SK, SK_LEN, PK_LEN>
759{
760    fn eq(&self, other: &Self) -> bool {
761        self.dk.eq(&other.dk)
762    }
763}
764
765impl<
766    const k: usize,
767    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
768    SK: MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
769        + MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN>,
770    const SK_LEN: usize,
771    const PK_LEN: usize,
772> Eq for MLKEMPrivateKeyExpanded<k, PK, SK, SK_LEN, PK_LEN>
773{
774}
775
776impl<
777    const k: usize,
778    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
779    SK: MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
780        + MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN>,
781    const SK_LEN: usize,
782    const PK_LEN: usize,
783> Debug for MLKEMPrivateKeyExpanded<k, PK, SK, SK_LEN, PK_LEN>
784{
785    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
786        let alg = match k {
787            2 => ML_KEM_512_NAME,
788            3 => ML_KEM_768_NAME,
789            4 => ML_KEM_1024_NAME,
790            _ => panic!("Unsupported key length"),
791        };
792        write!(
793            f,
794            "MLKEMPrivateKeyExpanded {{ alg: {}, pub_key_hash: {:x?}, has_seed: {} }}",
795            alg,
796            self.dk.pk().compute_hash(),
797            self.dk.seed().is_some(),
798        )
799    }
800}
801
802impl<
803    const k: usize,
804    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
805    SK: MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
806        + MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN>,
807    const SK_LEN: usize,
808    const PK_LEN: usize,
809> Display for MLKEMPrivateKeyExpanded<k, PK, SK, SK_LEN, PK_LEN>
810{
811    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
812        let alg = match k {
813            2 => ML_KEM_512_NAME,
814            3 => ML_KEM_768_NAME,
815            4 => ML_KEM_1024_NAME,
816            _ => panic!("Unsupported key length"),
817        };
818        write!(
819            f,
820            "MLKEMPrivateKeyExpanded {{ alg: {}, pub_key_hash: {:x?}, has_seed: {} }}",
821            alg,
822            self.dk.pk().compute_hash(),
823            self.dk.seed().is_some(),
824        )
825    }
826}
827
828impl<
829    const k: usize,
830    PK: MLKEMPublicKeyInternalTrait<k, PK_LEN>,
831    SK: MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
832        + MLKEMPrivateKeyInternalTrait<k, PK, SK_LEN, PK_LEN>,
833    const SK_LEN: usize,
834    const PK_LEN: usize,
835> MLKEMPrivateKeyTrait<k, PK, SK_LEN, PK_LEN>
836    for MLKEMPrivateKeyExpanded<k, PK, SK, SK_LEN, PK_LEN>
837{
838    fn seed(&self) -> Option<KeyMaterial<64>> {
839        self.dk.seed()
840    }
841
842    fn pk(&self) -> &PK {
843        self.dk.pk()
844    }
845
846    fn pk_hash(&self) -> &[u8; 32] {
847        &self.dk.pk_hash()
848    }
849
850    fn sk_decode(sk: &[u8; SK_LEN]) -> Result<Self, KEMError> {
851        let dk = SK::sk_decode(sk)?;
852        let A_hat = dk.pk().A_hat();
853
854        Ok(Self { _phantom: core::marker::PhantomData, dk: dk.clone(), A_hat })
855    }
856}
bouncycastle_mlkem/mlkem_keys.rs

bouncycastle_mlkem/
mlkem_keys.rs
