bouncycastle_mldsa/
mldsa_keys.rs

1use crate::aux_functions::{
2    bit_pack_eta, bit_pack_t0, bit_unpack_eta, bit_unpack_t0, bitlen_eta, expandA,
3    power_2_round_vec, simple_bit_pack_t1, simple_bit_unpack_t1,
4};
5use crate::matrix::{Matrix, Vector};
6use crate::mldsa::H;
7use crate::mldsa::{MLDSA44_ETA, MLDSA44_PK_LEN, MLDSA44_SK_LEN, MLDSA44_k, MLDSA44_l};
8use crate::mldsa::{MLDSA65_ETA, MLDSA65_PK_LEN, MLDSA65_SK_LEN, MLDSA65_k, MLDSA65_l};
9use crate::mldsa::{MLDSA87_ETA, MLDSA87_PK_LEN, MLDSA87_SK_LEN, MLDSA87_k, MLDSA87_l};
10use crate::mldsa::{POLY_T0PACKED_LEN, POLY_T1PACKED_LEN};
11use crate::{ML_DSA_44_NAME, ML_DSA_65_NAME, ML_DSA_87_NAME};
12use bouncycastle_core::errors::SignatureError;
13use bouncycastle_core::key_material::KeyMaterial;
14use bouncycastle_core::traits::{SignaturePrivateKey, SignaturePublicKey, XOF};
15use bouncycastle_utils::secret::Secret;
16use core::fmt;
17use core::fmt::{Debug, Display, Formatter};
18
19// imports just for docs
20#[allow(unused_imports)]
21use crate::mldsa::MLDSATrait;
22#[allow(unused_imports)]
23use crate::polynomial::Polynomial;
24
25/* Pub Types */
26
27/// ML-DSA-44 Public Key
28pub type MLDSA44PublicKey = MLDSAPublicKey<MLDSA44_k, MLDSA44_l, MLDSA44_PK_LEN>;
29/// ML-DSA-44 Private Key
30pub type MLDSA44PrivateKey =
31    MLDSAPrivateKey<MLDSA44_k, MLDSA44_l, MLDSA44_ETA, MLDSA44_SK_LEN, MLDSA44_PK_LEN>;
32/// ML-DSA-65 Public Key
33pub type MLDSA65PublicKey = MLDSAPublicKey<MLDSA65_k, MLDSA65_l, MLDSA65_PK_LEN>;
34/// ML-DSA-65 Private Key
35pub type MLDSA65PrivateKey =
36    MLDSAPrivateKey<MLDSA65_k, MLDSA65_l, MLDSA65_ETA, MLDSA65_SK_LEN, MLDSA65_PK_LEN>;
37/// ML-DSA-87 Public Key
38pub type MLDSA87PublicKey = MLDSAPublicKey<MLDSA87_k, MLDSA87_l, MLDSA87_PK_LEN>;
39/// ML-DSA-87 Private Key
40pub type MLDSA87PrivateKey =
41    MLDSAPrivateKey<MLDSA87_k, MLDSA87_l, MLDSA87_ETA, MLDSA87_SK_LEN, MLDSA87_PK_LEN>;
42
43/* Pre-expanded keys for repeated operations */
44
45/// ML-DSA-44 Public Key with a pre-expanded public matrix A for repeated encaps operations.
46pub type MLDSA44PublicKeyExpanded =
47    MLDSAPublicKeyExpanded<MLDSA44_k, MLDSA44_l, MLDSA44PublicKey, MLDSA44_PK_LEN>;
48/// ML-DSA-44 Private Key with a pre-expanded public matrix A for repeated decaps operations.
49pub type MLDSA44PrivateKeyExpanded = MLDSAPrivateKeyExpanded<
50    MLDSA44_k,
51    MLDSA44_l,
52    MLDSA44_ETA,
53    MLDSA44PublicKey,
54    MLDSA44PrivateKey,
55    MLDSA44_SK_LEN,
56    MLDSA44_PK_LEN,
57>;
58/// ML-DSA-65 Public Key with a pre-expanded public matrix A for repeated encaps operations.
59pub type MLDSA65PublicKeyExpanded =
60    MLDSAPublicKeyExpanded<MLDSA65_k, MLDSA65_l, MLDSA65PublicKey, MLDSA65_PK_LEN>;
61/// ML-DSA-65 Private Key with a pre-expanded public matrix A for repeated decaps operations.
62pub type MLDSA65PrivateKeyExpanded = MLDSAPrivateKeyExpanded<
63    MLDSA65_k,
64    MLDSA65_l,
65    MLDSA65_ETA,
66    MLDSA65PublicKey,
67    MLDSA65PrivateKey,
68    MLDSA65_SK_LEN,
69    MLDSA65_PK_LEN,
70>;
71/// ML-DSA-87 Public Key with a pre-expanded public matrix A for repeated encaps operations.
72pub type MLDSA87PublicKeyExpanded =
73    MLDSAPublicKeyExpanded<MLDSA87_k, MLDSA87_l, MLDSA87PublicKey, MLDSA87_PK_LEN>;
74/// ML-DSA-87 Private Key with a pre-expanded public matrix A for repeated decaps operations.
75pub type MLDSA87PrivateKeyExpanded = MLDSAPrivateKeyExpanded<
76    MLDSA87_k,
77    MLDSA87_l,
78    MLDSA87_ETA,
79    MLDSA87PublicKey,
80    MLDSA87PrivateKey,
81    MLDSA87_SK_LEN,
82    MLDSA87_PK_LEN,
83>;
84
85/// An ML-DSA public key.
86#[derive(Clone)]
87pub struct MLDSAPublicKey<const k: usize, const l: usize, const PK_LEN: usize> {
88    rho: [u8; 32],
89    t1: Vector<k>,
90}
91
92impl<const k: usize, const l: usize, const PK_LEN: usize> MLDSAPublicKey<k, l, PK_LEN> {
93    /// Algorithm 22 pkEncode(𝜌, 𝐭1)
94    /// Encodes a public key for ML-DSA into a byte string.
95    /// Input:𝜌 ∈ 𝔹32, 𝐭1 ∈ 𝑅𝑘 with coefficients in [0, 2bitlen (𝑞−1)−𝑑 − 1].
96    /// Output: Public key 𝑝𝑘 ∈ 𝔹32+32𝑘(bitlen (𝑞−1)−𝑑).
97    fn pk_encode_out(&self, out: &mut [u8; PK_LEN]) -> usize {
98        out.fill(0);
99
100        out[0..32].copy_from_slice(&self.rho);
101
102        let (pk_chunks, last_chunk) = out[32..].as_chunks_mut::<POLY_T1PACKED_LEN>();
103
104        // that should divide evenly the remainder of the array
105        debug_assert_eq!(pk_chunks.len(), k);
106        debug_assert_eq!(last_chunk.len(), 0);
107
108        for (pk_chunk, t1_i) in pk_chunks.into_iter().zip(&self.t1.vec) {
109            pk_chunk.copy_from_slice(&simple_bit_pack_t1(&t1_i));
110        }
111
112        PK_LEN
113    }
114}
115
116/// General trait for all ML-DSA public keys types.
117pub trait MLDSAPublicKeyTrait<const k: usize, const l: usize, const PK_LEN: usize>:
118    SignaturePublicKey<PK_LEN>
119{
120    /// Algorithm 23 pkDecode(𝑝𝑘)
121    /// Reverses the procedure pkEncode.
122    /// Input: Public key 𝑝𝑘 ∈ 𝔹32+32𝑘(bitlen (𝑞−1)−𝑑).
123    /// Output: 𝜌 ∈ 𝔹32, 𝐭1 ∈ 𝑅𝑘 with coefficients in [0, 2bitlen (𝑞−1)−𝑑 − 1].
124    fn pk_decode(pk: &[u8; PK_LEN]) -> Self;
125
126    /// Get a copy of the expanded public matrix A_hat
127    fn A_hat(&self) -> Matrix<k, l>;
128
129    /// Compute the public key hash (tr) from the public key.
130    ///
131    /// This is exposed as a public API for a few reasons:
132    /// 1. `tr` is required for some external-prehashing schemes such as the so-called "external mu" signing mode.
133    /// 2. `tr` is the canonical fingerprint of an ML-DSA public key, so would be an appropriate value
134    ///     to use, for example, to build a public key lookup or deny-listing table.
135    fn compute_tr(&self) -> [u8; 64];
136}
137
138pub(crate) trait MLDSAPublicKeyInternalTrait<const k: usize, const PK_LEN: usize>:
139    SignaturePublicKey<PK_LEN>
140{
141    /// Not exposing a constructor publicly because you should have to get an instance either by
142    /// running a keygen, or by decoding an existing key.
143    fn new(rho: [u8; 32], t1: Vector<k>) -> Self;
144
145    /// Get a ref to t1
146    fn t1(&self) -> &Vector<k>;
147}
148
149impl<const k: usize, const l: usize, const PK_LEN: usize> MLDSAPublicKeyTrait<k, l, PK_LEN>
150    for MLDSAPublicKey<k, l, PK_LEN>
151{
152    // todo: block a t1 of all zeros? Maybe add to consistency_check() ?
153    fn pk_decode(pk: &[u8; PK_LEN]) -> Self {
154        let rho = pk[0..32].try_into().unwrap();
155        let mut t1 = Vector::<k>::new();
156
157        let (pk_chunks, last_chunk) = pk[32..].as_chunks::<POLY_T1PACKED_LEN>();
158
159        // that should divide evenly the remainder of the array
160        debug_assert_eq!(pk_chunks.len(), k);
161        debug_assert_eq!(last_chunk.len(), 0);
162
163        for (t1_i, pk_chunk) in t1.vec.iter_mut().zip(pk_chunks) {
164            // 3: 𝐭1[𝑖] ← SimpleBitUnpack(𝑧𝑖, 2bitlen (𝑞−1)−𝑑 − 1)
165            //  ▷ This is always in the correct range
166            //  Therefore, we don't need to check that the coeeffs are in range
167            t1_i.coeffs.copy_from_slice(&simple_bit_unpack_t1(pk_chunk).coeffs);
168        }
169
170        Self::new(rho, t1)
171    }
172
173    fn A_hat(&self) -> Matrix<k, l> {
174        expandA::<k, l>(&self.rho)
175    }
176
177    fn compute_tr(&self) -> [u8; 64] {
178        let mut tr = [0u8; 64];
179        H::new().hash_xof_out(&self.encode(), &mut tr);
180
181        tr
182    }
183}
184
185impl<const k: usize, const l: usize, const PK_LEN: usize> MLDSAPublicKeyInternalTrait<k, PK_LEN>
186    for MLDSAPublicKey<k, l, PK_LEN>
187{
188    fn new(rho: [u8; 32], t1: Vector<k>) -> Self {
189        Self { rho, t1 }
190    }
191
192    fn t1(&self) -> &Vector<k> {
193        &self.t1
194    }
195}
196
197impl<const k: usize, const l: usize, const PK_LEN: usize> SignaturePublicKey<PK_LEN>
198    for MLDSAPublicKey<k, l, PK_LEN>
199{
200    fn encode(&self) -> [u8; PK_LEN] {
201        let mut pk = [0u8; PK_LEN];
202        let bytes_written = self.encode_out(&mut pk);
203        debug_assert_eq!(bytes_written, PK_LEN);
204
205        pk
206    }
207
208    fn encode_out(&self, out: &mut [u8; PK_LEN]) -> usize {
209        out.fill(0);
210
211        self.pk_encode_out(out)
212    }
213
214    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
215        if bytes.len() != PK_LEN {
216            return Err(SignatureError::DecodingError(
217                "Provided key bytes are the incorrect length",
218            ));
219        }
220        let bytes_sized: [u8; PK_LEN] = bytes[..PK_LEN].try_into().unwrap();
221        Ok(Self::pk_decode(&bytes_sized))
222    }
223}
224
225impl<const k: usize, const l: usize, const PK_LEN: usize> Eq for MLDSAPublicKey<k, l, PK_LEN> {}
226
227impl<const k: usize, const l: usize, const PK_LEN: usize> PartialEq
228    for MLDSAPublicKey<k, l, PK_LEN>
229{
230    fn eq(&self, other: &Self) -> bool {
231        let self_encoded = self.encode();
232        let other_encoded = other.encode();
233        bouncycastle_utils::ct::ct_eq_bytes(self_encoded.as_ref(), other_encoded.as_ref())
234    }
235}
236
237impl<const k: usize, const l: usize, const PK_LEN: usize> Debug for MLDSAPublicKey<k, l, PK_LEN> {
238    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
239        let alg = match k {
240            4 => ML_DSA_44_NAME,
241            6 => ML_DSA_65_NAME,
242            8 => ML_DSA_87_NAME,
243            _ => panic!("Unsupported key length"),
244        };
245        write!(f, "MLDSAPublicKey {{ alg: {}, pub_key_hash (tr): {:x?} }}", alg, self.compute_tr(),)
246    }
247}
248
249impl<const k: usize, const l: usize, const PK_LEN: usize> Display for MLDSAPublicKey<k, l, PK_LEN> {
250    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
251        let alg = match k {
252            4 => ML_DSA_44_NAME,
253            6 => ML_DSA_65_NAME,
254            8 => ML_DSA_87_NAME,
255            _ => panic!("Unsupported key length"),
256        };
257        write!(f, "MLDSAPublicKey {{ alg: {}, pub_key_hash (tr): {:x?} }}", alg, self.compute_tr(),)
258    }
259}
260
261/// A fully expanded ML-DSA public key that includes the intermediate values needed for performing
262/// multiple verification operations against the same public key, which causes the public key struct
263/// to take up more memory, but results in more efficient repeated verify() operations.
264#[derive(Clone)]
265pub struct MLDSAPublicKeyExpanded<
266    const k: usize,
267    const l: usize,
268    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
269    const PK_LEN: usize,
270> {
271    pub(crate) pk: PK,
272    pub(crate) A_hat: Matrix<k, l>,
273}
274
275impl<
276    const k: usize,
277    const l: usize,
278    PK: MLDSAPublicKeyTrait<k, l, PK_LEN> + MLDSAPublicKeyInternalTrait<k, PK_LEN>,
279    const PK_LEN: usize,
280> SignaturePublicKey<PK_LEN> for MLDSAPublicKeyExpanded<k, l, PK, PK_LEN>
281{
282    fn encode(&self) -> [u8; PK_LEN] {
283        self.pk.encode()
284    }
285
286    fn encode_out(&self, out: &mut [u8; PK_LEN]) -> usize {
287        out.fill(0);
288
289        self.pk.encode_out(out)
290    }
291
292    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
293        if bytes.len() != PK_LEN {
294            return Err(SignatureError::DecodingError(
295                "Provided key bytes are the incorrect length",
296            ));
297        }
298        let bytes_sized: [u8; PK_LEN] = bytes[..PK_LEN].try_into().unwrap();
299        Ok(Self::pk_decode(&bytes_sized))
300    }
301}
302
303impl<
304    const k: usize,
305    const l: usize,
306    PK: MLDSAPublicKeyTrait<k, l, PK_LEN> + MLDSAPublicKeyInternalTrait<k, PK_LEN>,
307    const PK_LEN: usize,
308> PartialEq for MLDSAPublicKeyExpanded<k, l, PK, PK_LEN>
309{
310    fn eq(&self, other: &Self) -> bool {
311        self.pk.eq(&other.pk)
312    }
313}
314
315impl<
316    const k: usize,
317    const l: usize,
318    PK: MLDSAPublicKeyTrait<k, l, PK_LEN> + MLDSAPublicKeyInternalTrait<k, PK_LEN>,
319    const PK_LEN: usize,
320> Eq for MLDSAPublicKeyExpanded<k, l, PK, PK_LEN>
321{
322}
323
324impl<
325    const k: usize,
326    const l: usize,
327    PK: MLDSAPublicKeyTrait<k, l, PK_LEN> + MLDSAPublicKeyInternalTrait<k, PK_LEN>,
328    const PK_LEN: usize,
329> Debug for MLDSAPublicKeyExpanded<k, l, PK, PK_LEN>
330{
331    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
332        let alg = match k {
333            4 => ML_DSA_44_NAME,
334            6 => ML_DSA_65_NAME,
335            8 => ML_DSA_87_NAME,
336            _ => panic!("Unsupported key length"),
337        };
338        write!(
339            f,
340            "MLDSAPublicKeyExpanded {{ alg: {}, pub_key_hash (tr): {:x?} }}",
341            alg,
342            self.compute_tr(),
343        )
344    }
345}
346
347impl<
348    const k: usize,
349    const l: usize,
350    PK: MLDSAPublicKeyTrait<k, l, PK_LEN> + MLDSAPublicKeyInternalTrait<k, PK_LEN>,
351    const PK_LEN: usize,
352> Display for MLDSAPublicKeyExpanded<k, l, PK, PK_LEN>
353{
354    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
355        let alg = match k {
356            4 => ML_DSA_44_NAME,
357            6 => ML_DSA_65_NAME,
358            8 => ML_DSA_87_NAME,
359            _ => panic!("Unsupported key length"),
360        };
361        write!(
362            f,
363            "MLDSAPublicKeyExpanded {{ alg: {}, pub_key_hash (tr): {:x?} }}",
364            alg,
365            self.compute_tr(),
366        )
367    }
368}
369
370impl<
371    const k: usize,
372    const l: usize,
373    PK: MLDSAPublicKeyTrait<k, l, PK_LEN> + MLDSAPublicKeyInternalTrait<k, PK_LEN>,
374    const PK_LEN: usize,
375> From<&PK> for MLDSAPublicKeyExpanded<k, l, PK, PK_LEN>
376{
377    /// Fully expands the intermediate values needed for performing multiple encaps operations
378    /// against the same public key, which causes the MLKEMPublicKey struct to take up
379    fn from(pk: &PK) -> Self {
380        let A_hat = pk.A_hat();
381
382        Self { pk: pk.clone(), A_hat }
383    }
384}
385
386impl<
387    const k: usize,
388    const l: usize,
389    PK: MLDSAPublicKeyTrait<k, l, PK_LEN> + MLDSAPublicKeyInternalTrait<k, PK_LEN>,
390    const PK_LEN: usize,
391> MLDSAPublicKeyTrait<k, l, PK_LEN> for MLDSAPublicKeyExpanded<k, l, PK, PK_LEN>
392{
393    fn pk_decode(pk: &[u8; PK_LEN]) -> Self {
394        let pk1 = PK::pk_decode(pk);
395        let A_hat = pk1.A_hat();
396        Self { pk: pk1, A_hat }
397    }
398
399    fn A_hat(&self) -> Matrix<k, l> {
400        self.A_hat.clone()
401    }
402
403    fn compute_tr(&self) -> [u8; 64] {
404        self.pk.compute_tr()
405    }
406}
407
408/// An ML-DSA private key.
409///
410/// This will automatically inherit the [`Secret`] protections because [`Polynomial`] wraps the underlying data with [`Secret`].
411#[derive(Clone)]
412pub struct MLDSAPrivateKey<
413    const k: usize,
414    const l: usize,
415    const eta: usize,
416    const SK_LEN: usize,
417    const PK_LEN: usize,
418> {
419    rho: [u8; 32],
420    K: Secret<[u8; 32]>,
421    tr: [u8; 64],
422    // Deviation from the FIPS:
423    //  s1, s2, and t0 are only ever used in their ntt form; the only time they need to be in their
424    //  natural domain form is when encoding or decoding to the standardized byte representation.
425    //  So we are going to hold them as s1_hat, s2_hat, and t0_hat.
426    //  Note: these are not necessarily in their reduced form; so you'll need to reduce them before
427    //  inv_ntt()'ing them or hashing them.
428    s1_hat: Secret<Vector<l>>,
429    s2_hat: Secret<Vector<k>>,
430    t0_hat: Vector<k>,
431    // note: KeyMaterial is inherently Secret
432    seed: Option<KeyMaterial<32>>,
433}
434
435impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
436    MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
437{
438    /// Algorithm 24 skEncode(𝜌, 𝐾, 𝑡𝑟, 𝐬1, 𝐬2, 𝐭0)
439    /// Encodes a secret key for ML-DSA into a byte string.
440    /// Input: 𝜌 ∈ 𝔹32, 𝐾 ∈ 𝔹32, 𝑡𝑟 ∈ 𝔹64 , 𝐬1 ∈ 𝑅ℓ with coefficients in [−𝜂, 𝜂], 𝐬2 ∈ 𝑅𝑘 with
441    /// coefficients in [−𝜂, 𝜂], 𝐭0 ∈ 𝑅𝑘 with coefficients in [−2𝑑−1 + 1, 2𝑑−1].
442    /// Output: Private key 𝑠𝑘 ∈ 𝔹32+32+64+32⋅((𝑘+ℓ)⋅bitlen (2𝜂)+𝑑𝑘).
443    fn sk_encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
444        out.fill(0);
445
446        // counter of progress along the output buffer
447        let mut off: usize = 0;
448
449        out[0..32].copy_from_slice(&self.rho);
450        out[32..64].copy_from_slice(&*self.K);
451        out[64..128].copy_from_slice(&self.tr);
452        off += 128;
453
454        let mut buf = [0u8; 32 * 4]; // largest possible buffer
455        let eta_pack_len = bitlen_eta(eta);
456
457        let sk_chunks = out[off..off + l * bitlen_eta(eta)].chunks_mut(bitlen_eta(eta));
458        debug_assert_eq!(sk_chunks.len(), l);
459        for (sk_chunk, s1_hat_i) in sk_chunks.into_iter().zip(&self.s1_hat.vec) {
460            // Deviation from the FIPS:
461            //   We are holding these in ntt form, so need to convert back to standard form
462            let mut s1_hat_i = s1_hat_i.clone();
463            s1_hat_i.reduce();
464            s1_hat_i.inv_ntt();
465            let s1_i = s1_hat_i;
466
467            bit_pack_eta::<eta>(&s1_i, &mut buf);
468            sk_chunk.copy_from_slice(&buf[..eta_pack_len]);
469        }
470        off += l * bitlen_eta(eta);
471
472        let sk_chunks = out[off..off + k * bitlen_eta(eta)].chunks_mut(bitlen_eta(eta));
473        debug_assert_eq!(sk_chunks.len(), k);
474        for (sk_chunk, s2_hat_i) in sk_chunks.into_iter().zip(&self.s2_hat.vec) {
475            // Deviation from the FIPS:
476            //   We are holding these in ntt form, so need to convert back to standard form
477            let mut s2_hat_i = s2_hat_i.clone();
478            s2_hat_i.reduce();
479            s2_hat_i.inv_ntt();
480            let s2_i = s2_hat_i;
481
482            bit_pack_eta::<eta>(&s2_i, &mut buf);
483            sk_chunk.copy_from_slice(&buf[..eta_pack_len]);
484        }
485        off += k * bitlen_eta(eta);
486
487        let sk_chunks = out[off..off + k * POLY_T0PACKED_LEN].chunks_mut(POLY_T0PACKED_LEN);
488        debug_assert_eq!(sk_chunks.len(), k);
489        for (sk_chunk, t0_hat_i) in sk_chunks.into_iter().zip(&self.t0_hat.vec) {
490            // Deviation from the FIPS:
491            //   We are holding these in ntt form, so need to convert back to standard form
492            let mut t0_hat_i = t0_hat_i.clone();
493            t0_hat_i.reduce();
494            t0_hat_i.inv_ntt();
495            let t0_i = t0_hat_i;
496
497            sk_chunk.copy_from_slice(&bit_pack_t0(&t0_i));
498        }
499
500        SK_LEN
501    }
502}
503
504/// General trait for all ML-DSA private keys types.
505pub trait MLDSAPrivateKeyTrait<
506    const k: usize,
507    const l: usize,
508    const eta: usize,
509    const SK_LEN: usize,
510    const PK_LEN: usize,
511>: SignaturePrivateKey<SK_LEN>
512{
513    /// Get a ref to the seed, if there is one stored with this private key
514    fn seed(&self) -> Option<&KeyMaterial<32>>;
515
516    /// Get a ref to the key hash `tr`.
517    fn tr(&self) -> &[u8; 64];
518
519    /// Get the public matrix A_hat.
520    fn A_hat(&self) -> Matrix<k, l>;
521
522    /// This is a partial implementation of keygen_internal(), and probably not allowed in FIPS mode.
523    fn derive_pk(&self) -> MLDSAPublicKey<k, l, PK_LEN>;
524    /// Algorithm 25 skDecode(𝑠𝑘)
525    /// Reverses the procedure skEncode.
526    /// Input: Private key 𝑠𝑘 ∈ 𝔹32+32+64+32⋅((ℓ+𝑘)⋅bitlen (2𝜂)+𝑑𝑘).
527    /// Output: 𝜌 ∈ 𝔹32, 𝐾 ∈ 𝔹32, 𝑡𝑟 ∈ 𝔹64 ,
528    /// 𝐬1 ∈ 𝑅ℓ, 𝐬2 ∈ 𝑅𝑘, 𝐭0 ∈ 𝑅𝑘 with coefficients in [−2𝑑−1 + 1, 2𝑑−1].
529    ///
530    /// Note: this object contains only the simple decoding routine to unpack a semi-expanded key.
531    /// See [`MLDSATrait`] for key generation functions, including derive-from-seed and consistency-check functions.
532    fn sk_decode(sk: &[u8; SK_LEN]) -> Result<Self, SignatureError>;
533}
534
535pub(crate) trait MLDSAPrivateKeyInternalTrait<
536    const k: usize,
537    const l: usize,
538    const eta: usize,
539    const SK_LEN: usize,
540    const PK_LEN: usize,
541>
542{
543    /// Not exposing a constructor publicly because you should have to get an instance either by
544    /// running a keygen, or by decoding an existing key.
545    fn new(
546        rho: [u8; 32],
547        K: Secret<[u8; 32]>,
548        tr: [u8; 64],
549        s1_hat: Secret<Vector<l>>,
550        s2_hat: Secret<Vector<k>>,
551        t0_hat: Vector<k>,
552        seed: Option<KeyMaterial<32>>,
553    ) -> Self;
554    /// Get a ref to K
555    fn K(&self) -> &Secret<[u8; 32]>;
556    /// Get a ref to s1
557    fn s1_hat(&self) -> &Vector<l>;
558    /// Get a ref to s2
559    fn s2_hat(&self) -> &Vector<k>;
560    /// Get a ref to t0
561    fn t0_hat(&self) -> &Vector<k>;
562}
563
564impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
565    MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN> for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
566{
567    fn seed(&self) -> Option<&KeyMaterial<32>> {
568        match self.seed {
569            Some(_) => self.seed.as_ref(),
570            None => None,
571        }
572    }
573
574    fn tr(&self) -> &[u8; 64] {
575        &self.tr
576    }
577
578    fn A_hat(&self) -> Matrix<k, l> {
579        expandA::<k, l>(&self.rho)
580    }
581
582    fn derive_pk(&self) -> MLDSAPublicKey<k, l, PK_LEN> {
583        // 5: 𝐭 ← NTT−1(𝐀 ∘ NTT(𝐬1)) + 𝐬2
584        //   ▷ compute 𝐭 = 𝐀𝐬1 + 𝐬2
585        let mut t = {
586            // scope for A_hat
587            // 3: 𝐀 ← ExpandA(𝜌)
588            //   ▷ 𝐀 is generated and stored in NTT representation as 𝐀
589            let A_hat = expandA::<k, l>(&self.rho);
590
591            let mut t_ntt = A_hat.matrix_vector_ntt(&self.s1_hat);
592            t_ntt.inv_ntt();
593            t_ntt
594        };
595
596        {
597            // Deviation from the FIPS:
598            // Because s2 is in ntt form, it is necessary to reverse that here before adding it to t
599            let mut s2 = self.s2_hat.clone();
600            s2.reduce();
601            s2.inv_ntt();
602
603            t.add_vector_ntt(&s2);
604            t.conditional_add_q();
605        }
606        // 6: (𝐭1, 𝐭0) ← Power2Round(𝐭)
607        //   ▷ compress 𝐭
608        //   ▷ PowerTwoRound is applied componentwise (see explanatory text in Section 7.4)
609        let (t1, _) = power_2_round_vec::<k>(&t);
610
611        MLDSAPublicKey::<k, l, PK_LEN>::new(self.rho.clone(), t1)
612    }
613    fn sk_decode(sk: &[u8; SK_LEN]) -> Result<Self, SignatureError> {
614        // Construct the (Secret-protected) key up front and unpack each field directly into it,
615        // rather than decoding into unprotected temporaries and copying them in at the end. This
616        // way the secret material is written straight into its protected home; and if a range
617        // check below fails, `key` is dropped and its `Secret` fields are zeroized on the way out.
618        let mut key = Self {
619            rho: sk[0..32].try_into().unwrap(),
620            K: Secret::new(),
621            tr: sk[64..128].try_into().unwrap(),
622            s1_hat: Secret::new(),
623            s2_hat: Secret::new(),
624            t0_hat: Vector::<k>::new(),
625            seed: None,
626        };
627        key.K.copy_from_slice(&sk[32..64]);
628        let mut off = 128;
629
630        // unpack s1 directly into key.s1_hat so that we don't make additional non-Secret copies.
631        let sk_chunks = sk[off..off + (l * bitlen_eta(eta))].chunks(bitlen_eta(eta));
632        debug_assert_eq!(sk_chunks.len(), l);
633        for (s1_i, sk_chunk) in key.s1_hat.vec.iter_mut().zip(sk_chunks) {
634            // 3: 𝐬1[𝑖] ← BitUnpack(𝑦𝑖, 𝜂, 𝜂)
635            //  ▷ this may lie outside [−𝜂, 𝜂] if input is malformed
636            s1_i.coeffs.copy_from_slice(&bit_unpack_eta::<eta>(&sk_chunk).coeffs);
637
638            // check that the coefficients are within the expected range
639            for coeff in s1_i.coeffs.iter() {
640                if *coeff < -(eta as i32) || *coeff > (eta as i32) {
641                    return Err(SignatureError::DecodingError("Invalid or corrupted key"));
642                }
643            }
644        }
645        // Deviation from the FIPS:
646        //   Convert this to ntt form as part of decode
647        key.s1_hat.ntt();
648        off += l * bitlen_eta(eta);
649
650        // unpack s2 directly into key.s2_hat so that we don't make additional non-Secret copies.
651        let sk_chunks = sk[off..off + (k * bitlen_eta(eta))].chunks(bitlen_eta(eta));
652        debug_assert_eq!(sk_chunks.len(), k);
653        for (s2_i, sk_chunk) in key.s2_hat.vec.iter_mut().zip(sk_chunks) {
654            // 6: 𝐬2[𝑖] ← BitUnpack(𝑧𝑖, 𝜂, 𝜂)
655            //  ▷ this may lie outside [−𝜂, 𝜂] if input is malformed
656            s2_i.coeffs.copy_from_slice(&bit_unpack_eta::<eta>(&sk_chunk).coeffs);
657
658            // check that the coefficients are within the expected range
659            for coeff in s2_i.coeffs.iter() {
660                if *coeff < -(eta as i32) || *coeff > (eta as i32) {
661                    return Err(SignatureError::DecodingError("Invalid or corrupted key"));
662                }
663            }
664        }
665        // Deviation from the FIPS:
666        // Convert this to ntt form as part of decode
667        key.s2_hat.ntt();
668        off += k * bitlen_eta(eta);
669
670        // unpack t0 directly into key.t0_hat
671        let (sk_chunks, last_chunk) =
672            sk[off..off + (k * POLY_T0PACKED_LEN)].as_chunks::<POLY_T0PACKED_LEN>();
673
674        // that should divide evenly the remainder of the array
675        debug_assert_eq!(sk_chunks.len(), k);
676        debug_assert_eq!(last_chunk.len(), 0);
677
678        for (t0_i, sk_chunk) in key.t0_hat.vec.iter_mut().zip(sk_chunks) {
679            t0_i.coeffs.copy_from_slice(&bit_unpack_t0(sk_chunk).coeffs);
680        }
681        // Deviation from the FIPS:
682        // Convert this to ntt form as part of decode
683        key.t0_hat.ntt();
684
685        Ok(key)
686    }
687}
688
689impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
690    MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>
691    for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
692{
693    fn new(
694        rho: [u8; 32],
695        K: Secret<[u8; 32]>,
696        tr: [u8; 64],
697        s1_hat: Secret<Vector<l>>,
698        s2_hat: Secret<Vector<k>>,
699        t0_hat: Vector<k>,
700        seed: Option<KeyMaterial<32>>,
701    ) -> Self {
702        Self {
703            rho: rho.clone(),
704            K: K.clone(),
705            tr: tr.clone(),
706            s1_hat: s1_hat.clone(),
707            s2_hat: s2_hat.clone(),
708            t0_hat: t0_hat.clone(),
709            seed: seed.clone(),
710        }
711    }
712
713    fn K(&self) -> &Secret<[u8; 32]> {
714        &self.K
715    }
716
717    fn s1_hat(&self) -> &Vector<l> {
718        &self.s1_hat
719    }
720
721    fn s2_hat(&self) -> &Vector<k> {
722        &self.s2_hat
723    }
724
725    fn t0_hat(&self) -> &Vector<k> {
726        &self.t0_hat
727    }
728}
729
730impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
731    SignaturePrivateKey<SK_LEN> for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
732{
733    fn encode(&self) -> [u8; SK_LEN] {
734        let mut out = [0u8; SK_LEN];
735        let bytes_written = self.sk_encode_out(&mut out);
736        debug_assert_eq!(bytes_written, SK_LEN);
737
738        out
739    }
740
741    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
742        out.fill(0);
743
744        self.sk_encode_out(out)
745    }
746
747    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
748        if bytes.len() != SK_LEN {
749            return Err(SignatureError::DecodingError(
750                "Provided key bytes are the incorrect length",
751            ));
752        }
753        let bytes_sized: [u8; SK_LEN] = bytes[..SK_LEN].try_into().unwrap();
754
755        Ok(Self::sk_decode(&bytes_sized)?)
756    }
757}
758
759impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize> Eq
760    for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
761{
762}
763
764impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
765    PartialEq for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
766{
767    fn eq(&self, other: &Self) -> bool {
768        let self_encoded = self.encode();
769        let other_encoded = other.encode();
770        bouncycastle_utils::ct::ct_eq_bytes(self_encoded.as_ref(), other_encoded.as_ref())
771    }
772}
773
774/// Debug impl mainly to prevent the secret key from being printed in logs.
775impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
776    fmt::Debug for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
777{
778    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
779        let alg = match k {
780            4 => ML_DSA_44_NAME,
781            6 => ML_DSA_65_NAME,
782            8 => ML_DSA_87_NAME,
783            _ => panic!("Unsupported key length"),
784        };
785        write!(
786            f,
787            "MLDSAPrivateKey {{ alg: {}, pub_key_hash (tr): {:x?}, has_seed: {} }}",
788            alg,
789            self.tr,
790            self.seed.is_some(),
791        )
792    }
793}
794
795/// Display impl mainly to prevent the secret key from being printed in logs.
796impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
797    Display for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
798{
799    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
800        let alg = match k {
801            4 => ML_DSA_44_NAME,
802            6 => ML_DSA_65_NAME,
803            8 => ML_DSA_87_NAME,
804            _ => panic!("Unsupported key length"),
805        };
806        write!(
807            f,
808            "MLDSAPrivateKey {{ alg: {}, pub_key_hash (tr): {:x?}, has_seed: {} }}",
809            alg,
810            self.tr,
811            self.seed.is_some(),
812        )
813    }
814}
815
816/// A fully expanded ML-DSA private key that includes the intermediate values needed for performing
817/// multiple sign operations with the same private key, which causes the private ey struct to take up
818/// more memory, but results in more efficient repeated sign() operations.
819#[derive(Clone)]
820pub struct MLDSAPrivateKeyExpanded<
821    const k: usize,
822    const l: usize,
823    const eta: usize,
824    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
825    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
826        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
827    const SK_LEN: usize,
828    const PK_LEN: usize,
829> {
830    _phantom: core::marker::PhantomData<PK>,
831    pub(crate) sk: SK,
832    pub(crate) A_hat: Matrix<k, l>,
833}
834
835impl<
836    const k: usize,
837    const l: usize,
838    const eta: usize,
839    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
840    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
841        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
842    const SK_LEN: usize,
843    const PK_LEN: usize,
844> PartialEq for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
845{
846    fn eq(&self, other: &Self) -> bool {
847        self.sk.eq(&other.sk)
848    }
849}
850
851impl<
852    const k: usize,
853    const l: usize,
854    const eta: usize,
855    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
856    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
857        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
858    const SK_LEN: usize,
859    const PK_LEN: usize,
860> Eq for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
861{
862}
863
864impl<
865    const k: usize,
866    const l: usize,
867    const eta: usize,
868    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
869    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
870        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
871    const SK_LEN: usize,
872    const PK_LEN: usize,
873> Debug for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
874{
875    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
876        let alg = match k {
877            4 => ML_DSA_44_NAME,
878            6 => ML_DSA_65_NAME,
879            8 => ML_DSA_87_NAME,
880            _ => panic!("Unsupported key length"),
881        };
882        write!(
883            f,
884            "MLDSAPrivateKeyExpanded {{ alg: {}, pub_key_hash (tr): {:x?}, has_seed: {} }}",
885            alg,
886            self.sk.tr(),
887            self.sk.seed().is_some(),
888        )
889    }
890}
891
892impl<
893    const k: usize,
894    const l: usize,
895    const eta: usize,
896    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
897    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
898        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
899    const SK_LEN: usize,
900    const PK_LEN: usize,
901> Display for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
902{
903    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
904        let alg = match k {
905            4 => ML_DSA_44_NAME,
906            6 => ML_DSA_65_NAME,
907            8 => ML_DSA_87_NAME,
908            _ => panic!("Unsupported key length"),
909        };
910        write!(
911            f,
912            "MLDSAPrivateKeyExpanded {{ alg: {}, pub_key_hash (tr): {:x?}, has_seed: {} }}",
913            alg,
914            self.sk.tr(),
915            self.sk.seed().is_some(),
916        )
917    }
918}
919
920impl<
921    const k: usize,
922    const l: usize,
923    const eta: usize,
924    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
925    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
926        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
927    const SK_LEN: usize,
928    const PK_LEN: usize,
929> From<&SK> for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
930{
931    /// Fully expands the intermediate values needed for performing multiple encaps operations
932    /// against the same public key, which causes the MLKEMPublicKey struct to take up
933    fn from(sk: &SK) -> Self {
934        let A_hat = sk.derive_pk().A_hat();
935
936        Self { _phantom: core::marker::PhantomData, sk: sk.clone(), A_hat }
937    }
938}
939
940impl<
941    const k: usize,
942    const l: usize,
943    const eta: usize,
944    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
945    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
946        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
947    const SK_LEN: usize,
948    const PK_LEN: usize,
949> SignaturePrivateKey<SK_LEN> for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
950{
951    fn encode(&self) -> [u8; SK_LEN] {
952        self.sk.encode()
953    }
954
955    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
956        out.fill(0);
957
958        self.sk.encode_out(out)
959    }
960
961    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
962        let sk = SK::from_bytes(bytes)?;
963        Ok(Self::from(&sk))
964    }
965}
966
967impl<
968    const k: usize,
969    const l: usize,
970    const eta: usize,
971    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
972    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
973        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
974    const SK_LEN: usize,
975    const PK_LEN: usize,
976> MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
977    for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
978{
979    fn seed(&self) -> Option<&KeyMaterial<32>> {
980        self.sk.seed()
981    }
982
983    fn tr(&self) -> &[u8; 64] {
984        self.sk.tr()
985    }
986
987    fn A_hat(&self) -> Matrix<k, l> {
988        self.sk.A_hat()
989    }
990
991    fn derive_pk(&self) -> MLDSAPublicKey<k, l, PK_LEN> {
992        self.sk.derive_pk()
993    }
994
995    fn sk_decode(sk: &[u8; SK_LEN]) -> Result<Self, SignatureError> {
996        let sk1 = SK::sk_decode(sk)?;
997        let A_hat = sk1.derive_pk().A_hat();
998
999        Ok(Self { _phantom: core::marker::PhantomData, sk: sk1, A_hat })
1000    }
1001}
bouncycastle_mldsa/mldsa_keys.rs

bouncycastle_mldsa/
mldsa_keys.rs
