Skip to main content

bouncycastle_mldsa_lowmemory/
mldsa_keys.rs

1use crate::aux_functions::{
2    bit_pack_eta, bit_pack_t0, bitlen_eta, power_2_round, rej_bounded_poly, simple_bit_pack_t1,
3    simple_bit_unpack_t1,
4};
5use crate::low_memory_helpers::{expandA_elem, s_unpack};
6use crate::mldsa::{H, N, POLY_T0PACKED_LEN};
7use crate::mldsa::{
8    MLDSA44_ETA, MLDSA44_FULL_SK_LEN, MLDSA44_GAMMA2, MLDSA44_LAMBDA, MLDSA44_PK_LEN,
9    MLDSA44_S1_PACKED_LEN, MLDSA44_S2_PACKED_LEN, MLDSA44_SK_LEN, MLDSA44_k, MLDSA44_l,
10};
11use crate::mldsa::{
12    MLDSA44_T1_PACKED_LEN, MLDSA65_T1_PACKED_LEN, MLDSA87_T1_PACKED_LEN, POLY_T1PACKED_LEN,
13};
14use crate::mldsa::{
15    MLDSA65_ETA, MLDSA65_FULL_SK_LEN, MLDSA65_GAMMA2, MLDSA65_LAMBDA, MLDSA65_PK_LEN,
16    MLDSA65_S1_PACKED_LEN, MLDSA65_S2_PACKED_LEN, MLDSA65_SK_LEN, MLDSA65_k, MLDSA65_l,
17};
18use crate::mldsa::{
19    MLDSA87_ETA, MLDSA87_FULL_SK_LEN, MLDSA87_GAMMA2, MLDSA87_LAMBDA, MLDSA87_PK_LEN,
20    MLDSA87_S1_PACKED_LEN, MLDSA87_S2_PACKED_LEN, MLDSA87_SK_LEN, MLDSA87_k, MLDSA87_l,
21};
22use crate::{ML_DSA_44_NAME, ML_DSA_65_NAME, ML_DSA_87_NAME};
23use bouncycastle_core::errors::SignatureError;
24use bouncycastle_core::key_material;
25use bouncycastle_core::key_material::{KeyMaterial, KeyMaterialTrait, KeyType};
26use bouncycastle_core::traits::{SecurityStrength, SignaturePrivateKey, SignaturePublicKey, XOF};
27use bouncycastle_utils::secret::Secret;
28use core::fmt;
29use core::fmt::{Debug, Display, Formatter};
30use core::ops::DerefMut;
31// imports just for docs
32#[allow(unused_imports)]
33use crate::mldsa::MLDSATrait;
34use crate::polynomial::Polynomial;
35
36/* Pub Types */
37
38/// ML-DSA-44 Public Key
39pub type MLDSA44PublicKey = MLDSAPublicKey<MLDSA44_k, MLDSA44_T1_PACKED_LEN, MLDSA44_PK_LEN>;
40/// ML-DSA-44 Private Key
41pub type MLDSA44PrivateKey = MLDSASeedPrivateKey<
42    MLDSA44_LAMBDA,
43    MLDSA44_GAMMA2,
44    MLDSA44_k,
45    MLDSA44_l,
46    MLDSA44_ETA,
47    MLDSA44_S1_PACKED_LEN,
48    MLDSA44_S2_PACKED_LEN,
49    MLDSA44_T1_PACKED_LEN,
50    MLDSA44_PK_LEN,
51    MLDSA44_SK_LEN,
52    MLDSA44_FULL_SK_LEN,
53>;
54/// ML-DSA-65 Public Key
55pub type MLDSA65PublicKey = MLDSAPublicKey<MLDSA65_k, MLDSA65_T1_PACKED_LEN, MLDSA65_PK_LEN>;
56/// ML-DSA-65 Private Key
57pub type MLDSA65PrivateKey = MLDSASeedPrivateKey<
58    MLDSA65_LAMBDA,
59    MLDSA65_GAMMA2,
60    MLDSA65_k,
61    MLDSA65_l,
62    MLDSA65_ETA,
63    MLDSA65_S1_PACKED_LEN,
64    MLDSA65_S2_PACKED_LEN,
65    MLDSA65_T1_PACKED_LEN,
66    MLDSA65_PK_LEN,
67    MLDSA65_SK_LEN,
68    MLDSA65_FULL_SK_LEN,
69>;
70/// ML-DSA-87 Public Key
71pub type MLDSA87PublicKey = MLDSAPublicKey<MLDSA87_k, MLDSA87_T1_PACKED_LEN, MLDSA87_PK_LEN>;
72/// ML-DSA-87 Private Key
73pub type MLDSA87PrivateKey = MLDSASeedPrivateKey<
74    MLDSA87_LAMBDA,
75    MLDSA87_GAMMA2,
76    MLDSA87_k,
77    MLDSA87_l,
78    MLDSA87_ETA,
79    MLDSA87_S1_PACKED_LEN,
80    MLDSA87_S2_PACKED_LEN,
81    MLDSA87_T1_PACKED_LEN,
82    MLDSA87_PK_LEN,
83    MLDSA87_SK_LEN,
84    MLDSA87_FULL_SK_LEN,
85>;
86
87/// An ML-DSA public key.
88#[derive(Clone)]
89pub struct MLDSAPublicKey<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize> {
90    pub(crate) rho: [u8; 32],
91    pub(crate) t1_packed: [u8; T1_PACKED_LEN],
92}
93
94/// General trait for all ML-DSA public keys types.
95pub trait MLDSAPublicKeyTrait<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize>:
96    SignaturePublicKey<PK_LEN>
97{
98    /// Algorithm 23 pkDecode(π‘π‘˜)
99    /// Reverses the procedure pkEncode.
100    /// Input: Public key π‘π‘˜ ∈ 𝔹32+32π‘˜(bitlen (π‘žβˆ’1)βˆ’π‘‘).
101    /// Output: 𝜌 ∈ 𝔹32, 𝐭1 ∈ π‘…π‘˜ with coefficients in [0, 2bitlen (π‘žβˆ’1)βˆ’π‘‘ βˆ’ 1].
102    fn pk_decode(pk: &[u8; PK_LEN]) -> Self;
103
104    /// Compute the public key hash (tr) from the public key.
105    ///
106    /// This is exposed as a public API for a few reasons:
107    /// 1. `tr` is required for some external-prehashing schemes such as the so-called "external mu" signing mode.
108    /// 2. `tr` is the canonical fingerprint of an ML-DSA public key, so would be an appropriate value
109    ///     to use, for example, to build a public key lookup or deny-listing table.
110    fn compute_tr(&self) -> [u8; 64];
111}
112
113pub(crate) trait MLDSAPublicKeyInternalTrait<
114    const k: usize,
115    const T1_PACKED_LEN: usize,
116    const PK_LEN: usize,
117>
118{
119    /// Not exposing a constructor publicly because the user should get an instance either by
120    /// running a keygen, or by decoding an existing key.
121    fn new(rho: [u8; 32], t1_packed: [u8; T1_PACKED_LEN]) -> Self;
122
123    /// Get a ref to rho
124    fn rho(&self) -> &[u8; 32];
125
126    /// Get a ref to t1
127    fn unpack_t1_row(&self, row: usize) -> Polynomial;
128}
129
130impl<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize>
131    MLDSAPublicKeyTrait<k, T1_PACKED_LEN, PK_LEN> for MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN>
132{
133    fn pk_decode(pk: &[u8; PK_LEN]) -> Self {
134        Self { rho: pk[..32].try_into().unwrap(), t1_packed: pk[32..].try_into().unwrap() }
135    }
136
137    fn compute_tr(&self) -> [u8; 64] {
138        let mut tr = [0u8; 64];
139        H::new().hash_xof_out(&self.encode(), &mut tr);
140
141        tr
142    }
143}
144
145impl<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize>
146    MLDSAPublicKeyInternalTrait<k, T1_PACKED_LEN, PK_LEN>
147    for MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN>
148{
149    fn new(rho: [u8; 32], t1_packed: [u8; T1_PACKED_LEN]) -> Self {
150        Self { rho, t1_packed }
151    }
152
153    fn rho(&self) -> &[u8; 32] {
154        &self.rho
155    }
156
157    fn unpack_t1_row(&self, row: usize) -> Polynomial {
158        simple_bit_unpack_t1(
159            &self.t1_packed[row * POLY_T1PACKED_LEN..(row + 1) * POLY_T1PACKED_LEN]
160                .try_into()
161                .unwrap(),
162        )
163    }
164}
165
166impl<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize> SignaturePublicKey<PK_LEN>
167    for MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN>
168{
169    /// Algorithm 22 pkEncode(𝜌, 𝐭1)
170    /// Encodes a public key for ML-DSA into a byte string.
171    /// Input:𝜌 ∈ 𝔹32, 𝐭1 ∈ π‘…π‘˜ with coefficients in [0, 2bitlen (π‘žβˆ’1)βˆ’π‘‘ βˆ’ 1].
172    /// Output: Public key π‘π‘˜ ∈ 𝔹32+32π‘˜(bitlen (π‘žβˆ’1)βˆ’π‘‘).
173    fn encode(&self) -> [u8; PK_LEN] {
174        let mut pk = [0u8; PK_LEN];
175        self.encode_out(&mut pk);
176
177        pk
178    }
179    /// Algorithm 22 pkEncode(𝜌, 𝐭1)
180    /// Encodes a public key for ML-DSA into a byte string.
181    /// Input:𝜌 ∈ 𝔹32, 𝐭1 ∈ π‘…π‘˜ with coefficients in [0, 2bitlen (π‘žβˆ’1)βˆ’π‘‘ βˆ’ 1].
182    /// Output: Public key π‘π‘˜ ∈ 𝔹32+32π‘˜(bitlen (π‘žβˆ’1)βˆ’π‘‘).
183    fn encode_out(&self, out: &mut [u8; PK_LEN]) -> usize {
184        debug_assert_eq!(out.len(), PK_LEN);
185
186        out.fill(0);
187
188        out[..32].copy_from_slice(&self.rho);
189        out[32..].copy_from_slice(&self.t1_packed);
190
191        PK_LEN
192    }
193
194    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
195        if bytes.len() != PK_LEN {
196            return Err(SignatureError::DecodingError(
197                "Provided key bytes are the incorrect length",
198            ));
199        }
200        let sized_bytes: [u8; PK_LEN] = bytes[..PK_LEN].try_into().unwrap();
201        Ok(Self::pk_decode(&sized_bytes))
202    }
203}
204
205impl<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize> Eq
206    for MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN>
207{
208}
209
210impl<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize> PartialEq
211    for MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN>
212{
213    fn eq(&self, other: &Self) -> bool {
214        let self_encoded = self.encode();
215        let other_encoded = other.encode();
216        bouncycastle_utils::ct::ct_eq_bytes(self_encoded.as_ref(), other_encoded.as_ref())
217    }
218}
219
220impl<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize> fmt::Debug
221    for MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN>
222{
223    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
224        let alg = match k {
225            4 => ML_DSA_44_NAME,
226            6 => ML_DSA_65_NAME,
227            8 => ML_DSA_87_NAME,
228            _ => panic!("Unsupported key length"),
229        };
230        write!(f, "MLDSAPublicKey {{ alg: {}, pub_key_hash (tr): {:x?} }}", alg, self.compute_tr(),)
231    }
232}
233
234impl<const k: usize, const T1_PACKED_LEN: usize, const PK_LEN: usize> Display
235    for MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN>
236{
237    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
238        let alg = match k {
239            4 => ML_DSA_44_NAME,
240            6 => ML_DSA_65_NAME,
241            8 => ML_DSA_87_NAME,
242            _ => panic!("Unsupported key length"),
243        };
244        write!(f, "MLDSAPublicKey {{ alg: {}, pub_key_hash (tr): {:x?} }}", alg, self.compute_tr(),)
245    }
246}
247
248/// General trait for all ML-DSA private keys types.
249pub trait MLDSAPrivateKeyTrait<
250    const k: usize,
251    const l: usize,
252    const S1_PACKED_LEN: usize,
253    const S2_PACKED_LEN: usize,
254    const T1_PACKED_LEN: usize,
255    const PK_LEN: usize,
256    const SK_LEN: usize,
257    const FULL_SK_LEN: usize,
258>: SignaturePrivateKey<SK_LEN>
259{
260    /// New from KeyMaterial. Can throw a SignatureError if the KeyMaterial does not contain sufficient entropy.
261    fn from_keymaterial(seed: &KeyMaterial<32>) -> Result<Self, SignatureError>;
262
263    /// Get a ref to the seed, if there is one stored with this private key
264    fn seed(&self) -> Option<&KeyMaterial<32>>;
265
266    /// Get a copy of the key hash `tr`.
267    /// This is computationally intensive as it requires fully re-computing the public key (and then discarding it).
268    /// It is highly recommended that, if the user already has a copy of the public key, they should get `tr` from that,
269    /// or else compute `tr` once and store it.
270    fn tr(&self) -> [u8; 64];
271    /// Returns the full public key, and has the side-effect of setting the public key hash tr in this MLDSASeedSK object.
272    fn derive_pk(&self) -> MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN>;
273    /// This produces the full private key in the encoding specified in FIPS 204 Algorithm 24 skEncode()
274    /// so that it is compatible with other implementations.
275    ///
276    /// Note that since this encoding does not include the seed, this is a one-way operation;
277    /// after exporting in this encoding, it will be impossible to re-import it into a [`MLDSASeedPrivateKey`].
278    fn encode_full_sk(&self) -> [u8; FULL_SK_LEN];
279    /// This produces the full private key in the encoding specified in FIPS 204 Algorithm 24 skEncode()
280    /// so that it is compatible with other implementations.
281    ///
282    /// Note that since this encoding does not include the seed, this is a one-way operation;
283    /// after exporting in this encoding, it will be impossible to re-import it into a [`MLDSASeedPrivateKey`].
284    fn encode_full_sk_out(&self, out: &mut [u8; FULL_SK_LEN]) -> usize;
285    /// Algorithm 25 skDecode(π‘ π‘˜)
286    /// Reverses the procedure skEncode.
287    /// Input: Private key π‘ π‘˜ ∈ 𝔹32+32+64+32β‹…((β„“+π‘˜)β‹…bitlen (2πœ‚)+π‘‘π‘˜).
288    /// Output: 𝜌 ∈ 𝔹32, 𝐾 ∈ 𝔹32, π‘‘π‘Ÿ ∈ 𝔹64 ,
289    /// 𝐬1 ∈ 𝑅ℓ , 𝐬2 ∈ π‘…π‘˜ , 𝐭0 ∈ π‘…π‘˜ with coefficients in [βˆ’2π‘‘βˆ’1 + 1, 2π‘‘βˆ’1].
290    ///
291    /// Note: this object contains only the simple decoding routine to unpack a semi-expanded key.
292    /// See [`MLDSATrait`] for key generation functions, including derive-from-seed and consistency-check functions.
293    fn sk_decode(sk: &[u8; SK_LEN]) -> Self;
294}
295
296/// Internal structure for holding a seed-based private key for ML-DSA.
297#[derive(Clone, PartialEq, Eq)]
298pub struct MLDSASeedPrivateKey<
299    const LAMBDA: i32,
300    const GAMMA2: i32,
301    const k: usize,
302    const l: usize,
303    const eta: usize,
304    const S1_PACKED_LEN: usize,
305    const S2_PACKED_LEN: usize,
306    const T1_PACKED_LEN: usize,
307    const PK_LEN: usize,
308    const SK_LEN: usize,
309    const FULL_SK_LEN: usize,
310> {
311    // note: KeyMaterial is inherently Secret
312    seed: KeyMaterial<32>,
313    // public seed rho does not need to be secret
314    rho: [u8; 32],
315    rho_prime: Secret<[u8; 64]>,
316    K: Secret<[u8; 32]>,
317}
318impl<
319    const LAMBDA: i32,
320    const GAMMA2: i32,
321    const k: usize,
322    const l: usize,
323    const eta: usize,
324    const S1_PACKED_LEN: usize,
325    const S2_PACKED_LEN: usize,
326    const T1_PACKED_LEN: usize,
327    const PK_LEN: usize,
328    const SK_LEN: usize,
329    const FULL_SK_LEN: usize,
330> Debug
331    for MLDSASeedPrivateKey<
332        LAMBDA,
333        GAMMA2,
334        k,
335        l,
336        eta,
337        S1_PACKED_LEN,
338        S2_PACKED_LEN,
339        T1_PACKED_LEN,
340        PK_LEN,
341        SK_LEN,
342        FULL_SK_LEN,
343    >
344{
345    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
346        let alg = match k {
347            4 => ML_DSA_44_NAME,
348            6 => ML_DSA_65_NAME,
349            8 => ML_DSA_87_NAME,
350            _ => panic!("Unsupported key length"),
351        };
352        write!(f, "MLDSASeedPrivateKey {{ alg: {}, pub_key_hash (tr): {:x?} }}", alg, self.tr(),)
353    }
354}
355
356impl<
357    const LAMBDA: i32,
358    const GAMMA2: i32,
359    const k: usize,
360    const l: usize,
361    const eta: usize,
362    const S1_PACKED_LEN: usize,
363    const S2_PACKED_LEN: usize,
364    const T1_PACKED_LEN: usize,
365    const PK_LEN: usize,
366    const SK_LEN: usize,
367    const FULL_SK_LEN: usize,
368> Display
369    for MLDSASeedPrivateKey<
370        LAMBDA,
371        GAMMA2,
372        k,
373        l,
374        eta,
375        S1_PACKED_LEN,
376        S2_PACKED_LEN,
377        T1_PACKED_LEN,
378        PK_LEN,
379        SK_LEN,
380        FULL_SK_LEN,
381    >
382{
383    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
384        let alg = match k {
385            4 => ML_DSA_44_NAME,
386            6 => ML_DSA_65_NAME,
387            8 => ML_DSA_87_NAME,
388            _ => panic!("Unsupported key length"),
389        };
390        write!(f, "MLDSASeedPrivateKey {{ alg: {}, pub_key_hash (tr): {:x?} }}", alg, self.tr(),)
391    }
392}
393
394impl<
395    const LAMBDA: i32,
396    const GAMMA2: i32,
397    const k: usize,
398    const l: usize,
399    const eta: usize,
400    const S1_PACKED_LEN: usize,
401    const S2_PACKED_LEN: usize,
402    const T1_PACKED_LEN: usize,
403    const PK_LEN: usize,
404    const SK_LEN: usize,
405    const FULL_SK_LEN: usize,
406>
407    MLDSASeedPrivateKey<
408        LAMBDA,
409        GAMMA2,
410        k,
411        l,
412        eta,
413        S1_PACKED_LEN,
414        S2_PACKED_LEN,
415        T1_PACKED_LEN,
416        PK_LEN,
417        SK_LEN,
418        FULL_SK_LEN,
419    >
420{
421    /// Create a new MLDSASeedPrivateKey from a 32-byte KeyMaterial.
422    /// Seed SecurityStrength must match algorithm security strength: 128-bit (ML-DSA-44), 192-bit (ML-DSA-65), or 256-bit (ML-DSA-87),
423    /// otherwise it throws a SignatureError::KeyGenError("SecurityStrength".
424    pub fn new(seed: &KeyMaterial<32>) -> Result<Self, SignatureError> {
425        if !(seed.key_type() == KeyType::Seed || seed.key_type() == KeyType::CryptographicRandom)
426            || seed.key_len() != 32
427        {
428            return Err(SignatureError::KeyGenError(
429                "Seed must be 32 bytes and KeyType::Seed or KeyType::BytesFullEntropy.",
430            ));
431        }
432
433        if seed.security_strength() < SecurityStrength::from_bits(LAMBDA as usize) {
434            return Err(SignatureError::KeyGenError("SecurityStrength"));
435        }
436
437        let (rho, rho_prime, K) = Self::compute_rhos_and_K(&seed);
438
439        Ok(Self { seed: seed.clone(), rho, rho_prime, K })
440    }
441
442    fn compute_rhos_and_K(
443        seed: &KeyMaterial<32>,
444    ) -> ([u8; 32], Secret<[u8; 64]>, Secret<[u8; 32]>) {
445        // derive sk.K
446        // Alg 6; 1: (rho, rho_prime, K) <- H(πœ‰||IntegerToBytes(π‘˜, 1)||IntegerToBytes(β„“, 1), 128)
447        //   β–· expand seed
448        let mut rho = [0u8; 32];
449        let mut rho_prime: Secret<[u8; 64]> = Secret::new();
450        let mut K: Secret<[u8; 32]> = Secret::new();
451
452        let mut h = H::default();
453        h.absorb(seed.ref_to_bytes()).expect("absorb before squeeze is infallible");
454        h.absorb(&(k as u8).to_le_bytes()).expect("absorb before squeeze is infallible");
455        h.absorb(&(l as u8).to_le_bytes()).expect("absorb before squeeze is infallible");
456        let bytes_written = h.squeeze_out(&mut rho);
457        debug_assert_eq!(bytes_written, 32);
458        let bytes_written = h.squeeze_out(rho_prime.deref_mut());
459        debug_assert_eq!(bytes_written, 64);
460        let bytes_written = h.squeeze_out(K.deref_mut());
461        debug_assert_eq!(bytes_written, 32);
462
463        (rho, rho_prime, K)
464    }
465
466    fn compute_t_row(
467        &self,
468        idx: usize,
469        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
470        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
471    ) -> Polynomial {
472        debug_assert!(idx < k);
473
474        // [Optimization Note]:
475        // This is one of the places that a row of s1 can be re-computed instead of expanded from the compressed form.
476        // let mut s1 = self.compute_s1_row(0);
477        let mut s1_hat_i = s_unpack::<eta, S1_PACKED_LEN>(s1_packed, 0);
478        s1_hat_i.ntt();
479
480        let mut t_i = {
481            let mut t_hat_i = expandA_elem(&self.rho, idx, 0);
482            t_hat_i.multiply_ntt(&s1_hat_i);
483
484            for col in 1..l {
485                // [Optimization Note]:
486                // This is one of the places that a row of s1 can be re-computed instead of expanded from the compressed form.
487                // s1 = self.compute_s1_row(col);
488                let mut s1_hat = s_unpack::<eta, S1_PACKED_LEN>(s1_packed, col);
489                s1_hat.ntt();
490                let mut A_elem = expandA_elem(&self.rho, idx, col);
491                A_elem.multiply_ntt(&s1_hat);
492                t_hat_i.add_ntt(&A_elem);
493            }
494            t_hat_i.inv_ntt();
495
496            t_hat_i
497        };
498
499        // [Optimization Note]:
500        // This is one of the places that a row of s2 can be re-computed instead of unpacked from the compressed form.
501        // let s2 = self.compute_s2_row(idx);
502        let s2 = s_unpack::<eta, S2_PACKED_LEN>(s2_packed, idx);
503        t_i.add_ntt(&s2);
504        t_i.conditional_add_q();
505
506        t_i
507    }
508}
509
510impl<
511    const LAMBDA: i32,
512    const GAMMA2: i32,
513    const k: usize,
514    const l: usize,
515    const eta: usize,
516    const S1_PACKED_LEN: usize,
517    const S2_PACKED_LEN: usize,
518    const T1_PACKED_LEN: usize,
519    const PK_LEN: usize,
520    const SK_LEN: usize,
521    const FULL_SK_LEN: usize,
522> SignaturePrivateKey<SK_LEN>
523    for MLDSASeedPrivateKey<
524        LAMBDA,
525        GAMMA2,
526        k,
527        l,
528        eta,
529        S1_PACKED_LEN,
530        S2_PACKED_LEN,
531        T1_PACKED_LEN,
532        PK_LEN,
533        SK_LEN,
534        FULL_SK_LEN,
535    >
536{
537    /// Encodes the private key seed.
538    fn encode(&self) -> [u8; SK_LEN] {
539        debug_assert_eq!(SK_LEN, /* seed */ 32);
540
541        self.seed.ref_to_bytes().try_into().unwrap()
542    }
543
544    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
545        out.fill(0);
546
547        out.copy_from_slice(self.seed.ref_to_bytes());
548
549        debug_assert_eq!(self.seed.ref_to_bytes().len(), SK_LEN);
550        SK_LEN
551    }
552
553    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
554        if bytes.len() != 32 {
555            return Err(SignatureError::DecodingError("Invalid seed length"));
556        }
557        let mut keymat = KeyMaterial::<32>::from_bytes(bytes)?;
558        key_material::do_hazardous_operations(&mut keymat, |keymat| {
559            keymat.set_key_type(KeyType::Seed)?;
560            keymat.set_security_strength(SecurityStrength::_256bit)
561        })?;
562
563        Self::new(&keymat)
564    }
565}
566
567impl<
568    const LAMBDA: i32,
569    const GAMMA2: i32,
570    const k: usize,
571    const l: usize,
572    const eta: usize,
573    const S1_PACKED_LEN: usize,
574    const S2_PACKED_LEN: usize,
575    const T1_PACKED_LEN: usize,
576    const PK_LEN: usize,
577    const SK_LEN: usize,
578    const FULL_SK_LEN: usize,
579>
580    MLDSAPrivateKeyTrait<
581        k,
582        l,
583        S1_PACKED_LEN,
584        S2_PACKED_LEN,
585        T1_PACKED_LEN,
586        PK_LEN,
587        SK_LEN,
588        FULL_SK_LEN,
589    >
590    for MLDSASeedPrivateKey<
591        LAMBDA,
592        GAMMA2,
593        k,
594        l,
595        eta,
596        S1_PACKED_LEN,
597        S2_PACKED_LEN,
598        T1_PACKED_LEN,
599        PK_LEN,
600        SK_LEN,
601        FULL_SK_LEN,
602    >
603{
604    fn from_keymaterial(seed: &KeyMaterial<32>) -> Result<Self, SignatureError> {
605        Self::new(seed)
606    }
607
608    fn seed(&self) -> Option<&KeyMaterial<32>> {
609        Some(&self.seed)
610    }
611
612    fn tr(&self) -> [u8; 64] {
613        let pk: MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN> = self.derive_pk();
614        pk.compute_tr()
615    }
616
617    fn derive_pk(&self) -> MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN> {
618        // The goal here is to get t1, which is built and compressed one row at a time.
619
620        let s1_packed: Secret<[u8; S1_PACKED_LEN]> = self.compute_s1_packed();
621        let s2_packed: Secret<[u8; S2_PACKED_LEN]> = self.compute_s2_packed();
622
623        let mut t1_packed = [0u8; T1_PACKED_LEN];
624        debug_assert_eq!(T1_PACKED_LEN, POLY_T1PACKED_LEN * k);
625
626        for i in 0..k {
627            t1_packed[i * POLY_T1PACKED_LEN..(i + 1) * POLY_T1PACKED_LEN].copy_from_slice(
628                &simple_bit_pack_t1(&self.compute_t1_row(i, &s1_packed, &s2_packed)),
629            );
630        }
631
632        MLDSAPublicKey::<k, T1_PACKED_LEN, PK_LEN>::new(self.rho.clone(), t1_packed)
633    }
634    fn encode_full_sk(&self) -> [u8; FULL_SK_LEN] {
635        let mut out = [0; FULL_SK_LEN];
636        _ = self.encode_full_sk_out(&mut out);
637
638        out
639    }
640    fn encode_full_sk_out(&self, out: &mut [u8; FULL_SK_LEN]) -> usize {
641        out.fill(0);
642
643        // Algorithm 24 skEncode(𝜌, 𝐾, π‘‘π‘Ÿ, 𝐬1, 𝐬2, 𝐭0)
644
645        let mut off: usize = 0;
646
647        // 1: π‘ π‘˜ ← 𝜌||𝐾||π‘‘π‘Ÿ
648        out[0..32].copy_from_slice(&self.rho);
649        // K is protected inside a Secret<[u8]>, but here we really do need to deref and copy it out.
650        out[32..64].copy_from_slice(&*self.K);
651        out[64..128].copy_from_slice(&self.tr());
652        off += 128;
653
654        // 2: for 𝑖 from 0 to β„“ βˆ’ 1 do
655        // 3:   π‘ π‘˜ ← π‘ π‘˜ || BitPack (𝐬1[𝑖], πœ‚, πœ‚)
656        // 4: end for
657        let s1_packed = self.compute_s1_packed();
658        out[off..off + S1_PACKED_LEN].copy_from_slice(&*s1_packed);
659        off += S1_PACKED_LEN;
660
661        // 5: for 𝑖 from 0 to π‘˜ βˆ’ 1 do
662        // 6:   π‘ π‘˜ ← π‘ π‘˜ || BitPack (𝐬2[𝑖], πœ‚, πœ‚)
663        // 7: end for
664        let s2_packed = self.compute_s2_packed();
665        out[off..off + S2_PACKED_LEN].copy_from_slice(&*s2_packed);
666        off += S2_PACKED_LEN;
667
668        // 8: for 𝑖 from 0 to π‘˜ βˆ’ 1 do
669        // 9:   π‘ π‘˜ ← π‘ π‘˜ || BitPack (𝐭0[𝑖], 2π‘‘βˆ’1 βˆ’ 1, 2π‘‘βˆ’1)
670        // 10: end for
671        debug_assert_eq!(off + k * POLY_T0PACKED_LEN, FULL_SK_LEN);
672        for row in 0..k {
673            let t0_i = self.compute_t0_row(row, &s1_packed, &s2_packed);
674            out[off..off + POLY_T0PACKED_LEN].copy_from_slice(&bit_pack_t0(&t0_i));
675            off += POLY_T0PACKED_LEN;
676        }
677        debug_assert_eq!(off, FULL_SK_LEN);
678
679        FULL_SK_LEN
680    }
681
682    fn sk_decode(sk: &[u8; SK_LEN]) -> Self {
683        Self::from_bytes(sk).unwrap()
684    }
685}
686
687pub(crate) trait MLDSAPrivateKeyInternalTrait<
688    const LAMBDA: i32,
689    const GAMMA2: i32,
690    const k: usize,
691    const l: usize,
692    const eta: usize,
693    const S1_PACKED_LEN: usize,
694    const S2_PACKED_LEN: usize,
695    const PK_LEN: usize,
696    const SK_LEN: usize,
697>: Sized
698{
699    fn rho(&self) -> &[u8; 32];
700    fn K(&self) -> &[u8; 32];
701
702    /// A single entry of a privacy key vector.
703    /// These tend to be used very transiently, so we won't bother wrapping it as a Secret.
704    fn compute_s1_row(&self, idx: usize) -> Polynomial;
705
706    /// Private key component.
707    /// The packed representation sticks around for the whole computation, so
708    /// we'll wrap in as a Secret.
709    fn compute_s1_packed(&self) -> Secret<[u8; S1_PACKED_LEN]>;
710
711    /// A single entry of a privacy key vector.
712    /// These tend to be used very transiently, so we won't bother wrapping it as a Secret.
713    fn compute_s2_row(&self, idx: usize) -> Polynomial;
714
715    /// Private key component.
716    /// The packed representation sticks around for the whole computation, so
717    /// we'll wrap in as a Secret.
718    fn compute_s2_packed(&self) -> Secret<[u8; S2_PACKED_LEN]>;
719
720    /// Public key component.
721    fn compute_t0_row(
722        &self,
723        idx: usize,
724        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
725        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
726    ) -> Polynomial;
727
728    /// Public key component.
729    fn compute_t1_row(
730        &self,
731        idx: usize,
732        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
733        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
734    ) -> Polynomial;
735}
736
737impl<
738    const LAMBDA: i32,
739    const GAMMA2: i32,
740    const k: usize,
741    const l: usize,
742    const eta: usize,
743    const S1_PACKED_LEN: usize,
744    const S2_PACKED_LEN: usize,
745    const T1_PACKED_LEN: usize,
746    const PK_LEN: usize,
747    const SK_LEN: usize,
748    const FULL_SK_LEN: usize,
749>
750    MLDSAPrivateKeyInternalTrait<
751        LAMBDA,
752        GAMMA2,
753        k,
754        l,
755        eta,
756        S1_PACKED_LEN,
757        S2_PACKED_LEN,
758        PK_LEN,
759        SK_LEN,
760    >
761    for MLDSASeedPrivateKey<
762        LAMBDA,
763        GAMMA2,
764        k,
765        l,
766        eta,
767        S1_PACKED_LEN,
768        S2_PACKED_LEN,
769        T1_PACKED_LEN,
770        PK_LEN,
771        SK_LEN,
772        FULL_SK_LEN,
773    >
774{
775    fn rho(&self) -> &[u8; 32] {
776        &self.rho
777    }
778
779    fn K(&self) -> &[u8; 32] {
780        &self.K
781    }
782
783    fn compute_s1_row(&self, idx: usize) -> Polynomial {
784        debug_assert!(idx < l);
785        rej_bounded_poly::<eta>(&self.rho_prime, &(idx as u16).to_le_bytes())
786    }
787
788    fn compute_s1_packed(&self) -> Secret<[u8; S1_PACKED_LEN]> {
789        let mut s1_packed: Secret<[u8; S1_PACKED_LEN]> = Secret::new();
790        for idx in 0..l {
791            let s1_i = self.compute_s1_row(idx);
792            bit_pack_eta::<eta>(
793                &s1_i,
794                &mut s1_packed[idx * bitlen_eta(eta)..(idx + 1) * bitlen_eta(eta)],
795            );
796        }
797        s1_packed
798    }
799
800    fn compute_s2_row(&self, idx: usize) -> Polynomial {
801        debug_assert!(idx < k);
802        rej_bounded_poly::<eta>(&self.rho_prime, &((idx + l) as u16).to_le_bytes())
803    }
804
805    fn compute_s2_packed(&self) -> Secret<[u8; S2_PACKED_LEN]> {
806        let mut s2_packed: Secret<[u8; S2_PACKED_LEN]> = Secret::new();
807        for idx in 0..k {
808            let s2_i = self.compute_s2_row(idx);
809            bit_pack_eta::<eta>(
810                &s2_i,
811                &mut s2_packed[idx * bitlen_eta(eta)..(idx + 1) * bitlen_eta(eta)],
812            );
813        }
814        s2_packed
815    }
816
817    fn compute_t0_row(
818        &self,
819        idx: usize,
820        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
821        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
822    ) -> Polynomial {
823        let mut t0 = self.compute_t_row(idx, s1_packed, s2_packed);
824        for j in 0..N {
825            (_, t0[j]) = power_2_round(t0[j]);
826        }
827
828        t0
829    }
830
831    fn compute_t1_row(
832        &self,
833        idx: usize,
834        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
835        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
836    ) -> Polynomial {
837        let mut t1 = self.compute_t_row(idx, s1_packed, s2_packed);
838        for j in 0..N {
839            (t1[j], _) = power_2_round(t1[j]);
840        }
841
842        t1
843    }
844}