add sanity checks and test non-prime input

andrewwhitehead · andrewwhitehead · commit a1f122a346d6 · 2026-01-25T19:39:48.000-07:00
Signed-off-by: Andrew Whitehead &lt;cywolf@gmail.com&gt;
diff --git a/src/modular/prime_params.rs b/src/modular/prime_params.rs
@@ -27,11 +27,12 @@ pub struct PrimeParams<const LIMBS: usize> {
 
 impl<const LIMBS: usize> PrimeParams<LIMBS> {
     /// Instantiates a new set of [`PrimeParams`] given [`FixedMontyParams`] for a prime modulus.
+    ///
+    /// This method will return `None` if the modulus is determined to be non-prime, however
+    /// this is not an exhaustive check and non-prime values can be accepted.
     #[must_use]
     #[allow(clippy::unwrap_in_result, clippy::missing_panics_doc)]
     pub const fn new_vartime(params: &FixedMontyParams<LIMBS>) -> Option<Self> {
-        // A primitive root exists if and only if n is 1, 2, 4, p^k or 2p^k, k > 0, p is an odd prime
-
         let p = params.modulus();
         let p_minus_one = p.as_ref().set_bit_vartime(0, false);
         let s = NonZeroU32::new(p_minus_one.trailing_zeros_vartime()).expect("ensured non-zero");
@@ -53,6 +54,11 @@ impl<const LIMBS: usize> PrimeParams<LIMBS> {
             let exp = t.shr_vartime(1);
             // the s'th root of unity is calculated as `generator^t`
             let root = FixedMontyForm::new(&gen_uint, params).pow_vartime(&t);
+            // root^(2^(s-1)) must be equal to -1
+            let check = root.square_repeat_vartime(s.get() - 1);
+            if !Uint::eq(&check.retrieve(), &p_minus_one).to_bool_vartime() {
+                return None;
+            }
             (exp, root)
         };
 
@@ -125,6 +131,7 @@ impl<const LIMBS: usize> subtle::ConditionallySelectable for PrimeParams<LIMBS>
 const fn find_primitive_root<const LIMBS: usize>(
     p: &Odd<Uint<LIMBS>>,
 ) -> Option<(NonZeroU32, Uint<LIMBS>)> {
+    // A primitive root exists iff p is 1, 2, 4, q^k or 2q^k, k > 0, q is an odd prime.
     // Find a quadratic non-residue (primitive roots are non-residue for powers of a prime)
     let mut g = NonZeroU32::new(2u32).expect("ensured non-zero");
     loop {
@@ -148,3 +155,55 @@ const fn find_primitive_root<const LIMBS: usize>(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::PrimeParams;
+    use crate::{Choice, CtEq, CtSelect, Odd, U128, modular::MontyParams};
+
+    #[test]
+    fn check_expected() {
+        let monty_params =
+            MontyParams::new_vartime(Odd::<U128>::from_be_hex("e38af050d74b8567f73c8713cbc7bc47"));
+        let prime_params = PrimeParams::new_vartime(&monty_params).expect("failed creating params");
+        assert_eq!(prime_params.s.get(), 1);
+        assert_eq!(prime_params.generator.get(), 5);
+    }
+
+    #[test]
+    fn check_non_prime() {
+        let monty_params =
+            MontyParams::new_vartime(Odd::<U128>::from_be_hex("e38af050d74b8567f73c8713cbc7bc01"));
+        assert!(PrimeParams::new_vartime(&monty_params).is_none());
+    }
+
+    #[test]
+    fn check_equality() {
+        let monty_params_1 =
+            MontyParams::new_vartime(Odd::<U128>::from_be_hex("e38af050d74b8567f73c8713cbc7bc47"));
+        let prime_params_1 =
+            PrimeParams::new_vartime(&monty_params_1).expect("failed creating params");
+
+        let monty_params_2 =
+            MontyParams::new_vartime(Odd::<U128>::from_be_hex("f2799d643ab7ff983437c3a86cdb1beb"));
+        let prime_params_2 =
+            PrimeParams::new_vartime(&monty_params_2).expect("failed creating params");
+
+        assert!(CtEq::ct_eq(&prime_params_1, &prime_params_1).to_bool_vartime());
+        assert!(CtEq::ct_ne(&prime_params_1, &prime_params_2).to_bool_vartime());
+        assert!(
+            CtEq::ct_eq(
+                &CtSelect::ct_select(&prime_params_1, &prime_params_2, Choice::FALSE),
+                &prime_params_1,
+            )
+            .to_bool_vartime()
+        );
+        assert!(
+            CtEq::ct_eq(
+                &CtSelect::ct_select(&prime_params_1, &prime_params_2, Choice::TRUE),
+                &prime_params_2,
+            )
+            .to_bool_vartime()
+        );
+    }
+}
diff --git a/src/modular/sqrt.rs b/src/modular/sqrt.rs
@@ -1,17 +1,17 @@
 use crate::{
     Choice, CtOption, Uint,
-    modular::{MontyForm, MontyParams, prime_params::PrimeParams},
+    modular::{FixedMontyForm, FixedMontyParams, prime_params::PrimeParams},
 };
 
 /// Compute a modular square root (if it exists) given [`MontyParams`]
 /// and [`PrimeParams`] corresponding to `monty_value`, in Montgomery form.
 #[must_use]
 pub const fn sqrt_montgomery_form<const LIMBS: usize>(
     monty_value: &Uint<LIMBS>,
-    monty_params: &MontyParams<LIMBS>,
+    monty_params: &FixedMontyParams<LIMBS>,
     prime_params: &PrimeParams<LIMBS>,
 ) -> CtOption<Uint<LIMBS>> {
-    let value = MontyForm::from_montgomery(*monty_value, *monty_params);
+    let value = FixedMontyForm::from_montgomery(*monty_value, monty_params);
     let b = value.pow_vartime(&prime_params.sqrt_exp);
 
     // Constant-time versions of modular square root algorithms based on:
@@ -25,23 +25,24 @@ pub const fn sqrt_montgomery_form<const LIMBS: usize>(
         }
         2 => {
             // Algorithm 3: p = 5 mod 8 (Atkins variant)
-            let ru = MontyForm::from_montgomery(prime_params.monty_root_unity, *monty_params);
+            let ru = FixedMontyForm::from_montgomery(prime_params.monty_root_unity, monty_params);
             let cb = value.mul(&b);
             let zeta = cb.mul(&b);
             let is_one = Uint::eq(zeta.as_montgomery(), monty_params.one());
             monty_select(&cb.mul(&ru), &cb, is_one)
         }
         3 => {
             // Algorithm 4: p = 9 mod 16
-            let ru = MontyForm::from_montgomery(prime_params.monty_root_unity, *monty_params);
-            let ru_2 = MontyForm::from_montgomery(prime_params.monty_root_unity_p2, *monty_params);
+            let ru = FixedMontyForm::from_montgomery(prime_params.monty_root_unity, monty_params);
+            let ru_2 =
+                FixedMontyForm::from_montgomery(prime_params.monty_root_unity_p2, monty_params);
             let ru_3 = ru.mul(&ru_2);
             let cb = value.mul(&b);
             let zeta = cb.mul(&b);
 
             let mut m = monty_select(
                 &ru,
-                &MontyForm::one(*ru.params()),
+                &FixedMontyForm::one(ru.params()),
                 Uint::eq(zeta.as_montgomery(), monty_params.one()),
             );
             // m = ru^2 if zeta = -1
@@ -57,8 +58,9 @@ pub const fn sqrt_montgomery_form<const LIMBS: usize>(
         }
         4 => {
             // Algorithm 5: p = 17 mod 32
-            let ru = MontyForm::from_montgomery(prime_params.monty_root_unity, *monty_params);
-            let ru_2 = MontyForm::from_montgomery(prime_params.monty_root_unity_p2, *monty_params);
+            let ru = FixedMontyForm::from_montgomery(prime_params.monty_root_unity, monty_params);
+            let ru_2 =
+                FixedMontyForm::from_montgomery(prime_params.monty_root_unity_p2, monty_params);
             let ru_4 = ru_2.square();
             let ru_6 = ru_2.mul(&ru_4);
             let cb = value.mul(&b);
@@ -71,7 +73,7 @@ pub const fn sqrt_montgomery_form<const LIMBS: usize>(
 
             // m = B if -zeta in (B, C), else 1
             let mut m = monty_select(
-                &MontyForm::one(*ru.params()),
+                &FixedMontyForm::one(ru.params()),
                 &ru_2,
                 neg_zeta_b.or(monty_eq(&neg_zeta, &ru_4)),
             );
@@ -99,7 +101,8 @@ pub const fn sqrt_montgomery_form<const LIMBS: usize>(
             // Tonelli-Shanks
             let mut x = value.mul(&b);
             let mut d = x.mul(&b);
-            let mut z = MontyForm::from_montgomery(prime_params.monty_root_unity, *monty_params);
+            let mut z =
+                FixedMontyForm::from_montgomery(prime_params.monty_root_unity, monty_params);
             let mut v = prime_params.s.get();
             let mut max_v = v;
 
@@ -134,18 +137,21 @@ pub const fn sqrt_montgomery_form<const LIMBS: usize>(
     CtOption::new(x.to_montgomery(), monty_eq(&x.square(), &value))
 }
 
-const fn monty_eq<const LIMBS: usize>(a: &MontyForm<LIMBS>, b: &MontyForm<LIMBS>) -> Choice {
+const fn monty_eq<const LIMBS: usize>(
+    a: &FixedMontyForm<LIMBS>,
+    b: &FixedMontyForm<LIMBS>,
+) -> Choice {
     Uint::eq(a.as_montgomery(), b.as_montgomery())
 }
 
 const fn monty_select<const LIMBS: usize>(
-    a: &MontyForm<LIMBS>,
-    b: &MontyForm<LIMBS>,
+    a: &FixedMontyForm<LIMBS>,
+    b: &FixedMontyForm<LIMBS>,
     c: Choice,
-) -> MontyForm<LIMBS> {
-    MontyForm::from_montgomery(
+) -> FixedMontyForm<LIMBS> {
+    FixedMontyForm::from_montgomery(
         Uint::select(a.as_montgomery(), b.as_montgomery(), c),
-        *a.params(),
+        a.params(),
     )
 }
 
@@ -154,29 +160,29 @@ mod tests {
     use super::sqrt_montgomery_form;
     use crate::{
         Odd, U256, U576, Uint,
-        modular::{MontyForm, MontyParams, PrimeParams},
+        modular::{FixedMontyForm, FixedMontyParams, PrimeParams},
     };
 
     fn root_of_unity<const LIMBS: usize>(
-        monty_params: &MontyParams<LIMBS>,
+        monty_params: &FixedMontyParams<LIMBS>,
         prime_params: &PrimeParams<LIMBS>,
     ) -> Uint<LIMBS> {
-        MontyForm::from_montgomery(prime_params.monty_root_unity, *monty_params).retrieve()
+        FixedMontyForm::from_montgomery(prime_params.monty_root_unity, monty_params).retrieve()
     }
 
     fn test_monty_sqrt<const LIMBS: usize>(
-        monty_params: MontyParams<LIMBS>,
+        monty_params: FixedMontyParams<LIMBS>,
         prime_params: PrimeParams<LIMBS>,
     ) {
         let modulus = monty_params.modulus.get();
         let rounds = if cfg!(miri) { 1..=2 } else { 0..=256 };
         for i in rounds {
             let s = i * i;
-            let s_monty = MontyForm::new(&Uint::from_u32(s), &monty_params);
+            let s_monty = FixedMontyForm::new(&Uint::from_u32(s), &monty_params);
             let rt_monty =
                 sqrt_montgomery_form(s_monty.as_montgomery(), &monty_params, &prime_params)
                     .expect("no sqrt found");
-            let rt = MontyForm::from_montgomery(rt_monty, monty_params).retrieve();
+            let rt = FixedMontyForm::from_montgomery(rt_monty, &monty_params).retrieve();
             let i = Uint::from_u32(i);
             assert!(
                 Uint::eq(&rt, &i)
@@ -187,7 +193,7 @@ mod tests {
 
         // generator must be non-residue
         let generator = Uint::from_u32(prime_params.generator.get());
-        let gen_monty = MontyForm::new(&generator, &monty_params);
+        let gen_monty = FixedMontyForm::new(&generator, &monty_params);
         assert!(
             sqrt_montgomery_form(gen_monty.as_montgomery(), &monty_params, &prime_params)
                 .is_none()
@@ -199,7 +205,7 @@ mod tests {
     fn mod_sqrt_s_1() {
         // p = 3 mod 4, s = 1
         // P-256 field modulus
-        let monty_params = MontyParams::new_vartime(Odd::<U256>::from_be_hex(
+        let monty_params = FixedMontyParams::new_vartime(Odd::<U256>::from_be_hex(
             "ffffffff00000001000000000000000000000000ffffffffffffffffffffffff",
         ));
         let prime_params = PrimeParams::new_vartime(&monty_params).expect("failed creating params");
@@ -217,7 +223,7 @@ mod tests {
     fn mod_sqrt_s_2() {
         // p = 5 mod 8, s = 2
         // ed25519 base field
-        let monty_params = MontyParams::new_vartime(Odd::<U256>::from_be_hex(
+        let monty_params = FixedMontyParams::new_vartime(Odd::<U256>::from_be_hex(
             "7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffed",
         ));
         let prime_params = PrimeParams::new_vartime(&monty_params).expect("failed creating params");
@@ -235,7 +241,7 @@ mod tests {
     fn mod_sqrt_s_3() {
         // p = 9 mod 16, s = 3
         // brainpoolP384 scalar field
-        let monty_params = MontyParams::new_vartime(Odd::<U576>::from_be_hex(
+        let monty_params = FixedMontyParams::new_vartime(Odd::<U576>::from_be_hex(
             "00000000000001fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffa51868783bf2f966b7fcc0148f709a5d03bb5c9b8899c47aebb6fb71e91386409",
         ));
         let prime_params = PrimeParams::new_vartime(&monty_params).expect("failed creating params");
@@ -255,7 +261,7 @@ mod tests {
     fn mod_sqrt_s_4() {
         // p = 17 mod 32, s = 4
         // P-256 scalar field
-        let monty_params = MontyParams::new_vartime(Odd::<U256>::from_be_hex(
+        let monty_params = FixedMontyParams::new_vartime(Odd::<U256>::from_be_hex(
             "ffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551",
         ));
         let prime_params = PrimeParams::new_vartime(&monty_params).expect("failed creating params");
@@ -273,7 +279,7 @@ mod tests {
     fn mod_sqrt_s_6() {
         // s = 6
         // K-256 scalar field
-        let monty_params = MontyParams::new_vartime(Odd::<U256>::from_be_hex(
+        let monty_params = FixedMontyParams::new_vartime(Odd::<U256>::from_be_hex(
             "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141",
         ));
         let prime_params = PrimeParams::new_vartime(&monty_params).expect("failed creating params");
