cnfy-uint 0.2.3

//! Constant-time modular exponentiation for [`U256`].
use super::U256;

impl U256 {
    /// Computes `self^exp mod modulus` in constant time.
    ///
    /// Always performs exactly 256 iterations (one per bit of the exponent)
    /// regardless of the actual exponent value. Each iteration unconditionally
    /// squares and multiplies, then uses `ct_select` to discard the multiply
    /// result when the exponent bit is zero. This prevents leaking the
    /// exponent through timing or branch prediction.
    ///
    /// Returns 1 when `exp` is zero (for any base, including zero), matching
    /// the mathematical convention `0^0 = 1`. This falls out naturally: no
    /// exponent bits are set, so every multiply result is discarded.
    #[inline]
    pub fn pow_mod_ct(&self, exp: &U256, modulus: &U256) -> U256 {
        let base = self.modulo(modulus);
        let mut result = U256::ONE;

        let mut i = 256u32;
        while i > 0 {
            i -= 1;
            result = result.square_mod(modulus);
            let multiplied = result.mul_mod(&base, modulus);
            result = multiplied.ct_select(&result, exp.bit(i));
        }

        result
    }
}

#[cfg(test)]
mod ai_tests {
    use super::*;

    const P: U256 = U256::from_be_limbs([
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    ]);

    /// x^0 = 1 for any x.
    #[test]
    fn exp_zero() {
        let base = U256::from_be_limbs([0, 0, 0, 42]);
        assert_eq!(base.pow_mod_ct(&U256::ZERO, &P), U256::ONE);
    }

    /// x^1 = x.
    #[test]
    fn exp_one() {
        let base = U256::from_be_limbs([0, 0, 0, 42]);
        assert_eq!(base.pow_mod_ct(&U256::ONE, &P), base);
    }

    /// 2^10 = 1024.
    #[test]
    fn two_pow_ten() {
        let base = U256::from_be_limbs([0, 0, 0, 2]);
        let exp = U256::from_be_limbs([0, 0, 0, 10]);
        assert_eq!(base.pow_mod_ct(&exp, &P), U256::from_be_limbs([0, 0, 0, 1024]));
    }

    /// 3^5 = 243.
    #[test]
    fn three_pow_five() {
        let base = U256::from_be_limbs([0, 0, 0, 3]);
        let exp = U256::from_be_limbs([0, 0, 0, 5]);
        assert_eq!(base.pow_mod_ct(&exp, &P), U256::from_be_limbs([0, 0, 0, 243]));
    }

    /// Fermat's little theorem: a^(P-1) ≡ 1 (mod P).
    #[test]
    fn fermats_little_theorem() {
        let base = U256::from_be_limbs([0, 0, 0, 7]);
        let p_minus_1 = P - U256::ONE;
        assert_eq!(base.pow_mod_ct(&p_minus_1, &P), U256::ONE);
    }

    /// 0^0 = 1 (mathematical convention).
    #[test]
    fn zero_pow_zero() {
        assert_eq!(U256::ZERO.pow_mod_ct(&U256::ZERO, &P), U256::ONE);
    }

    /// 0^n = 0 for n > 0.
    #[test]
    fn zero_base() {
        let exp = U256::from_be_limbs([0, 0, 0, 5]);
        assert_eq!(U256::ZERO.pow_mod_ct(&exp, &P), U256::ZERO);
    }

    /// CT matches VT for various inputs.
    #[test]
    fn matches_vt() {
        let cases = [
            (U256::from_be_limbs([0, 0, 0, 2]), U256::from_be_limbs([0, 0, 0, 10])),
            (U256::from_be_limbs([0, 0, 0, 7]), P - U256::ONE),
            (U256::ZERO, U256::ZERO),
            (U256::from_be_limbs([0, 0, 0, 13]), U256::from_be_limbs([0, 0, 0, 7])),
        ];
        for (base, exp) in &cases {
            assert_eq!(
                base.pow_mod_ct(exp, &P),
                base.pow_mod_vt(exp, &P),
                "CT/VT mismatch for {base:?}^{exp:?}"
            );
        }
    }

    /// Matches iterative multiplication for small cases.
    #[test]
    fn matches_iterative_mul() {
        let base = U256::from_be_limbs([0, 0, 0, 13]);
        let mut expected = U256::ONE;
        for _ in 0..7 {
            expected = expected.mul_mod(&base, &P);
        }
        let exp = U256::from_be_limbs([0, 0, 0, 7]);
        assert_eq!(base.pow_mod_ct(&exp, &P), expected);
    }
}