dcrypt_sign/dilithium/
mod.rs

1// File: crates/sign/src/pq/dilithium/mod.rs
2//! Dilithium Digital Signature Algorithm (as per FIPS 204)
3//!
4//! This module provides high-level implementations for Dilithium2, Dilithium3, and Dilithium5,
5//! which are lattice-based digital signature schemes standardized by NIST.
6//!
7//! The core cryptographic logic relies on polynomial arithmetic over rings, specific sampling
8//! distributions (Centered Binomial Distribution, uniform bounded for `y`, sparse ternary for `c`),
9//! and cryptographic hash functions (SHA3, SHAKE) provided by the `dcrypt-algorithms` crate.
10//! The security of Dilithium is based on the hardness of the Module Learning With Errors (MLWE)
11//! and Module Short Integer Solution (MSIS) problems over polynomial rings.
12//!
13//! The signing process employs the Fiat-Shamir with Aborts paradigm to achieve security
14//! against chosen message attacks.
15//!
16//! This module defines the public API for Dilithium, conforming to the `dcrypt-api::Signature` trait.
17//! Detailed implementations of internal operations are found in submodules:
18//! - `polyvec.rs`: Defines `PolyVecL`, `PolyVecK` and Dilithium-specific polynomial vector operations.
19//! - `arithmetic.rs`: Implements crucial arithmetic functions like `Power2Round`, `Decompose`,
20//!   `MakeHint`, `UseHint`, and coefficient norm checking.
21//! - `sampling.rs`: Implements Dilithium-specific sampling procedures for secret polynomials,
22//!   the masking vector `y`, and the challenge polynomial `c`.
23//! - `encoding.rs`: Handles the precise serialization and deserialization formats for public keys,
24//!   secret keys, and signatures as specified by FIPS 204.
25//! - `sign.rs`: Contains the core `keypair_internal`, `sign_internal`, and `verify_internal` logic.
26
27use crate::error::Error as SignError;
28use core::marker::PhantomData;
29use dcrypt_api::{Result as ApiResult, Signature as SignatureTrait};
30use rand::{CryptoRng, RngCore};
31use zeroize::{Zeroize, ZeroizeOnDrop};
32
33// Internal modules for Dilithium logic
34mod arithmetic;
35mod encoding;
36mod polyvec;
37mod sampling;
38mod sign;
39
40// Import what we need for public key reconstruction
41use arithmetic::power2round_polyvec;
42use polyvec::{expand_matrix_a, matrix_polyvecl_mul};
43
44// Make encoding functions accessible for serialization
45use encoding::{unpack_public_key, unpack_secret_key, unpack_signature};
46
47// Re-export from params crate for easy access to DilithiumNParams structs.
48// These structs from `dcrypt-params` hold the specific numerical parameters (K, L, eta, gamma1, etc.)
49// that define each Dilithium security level.
50use dcrypt_params::pqc::dilithium::{
51    Dilithium2Params, Dilithium3Params, Dilithium5Params, DilithiumSchemeParams,
52};
53
54// --- Public Key, Secret Key, Signature Data Wrapper Structs ---
55// These structs wrap byte vectors (`Vec<u8>`) that store the serialized representations
56// of the cryptographic objects. They provide a type-safe interface at the API boundary.
57
58/// Dilithium Public Key.
59///
60/// Stores the packed representation of `(rho, t1)`.
61/// - `rho`: A 32-byte seed used to deterministically generate the matrix A.
62/// - `t1`: A vector of K polynomials, where each coefficient is the high-order bits
63///   of `t_i = (A*s1)_i + (s2)_i`. Packed according to `P::D_PARAM` bits.
64#[derive(Clone, Debug, Zeroize)]
65pub struct DilithiumPublicKey(pub(crate) Vec<u8>);
66
67/// Dilithium Secret Key.
68///
69/// Stores the FIPS 204 compliant packed representation of `(rho, K, tr, s1, s2, t0)`.
70/// This implementation follows the standard FIPS 204 format exclusively.
71#[derive(Clone, Debug, Zeroize, ZeroizeOnDrop)]
72pub struct DilithiumSecretKey(Vec<u8>);
73
74/// Dilithium Signature Data.
75///
76/// Stores the packed representation of `(c_tilde, z, h)`.
77/// - `c_tilde`: A short seed from which the challenge polynomial `c` is derived.
78/// - `z`: A vector of L polynomials, `z = y + c*s1`.
79/// - `h`: A hint vector indicating which coefficients required correction during verification.
80#[derive(Clone, Debug)]
81pub struct DilithiumSignatureData(pub(crate) Vec<u8>);
82
83// AsRef/AsMut implementations allow access to the raw byte data.
84impl AsRef<[u8]> for DilithiumPublicKey {
85    fn as_ref(&self) -> &[u8] {
86        &self.0
87    }
88}
89impl AsMut<[u8]> for DilithiumPublicKey {
90    fn as_mut(&mut self) -> &mut [u8] {
91        &mut self.0
92    }
93}
94impl AsRef<[u8]> for DilithiumSignatureData {
95    fn as_ref(&self) -> &[u8] {
96        &self.0
97    }
98}
99impl AsMut<[u8]> for DilithiumSignatureData {
100    fn as_mut(&mut self) -> &mut [u8] {
101        &mut self.0
102    }
103}
104
105// --- DilithiumSecretKey Implementation ---
106
107impl AsRef<[u8]> for DilithiumSecretKey {
108    fn as_ref(&self) -> &[u8] {
109        &self.0
110    }
111}
112
113// NOTE: AsMut<[u8]> implementation removed for security reasons.
114// Use from_bytes() and to_bytes() for safe secret key manipulation.
115
116impl DilithiumSecretKey {
117    /// Create from FIPS 204 format bytes
118    ///
119    /// The secret key must be in the standard FIPS 204 format which includes
120    /// the tr component and appropriate padding.
121    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SignError> {
122        // Validate that the size matches one of the standard FIPS 204 sizes
123        match bytes.len() {
124            2560 => {} // Dilithium2 FIPS 204 format
125            4032 => {} // Dilithium3 FIPS 204 format
126            4896 => {} // Dilithium5 FIPS 204 format
127            _ => {
128                return Err(SignError::Deserialization(format!(
129                    "Invalid FIPS 204 secret key size: {} bytes",
130                    bytes.len()
131                )))
132            }
133        };
134
135        // Basic validation by attempting to unpack
136        match bytes.len() {
137            2560 => {
138                let _ = unpack_secret_key::<Dilithium2Params>(bytes)?;
139            }
140            4032 => {
141                let _ = unpack_secret_key::<Dilithium3Params>(bytes)?;
142            }
143            4896 => {
144                let _ = unpack_secret_key::<Dilithium5Params>(bytes)?;
145            }
146            _ => unreachable!(),
147        }
148
149        Ok(Self(bytes.to_vec()))
150    }
151
152    /// Get the serialized bytes of this secret key (FIPS 204 format)
153    pub fn to_bytes(&self) -> &[u8] {
154        &self.0
155    }
156
157    /// Extract the public key from this secret key
158    pub fn public_key(&self) -> Result<DilithiumPublicKey, SignError> {
159        match self.0.len() {
160            2560 => {
161                let (rho, _, _, s1, s2, _) = unpack_secret_key::<Dilithium2Params>(&self.0)?;
162                let pk_bytes = reconstruct_public_key::<Dilithium2Params>(&rho, &s1, &s2)?;
163                Ok(DilithiumPublicKey(pk_bytes))
164            }
165            4032 => {
166                let (rho, _, _, s1, s2, _) = unpack_secret_key::<Dilithium3Params>(&self.0)?;
167                let pk_bytes = reconstruct_public_key::<Dilithium3Params>(&rho, &s1, &s2)?;
168                Ok(DilithiumPublicKey(pk_bytes))
169            }
170            4896 => {
171                let (rho, _, _, s1, s2, _) = unpack_secret_key::<Dilithium5Params>(&self.0)?;
172                let pk_bytes = reconstruct_public_key::<Dilithium5Params>(&rho, &s1, &s2)?;
173                Ok(DilithiumPublicKey(pk_bytes))
174            }
175            _ => unreachable!(),
176        }
177    }
178}
179
180// Helper function to reconstruct public key from secret key components
181fn reconstruct_public_key<P: DilithiumSchemeParams>(
182    rho: &[u8; 32],
183    s1: &polyvec::PolyVecL<P>,
184    s2: &polyvec::PolyVecK<P>,
185) -> Result<Vec<u8>, SignError> {
186    // Expand matrix A from rho
187    let matrix_a = expand_matrix_a::<P>(rho)?;
188
189    // Convert to NTT domain
190    let mut matrix_a_hat = Vec::with_capacity(P::K_DIM);
191    for row in matrix_a {
192        let mut row_ntt = row;
193        row_ntt.ntt_inplace().map_err(SignError::from_algo)?;
194        matrix_a_hat.push(row_ntt);
195    }
196
197    let mut s1_hat = s1.clone();
198    s1_hat.ntt_inplace().map_err(SignError::from_algo)?;
199
200    let mut s2_hat = s2.clone();
201    s2_hat.ntt_inplace().map_err(SignError::from_algo)?;
202
203    // t = As1 + s2
204    let mut t_hat = matrix_polyvecl_mul(&matrix_a_hat, &s1_hat);
205    t_hat = t_hat.add(&s2_hat);
206
207    // Convert back to standard domain
208    let mut t = t_hat;
209    t.inv_ntt_inplace().map_err(SignError::from_algo)?;
210
211    // Get t1 using Power2Round
212    let (_, t1) = power2round_polyvec(&t, P::D_PARAM);
213
214    // Pack public key
215    encoding::pack_public_key::<P>(rho, &t1)
216}
217
218// --- Serialization/Deserialization Methods ---
219
220impl DilithiumPublicKey {
221    /// Deserialize a public key from bytes with full validation
222    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SignError> {
223        // Determine parameter set from key size and validate
224        match bytes.len() {
225            n if n == Dilithium2Params::PUBLIC_KEY_BYTES => {
226                let _ = unpack_public_key::<Dilithium2Params>(bytes)?;
227            }
228            n if n == Dilithium3Params::PUBLIC_KEY_BYTES => {
229                let _ = unpack_public_key::<Dilithium3Params>(bytes)?;
230            }
231            n if n == Dilithium5Params::PUBLIC_KEY_BYTES => {
232                let _ = unpack_public_key::<Dilithium5Params>(bytes)?;
233            }
234            _ => {
235                return Err(SignError::Deserialization(format!(
236                    "Invalid public key size: {} bytes",
237                    bytes.len()
238                )))
239            }
240        }
241
242        Ok(Self(bytes.to_vec()))
243    }
244
245    /// Get the serialized bytes of this public key
246    pub fn to_bytes(&self) -> &[u8] {
247        &self.0
248    }
249}
250
251impl DilithiumSignatureData {
252    /// Deserialize a signature from bytes with full validation
253    pub fn from_bytes(bytes: &[u8]) -> Result<Self, SignError> {
254        // Determine parameter set from signature size and validate
255        match bytes.len() {
256            n if n == Dilithium2Params::SIGNATURE_SIZE => {
257                let _ = unpack_signature::<Dilithium2Params>(bytes)?;
258            }
259            n if n == Dilithium3Params::SIGNATURE_SIZE => {
260                let _ = unpack_signature::<Dilithium3Params>(bytes)?;
261            }
262            n if n == Dilithium5Params::SIGNATURE_SIZE => {
263                let _ = unpack_signature::<Dilithium5Params>(bytes)?;
264            }
265            _ => {
266                return Err(SignError::Deserialization(format!(
267                    "Invalid signature size: {} bytes",
268                    bytes.len()
269                )))
270            }
271        }
272
273        Ok(Self(bytes.to_vec()))
274    }
275
276    /// Get the serialized bytes of this signature
277    pub fn to_bytes(&self) -> &[u8] {
278        &self.0
279    }
280}
281
282/// Generic Dilithium signature structure parameterized by `P: DilithiumSchemeParams`.
283pub struct Dilithium<P: DilithiumSchemeParams + 'static> {
284    _params: PhantomData<P>,
285}
286
287// --- Implement api::Signature for Dilithium<P> ---
288impl<P: DilithiumSchemeParams + Send + Sync + 'static> SignatureTrait for Dilithium<P> {
289    type PublicKey = DilithiumPublicKey;
290    type SecretKey = DilithiumSecretKey;
291    type SignatureData = DilithiumSignatureData;
292    type KeyPair = (Self::PublicKey, Self::SecretKey);
293
294    fn name() -> &'static str {
295        P::NAME
296    }
297
298    fn keypair<R: CryptoRng + RngCore>(rng: &mut R) -> ApiResult<Self::KeyPair> {
299        let (pk_bytes, sk_bytes) =
300            sign::keypair_internal::<P, R>(rng).map_err(dcrypt_api::Error::from)?;
301        let sk = DilithiumSecretKey::from_bytes(&sk_bytes).map_err(dcrypt_api::Error::from)?;
302        Ok((DilithiumPublicKey(pk_bytes), sk))
303    }
304
305    fn public_key(keypair: &Self::KeyPair) -> Self::PublicKey {
306        keypair.0.clone()
307    }
308    fn secret_key(keypair: &Self::KeyPair) -> Self::SecretKey {
309        keypair.1.clone()
310    }
311
312    fn sign(message: &[u8], secret_key: &Self::SecretKey) -> ApiResult<Self::SignatureData> {
313        let mut rng = rand::rngs::OsRng;
314        let sig_bytes = sign::sign_internal::<P, _>(message, &secret_key.0, &mut rng)
315            .map_err(dcrypt_api::Error::from)?;
316        Ok(DilithiumSignatureData(sig_bytes))
317    }
318
319    fn verify(
320        message: &[u8],
321        signature: &Self::SignatureData,
322        public_key: &Self::PublicKey,
323    ) -> ApiResult<()> {
324        sign::verify_internal::<P>(message, &signature.0, &public_key.0)
325            .map_err(dcrypt_api::Error::from)
326    }
327}
328
329// Concrete types for different Dilithium levels
330pub type Dilithium2 = Dilithium<Dilithium2Params>;
331pub type Dilithium3 = Dilithium<Dilithium3Params>;
332pub type Dilithium5 = Dilithium<Dilithium5Params>;
333
334#[cfg(test)]
335mod tests;