bouncycastle_mlkem_lowmemory/
mlkem_keys.rs

1use crate::aux_functions::sample_poly_CBD;
2use crate::low_memory_helpers::{
3    compute_A_hat_dot_s_hat, pack_s_hat_row, pack_t_hat_row, unpack_t_hat_row,
4};
5use crate::mlkem::{G, H, POLY_BYTES, q};
6use crate::mlkem::{
7    MLKEM512_ETA1, MLKEM512_FULL_SK_LEN, MLKEM512_LAMBDA, MLKEM512_PK_LEN, MLKEM512_SK_LEN,
8    MLKEM512_T_PACKED_LEN, MLKEM512_k,
9};
10use crate::mlkem::{
11    MLKEM768_ETA1, MLKEM768_FULL_SK_LEN, MLKEM768_LAMBDA, MLKEM768_PK_LEN, MLKEM768_SK_LEN,
12    MLKEM768_T_PACKED_LEN, MLKEM768_k,
13};
14use crate::mlkem::{
15    MLKEM1024_ETA1, MLKEM1024_FULL_SK_LEN, MLKEM1024_LAMBDA, MLKEM1024_PK_LEN, MLKEM1024_SK_LEN,
16    MLKEM1024_T_PACKED_LEN, MLKEM1024_k,
17};
18use crate::polynomial::Polynomial;
19use crate::{ML_KEM_512_NAME, ML_KEM_768_NAME, ML_KEM_1024_NAME};
20use bouncycastle_core::errors::KEMError;
21use bouncycastle_core::key_material::{
22    KeyMaterial, KeyMaterialTrait, KeyType, do_hazardous_operations,
23};
24use bouncycastle_core::traits::{Hash, KEMPrivateKey, KEMPublicKey, SecurityStrength};
25use bouncycastle_sha3::SHA3_256;
26use bouncycastle_utils::secret::Secret;
27use core::fmt;
28use core::fmt::{Debug, Display, Formatter};
29// imports just for docs
30
31/* Pub Types */
32
33/// ML-KEM-512 Public Key
34pub type MLKEM512PublicKey = MLKEMPublicKey<MLKEM512_k, MLKEM512_PK_LEN, MLKEM512_T_PACKED_LEN>;
35/// ML-KEM-512 Private Key
36pub type MLKEM512PrivateKey = MLKEMSeedPrivateKey<
37    MLKEM512_k,
38    MLKEM512_ETA1,
39    MLKEM512_LAMBDA,
40    MLKEM512_SK_LEN,
41    MLKEM512_FULL_SK_LEN,
42    MLKEM512_PK_LEN,
43    MLKEM512_T_PACKED_LEN,
44>;
45/// ML-KEM-768 Public Key
46pub type MLKEM768PublicKey = MLKEMPublicKey<MLKEM768_k, MLKEM768_PK_LEN, MLKEM768_T_PACKED_LEN>;
47/// ML-KEM-768 Private Key
48pub type MLKEM768PrivateKey = MLKEMSeedPrivateKey<
49    MLKEM768_k,
50    MLKEM768_ETA1,
51    MLKEM768_LAMBDA,
52    MLKEM768_SK_LEN,
53    MLKEM768_FULL_SK_LEN,
54    MLKEM768_PK_LEN,
55    MLKEM768_T_PACKED_LEN,
56>;
57/// ML-KEM-1024 Public Key
58pub type MLKEM1024PublicKey = MLKEMPublicKey<MLKEM1024_k, MLKEM1024_PK_LEN, MLKEM1024_T_PACKED_LEN>;
59/// ML-KEM-1024 Private Key
60pub type MLKEM1024PrivateKey = MLKEMSeedPrivateKey<
61    MLKEM1024_k,
62    MLKEM1024_ETA1,
63    MLKEM1024_LAMBDA,
64    MLKEM1024_SK_LEN,
65    MLKEM1024_FULL_SK_LEN,
66    MLKEM1024_PK_LEN,
67    MLKEM1024_T_PACKED_LEN,
68>;
69
70/// An ML-KEM public key.
71#[derive(Clone)]
72pub struct MLKEMPublicKey<const k: usize, const PK_LEN: usize, const T_PACKED_LEN: usize> {
73    pub(crate) t_hat_packed: [u8; T_PACKED_LEN],
74    pub(crate) rho: [u8; 32],
75}
76
77/// General trait for all ML-KEM public keys types.
78pub trait MLKEMPublicKeyTrait<const k: usize, const PK_LEN: usize, const T_PACKED_LEN: usize>:
79    KEMPublicKey<PK_LEN>
80{
81    /// Algorithm 23 pkDecode(𝑝𝑘)
82    /// Reverses the procedure pkEncode.
83    /// Input: Public key 𝑝𝑘 ∈ 𝔹32+32𝑘(bitlen (𝑞−1)−𝑑).
84    /// Output: 𝜌 ∈ 𝔹32, 𝐭1 ∈ 𝑅𝑘 with coefficients in [0, 2bitlen (𝑞−1)−𝑑 − 1].
85    // todo: go make the equivalent thing also throw an error in the non-optimized impl
86    fn pk_decode(pk: &[u8; PK_LEN]) -> Result<Self, KEMError>;
87
88    /// Get a ref to t_hat_packed byte array
89    fn t_hat_packed(&self) -> &[u8; T_PACKED_LEN];
90
91    /// Get a ref to rho
92    fn rho(&self) -> &[u8; 32];
93
94    /// Get the hash of the public key
95    fn compute_hash(&self) -> [u8; 32];
96}
97
98pub(crate) trait MLKEMPublicKeyInternalTrait<
99    const k: usize,
100    const T_PACKED_LEN: usize,
101    const PK_LEN: usize,
102>: MLKEMPublicKeyTrait<k, PK_LEN, T_PACKED_LEN>
103{
104    /// Not exposing a constructor publicly because you should have to get an instance either by
105    /// running a keygen, or by decoding an existing key.
106    fn new(t_hat: [u8; T_PACKED_LEN], rho: [u8; 32]) -> Self;
107}
108
109impl<const k: usize, const PK_LEN: usize, const T_PACKED_LEN: usize>
110    MLKEMPublicKeyTrait<k, PK_LEN, T_PACKED_LEN> for MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN>
111{
112    fn pk_decode(pk: &[u8; PK_LEN]) -> Result<Self, KEMError> {
113        let pk = Self::new(
114            pk[..T_PACKED_LEN].try_into().unwrap(),
115            pk[T_PACKED_LEN..].try_into().unwrap(),
116        );
117
118        // check that all entries are in range
119        for i in 0..k {
120            let p = unpack_t_hat_row(&pk.t_hat_packed, i);
121            for w in p.coeffs.iter() {
122                if *w >= q {
123                    return Err(KEMError::DecodingError("Invalid public key"));
124                }
125            }
126        }
127
128        Ok(pk)
129    }
130
131    fn t_hat_packed(&self) -> &[u8; T_PACKED_LEN] {
132        &self.t_hat_packed
133    }
134
135    fn rho(&self) -> &[u8; 32] {
136        &self.rho
137    }
138
139    fn compute_hash(&self) -> [u8; 32] {
140        // The encoded public key is just t_hat and rho, so feed the elements of the public key into the hash one-by-one
141
142        let mut out = [0u8; 32];
143        let mut h = H::default();
144        h.do_update(&self.t_hat_packed);
145        h.do_update(&self.rho);
146        let bytes_written = h.do_final_out(&mut out);
147        debug_assert_eq!(bytes_written, 32);
148        out
149    }
150}
151
152impl<const k: usize, const T_PACKED_LEN: usize, const PK_LEN: usize>
153    MLKEMPublicKeyInternalTrait<k, T_PACKED_LEN, PK_LEN>
154    for MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN>
155{
156    fn new(t_hat_packed: [u8; T_PACKED_LEN], rho: [u8; 32]) -> Self {
157        Self { rho, t_hat_packed }
158    }
159}
160
161impl<const k: usize, const PK_LEN: usize, const T_PACKED_LEN: usize> KEMPublicKey<PK_LEN>
162    for MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN>
163{
164    /// Algorithm 22 pkEncode(𝜌, 𝐭1)
165    /// Encodes a public key for ML-DSA into a byte string.
166    /// Input:𝜌 ∈ 𝔹32, 𝐭1 ∈ 𝑅𝑘 with coefficients in [0, 2bitlen (𝑞−1)−𝑑 − 1].
167    /// Output: Public key 𝑝𝑘 ∈ 𝔹32+32𝑘(bitlen (𝑞−1)−𝑑).
168    fn encode(&self) -> [u8; PK_LEN] {
169        debug_assert_eq!(PK_LEN, 32 + 12 * k * 32);
170        let mut pk = [0u8; PK_LEN];
171        self.encode_out(&mut pk);
172
173        pk
174    }
175
176    fn encode_out(&self, out: &mut [u8; PK_LEN]) -> usize {
177        debug_assert_eq!(self.t_hat_packed.len(), T_PACKED_LEN);
178
179        out.fill(0);
180
181        out[..T_PACKED_LEN].copy_from_slice(&self.t_hat_packed);
182        debug_assert_eq!(out[T_PACKED_LEN..].len(), 32);
183        out[T_PACKED_LEN..].copy_from_slice(&self.rho);
184
185        PK_LEN
186    }
187
188    fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
189        if bytes.len() != PK_LEN {
190            return Err(KEMError::DecodingError("Provided key bytes are the incorrect length"));
191        }
192        let bytes_sized: [u8; PK_LEN] = bytes[..PK_LEN].try_into().unwrap();
193        Self::pk_decode(&bytes_sized)
194    }
195}
196
197impl<const k: usize, const PK_LEN: usize, const T_PACKED_LEN: usize> Eq
198    for MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN>
199{
200}
201
202impl<const k: usize, const PK_LEN: usize, const T_PACKED_LEN: usize> PartialEq
203    for MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN>
204{
205    fn eq(&self, other: &Self) -> bool {
206        bouncycastle_utils::ct::ct_eq_bytes(&self.encode(), &other.encode())
207    }
208}
209
210impl<const k: usize, const PK_LEN: usize, const T_PACKED_LEN: usize> Debug
211    for MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN>
212{
213    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
214        let alg = match k {
215            2 => ML_KEM_512_NAME,
216            3 => ML_KEM_768_NAME,
217            4 => ML_KEM_1024_NAME,
218            _ => panic!("Unsupported key length"),
219        };
220        let hash = SHA3_256::new().hash(&self.encode());
221        write!(f, "MLKEMPublicKey {{ alg: {}, pub_key_hash: {:x?} }}", alg, hash)
222    }
223}
224
225impl<const k: usize, const PK_LEN: usize, const T_PACKED_LEN: usize> Display
226    for MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN>
227{
228    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
229        let alg = match k {
230            2 => ML_KEM_512_NAME,
231            3 => ML_KEM_768_NAME,
232            4 => ML_KEM_1024_NAME,
233            _ => panic!("Unsupported key length"),
234        };
235        let hash = SHA3_256::new().hash(&self.encode());
236        write!(f, "MLKEMPublicKey {{ alg: {}, pub_key_hash: {:x?} }}", alg, hash)
237    }
238}
239
240/// An ML-KEM private key.
241#[derive(Clone)]
242pub struct MLKEMSeedPrivateKey<
243    const k: usize,
244    const eta1: i16,
245    const LAMBDA: i16,
246    const SK_LEN: usize,
247    const FULL_SK_LEN: usize,
248    const PK_LEN: usize,
249    const T_PACKED_LEN: usize,
250> {
251    rho: [u8; 32],
252    sigma: Secret<[u8; 32]>,
253    pk_hash: Option<[u8; 32]>,
254    z: Secret<[u8; 32]>,
255    seed_d: Secret<[u8; 32]>,
256}
257
258impl<
259    const k: usize,
260    const eta1: i16,
261    const LAMBDA: i16,
262    const SK_LEN: usize,
263    const FULL_SK_LEN: usize,
264    const PK_LEN: usize,
265    const T_PACKED_LEN: usize,
266> MLKEMSeedPrivateKey<k, eta1, LAMBDA, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
267{
268    /// Create a new MLKEMSeedPrivateKey from a 64-byte KeyMaterial.
269    /// Seed SecurityStrength must match algorithm security strength: 128-bit (ML-KEM-512), 192-bit (ML-KEM-768), or 256-bit (ML-KEM-1024).
270    pub fn new(seed: &KeyMaterial<64>) -> Result<Self, KEMError> {
271        if !(seed.key_type() == KeyType::Seed || seed.key_type() == KeyType::CryptographicRandom)
272            || seed.key_len() != 64
273        {
274            return Err(KEMError::KeyGenError(
275                "Seed must be 64 bytes and KeyType::Seed or KeyType::BytesFullEntropy.",
276            ));
277        }
278
279        if seed.security_strength() < SecurityStrength::from_bits(LAMBDA as usize) {
280            return Err(KEMError::KeyGenError("SecurityStrength"));
281        }
282
283        // These are Secret-safe because we're using .copy_from_slice directly out of one Secret<[u8]>
284        // into another and the contents are never touching a non-Secret buffer.
285        let mut seed_d = Secret::<[u8; 32]>::new();
286        seed_d.as_mut().copy_from_slice(seed.ref_to_bytes()[..32].try_into().unwrap());
287        let mut z = Secret::<[u8; 32]>::new();
288        z.as_mut().copy_from_slice(seed.ref_to_bytes()[32..].try_into().unwrap());
289
290        let (rho, sigma) = Self::compute_rho_and_sigma(&seed_d);
291
292        // Deviation from the FIPS: The implementation does not persist the hash of the public key H(ek) in the
293        // in-memory representation because it can be re-computed as needed.
294        Ok(Self { rho, sigma, pk_hash: None, z, seed_d })
295    }
296    /// Algorithm 13 K-PKE.KeyGen(𝑑)
297    /// 1: (𝜌, 𝜎) ← G(𝑑‖𝑘)
298    ///  ▷ expand 32+1 bytes to two pseudorandom 32-byte seeds1
299    /// rho: public seed
300    /// sigma: noise seed
301    fn compute_rho_and_sigma(seed_d: &[u8; 32]) -> ([u8; 32], Secret<[u8; 32]>) {
302        // Only the second half of the output is secret, but we'll wrap the whole thing
303        // so that the local copy gets zeriozed on drop.
304        let mut buf: Secret<[u8; 64]> = Secret::new();
305
306        let mut g = G::new();
307        g.do_update(seed_d);
308        g.do_update(&[k as u8]);
309        let bytes_written = g.do_final_out(buf.as_mut());
310        debug_assert_eq!(bytes_written, 64);
311
312        let mut sigma = Secret::<[u8; 32]>::new();
313        sigma.as_mut().copy_from_slice(buf[32..64].try_into().unwrap());
314
315        (buf[..32].try_into().unwrap(), sigma)
316    }
317}
318
319/// General trait for all ML-KEM private keys types.
320pub trait MLKEMPrivateKeyTrait<
321    const k: usize,
322    const SK_LEN: usize,
323    const FULL_SK_LEN: usize,
324    const PK_LEN: usize,
325    const T_PACKED_LEN: usize,
326>: KEMPrivateKey<SK_LEN>
327{
328    /// New from KeyMaterial. Can throw a KEMError if the KeyMaterial does not contain sufficient entropy.
329    fn from_keymaterial(seed: &KeyMaterial<64>) -> Result<Self, KEMError>;
330    /// Get a ref to the seed, which there always will be for a MLKEMSeedPrivateKey
331    /// In this implementation, we always have a seed, so will always return Some.
332    fn seed(&self) -> Option<KeyMaterial<64>>;
333    /// Runs essentially a full keygen according to Algorithm 13.
334    // Dev note: This is a partial implementation of keygen_internal(), and probably not allowed in FIPS mode.
335    fn pk(&self) -> MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN>;
336    /// Get a ref to the stored public key hash.
337    /// Since in this implementation, this requires running the full keygen, this is a lazy evaluation and
338    /// will only be computationally heavy the first time it is called for a given key.
339    /// This requires a mutable copy. If you don't have then, then you can compute the full public key via [`MLKEMPrivateKeyTrait::pk`]
340    /// and then get the hash of that.
341    fn pk_hash(&mut self) -> &[u8; 32];
342    /// This produces the full private key in the encoding specified in FIPS 203 so that it is
343    /// compatible with other implementations.
344    ///
345    /// Note that since this encoding does not include the seed, this is a one-way operation;
346    /// after exporting in this encoding, it will be impossible to re-import it into a [`MLKEMSeedPrivateKey`].
347    ///
348    /// As described on Algorithm 16 line
349    ///   3: dk ← (dkPKE ‖ ek ‖ H(ek) ‖ 𝑧)
350    fn encode_full_sk(&self) -> [u8; FULL_SK_LEN];
351    /// This produces the full private key in the encoding specified in FIPS 203 so that it is
352    /// compatible with other implementations.
353    ///
354    /// Note that since this encoding does not include the seed, this is a one-way operation;
355    /// after exporting in this encoding, it will be impossible to re-import it into a [`MLKEMSeedPrivateKey`].
356    ///
357    /// As described on Algorithm 16 line
358    ///   3: dk ← (dkPKE ‖ ek ‖ H(ek) ‖ 𝑧)
359    fn encode_full_sk_out(&self, out: &mut [u8; FULL_SK_LEN]) -> usize;
360    /// Decode the private key.
361    fn sk_decode(sk: &[u8; SK_LEN]) -> Self;
362}
363
364pub(crate) trait MLKEMPrivateKeyInternalTrait<
365    const k: usize,
366    const SK_LEN: usize,
367    const PK_LEN: usize,
368    const T_PACKED_LEN: usize,
369>
370{
371    fn z(&self) -> &[u8; 32];
372
373    fn compute_s_hat_row(&self, idx: usize) -> Polynomial;
374
375    fn rho(&self) -> &[u8; 32];
376
377    /// Note: this one is not a ref because the data does not exist in the private key.
378    fn t_hat_packed(&self) -> [u8; T_PACKED_LEN];
379}
380
381impl<
382    const k: usize,
383    const eta1: i16,
384    const LAMBDA: i16,
385    const SK_LEN: usize,
386    const FULL_SK_LEN: usize,
387    const PK_LEN: usize,
388    const T_PACKED_LEN: usize,
389> MLKEMPrivateKeyTrait<k, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
390    for MLKEMSeedPrivateKey<k, eta1, LAMBDA, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
391{
392    fn from_keymaterial(seed: &KeyMaterial<64>) -> Result<Self, KEMError> {
393        Self::new(seed)
394    }
395    fn seed(&self) -> Option<KeyMaterial<64>> {
396        let mut tmp = Secret::<[u8; 64]>::new();
397        tmp[..32].as_mut().copy_from_slice(&*self.seed_d);
398        tmp[32..].as_mut().copy_from_slice(&*self.z);
399        let mut seed = KeyMaterial::<64>::from_bytes_as_type(&*tmp, KeyType::Seed).unwrap();
400        do_hazardous_operations(&mut seed, |seed| {
401            seed.set_security_strength(match k {
402                2 => SecurityStrength::_128bit,
403                3 => SecurityStrength::_192bit,
404                4 => SecurityStrength::_256bit,
405                _ => unreachable!("Invalid mlkem param set"),
406            })
407        })
408        .unwrap();
409
410        Some(seed)
411    }
412    fn pk(&self) -> MLKEMPublicKey<k, PK_LEN, T_PACKED_LEN> {
413        MLKEMPublicKey::<k, PK_LEN, T_PACKED_LEN>::new(self.t_hat_packed(), self.rho)
414    }
415    fn pk_hash(&mut self) -> &[u8; 32] {
416        if self.pk_hash.is_none() {
417            self.pk_hash = Some(self.pk().compute_hash().clone());
418        }
419
420        &self.pk_hash.as_ref().unwrap()
421    }
422    /// This produces the full private key in the encoding specified in FIPS 203 so that it is
423    /// compatible with other implementations.
424    ///
425    /// Note that since this encoding does not include the seed, this is a one-way operation;
426    /// after exporting in this encoding, it will be impossible to re-import it into a [`MLKEMSeedPrivateKey`].
427    ///
428    /// As described on Algorithm 16 line
429    ///   3: dk ← (dkPKE ‖ ek ‖ H(ek) ‖ 𝑧)
430    fn encode_full_sk(&self) -> [u8; FULL_SK_LEN] {
431        let mut out = [0u8; FULL_SK_LEN];
432        self.encode_full_sk_out(&mut out);
433
434        out
435    }
436    /// This produces the full private key in the encoding specified in the FIPS so that it is
437    /// compatible with other implementations.
438    /// Note that this encoding does not include the seed, so if exporting in this encoding, it will
439    /// be impossible to re-import it into this implementation.
440    ///
441    /// As described on Algorithm 16 line
442    ///   3: dk ← (dkPKE ‖ ek ‖ H(ek) ‖ 𝑧)
443    fn encode_full_sk_out(&self, out: &mut [u8; FULL_SK_LEN]) -> usize {
444        out.fill(0);
445
446        let mut pos = 0usize;
447
448        /* dk_pke */
449        // Alg 13; line 20: dkPKE ← ByteEncode12(𝐬)
450        for i in 0..k {
451            pack_s_hat_row::<k>(&self.compute_s_hat_row(i), i, out);
452        }
453        pos += k * POLY_BYTES;
454
455        /* ek */
456        // Alg 13; line 19: ekPKE ← ByteEncode12(𝐭)‖𝜌
457        let pk = self.pk();
458        out[pos..pos + PK_LEN].copy_from_slice(&pk.encode());
459        pos += PK_LEN;
460
461        /* H(ek) */
462        out[pos..pos + 32].copy_from_slice(&pk.compute_hash());
463        pos += 32;
464
465        /* z */
466        out[pos..pos + 32].copy_from_slice(&*self.z);
467
468        FULL_SK_LEN
469    }
470    fn sk_decode(sk: &[u8; SK_LEN]) -> Self {
471        debug_assert_eq!(SK_LEN, /* seed*/ 64);
472        Self::from_bytes(sk).unwrap()
473    }
474}
475
476impl<
477    const k: usize,
478    const eta1: i16,
479    const LAMBDA: i16,
480    const SK_LEN: usize,
481    const FULL_SK_LEN: usize,
482    const PK_LEN: usize,
483    const T_PACKED_LEN: usize,
484> MLKEMPrivateKeyInternalTrait<k, SK_LEN, PK_LEN, T_PACKED_LEN>
485    for MLKEMSeedPrivateKey<k, eta1, LAMBDA, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
486{
487    fn z(&self) -> &[u8; 32] {
488        &self.z
489    }
490
491    fn compute_s_hat_row(&self, idx: usize) -> Polynomial {
492        debug_assert!(idx < k);
493
494        // We're doing just one row of this:
495        // 8: for (𝑖 ← 0; 𝑖 < 𝑘; 𝑖++)
496        //  ▷ generate 𝐬 ∈ (ℤ256)^k
497        // 9: 𝐬[𝑖] ← SamplePolyCBD𝜂1(PRF𝜂1 (𝜎, 𝑁 ))
498        //   ▷ 𝐬[𝑖] ∈ ℤ256 sampled from CBD
499        // 10: 𝑁 ← 𝑁 + 1
500        // Note: here n = 0
501        let mut s_i = sample_poly_CBD::<eta1>(&self.sigma, idx as u8);
502
503        // 16: 𝐬_hat ← NTT(𝐬)̂
504        s_i.ntt();
505        s_i
506    }
507
508    fn rho(&self) -> &[u8; 32] {
509        &self.rho
510    }
511    /// Runs essentially a full keygen according to Algorithm 13
512    /// Outputs t_hat in the packed encoding specified in FIPS 203
513    fn t_hat_packed(&self) -> [u8; T_PACKED_LEN] {
514        let mut t_hat_packed = [0u8; T_PACKED_LEN];
515
516        for i in 0..k {
517            // first half of
518            // 18: 𝐭_hat ← 𝐀_hat ∘ 𝐬_hat + 𝐞_hat
519            let mut t_hat_i = compute_A_hat_dot_s_hat::<k, eta1>(&self.rho, &self.sigma, i);
520
521            // second half of
522            // 18: 𝐭_hat ← 𝐀_hat ∘ 𝐬_hat + 𝐞_hat
523            {
524                // 12: for (𝑖 ← 0; 𝑖 < 𝑘; 𝑖++)
525                //  ▷ generate 𝐞 ∈ (ℤ256)^k
526                // 13: 𝐞[𝑖] ← SamplePolyCBD𝜂1(PRF𝜂1 (𝜎, 𝑁))
527                //   ▷ 𝐞[𝑖] ∈ ℤ256 sampled from CBD
528                // 14: 𝑁 ← 𝑁 + 1
529                // Note: here n = k
530                let mut e_i = sample_poly_CBD::<eta1>(&self.sigma, (k + i) as u8);
531
532                e_i.ntt(); // technically now e_hat_i
533                t_hat_i.add(&e_i);
534            }
535            t_hat_i.poly_reduce();
536
537            pack_t_hat_row::<T_PACKED_LEN>(&t_hat_i, i, &mut t_hat_packed);
538        }
539
540        t_hat_packed
541    }
542}
543
544impl<
545    const k: usize,
546    const eta1: i16,
547    const LAMBDA: i16,
548    const SK_LEN: usize,
549    const FULL_SK_LEN: usize,
550    const PK_LEN: usize,
551    const T_PACKED_LEN: usize,
552> KEMPrivateKey<SK_LEN>
553    for MLKEMSeedPrivateKey<k, eta1, LAMBDA, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
554{
555    /// Encode the private key as a 64-byte seed (d || z)
556    fn encode(&self) -> [u8; SK_LEN] {
557        let mut sk = [0u8; SK_LEN];
558        self.encode_out(&mut sk);
559
560        sk
561    }
562
563    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
564        debug_assert_eq!(SK_LEN, 64);
565
566        out.fill(0);
567
568        out[..32].copy_from_slice(&*self.seed_d);
569        out[32..].copy_from_slice(&*self.z);
570
571        SK_LEN
572    }
573
574    fn from_bytes(bytes: &[u8]) -> Result<Self, KEMError> {
575        if bytes.len() != 64 {
576            return Err(KEMError::DecodingError("Invalid seed length"));
577        }
578        let mut keymat = KeyMaterial::<64>::from_bytes(bytes)?;
579        do_hazardous_operations(&mut keymat, |keymat| {
580            keymat.set_key_type(KeyType::Seed)?;
581            keymat.set_security_strength(SecurityStrength::_256bit)
582        })?;
583
584        Self::new(&keymat)
585    }
586}
587
588impl<
589    const k: usize,
590    const eta1: i16,
591    const LAMBDA: i16,
592    const SK_LEN: usize,
593    const FULL_SK_LEN: usize,
594    const PK_LEN: usize,
595    const T_PACKED_LEN: usize,
596> Eq for MLKEMSeedPrivateKey<k, eta1, LAMBDA, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
597{
598}
599
600impl<
601    const k: usize,
602    const eta1: i16,
603    const LAMBDA: i16,
604    const SK_LEN: usize,
605    const FULL_SK_LEN: usize,
606    const PK_LEN: usize,
607    const T_PACKED_LEN: usize,
608> PartialEq for MLKEMSeedPrivateKey<k, eta1, LAMBDA, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
609{
610    fn eq(&self, other: &Self) -> bool {
611        let self_encoded = self.encode();
612        let other_encoded = other.encode();
613        bouncycastle_utils::ct::ct_eq_bytes(self_encoded.as_ref(), other_encoded.as_ref())
614    }
615}
616
617/// Debug impl mainly to prevent the secret key from being printed in logs.
618impl<
619    const k: usize,
620    const eta1: i16,
621    const LAMBDA: i16,
622    const SK_LEN: usize,
623    const FULL_SK_LEN: usize,
624    const PK_LEN: usize,
625    const T_PACKED_LEN: usize,
626> fmt::Debug for MLKEMSeedPrivateKey<k, eta1, LAMBDA, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
627{
628    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
629        let alg = match k {
630            2 => ML_KEM_512_NAME,
631            3 => ML_KEM_768_NAME,
632            4 => ML_KEM_1024_NAME,
633            _ => panic!("Unsupported key length"),
634        };
635        let pk_hash = self.pk().compute_hash();
636        write!(f, "MLKEMSeedPrivateKey {{ alg: {}, pub_key_hash: {:x?} }}", alg, &pk_hash,)
637    }
638}
639
640/// Display impl mainly to prevent the secret key from being printed in logs.
641impl<
642    const k: usize,
643    const eta1: i16,
644    const LAMBDA: i16,
645    const SK_LEN: usize,
646    const FULL_SK_LEN: usize,
647    const PK_LEN: usize,
648    const T_PACKED_LEN: usize,
649> Display for MLKEMSeedPrivateKey<k, eta1, LAMBDA, SK_LEN, FULL_SK_LEN, PK_LEN, T_PACKED_LEN>
650{
651    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
652        let alg = match k {
653            2 => ML_KEM_512_NAME,
654            3 => ML_KEM_768_NAME,
655            4 => ML_KEM_1024_NAME,
656            _ => panic!("Unsupported key length"),
657        };
658        let pk_hash = self.pk().compute_hash();
659        write!(f, "MLKEMSeedPrivateKey {{ alg: {}, pub_key_hash: {:x?} }}", alg, &pk_hash,)
660    }
661}
bouncycastle_mlkem_lowmemory/mlkem_keys.rs

bouncycastle_mlkem_lowmemory/
mlkem_keys.rs
