bouncycastle_mldsa_lowmemory/

polynomial.rs

//! Represents a polynomial over the ML-DSA ring.

use crate::aux_functions::{high_bits, low_bits, make_hint, use_hint};
use crate::mldsa::{MLDSA44_POLY_W1_PACKED_LEN, MLDSA65_POLY_W1_PACKED_LEN, N, q, q_inv};
use core::ops::{Index, IndexMut};

/// A polynomial over the ML-DSA ring.
///
/// Dev note: The following structure does not necessarily need to be declared as public.
/// There is no real scenario where this function needs to be called directly.
/// However, in order to test the Debug and Display traits, it is necessary to use STD, so those
/// can't be tested from inline tests in this file and the real unit tests are in a different crate.
/// That's the reason why pub is used.
///
/// # 🚨 Security 🚨
/// Polynomials themselves are not inherently secret since sometimes they are part of public keys
/// and sometimes private keys.
/// It is the responsibility of the caller to wrap sensitive instances in `Secret<Polynomial>`.
/// Note: at the moment, nothing in this crate uses `Secret<Polynomial>`, so I have left the `impl ZeroizablePrimitive` commented-out.
#[derive(Clone, Copy)]
pub struct Polynomial {
    pub(crate) coeffs: [i32; N],
}

/// Convenience function to avoid ".0" all over the place.
impl Index<usize> for Polynomial {
    type Output = i32;

    fn index(&self, index: usize) -> &Self::Output {
        &self.coeffs[index]
    }
}
/// Convenience function to avoid ".0" all over the place.
impl IndexMut<usize> for Polynomial {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.coeffs[index]
    }
}

// Turn this back on if we want to start tagging things as `Secret<Polynomial>`.
// impl ZeroizablePrimitive for Polynomial {
//     const ZEROED: Self = Self::new();
// }

impl Polynomial {
    /// Create a new polynomial with all coefficients set to zero.
    pub const fn new() -> Self {
        Self { coeffs: [0i32; N] }
    }

    pub(crate) fn conditional_add_q(&mut self) {
        for x in self.coeffs.iter_mut() {
            *x = conditional_add_q(*x);
        }
    }

    /// Reduces every coefficient to |𝑐| < 𝑞 without changing its residue. See [`reduce32`] for when
    /// this is required for security reasons; note that [`Self::reduce`] is a *Montgomery* reduction and is not a substitute.
    pub(crate) fn reduce32(&mut self) {
        for x in self.coeffs.iter_mut() {
            *x = reduce32(*x);
        }
    }

    /// Algorithm 44 AddNTT(𝑎, 𝑏)̂
    /// Computes the sum a + 𝑏 of two elements 𝑎, 𝑏 ∈ 𝑇𝑞.
    /// Note: result could be up to 2q.
    pub(crate) fn add_ntt(&mut self, w: &Self) {
        for i in 0..N {
            self[i] += w[i];
        }
    }

    pub(crate) fn sub(&mut self, w: &Self) {
        for i in 0..N {
            self[i] -= w[i];
        }
    }

    /// Algorithm 45 MultiplyNTT(𝑎, 𝑏)̂
    /// Computes the product 𝑎 ∘̂ 𝑏 of two elements 𝑎, 𝑏 ∈ 𝑇𝑞.
    /// Input: 𝑎, 𝑏 ∈ 𝑇𝑞.
    /// Output: 𝑐 ∈ 𝑇𝑞.
    /// Multiply the coefficients in this polynomial by those in another polynomial and perform montgomery reduction.
    /// Also called pointwise montgomery multiplication
    pub(crate) fn multiply_ntt(&mut self, b: &Polynomial) {
        for i in 0..N {
            self[i] = montgomery_reduce((self[i] as i64) * (b[i] as i64));
        }
    }

    pub(crate) fn high_bits<const GAMMA2: i32>(&mut self) {
        for i in 0..N {
            self[i] = high_bits::<GAMMA2>(self[i]);
        }
    }

    pub(crate) fn low_bits<const GAMMA2: i32>(&mut self) {
        for i in 0..N {
            self[i] = low_bits::<GAMMA2>(self[i]);
        }
    }

    pub(crate) fn check_norm<const BOUND: i32>(&self) -> bool {
        // Fine that this is not constant-time (returns true early) because it is used in a rejection loop.
        // IE the early quit here leads to rejection and continuing to the top of the rejection loop, or failing 
        // the signature validation.
        // So the i32 that was just checked in a non-constant-time manner is about to get thrown away.

        // Note: this formulation of the check_norm function usually requires this bounds check
        //  if bound > (q - 1) / 8 {
        //     return true;
        //  }
        // but since BOUND is a constant here, a debug_assert is done to ensure the value is what is expected.
        debug_assert!(BOUND <= (q - 1) / 8);

        let mut t: i32;
        for x in self.coeffs.iter() {
            t = *x >> 31;
            t = *x - (t & (2 * *x));

            if t >= BOUND {
                return true;
            }
        }
        false
    }

