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            // Bound-keeping step before NTT⁻¹, not in FIPS 204: see [`Polynomial::reduce32`].
495            t_hat_i.reduce32();
496            t_hat_i.inv_ntt();
497
498            t_hat_i
499        };
500
501        // [Optimization Note]:
502        // This is one of the places that a row of s2 can be re-computed instead of unpacked from the compressed form.
503        // let s2 = self.compute_s2_row(idx);
504        let s2 = s_unpack::<eta, S2_PACKED_LEN>(s2_packed, idx);
505        t_i.add_ntt(&s2);
506        t_i.conditional_add_q();
507
508        t_i
509    }
510}
511
512impl<
513    const LAMBDA: i32,
514    const GAMMA2: i32,
515    const k: usize,
516    const l: usize,
517    const eta: usize,
518    const S1_PACKED_LEN: usize,
519    const S2_PACKED_LEN: usize,
520    const T1_PACKED_LEN: usize,
521    const PK_LEN: usize,
522    const SK_LEN: usize,
523    const FULL_SK_LEN: usize,
524> SignaturePrivateKey<SK_LEN>
525    for MLDSASeedPrivateKey<
526        LAMBDA,
527        GAMMA2,
528        k,
529        l,
530        eta,
531        S1_PACKED_LEN,
532        S2_PACKED_LEN,
533        T1_PACKED_LEN,
534        PK_LEN,
535        SK_LEN,
536        FULL_SK_LEN,
537    >
538{
539    /// Encodes the private key seed.
540    fn encode(&self) -> [u8; SK_LEN] {
541        debug_assert_eq!(SK_LEN, /* seed */ 32);
542
543        self.seed.ref_to_bytes().try_into().unwrap()
544    }
545
546    fn encode_out(&self, out: &mut [u8; SK_LEN]) -> usize {
547        out.fill(0);
548
549        out.copy_from_slice(self.seed.ref_to_bytes());
550
551        debug_assert_eq!(self.seed.ref_to_bytes().len(), SK_LEN);
552        SK_LEN
553    }
554
555    fn from_bytes(bytes: &[u8]) -> Result<Self, SignatureError> {
556        if bytes.len() != 32 {
557            return Err(SignatureError::DecodingError("Invalid seed length"));
558        }
559        let mut keymat = KeyMaterial::<32>::from_bytes(bytes)?;
560        key_material::do_hazardous_operations(&mut keymat, |keymat| {
561            keymat.set_key_type(KeyType::Seed)?;
562            keymat.set_security_strength(SecurityStrength::_256bit)
563        })?;
564
565        Self::new(&keymat)
566    }
567}
568
569impl<
570    const LAMBDA: i32,
571    const GAMMA2: i32,
572    const k: usize,
573    const l: usize,
574    const eta: usize,
575    const S1_PACKED_LEN: usize,
576    const S2_PACKED_LEN: usize,
577    const T1_PACKED_LEN: usize,
578    const PK_LEN: usize,
579    const SK_LEN: usize,
580    const FULL_SK_LEN: usize,
581>
582    MLDSAPrivateKeyTrait<
583        k,
584        l,
585        S1_PACKED_LEN,
586        S2_PACKED_LEN,
587        T1_PACKED_LEN,
588        PK_LEN,
589        SK_LEN,
590        FULL_SK_LEN,
591    >
592    for MLDSASeedPrivateKey<
593        LAMBDA,
594        GAMMA2,
595        k,
596        l,
597        eta,
598        S1_PACKED_LEN,
599        S2_PACKED_LEN,
600        T1_PACKED_LEN,
601        PK_LEN,
602        SK_LEN,
603        FULL_SK_LEN,
604    >
605{
606    fn from_keymaterial(seed: &KeyMaterial<32>) -> Result<Self, SignatureError> {
607        Self::new(seed)
608    }
609
610    fn seed(&self) -> Option<&KeyMaterial<32>> {
611        Some(&self.seed)
612    }
613
614    fn tr(&self) -> [u8; 64] {
615        let pk: MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN> = self.derive_pk();
616        pk.compute_tr()
617    }
618
619    fn derive_pk(&self) -> MLDSAPublicKey<k, T1_PACKED_LEN, PK_LEN> {
620        // The goal here is to get t1, which is built and compressed one row at a time.
621
622        let s1_packed: Secret<[u8; S1_PACKED_LEN]> = self.compute_s1_packed();
623        let s2_packed: Secret<[u8; S2_PACKED_LEN]> = self.compute_s2_packed();
624
625        let mut t1_packed = [0u8; T1_PACKED_LEN];
626        debug_assert_eq!(T1_PACKED_LEN, POLY_T1PACKED_LEN * k);
627
628        for i in 0..k {
629            t1_packed[i * POLY_T1PACKED_LEN..(i + 1) * POLY_T1PACKED_LEN].copy_from_slice(
630                &simple_bit_pack_t1(&self.compute_t1_row(i, &s1_packed, &s2_packed)),
631            );
632        }
633
634        MLDSAPublicKey::<k, T1_PACKED_LEN, PK_LEN>::new(self.rho.clone(), t1_packed)
635    }
636    fn encode_full_sk(&self) -> [u8; FULL_SK_LEN] {
637        let mut out = [0; FULL_SK_LEN];
638        _ = self.encode_full_sk_out(&mut out);
639
640        out
641    }
642    fn encode_full_sk_out(&self, out: &mut [u8; FULL_SK_LEN]) -> usize {
643        out.fill(0);
644
645        // Algorithm 24 skEncode(𝜌, 𝐾, π‘‘π‘Ÿ, 𝐬1, 𝐬2, 𝐭0)
646
647        let mut off: usize = 0;
648
649        // 1: π‘ π‘˜ ← 𝜌||𝐾||π‘‘π‘Ÿ
650        out[0..32].copy_from_slice(&self.rho);
651        // K is protected inside a Secret<[u8]>, but here we really do need to deref and copy it out.
652        out[32..64].copy_from_slice(&*self.K);
653        out[64..128].copy_from_slice(&self.tr());
654        off += 128;
655
656        // 2: for 𝑖 from 0 to β„“ βˆ’ 1 do
657        // 3:   π‘ π‘˜ ← π‘ π‘˜ || BitPack (𝐬1[𝑖], πœ‚, πœ‚)
658        // 4: end for
659        let s1_packed = self.compute_s1_packed();
660        out[off..off + S1_PACKED_LEN].copy_from_slice(&*s1_packed);
661        off += S1_PACKED_LEN;
662
663        // 5: for 𝑖 from 0 to π‘˜ βˆ’ 1 do
664        // 6:   π‘ π‘˜ ← π‘ π‘˜ || BitPack (𝐬2[𝑖], πœ‚, πœ‚)
665        // 7: end for
666        let s2_packed = self.compute_s2_packed();
667        out[off..off + S2_PACKED_LEN].copy_from_slice(&*s2_packed);
668        off += S2_PACKED_LEN;
669
670        // 8: for 𝑖 from 0 to π‘˜ βˆ’ 1 do
671        // 9:   π‘ π‘˜ ← π‘ π‘˜ || BitPack (𝐭0[𝑖], 2π‘‘βˆ’1 βˆ’ 1, 2π‘‘βˆ’1)
672        // 10: end for
673        debug_assert_eq!(off + k * POLY_T0PACKED_LEN, FULL_SK_LEN);
674        for row in 0..k {
675            let t0_i = self.compute_t0_row(row, &s1_packed, &s2_packed);
676            out[off..off + POLY_T0PACKED_LEN].copy_from_slice(&bit_pack_t0(&t0_i));
677            off += POLY_T0PACKED_LEN;
678        }
679        debug_assert_eq!(off, FULL_SK_LEN);
680
681        FULL_SK_LEN
682    }
683
684    fn sk_decode(sk: &[u8; SK_LEN]) -> Self {
685        Self::from_bytes(sk).unwrap()
686    }
687}
688
689pub(crate) trait MLDSAPrivateKeyInternalTrait<
690    const LAMBDA: i32,
691    const GAMMA2: i32,
692    const k: usize,
693    const l: usize,
694    const eta: usize,
695    const S1_PACKED_LEN: usize,
696    const S2_PACKED_LEN: usize,
697    const PK_LEN: usize,
698    const SK_LEN: usize,
699>: Sized
700{
701    fn rho(&self) -> &[u8; 32];
702    fn K(&self) -> &[u8; 32];
703
704    /// A single entry of a privacy key vector.
705    /// These tend to be used very transiently, so we won't bother wrapping it as a Secret.
706    fn compute_s1_row(&self, idx: usize) -> Polynomial;
707
708    /// Private key component.
709    /// The packed representation sticks around for the whole computation, so
710    /// we'll wrap in as a Secret.
711    fn compute_s1_packed(&self) -> Secret<[u8; S1_PACKED_LEN]>;
712
713    /// A single entry of a privacy key vector.
714    /// These tend to be used very transiently, so we won't bother wrapping it as a Secret.
715    fn compute_s2_row(&self, idx: usize) -> Polynomial;
716
717    /// Private key component.
718    /// The packed representation sticks around for the whole computation, so
719    /// we'll wrap in as a Secret.
720    fn compute_s2_packed(&self) -> Secret<[u8; S2_PACKED_LEN]>;
721
722    /// Public key component.
723    fn compute_t0_row(
724        &self,
725        idx: usize,
726        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
727        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
728    ) -> Polynomial;
729
730    /// Public key component.
731    fn compute_t1_row(
732        &self,
733        idx: usize,
734        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
735        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
736    ) -> Polynomial;
737}
738
739impl<
740    const LAMBDA: i32,
741    const GAMMA2: i32,
742    const k: usize,
743    const l: usize,
744    const eta: usize,
745    const S1_PACKED_LEN: usize,
746    const S2_PACKED_LEN: usize,
747    const T1_PACKED_LEN: usize,
748    const PK_LEN: usize,
749    const SK_LEN: usize,
750    const FULL_SK_LEN: usize,
751>
752    MLDSAPrivateKeyInternalTrait<
753        LAMBDA,
754        GAMMA2,
755        k,
756        l,
757        eta,
758        S1_PACKED_LEN,
759        S2_PACKED_LEN,
760        PK_LEN,
761        SK_LEN,
762    >
763    for MLDSASeedPrivateKey<
764        LAMBDA,
765        GAMMA2,
766        k,
767        l,
768        eta,
769        S1_PACKED_LEN,
770        S2_PACKED_LEN,
771        T1_PACKED_LEN,
772        PK_LEN,
773        SK_LEN,
774        FULL_SK_LEN,
775    >
776{
777    fn rho(&self) -> &[u8; 32] {
778        &self.rho
779    }
780
781    fn K(&self) -> &[u8; 32] {
782        &self.K
783    }
784
785    fn compute_s1_row(&self, idx: usize) -> Polynomial {
786        debug_assert!(idx < l);
787        rej_bounded_poly::<eta>(&self.rho_prime, &(idx as u16).to_le_bytes())
788    }
789
790    fn compute_s1_packed(&self) -> Secret<[u8; S1_PACKED_LEN]> {
791        let mut s1_packed: Secret<[u8; S1_PACKED_LEN]> = Secret::new();
792        for idx in 0..l {
793            let s1_i = self.compute_s1_row(idx);
794            bit_pack_eta::<eta>(
795                &s1_i,
796                &mut s1_packed[idx * bitlen_eta(eta)..(idx + 1) * bitlen_eta(eta)],
797            );
798        }
799        s1_packed
800    }
801
802    fn compute_s2_row(&self, idx: usize) -> Polynomial {
803        debug_assert!(idx < k);
804        rej_bounded_poly::<eta>(&self.rho_prime, &((idx + l) as u16).to_le_bytes())
805    }
806
807    fn compute_s2_packed(&self) -> Secret<[u8; S2_PACKED_LEN]> {
808        let mut s2_packed: Secret<[u8; S2_PACKED_LEN]> = Secret::new();
809        for idx in 0..k {
810            let s2_i = self.compute_s2_row(idx);
811            bit_pack_eta::<eta>(
812                &s2_i,
813                &mut s2_packed[idx * bitlen_eta(eta)..(idx + 1) * bitlen_eta(eta)],
814            );
815        }
816        s2_packed
817    }
818
819    fn compute_t0_row(
820        &self,
821        idx: usize,
822        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
823        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
824    ) -> Polynomial {
825        let mut t0 = self.compute_t_row(idx, s1_packed, s2_packed);
826        for j in 0..N {
827            (_, t0[j]) = power_2_round(t0[j]);
828        }
829
830        t0
831    }
832
833    fn compute_t1_row(
834        &self,
835        idx: usize,
836        s1_packed: &Secret<[u8; S1_PACKED_LEN]>,
837        s2_packed: &Secret<[u8; S2_PACKED_LEN]>,
838    ) -> Polynomial {
839        let mut t1 = self.compute_t_row(idx, s1_packed, s2_packed);
840        for j in 0..N {
841            (t1[j], _) = power_2_round(t1[j]);
842        }
843
844        t1
845    }
846}