Find squre root modulo prime number or RSA modulo

tonelli_shanks.rs

use rand_core::RngCore;
use unknown_order::BigNumber;

pub fn tonelli_shanks(n: &BigNumber, p: &BigNumber, z: &BigNumber) -> Option<BigNumber> {
    let zero = BigNumber::zero();
    let one = BigNumber::one();
    let two = BigNumber::from(2);
    // find q, s such that n = q * 2^s
    let (q, s) = {
        let mut s = BigNumber::zero();
        let mut q = p.clone() - 1;
        while &q % 2 == zero {
            q = &q >> 1;
            s += 1;
        }
        (q, s)
    };
    // main search loop
    let mut m = s;
    let mut c = z.modpow(&q, &p);
    let mut t = n.modpow(&q, &p);
    let mut r = n.modpow(&((&q + 1) / 2), &p);
    loop {
        if t == zero {
            break Some(zero);
        } else if t == one {
            break Some(r);
        } else {
            // find smallest i such that 0 < i < m, t^(2^i) = 1
            // If it doesn't exist, n is non-residue and we exit without result
            // with `?`
            let i = {
                let mut i = BigNumber::one();
                let mut acc = t.modmul(&t, &p);
                loop {
                    if i == m {
                        break None;
                    } else if acc == one {
                        break Some(i);
                    } else {
                        acc = acc.modmul(&acc, &p);
                        i += 1;
                    }
                }
            }?;
            let e = two.modpow(&(m - &i - 1), &p);
            let b = c.modpow(&e, &p);
            m = i;
            c = b.modmul(&b, &p);
            t = t.modmul(&c, &p);
            r = r.modmul(&b, &p);
        }
    }
}


pub fn chinese_remainder_construction((mp, p): (&BigNumber, &BigNumber), (mq, q): (&BigNumber, &BigNumber)) -> BigNumber {
    let egcd = p.extended_gcd(q);
    let kp = egcd.x;
    let kq = egcd.y;
    (mp * kq * q).modadd(&(mq * kp * p), &(p * q))
}


pub fn square_roots(x: BigNumber, p: BigNumber, q: BigNumber) -> Option<impl Iterator<Item = BigNumber>> {
    let mut rng = rand_core::OsRng::default(); // XXX
    let zp = non_residue_in(&p, &mut rng);
    let zq = non_residue_in(&q, &mut rng);
    let rp = tonelli_shanks(&x, &p, &zp)?;
    let rq = tonelli_shanks(&x, &q, &zq)?;
    let nrp = &p - &rp;
    let nrq = &q - &rq;
    Some(TWO_BOOLS.iter().map(move |(sp, sq)| {
        let mp = if *sp { &rp } else { &nrp };
        let mq = if *sq { &rq } else { &nrq };
        let root = chinese_remainder_construction((&mp, &p), (&mq, &q));
        root
    }))
}

pub const TWO_BOOLS: [(bool, bool); 4] = [
    (false, false),
    (true, false),
    (false, true),
    (true, true),
];


pub fn non_residue_in<R: RngCore>(n: &BigNumber, mut rng: R) -> BigNumber {
    loop {
        let w = BigNumber::from_rng(&n, &mut rng);
        if jacobi(&w, &n) == -1 {
            break w;
        }
    }
}


#[cfg(test)]
mod test {
    #[test]
    fn tonelli_shanks() {
        let mut rng = rand_core::OsRng::default();
        let p = BigNumber::prime(256);
        let z = super::non_residue_in(&p, &mut rng);
        let mut fails = 0;
        for _ in 0..128 {
            let n = BigNumber::from_rng(&p, &mut rng);
            match super::tonelli_shanks(&n, &p, &z) {
                Some(r) => assert_eq!(r.modmul(&r, &p), n),
                None => fails += 1,
            }
        }
        assert!(fails <= 80, "Too many non residues");
    }

    #[test]
    fn chinese_remainder() {
        let mut rng = rand_core::OsRng::default();
        let p = BigNumber::prime(256);
        let q = BigNumber::prime(256);
        assert_ne!(p, q);
        let mp = BigNumber::from_rng(&p, &mut rng);
        let mq = BigNumber::from_rng(&q, &mut rng);
        let r = super::chinese_remainder_construction((&mp, &p), (&mq, &q));
        assert_eq!(&r % p, mp);
        assert_eq!(&r % q, mq);
    }
}