    pub(crate) fn shift_left<const d: i32>(&mut self) {
        for x in self.coeffs.iter_mut() {
            *x <<= d;
        }
    }

    /// Creates the hint vector, and also returns its hamming weight (ie the number of 1's).
    pub(crate) fn make_hint_row<const GAMMA2: i32>(&self, r: &Self) -> (Self, i32) {
        let mut out = Polynomial::new();
        let mut count = 0i32;
        for i in 0..N {
            let x = make_hint::<GAMMA2>(self[i], r[i]);
            out[i] = x;
            count += x;
        }

        (out, count)
    }

    pub(crate) fn w1_encode<const POLY_W1_PACKED_LEN: usize>(&self) -> [u8; POLY_W1_PACKED_LEN] {
        // It might seem counter-intuitive for a low-memory implementation to create a tmp buffer
        // rather than work in the provided buffer, but experimentation shows that
        // rust is roughly an order of magnitude faster working in a scope-local array than
        // in a referenced piece of memory.
        // This is possibly because, once the rust compiler understands that the intermediate values are scope-local,
        // it performs optimizations throughout all of the computation into CPU registers and skips, in this case,
        // several hundred physical memory writes.
        // So while it looks odd to use a scope variable in a low-memory implementation, it's way faster
        // while seemingly maintaining the same physical memory footprint.
        let mut r = [0u8; POLY_W1_PACKED_LEN];

        match POLY_W1_PACKED_LEN {
            MLDSA44_POLY_W1_PACKED_LEN => {
                for i in 0..N / 4 {
                    r[3 * i] = ((self[4 * i]) as u8) | ((self[4 * i + 1] << 6) as u8);
                    r[3 * i + 1] = ((self[4 * i + 1] >> 2) as u8) | ((self[4 * i + 2] << 4) as u8);
                    r[3 * i + 2] = ((self[4 * i + 2] >> 4) as u8) | ((self[4 * i + 3] << 2) as u8);
                }
            }
            // ML-DSA65 and 87 share a POLY_W1_PACKED_LEN value
            MLDSA65_POLY_W1_PACKED_LEN => {
                for i in 0..N / 2 {
                    r[i] = ((self[2 * i]) | (self[2 * i + 1] << 4)) as u8;
                }
            }
            _ => {
                unreachable!()
            }
        }

        r
    }

    /// Algorithm 41 NTT(𝑤)
    /// Computes the NTT.
    /// Input: Polynomial 𝑤(𝑋) = Σ_{j=0}^{255} 𝑤𝑗𝑋𝑗 ∈ 𝑅𝑞.
    /// Output: 𝑤_hat = (𝑤_hat\[0], ..., 𝑤_hat\[255]) ∈ 𝑇𝑞.
    ///
    /// Note: by convention, variables holding the output of the NTT function should be named "_hat"
    /// to indicate that they are in the NTT domain (sometimes called the frequency domain), not the natural domain.
    /// Usage of the rust type system to enforce this is arguably unnecessary, since that's what the NIST
    /// test vectors are for.
    ///
    /// Design choice: the NTT is not done in-place, but copy data to a new array.
    /// This uses slightly more memory and requires a copy, but makes the code easier to read
    /// and less likely to contain a bug. But this optimization could be considered in the future.
    pub(crate) fn ntt(&mut self) {
        let mut m: usize = 0;
        let mut len: usize = 128;

        while len >= 1 {
            let mut start: usize = 0;
            while start < N {
                m += 1;
                let z: i32 = ZETAS[m];

                for j in start..start + len {
                    let t = montgomery_reduce(z as i64 * self[j + len] as i64);
                    self[j + len] = self[j] - t; // '% q' not strictly needed cause it gets reduced at some point later. Removing it gave +5% in benchmarking
                    self[j] = self[j] + t; // '% q' not strictly needed
                }
                start = start + 2 * len;
            }
            len >>= 1;
        }
    }

