bouncycastle_mldsa_lowmemory/
polynomial.rs

1//! Represents a polynomial over the ML-DSA ring.
2
3use crate::aux_functions::{high_bits, low_bits, make_hint, use_hint};
4use crate::mldsa::{MLDSA44_POLY_W1_PACKED_LEN, MLDSA65_POLY_W1_PACKED_LEN, N, q, q_inv};
5use core::ops::{Index, IndexMut};
6
7/// A polynomial over the ML-DSA ring.
8///
9/// Dev note: The following structure does not necessarily need to be declared as public.
10/// There is no real scenario where this function needs to be called directly.
11/// However, in order to test the Debug and Display traits, it is necessary to use STD, so those
12/// can't be tested from inline tests in this file and the real unit tests are in a different crate.
13/// That's the reason why pub is used.
14///
15/// # 🚨 Security 🚨
16/// Polynomials themselves are not inherently secret since sometimes they are part of public keys
17/// and sometimes private keys.
18/// It is the responsibility of the caller to wrap sensitive instances in `Secret<Polynomial>`.
19/// Note: at the moment, nothing in this crate uses `Secret<Polynomial>`, so I have left the `impl ZeroizablePrimitive` commented-out.
20#[derive(Clone, Copy)]
21pub struct Polynomial {
22    pub(crate) coeffs: [i32; N],
23}
24
25/// Convenience function to avoid ".0" all over the place.
26impl Index<usize> for Polynomial {
27    type Output = i32;
28
29    fn index(&self, index: usize) -> &Self::Output {
30        &self.coeffs[index]
31    }
32}
33/// Convenience function to avoid ".0" all over the place.
34impl IndexMut<usize> for Polynomial {
35    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
36        &mut self.coeffs[index]
37    }
38}
39
40// Turn this back on if we want to start tagging things as `Secret<Polynomial>`.
41// impl ZeroizablePrimitive for Polynomial {
42//     const ZEROED: Self = Self::new();
43// }
44
45impl Polynomial {
46    /// Create a new polynomial with all coefficients set to zero.
47    pub const fn new() -> Self {
48        Self { coeffs: [0i32; N] }
49    }
50
51    pub(crate) fn conditional_add_q(&mut self) {
52        for x in self.coeffs.iter_mut() {
53            *x = conditional_add_q(*x);
54        }
55    }
56
57    /// Algorithm 44 AddNTT(𝑎, 𝑏)̂
58    /// Computes the sum a + 𝑏 of two elements 𝑎, 𝑏 ∈ 𝑇𝑞.
59    /// Note: result could be up to 2q.
60    pub(crate) fn add_ntt(&mut self, w: &Self) {
61        for i in 0..N {
62            self[i] += w[i];
63        }
64    }
65
66    pub(crate) fn sub(&mut self, w: &Self) {
67        for i in 0..N {
68            self[i] -= w[i];
69        }
70    }
71
72    /// Algorithm 45 MultiplyNTT(𝑎, 𝑏)̂
73    /// Computes the product 𝑎 ∘̂ 𝑏 of two elements 𝑎, 𝑏 ∈ 𝑇𝑞.
74    /// Input: 𝑎, 𝑏 ∈ 𝑇𝑞.
75    /// Output: 𝑐 ∈ 𝑇𝑞.
76    /// Multiply the coefficients in this polynomial by those in another polynomial and perform montgomery reduction.
77    /// Also called pointwise montgomery multiplication
78    pub(crate) fn multiply_ntt(&mut self, b: &Polynomial) {
79        for i in 0..N {
80            self[i] = montgomery_reduce((self[i] as i64) * (b[i] as i64));
81        }
82    }
83
84    pub(crate) fn high_bits<const GAMMA2: i32>(&mut self) {
85        for i in 0..N {
86            self[i] = high_bits::<GAMMA2>(self[i]);
87        }
88    }
89
90    pub(crate) fn low_bits<const GAMMA2: i32>(&mut self) {
91        for i in 0..N {
92            self[i] = low_bits::<GAMMA2>(self[i]);
93        }
94    }
95
96    pub(crate) fn check_norm<const BOUND: i32>(&self) -> bool {
97        // Fine that this is not constant-time (returns true early) because it is used in a rejection loop.
98        // IE the early quit here leads to rejection and continuing to the top of the rejection loop, or failing 
99        // the signature validation.
100        // So the i32 that was just checked in a non-constant-time manner is about to get thrown away.
101
102        // Note: this formulation of the check_norm function usually requires this bounds check
103        //  if bound > (q - 1) / 8 {
104        //     return true;
105        //  }
106        // but since BOUND is a constant here, a debug_assert is done to ensure the value is what is expected.
107        debug_assert!(BOUND <= (q - 1) / 8);
108
109        let mut t: i32;
110        for x in self.coeffs.iter() {
111            t = *x >> 31;
112            t = *x - (t & (2 * *x));
113
114            if t >= BOUND {
115                return true;
116            }
117        }
118        false
119    }
120
121    pub(crate) fn shift_left<const d: i32>(&mut self) {
122        for x in self.coeffs.iter_mut() {
123            *x <<= d;
124        }
125    }
126
127    /// Creates the hint vector, and also returns its hamming weight (ie the number of 1's).
128    pub(crate) fn make_hint_row<const GAMMA2: i32>(&self, r: &Self) -> (Self, i32) {
129        let mut out = Polynomial::new();
130        let mut count = 0i32;
131        for i in 0..N {
132            let x = make_hint::<GAMMA2>(self[i], r[i]);
133            out[i] = x;
134            count += x;
135        }
136
137        (out, count)
138    }
139
140    pub(crate) fn w1_encode<const POLY_W1_PACKED_LEN: usize>(&self) -> [u8; POLY_W1_PACKED_LEN] {
141        // It might seem counter-intuitive for a low-memory implementation to create a tmp buffer
142        // rather than work in the provided buffer, but experimentation shows that
143        // rust is roughly an order of magnitude faster working in a scope-local array than
144        // in a referenced piece of memory.
145        // This is possibly because, once the rust compiler understands that the intermediate values are scope-local,
146        // it performs optimizations throughout all of the computation into CPU registers and skips, in this case,
147        // several hundred physical memory writes.
148        // So while it looks odd to use a scope variable in a low-memory implementation, it's way faster
149        // while seemingly maintaining the same physical memory footprint.
150        let mut r = [0u8; POLY_W1_PACKED_LEN];
151
152        match POLY_W1_PACKED_LEN {
153            MLDSA44_POLY_W1_PACKED_LEN => {
154                for i in 0..N / 4 {
155                    r[3 * i] = ((self[4 * i]) as u8) | ((self[4 * i + 1] << 6) as u8);
156                    r[3 * i + 1] = ((self[4 * i + 1] >> 2) as u8) | ((self[4 * i + 2] << 4) as u8);
157                    r[3 * i + 2] = ((self[4 * i + 2] >> 4) as u8) | ((self[4 * i + 3] << 2) as u8);
158                }
159            }
160            // ML-DSA65 and 87 share a POLY_W1_PACKED_LEN value
161            MLDSA65_POLY_W1_PACKED_LEN => {
162                for i in 0..N / 2 {
163                    r[i] = ((self[2 * i]) | (self[2 * i + 1] << 4)) as u8;
164                }
165            }
166            _ => {
167                unreachable!()
168            }
169        }
170
171        r
172    }
173
174    /// Algorithm 41 NTT(𝑤)
175    /// Computes the NTT.
176    /// Input: Polynomial 𝑤(𝑋) = Σ_{j=0}^{255} 𝑤𝑗𝑋𝑗 ∈ 𝑅𝑞.
177    /// Output: 𝑤_hat = (𝑤_hat\[0], ..., 𝑤_hat\[255]) ∈ 𝑇𝑞.
178    ///
179    /// Note: by convention, variables holding the output of the NTT function should be named "_hat"
180    /// to indicate that they are in the NTT domain (sometimes called the frequency domain), not the natural domain.
181    /// Usage of the rust type system to enforce this is arguably unnecessary, since that's what the NIST
182    /// test vectors are for.
183    ///
184    /// Design choice: the NTT is not done in-place, but copy data to a new array.
185    /// This uses slightly more memory and requires a copy, but makes the code easier to read
186    /// and less likely to contain a bug. But this optimization could be considered in the future.
187    pub(crate) fn ntt(&mut self) {
188        let mut m: usize = 0;
189        let mut len: usize = 128;
190
191        while len >= 1 {
192            let mut start: usize = 0;
193            while start < N {
194                m += 1;
195                let z: i32 = ZETAS[m];
196
197                for j in start..start + len {
198                    let t = montgomery_reduce(z as i64 * self[j + len] as i64);
199                    self[j + len] = self[j] - t; // '% q' not strictly needed cause it gets reduced at some point later. Removing it gave +5% in benchmarking
200                    self[j] = self[j] + t; // '% q' not strictly needed
201                }
202                start = start + 2 * len;
203            }
204            len >>= 1;
205        }
206    }
207
208    /// Algorithm 42 NTT−1(𝑤)̂
209    /// Computes the inverse of the NTT.
210    /// Input: ̂̂ ̂ 𝑤 = (𝑤\[0], … , 𝑤\[255]) ∈ 𝑇𝑞.
211    /// Output: Polynomial 𝑤(𝑋) = ∑255
212    /// 𝑗=0 𝑤𝑗𝑋𝑗 ∈ 𝑅𝑞
213    pub(crate) fn inv_ntt(&mut self) {
214        let mut m: usize = N;
215        let mut len: usize = 1;
216
217        while len < N {
218            let mut start: usize = 0;
219            while start < N {
220                m -= 1;
221                let z = (-1) * ZETAS[m];
222
223                // j = start;
224                // while j < start + len {
225                for j in start..start + len {
226                    // 𝑡 ← 𝑤𝑗
227                    let t: i32 = self[j];
228
229                    // 𝑤𝑗 ← (𝑡 + 𝑤𝑗+𝑙𝑒𝑛) mod 𝑞
230                    self[j] = t + self[j + len];
231
232                    // 𝑤𝑗+𝑙𝑒𝑛 ← (𝑡 − 𝑤𝑗+𝑙𝑒𝑛) mod 𝑞
233                    self[j + len] = t - self[j + len];
234
235                    // 𝑤𝑗+𝑙𝑒𝑛 ← (𝑧 ⋅ 𝑤𝑗+𝑙𝑒𝑛) mod 𝑞
236                    self[j + len] = montgomery_reduce(z as i64 * self[j + len] as i64);
237                }
238                start = start + 2 * len; // could be optimized to save the multiply-by-two since j finishes as `start + len`. That said 2* is just << 1, which is basically free.
239            }
240            len <<= 1;
241        }
242
243        // f = 256^-1 mod q
244        // const f: i64 = 8347681;
245        // bc-java uses this value rather than the one in FIPS 204
246        const f: i64 = 41978;
247        for j in 0..N {
248            // equiv. to the global constant N
249            self[j] = montgomery_reduce(f * self[j] as i64);
250        }
251    }
252
253    pub(crate) fn use_hint<const GAMMA2: i32>(&mut self, h: &Polynomial) {
254        for i in 0..N {
255            self[i] = use_hint::<GAMMA2>(self[i], h[i]);
256        }
257    }
258}
259
260/// FIPS 204 Algorithm 49
261/// As described in FIPS 204 Appendix A, montgomery reduction allows for efficient computation
262/// of expressions of the form c = a * b (mod q).
263/// The output is not necessarily less than q in absolute value, but it is less than 2q in absolute value
264pub(crate) fn montgomery_reduce(a: i64) -> i32 {
265    debug_assert!(a > -((q as i64) << 31) && a < ((q as i64) << 31));
266
267    // 2: 𝑡 ← ((𝑎 mod 2^32) ⋅ QINV) mod 2^32
268    let t: i32 = (a as i32).wrapping_mul(q_inv);
269
270    // 3: 𝑟 ← (𝑎 − 𝑡 ⋅ 𝑞)/2^32
271    ((a - ((t as i64) * (q as i64))) >> 32) as i32
272}
273
274pub(crate) fn conditional_add_q(a: i32) -> i32 {
275    a + ((a >> 31) & q)
276}
277
278#[test]
279/// These are the results it's giving; I'm not sure if these are "correct" or not.
280fn test_conditional_add_q() {
281    assert_eq!(conditional_add_q(-q - 1), -1);
282    assert_eq!(conditional_add_q(-q), 0);
283    assert_eq!(conditional_add_q(-q - 2), -2);
284    assert_eq!(conditional_add_q(-q + 1), 1);
285    assert_eq!(conditional_add_q(-1), q - 1);
286    assert_eq!(conditional_add_q(0), 0);
287    assert_eq!(conditional_add_q(1), 1);
288    assert_eq!(conditional_add_q(q - 1), q - 1);
289    assert_eq!(conditional_add_q(q), q);
290    assert_eq!(conditional_add_q(q + 1), q + 1);
291}
292
293/// Constants for NTT
294const ZETAS: [i32; 256] = [
295    0, 25847, -2608894, -518909, 237124, -777960, -876248, 466468, 1826347, 2353451, -359251,
296    -2091905, 3119733, -2884855, 3111497, 2680103, 2725464, 1024112, -1079900, 3585928, -549488,
297    -1119584, 2619752, -2108549, -2118186, -3859737, -1399561, -3277672, 1757237, -19422, 4010497,
298    280005, 2706023, 95776, 3077325, 3530437, -1661693, -3592148, -2537516, 3915439, -3861115,
299    -3043716, 3574422, -2867647, 3539968, -300467, 2348700, -539299, -1699267, -1643818, 3505694,
300    -3821735, 3507263, -2140649, -1600420, 3699596, 811944, 531354, 954230, 3881043, 3900724,
301    -2556880, 2071892, -2797779, -3930395, -1528703, -3677745, -3041255, -1452451, 3475950,
302    2176455, -1585221, -1257611, 1939314, -4083598, -1000202, -3190144, -3157330, -3632928, 126922,
303    3412210, -983419, 2147896, 2715295, -2967645, -3693493, -411027, -2477047, -671102, -1228525,
304    -22981, -1308169, -381987, 1349076, 1852771, -1430430, -3343383, 264944, 508951, 3097992,
305    44288, -1100098, 904516, 3958618, -3724342, -8578, 1653064, -3249728, 2389356, -210977, 759969,
306    -1316856, 189548, -3553272, 3159746, -1851402, -2409325, -177440, 1315589, 1341330, 1285669,
307    -1584928, -812732, -1439742, -3019102, -3881060, -3628969, 3839961, 2091667, 3407706, 2316500,
308    3817976, -3342478, 2244091, -2446433, -3562462, 266997, 2434439, -1235728, 3513181, -3520352,
309    -3759364, -1197226, -3193378, 900702, 1859098, 909542, 819034, 495491, -1613174, -43260,
310    -522500, -655327, -3122442, 2031748, 3207046, -3556995, -525098, -768622, -3595838, 342297,
311    286988, -2437823, 4108315, 3437287, -3342277, 1735879, 203044, 2842341, 2691481, -2590150,
312    1265009, 4055324, 1247620, 2486353, 1595974, -3767016, 1250494, 2635921, -3548272, -2994039,
313    1869119, 1903435, -1050970, -1333058, 1237275, -3318210, -1430225, -451100, 1312455, 3306115,
314    -1962642, -1279661, 1917081, -2546312, -1374803, 1500165, 777191, 2235880, 3406031, -542412,
315    -2831860, -1671176, -1846953, -2584293, -3724270, 594136, -3776993, -2013608, 2432395, 2454455,
316    -164721, 1957272, 3369112, 185531, -1207385, -3183426, 162844, 1616392, 3014001, 810149,
317    1652634, -3694233, -1799107, -3038916, 3523897, 3866901, 269760, 2213111, -975884, 1717735,
318    472078, -426683, 1723600, -1803090, 1910376, -1667432, -1104333, -260646, -3833893, -2939036,
319    -2235985, -420899, -2286327, 183443, -976891, 1612842, -3545687, -554416, 3919660, -48306,
320    -1362209, 3937738, 1400424, -846154, 1976782,
321];
bouncycastle_mldsa_lowmemory/polynomial.rs

bouncycastle_mldsa_lowmemory/
polynomial.rs
