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            // Bound-keeping step before NTT⁻¹, not in FIPS 204: see [`Polynomial::reduce32`].
593            t_ntt.reduce32();
594            t_ntt.inv_ntt();
595            t_ntt
596        };
597
598        {
599            // Deviation from the FIPS:
600            // Because s2 is in ntt form, it is necessary to reverse that here before adding it to t
601            let mut s2 = self.s2_hat.clone();
602            s2.reduce();
603            s2.inv_ntt();
604
605            t.add_vector_ntt(&s2);
606            t.conditional_add_q();
607        }
608        // 6: (𝐭1, 𝐭0) ← Power2Round(𝐭)
609        //   ▷ compress 𝐭
610        //   ▷ PowerTwoRound is applied componentwise (see explanatory text in Section 7.4)
611        let (t1, _) = power_2_round_vec::<k>(&t);
612
613        MLDSAPublicKey::<k, l, PK_LEN>::new(self.rho.clone(), t1)
614    }
615    fn sk_decode(sk: &[u8; SK_LEN]) -> Result<Self, SignatureError> {
616        // Construct the (Secret-protected) key up front and unpack each field directly into it,
617        // rather than decoding into unprotected temporaries and copying them in at the end. This
618        // way the secret material is written straight into its protected home; and if a range
619        // check below fails, `key` is dropped and its `Secret` fields are zeroized on the way out.
620        let mut key = Self {
621            rho: sk[0..32].try_into().unwrap(),
622            K: Secret::new(),
623            tr: sk[64..128].try_into().unwrap(),
624            s1_hat: Secret::new(),
625            s2_hat: Secret::new(),
626            t0_hat: Vector::<k>::new(),
627            seed: None,
628        };
629        key.K.copy_from_slice(&sk[32..64]);
630        let mut off = 128;
631
632        // unpack s1 directly into key.s1_hat so that we don't make additional non-Secret copies.
633        let sk_chunks = sk[off..off + (l * bitlen_eta(eta))].chunks(bitlen_eta(eta));
634        debug_assert_eq!(sk_chunks.len(), l);
635        for (s1_i, sk_chunk) in key.s1_hat.vec.iter_mut().zip(sk_chunks) {
636            // 3: 𝐬1[𝑖] ← BitUnpack(𝑦𝑖, 𝜂, 𝜂)
637            //  ▷ this may lie outside [−𝜂, 𝜂] if input is malformed
638            s1_i.coeffs.copy_from_slice(&bit_unpack_eta::<eta>(&sk_chunk).coeffs);
639
640            // check that the coefficients are within the expected range
641            for coeff in s1_i.coeffs.iter() {
642                if *coeff < -(eta as i32) || *coeff > (eta as i32) {
643                    return Err(SignatureError::DecodingError("Invalid or corrupted key"));
644                }
645            }
646        }
647        // Deviation from the FIPS:
648        //   Convert this to ntt form as part of decode
649        key.s1_hat.ntt();
650        off += l * bitlen_eta(eta);
651
652        // unpack s2 directly into key.s2_hat so that we don't make additional non-Secret copies.
653        let sk_chunks = sk[off..off + (k * bitlen_eta(eta))].chunks(bitlen_eta(eta));
654        debug_assert_eq!(sk_chunks.len(), k);
655        for (s2_i, sk_chunk) in key.s2_hat.vec.iter_mut().zip(sk_chunks) {
656            // 6: 𝐬2[𝑖] ← BitUnpack(𝑧𝑖, 𝜂, 𝜂)
657            //  ▷ this may lie outside [−𝜂, 𝜂] if input is malformed
658            s2_i.coeffs.copy_from_slice(&bit_unpack_eta::<eta>(&sk_chunk).coeffs);
659
660            // check that the coefficients are within the expected range
661            for coeff in s2_i.coeffs.iter() {
662                if *coeff < -(eta as i32) || *coeff > (eta as i32) {
663                    return Err(SignatureError::DecodingError("Invalid or corrupted key"));
664                }
665            }
666        }
667        // Deviation from the FIPS:
668        // Convert this to ntt form as part of decode
669        key.s2_hat.ntt();
670        off += k * bitlen_eta(eta);
671
672        // unpack t0 directly into key.t0_hat
673        let (sk_chunks, last_chunk) =
674            sk[off..off + (k * POLY_T0PACKED_LEN)].as_chunks::<POLY_T0PACKED_LEN>();
675
676        // that should divide evenly the remainder of the array
677        debug_assert_eq!(sk_chunks.len(), k);
678        debug_assert_eq!(last_chunk.len(), 0);
679
680        for (t0_i, sk_chunk) in key.t0_hat.vec.iter_mut().zip(sk_chunks) {
681            t0_i.coeffs.copy_from_slice(&bit_unpack_t0(sk_chunk).coeffs);
682        }
683        // Deviation from the FIPS:
684        // Convert this to ntt form as part of decode
685        key.t0_hat.ntt();
686
687        Ok(key)
688    }
689}
690
691impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
692    MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>
693    for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
694{
695    fn new(
696        rho: [u8; 32],
697        K: Secret<[u8; 32]>,
698        tr: [u8; 64],
699        s1_hat: Secret<Vector<l>>,
700        s2_hat: Secret<Vector<k>>,
701        t0_hat: Vector<k>,
702        seed: Option<KeyMaterial<32>>,
703    ) -> Self {
704        Self {
705            rho: rho.clone(),
706            K: K.clone(),
707            tr: tr.clone(),
708            s1_hat: s1_hat.clone(),
709            s2_hat: s2_hat.clone(),
710            t0_hat: t0_hat.clone(),
711            seed: seed.clone(),
712        }
713    }
714
715    fn K(&self) -> &Secret<[u8; 32]> {
716        &self.K
717    }
718
719    fn s1_hat(&self) -> &Vector<l> {
720        &self.s1_hat
721    }
722
723    fn s2_hat(&self) -> &Vector<k> {
724        &self.s2_hat
725    }
726
727    fn t0_hat(&self) -> &Vector<k> {
728        &self.t0_hat
729    }
730}
731
732impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
733    SignaturePrivateKey<SK_LEN> for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
734{
735    fn encode(&self) -> [u8; SK_LEN] {
736        let mut out = [0u8; SK_LEN];
737        let bytes_written = self.sk_encode_out(&mut out);
738        debug_assert_eq!(bytes_written, SK_LEN);
739
740        out
741    }
742
743    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
744        out.fill(0);
745
746        self.sk_encode_out(out)
747    }
748
749    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
750        if bytes.len() != SK_LEN {
751            return Err(SignatureError::DecodingError(
752                "Provided key bytes are the incorrect length",
753            ));
754        }
755        let bytes_sized: [u8; SK_LEN] = bytes[..SK_LEN].try_into().unwrap();
756
757        Ok(Self::sk_decode(&bytes_sized)?)
758    }
759}
760
761impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize> Eq
762    for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
763{
764}
765
766impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
767    PartialEq for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
768{
769    fn eq(&self, other: &Self) -> bool {
770        let self_encoded = self.encode();
771        let other_encoded = other.encode();
772        bouncycastle_utils::ct::ct_eq_bytes(self_encoded.as_ref(), other_encoded.as_ref())
773    }
774}
775
776/// Debug impl mainly to prevent the secret key from being printed in logs.
777impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
778    fmt::Debug for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
779{
780    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
781        let alg = match k {
782            4 => ML_DSA_44_NAME,
783            6 => ML_DSA_65_NAME,
784            8 => ML_DSA_87_NAME,
785            _ => panic!("Unsupported key length"),
786        };
787        write!(
788            f,
789            "MLDSAPrivateKey {{ alg: {}, pub_key_hash (tr): {:x?}, has_seed: {} }}",
790            alg,
791            self.tr,
792            self.seed.is_some(),
793        )
794    }
795}
796
797/// Display impl mainly to prevent the secret key from being printed in logs.
798impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, const PK_LEN: usize>
799    Display for MLDSAPrivateKey<k, l, eta, SK_LEN, PK_LEN>
800{
801    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
802        let alg = match k {
803            4 => ML_DSA_44_NAME,
804            6 => ML_DSA_65_NAME,
805            8 => ML_DSA_87_NAME,
806            _ => panic!("Unsupported key length"),
807        };
808        write!(
809            f,
810            "MLDSAPrivateKey {{ alg: {}, pub_key_hash (tr): {:x?}, has_seed: {} }}",
811            alg,
812            self.tr,
813            self.seed.is_some(),
814        )
815    }
816}
817
818/// A fully expanded ML-DSA private key that includes the intermediate values needed for performing
819/// multiple sign operations with the same private key, which causes the private ey struct to take up
820/// more memory, but results in more efficient repeated sign() operations.
821#[derive(Clone)]
822pub struct MLDSAPrivateKeyExpanded<
823    const k: usize,
824    const l: usize,
825    const eta: usize,
826    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
827    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
828        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
829    const SK_LEN: usize,
830    const PK_LEN: usize,
831> {
832    _phantom: core::marker::PhantomData<PK>,
833    pub(crate) sk: SK,
834    pub(crate) A_hat: Matrix<k, l>,
835}
836
837impl<
838    const k: usize,
839    const l: usize,
840    const eta: usize,
841    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
842    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
843        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
844    const SK_LEN: usize,
845    const PK_LEN: usize,
846> PartialEq for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
847{
848    fn eq(&self, other: &Self) -> bool {
849        self.sk.eq(&other.sk)
850    }
851}
852
853impl<
854    const k: usize,
855    const l: usize,
856    const eta: usize,
857    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
858    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
859        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
860    const SK_LEN: usize,
861    const PK_LEN: usize,
862> Eq for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
863{
864}
865
866impl<
867    const k: usize,
868    const l: usize,
869    const eta: usize,
870    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
871    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
872        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
873    const SK_LEN: usize,
874    const PK_LEN: usize,
875> Debug for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
876{
877    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
878        let alg = match k {
879            4 => ML_DSA_44_NAME,
880            6 => ML_DSA_65_NAME,
881            8 => ML_DSA_87_NAME,
882            _ => panic!("Unsupported key length"),
883        };
884        write!(
885            f,
886            "MLDSAPrivateKeyExpanded {{ alg: {}, pub_key_hash (tr): {:x?}, has_seed: {} }}",
887            alg,
888            self.sk.tr(),
889            self.sk.seed().is_some(),
890        )
891    }
892}
893
894impl<
895    const k: usize,
896    const l: usize,
897    const eta: usize,
898    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
899    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
900        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
901    const SK_LEN: usize,
902    const PK_LEN: usize,
903> Display for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
904{
905    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
906        let alg = match k {
907            4 => ML_DSA_44_NAME,
908            6 => ML_DSA_65_NAME,
909            8 => ML_DSA_87_NAME,
910            _ => panic!("Unsupported key length"),
911        };
912        write!(
913            f,
914            "MLDSAPrivateKeyExpanded {{ alg: {}, pub_key_hash (tr): {:x?}, has_seed: {} }}",
915            alg,
916            self.sk.tr(),
917            self.sk.seed().is_some(),
918        )
919    }
920}
921
922impl<
923    const k: usize,
924    const l: usize,
925    const eta: usize,
926    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
927    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
928        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
929    const SK_LEN: usize,
930    const PK_LEN: usize,
931> From<&SK> for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
932{
933    /// Fully expands the intermediate values needed for performing multiple encaps operations
934    /// against the same public key, which causes the MLKEMPublicKey struct to take up
935    fn from(sk: &SK) -> Self {
936        let A_hat = sk.derive_pk().A_hat();
937
938        Self { _phantom: core::marker::PhantomData, sk: sk.clone(), A_hat }
939    }
940}
941
942impl<
943    const k: usize,
944    const l: usize,
945    const eta: usize,
946    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
947    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
948        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
949    const SK_LEN: usize,
950    const PK_LEN: usize,
951> SignaturePrivateKey<SK_LEN> for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
952{
953    fn encode(&self) -> [u8; SK_LEN] {
954        self.sk.encode()
955    }
956
957    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
958        out.fill(0);
959
960        self.sk.encode_out(out)
961    }
962
963    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
964        let sk = SK::from_bytes(bytes)?;
965        Ok(Self::from(&sk))
966    }
967}
968
969impl<
970    const k: usize,
971    const l: usize,
972    const eta: usize,
973    PK: MLDSAPublicKeyInternalTrait<k, PK_LEN>,
974    SK: MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
975        + MLDSAPrivateKeyInternalTrait<k, l, eta, SK_LEN, PK_LEN>,
976    const SK_LEN: usize,
977    const PK_LEN: usize,
978> MLDSAPrivateKeyTrait<k, l, eta, SK_LEN, PK_LEN>
979    for MLDSAPrivateKeyExpanded<k, l, eta, PK, SK, SK_LEN, PK_LEN>
980{
981    fn seed(&self) -> Option<&KeyMaterial<32>> {
982        self.sk.seed()
983    }
984
985    fn tr(&self) -> &[u8; 64] {
986        self.sk.tr()
987    }
988
989    fn A_hat(&self) -> Matrix<k, l> {
990        self.sk.A_hat()
991    }
992
993    fn derive_pk(&self) -> MLDSAPublicKey<k, l, PK_LEN> {
994        self.sk.derive_pk()
995    }
996
997    fn sk_decode(sk: &[u8; SK_LEN]) -> Result<Self, SignatureError> {
998        let sk1 = SK::sk_decode(sk)?;
999        let A_hat = sk1.derive_pk().A_hat();
1000
1001        Ok(Self { _phantom: core::marker::PhantomData, sk: sk1, A_hat })
1002    }
1003}
bouncycastle_mldsa/mldsa_keys.rs

bouncycastle_mldsa/
mldsa_keys.rs