    /// Algorithm 42 NTT−1(𝑤)̂
    /// Computes the inverse of the NTT.
    /// Input: ̂̂ ̂ 𝑤 = (𝑤\[0], … , 𝑤\[255]) ∈ 𝑇𝑞.
    /// Output: Polynomial 𝑤(𝑋) = ∑255
    /// 𝑗=0 𝑤𝑗𝑋𝑗 ∈ 𝑅𝑞
    pub(crate) fn inv_ntt(&mut self) {
        // A core input condition on the InverseNTT is that every input
        // coefficient must satisfy |𝑤| < 𝑞 for the sums to stay within an i32.
        // The reason this is required is because Algorithm 42's butterflies (steps 13-14)
        // run 8 levels without reducing, leading to an i32 overflow if the input was out-of-range.
        // Callers must pass them through `reduce32` first.
        // Note that only a crafted input, such as (Wycheproof mldsa_87_verify tcId 240/241) will trigger this.
        debug_assert!(
            self.coeffs.iter().all(|c| c.abs() < q),
            "inv_ntt input coefficient not reduced below q; call reduce32() first"
        );

        let mut m: usize = N;
        let mut len: usize = 1;

        while len < N {
            let mut start: usize = 0;
            while start < N {
                m -= 1;
                let z = (-1) * ZETAS[m];

                // j = start;
                // while j < start + len {
                for j in start..start + len {
                    // 𝑡 ← 𝑤𝑗
                    let t: i32 = self[j];

                    // 𝑤𝑗 ← (𝑡 + 𝑤𝑗+𝑙𝑒𝑛) mod 𝑞
                    self[j] = t + self[j + len];

                    // 𝑤𝑗+𝑙𝑒𝑛 ← (𝑡 − 𝑤𝑗+𝑙𝑒𝑛) mod 𝑞
                    self[j + len] = t - self[j + len];

                    // 𝑤𝑗+𝑙𝑒𝑛 ← (𝑧 ⋅ 𝑤𝑗+𝑙𝑒𝑛) mod 𝑞
                    self[j + len] = montgomery_reduce(z as i64 * self[j + len] as i64);
                }
                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.
            }
            len <<= 1;
        }

        // f = 256^-1 mod q
        // const f: i64 = 8347681;
        // bc-java uses this value rather than the one in FIPS 204
        const f: i64 = 41978;
        for j in 0..N {
            // equiv. to the global constant N
            self[j] = montgomery_reduce(f * self[j] as i64);
        }
    }

    pub(crate) fn use_hint<const GAMMA2: i32>(&mut self, h: &Polynomial) {
        for i in 0..N {
            self[i] = use_hint::<GAMMA2>(self[i], h[i]);
        }
    }
}

/// FIPS 204 Algorithm 49
/// As described in FIPS 204 Appendix A, montgomery reduction allows for efficient computation
/// of expressions of the form c = a * b (mod q).
/// The output is not necessarily less than q in absolute value, but it is less than 2q in absolute value
pub(crate) fn montgomery_reduce(a: i64) -> i32 {
    debug_assert!(a > -((q as i64) << 31) && a < ((q as i64) << 31));

    // 2: 𝑡 ← ((𝑎 mod 2^32) ⋅ QINV) mod 2^32
    let t: i32 = (a as i32).wrapping_mul(q_inv);

    // 3: 𝑟 ← (𝑎 − 𝑡 ⋅ 𝑞)/2^32
    ((a - ((t as i64) * (q as i64))) >> 32) as i32
}

pub(crate) fn conditional_add_q(a: i32) -> i32 {
    a + ((a >> 31) & q)
}

/// Plain (non-Montgomery) reduction: for 𝑎 ≤ 2^31 − 2^22 − 1 returns 𝑟 ≡ 𝑎 (mod 𝑞) with |𝑟| ≤ 6283008 < 𝑞.
/// `reduce32` in the reference implementation (pq-crystals/dilithium, ref/reduce.c).
///
/// This is not in FIPS 204 since it assumes all arithmetic is done mod q.
/// In this implementation, sums of Montgomery products are left unreduced, but `inv_ntt` (Algorithm 42)
/// needs |input| < 𝑞 to stay within an `i32`, so this is applied to those sums first. Omitting it before
/// the verifier's NTT⁻¹ is exploitable (<https://eprint.iacr.org/2026/1032>, Wycheproof mldsa_87_verify tcId 240/241).
pub(crate) fn reduce32(a: i32) -> i32 {
    // The reference implementation's stated input bound; above it `a + 2^22` overflows.
    debug_assert!(a <= i32::MAX - (1 << 22));
    let t = (a + (1 << 22)) >> 23;
    a - t * q
}

#[test]
/// These are the results it's giving; I'm not sure if these are "correct" or not.
fn test_conditional_add_q() {
    assert_eq!(conditional_add_q(-q - 1), -1);
    assert_eq!(conditional_add_q(-q), 0);
    assert_eq!(conditional_add_q(-q - 2), -2);
    assert_eq!(conditional_add_q(-q + 1), 1);
    assert_eq!(conditional_add_q(-1), q - 1);
    assert_eq!(conditional_add_q(0), 0);
    assert_eq!(conditional_add_q(1), 1);
    assert_eq!(conditional_add_q(q - 1), q - 1);
    assert_eq!(conditional_add_q(q), q);
    assert_eq!(conditional_add_q(q + 1), q + 1);
}

/// Constants for NTT
const ZETAS: [i32; 256] = [
    0, 25847, -2608894, -518909, 237124, -777960, -876248, 466468, 1826347, 2353451, -359251,
    -2091905, 3119733, -2884855, 3111497, 2680103, 2725464, 1024112, -1079900, 3585928, -549488,
    -1119584, 2619752, -2108549, -2118186, -3859737, -1399561, -3277672, 1757237, -19422, 4010497,
    280005, 2706023, 95776, 3077325, 3530437, -1661693, -3592148, -2537516, 3915439, -3861115,
    -3043716, 3574422, -2867647, 3539968, -300467, 2348700, -539299, -1699267, -1643818, 3505694,
    -3821735, 3507263, -2140649, -1600420, 3699596, 811944, 531354, 954230, 3881043, 3900724,
    -2556880, 2071892, -2797779, -3930395, -1528703, -3677745, -3041255, -1452451, 3475950,
    2176455, -1585221, -1257611, 1939314, -4083598, -1000202, -3190144, -3157330, -3632928, 126922,
    3412210, -983419, 2147896, 2715295, -2967645, -3693493, -411027, -2477047, -671102, -1228525,
    -22981, -1308169, -381987, 1349076, 1852771, -1430430, -3343383, 264944, 508951, 3097992,
    44288, -1100098, 904516, 3958618, -3724342, -8578, 1653064, -3249728, 2389356, -210977, 759969,
    -1316856, 189548, -3553272, 3159746, -1851402, -2409325, -177440, 1315589, 1341330, 1285669,
    -1584928, -812732, -1439742, -3019102, -3881060, -3628969, 3839961, 2091667, 3407706, 2316500,
    3817976, -3342478, 2244091, -2446433, -3562462, 266997, 2434439, -1235728, 3513181, -3520352,
    -3759364, -1197226, -3193378, 900702, 1859098, 909542, 819034, 495491, -1613174, -43260,
    -522500, -655327, -3122442, 2031748, 3207046, -3556995, -525098, -768622, -3595838, 342297,
    286988, -2437823, 4108315, 3437287, -3342277, 1735879, 203044, 2842341, 2691481, -2590150,
    1265009, 4055324, 1247620, 2486353, 1595974, -3767016, 1250494, 2635921, -3548272, -2994039,
    1869119, 1903435, -1050970, -1333058, 1237275, -3318210, -1430225, -451100, 1312455, 3306115,
    -1962642, -1279661, 1917081, -2546312, -1374803, 1500165, 777191, 2235880, 3406031, -542412,
    -2831860, -1671176, -1846953, -2584293, -3724270, 594136, -3776993, -2013608, 2432395, 2454455,
    -164721, 1957272, 3369112, 185531, -1207385, -3183426, 162844, 1616392, 3014001, 810149,
    1652634, -3694233, -1799107, -3038916, 3523897, 3866901, 269760, 2213111, -975884, 1717735,
    472078, -426683, 1723600, -1803090, 1910376, -1667432, -1104333, -260646, -3833893, -2939036,
    -2235985, -420899, -2286327, 183443, -976891, 1612842, -3545687, -554416, 3919660, -48306,
    -1362209, 3937738, 1400424, -846154, 1976782,
];

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_reduce32() {
        // congruent to the input mod q, and within the reference implementation's stated output range
        for &a in &[
            0,
            1,
            -1,
            q - 1,
            q,
            q + 1,
            -q,
            -q - 1,
            8 * q,
            -8 * q,
            6283008,
            -6283008,
            i32::MIN,
            i32::MAX - (1 << 22),
        ] {
            let r = reduce32(a);
            assert!((-6283008..=6283008).contains(&r), "reduce32({a}) = {r} out of range");
            assert_eq!(
                (r as i64 - a as i64).rem_euclid(q as i64),
                0,
                "reduce32({a}) = {r} not congruent"
            );
            assert!(r.abs() < q);
        }
        // the largest sum inv_ntt may see: (l + 1) products each in (-q, q), for l = 7
        assert!(reduce32(8 * q - 8).abs() < q);
        assert!(reduce32(-(8 * q - 8)).abs() < q);
    }
}