composite_modulus_proofs 0.1.0

//! Proof that `gcd(x, phi(N)) = 1` where `x` is coprime to a composite `N` and `phi` is the Euler's totient
//! Described in section 3.1 of the paper [Efficient Noninteractive Certification of RSA Moduli and Beyond](https://eprint.iacr.org/2018/057)

use crate::{
    error::Error,
    join,
    math::{
        misc::{crt_combine, euler_totient},
        prime_check::is_divisible_by_prime_upto,
    },
    setup::{Modulus, Primes, PrimesWithPrecomp},
    util::{log_base_2, uint_le_bytes},
};
use alloc::vec;
use ark_std::{cfg_into_iter, cfg_iter_mut, io::Write};
use crypto_bigint::{
    modular::{MontyForm, MontyParams, SafeGcdInverter},
    Concat, Odd, PrecomputeInverter, Split, Uint,
};
use digest::{ExtendableOutput, Update};

use crate::util::get_inv_mod_p_minus_1_and_q_minus_1;
#[cfg(feature = "parallel")]
use rayon::prelude::*;

/// Proof that `gcd(x, phi(N)) = 1`. `NUM_CHALLENGES` depends on `ALPHA` and `KAPPA` as per the paper but can also
/// be set independently as done in another protocol
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProofGcdIsOne<
    const MODULUS_LIMBS: usize,
    const NUM_CHALLENGES: usize,
    const ALPHA: usize,
    const KAPPA: usize,
>(
    /// Responses to the challenges
    pub [Uint<MODULUS_LIMBS>; NUM_CHALLENGES],
);

impl<
        const MODULUS_LIMBS: usize,
        const NUM_CHALLENGES: usize,
        const ALPHA: usize,
        const KAPPA: usize,
    > ProofGcdIsOne<MODULUS_LIMBS, NUM_CHALLENGES, ALPHA, KAPPA>
{
    /// Create a proof that `gcd(x, phi(N)) = 1`. `transcript` is the proof transcript that is shared by
    /// other sub-protocols as well which this is running as a sub-protocol of a larger protocol.
    pub fn new<
        D: Default + Update + ExtendableOutput,
        W: Write + AsRef<[u8]>,
        const PRIME_LIMBS: usize,
        const MODULUS_UNSAT_LIMBS: usize,
    >(
        x: &Uint<MODULUS_LIMBS>,
        primes: Primes<PRIME_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> Result<Self, Error>
    where
        Uint<PRIME_LIMBS>: Concat<Output = Uint<MODULUS_LIMBS>>,
        Odd<Uint<MODULUS_LIMBS>>:
            PrecomputeInverter<Inverter = SafeGcdInverter<MODULUS_LIMBS, MODULUS_UNSAT_LIMBS>>,
    {
        let (proof, _) = Self::new_and_return_challenges::<D, W, PRIME_LIMBS, MODULUS_UNSAT_LIMBS>(
            x, primes, modulus, nonce, transcript,
        )?;
        Ok(proof)
    }

    /// Same as `Self::new` but returns the challenges used to create the proof. Useful when used with certain other protocols
    pub fn new_and_return_challenges<
        D: Default + Update + ExtendableOutput,
        W: Write + AsRef<[u8]>,
        const PRIME_LIMBS: usize,
        const MODULUS_UNSAT_LIMBS: usize,
    >(
        x: &Uint<MODULUS_LIMBS>,
        primes: Primes<PRIME_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> Result<(Self, [Uint<MODULUS_LIMBS>; NUM_CHALLENGES]), Error>
    where
        Uint<PRIME_LIMBS>: Concat<Output = Uint<MODULUS_LIMBS>>,
        Odd<Uint<MODULUS_LIMBS>>:
            PrecomputeInverter<Inverter = SafeGcdInverter<MODULUS_LIMBS, MODULUS_UNSAT_LIMBS>>,
    {
        let phi = euler_totient(&primes.p, &primes.q);
        let x_inv = x.inv_mod(&phi);
        if x_inv.is_none().into() {
            // gcd(x, phi(N)) isn't 1
            return Err(Error::NotInvertible);
        }
        let x_inv = x_inv.unwrap();
        let challenges = Self::challenges::<D, W>(x, modulus, nonce, transcript);
        let mut responses = [Uint::<MODULUS_LIMBS>::ZERO; NUM_CHALLENGES];
        let params = MontyParams::new_vartime(modulus.0);
        cfg_iter_mut!(responses).enumerate().for_each(|(i, s)| {
            *s = MontyForm::new(&challenges[i], params)
                .pow(&x_inv)
                .retrieve();
        });
        Ok((Self(responses), challenges))
    }

    /// Same as `Self::new` but takes the primes and some precomputation to make the proof generation is faster.
    /// Use it if creating frequent proofs.
    pub fn new_given_precomputation<
        D: Default + Update + ExtendableOutput,
        W: Write + AsRef<[u8]>,
        const PRIME_LIMBS: usize,
        const PRIME_UNSAT_LIMBS: usize,
    >(
        x: &Uint<MODULUS_LIMBS>,
        primes: PrimesWithPrecomp<PRIME_LIMBS, MODULUS_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> Result<Self, Error>
    where
        Uint<PRIME_LIMBS>: Concat<Output = Uint<MODULUS_LIMBS>>,
        Uint<MODULUS_LIMBS>: Split<Output = Uint<PRIME_LIMBS>>,
        Odd<Uint<PRIME_LIMBS>>: PrecomputeInverter<
            Inverter = SafeGcdInverter<PRIME_LIMBS, PRIME_UNSAT_LIMBS>,
            Output = Uint<PRIME_LIMBS>,
        >,
    {
        let (proof, _) = Self::new_given_precomputation_and_return_challenges::<
            D,
            W,
            PRIME_LIMBS,
            PRIME_UNSAT_LIMBS,
        >(x, primes, modulus, nonce, transcript)?;
        Ok(proof)
    }

    /// Same as `Self::new_given_precomputation` but takes `x^-1 mod p-1` and `x^-1 mod q-1` as well.
    pub fn new_given_x_inv_and_precomputation<
        D: Default + Update + ExtendableOutput,
        W: Write + AsRef<[u8]>,
        const PRIME_LIMBS: usize,
    >(
        x: &Uint<MODULUS_LIMBS>,
        x_inv_p_minus_1: &Uint<PRIME_LIMBS>,
        x_inv_q_minus_1: &Uint<PRIME_LIMBS>,
        primes: PrimesWithPrecomp<PRIME_LIMBS, MODULUS_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> Result<Self, Error>
    where
        Uint<PRIME_LIMBS>: Concat<Output = Uint<MODULUS_LIMBS>>,
        Uint<MODULUS_LIMBS>: Split<Output = Uint<PRIME_LIMBS>>,
    {
        let (proof, _) =
            Self::new_given_x_inv_and_precomputation_and_return_challenges::<D, W, PRIME_LIMBS>(
                x,
                x_inv_p_minus_1,
                x_inv_q_minus_1,
                primes,
                modulus,
                nonce,
                transcript,
            )?;
        Ok(proof)
    }

    /// Same as `Self::new_given_precomputation` but returns the challenges used to create the proof. Useful when used with certain other protocols
    pub fn new_given_precomputation_and_return_challenges<
        D: Default + Update + ExtendableOutput,
        W: Write + AsRef<[u8]>,
        const PRIME_LIMBS: usize,
        const PRIME_UNSAT_LIMBS: usize,
    >(
        x: &Uint<MODULUS_LIMBS>,
        primes: PrimesWithPrecomp<PRIME_LIMBS, MODULUS_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> Result<(Self, [Uint<MODULUS_LIMBS>; NUM_CHALLENGES]), Error>
    where
        Uint<PRIME_LIMBS>: Concat<Output = Uint<MODULUS_LIMBS>>,
        Uint<MODULUS_LIMBS>: Split<Output = Uint<PRIME_LIMBS>>,
        Odd<Uint<PRIME_LIMBS>>: PrecomputeInverter<
            Inverter = SafeGcdInverter<PRIME_LIMBS, PRIME_UNSAT_LIMBS>,
            Output = Uint<PRIME_LIMBS>,
        >,
    {
        // As x^-1 will be the exponent when computing mod p and mod q using CRT, reduce x^-1 mod p-1 and mod q-1 as per Fermat's little theorem
        // x_inv_p_minus_1 = x^-1 mod p-1, x_inv_q_minus_1 = x^-1 mod q-1
        let (x_inv_p_minus_1, x_inv_q_minus_1) = get_inv_mod_p_minus_1_and_q_minus_1(
            &x,
            primes.p_mtg.modulus(),
            primes.q_mtg.modulus(),
        )?;
        Self::new_given_x_inv_and_precomputation_and_return_challenges::<D, W, PRIME_LIMBS>(
            x,
            &x_inv_p_minus_1,
            &x_inv_q_minus_1,
            primes,
            modulus,
            nonce,
            transcript,
        )
    }

    /// Same as `Self::new_given_precomputation_and_return_challenges` but takes `x^-1 mod p-1` and `x^-1 mod q-1` as well.
    pub fn new_given_x_inv_and_precomputation_and_return_challenges<
        D: Default + Update + ExtendableOutput,
        W: Write + AsRef<[u8]>,
        const PRIME_LIMBS: usize,
    >(
        x: &Uint<MODULUS_LIMBS>,
        x_inv_p_minus_1: &Uint<PRIME_LIMBS>,
        x_inv_q_minus_1: &Uint<PRIME_LIMBS>,
        primes: PrimesWithPrecomp<PRIME_LIMBS, MODULUS_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> Result<(Self, [Uint<MODULUS_LIMBS>; NUM_CHALLENGES]), Error>
    where
        Uint<PRIME_LIMBS>: Concat<Output = Uint<MODULUS_LIMBS>>,
        Uint<MODULUS_LIMBS>: Split<Output = Uint<PRIME_LIMBS>>,
    {
        let challenges = Self::challenges::<D, W>(x, modulus, nonce, transcript);
        let mut responses = [Uint::<MODULUS_LIMBS>::ZERO; NUM_CHALLENGES];

        cfg_iter_mut!(responses).enumerate().for_each(|(i, s)| {
            let (resp_p, resp_q) = join!(
                {
                    let c_p = MontyForm::new(&challenges[i], primes.p_mtg_1)
                        .retrieve()
                        .resize();
                    MontyForm::new(&c_p, primes.p_mtg)
                        .pow(x_inv_p_minus_1)
                        .retrieve()
                },
                {
                    let c_q = MontyForm::new(&challenges[i], primes.q_mtg_1)
                        .retrieve()
                        .resize();
                    MontyForm::new(&c_q, primes.q_mtg)
                        .pow(x_inv_q_minus_1)
                        .retrieve()
                }
            );
            *s = crt_combine::<PRIME_LIMBS, MODULUS_LIMBS>(
                &resp_p,
                &resp_q,
                primes.p_inv,
                primes.p_mtg.modulus(),
                primes.q_mtg,
            );
        });
        Ok((Self(responses), challenges))
    }

    pub fn num_responses(&self) -> usize {
        self.0.len()
    }

    /// Verify the proof that `gcd(x, phi(N)) = 1`. `transcript` is the proof transcript that is shared by
    /// other sub-protocols as well which this is running as a sub-protocol of a larger protocol.
    pub fn verify<D: Default + Update + ExtendableOutput, W: Write + AsRef<[u8]>>(
        &self,
        x: &Uint<MODULUS_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> Result<(), Error> {
        let (res, _) = self.verify_and_return_challenges::<D, W>(x, modulus, nonce, transcript)?;
        Ok(res)
    }

    /// Same as `Self::verify` but returns the challenges used to create the proof. Useful when used with certain other protocols
    pub fn verify_and_return_challenges<
        D: Default + Update + ExtendableOutput,
        W: Write + AsRef<[u8]>,
    >(
        &self,
        x: &Uint<MODULUS_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> Result<((), [Uint<MODULUS_LIMBS>; NUM_CHALLENGES]), Error> {
        if let Some(i) = is_divisible_by_prime_upto(&modulus.0, ALPHA as u32) {
            return Err(Error::DivisibleByPrime(i));
        }
        let challenges = Self::challenges::<D, W>(x, modulus, nonce, transcript);
        let params = MontyParams::new_vartime(modulus.0);
        #[allow(unused_mut)]
        let mut res = cfg_into_iter!(0..NUM_CHALLENGES).map(|i| {
            if !modulus.is_greater_than(&self.0[i]) {
                return Err(Error::GreaterThanModulus(i as u32));
            }
            if MontyForm::new(&self.0[i], params).pow(x).retrieve() != challenges[i] {
                return Err(Error::InvalidProofForIndex(i as u32));
            }
            Ok(())
        });

        let res = report_error_if_any!(res)?;
        Ok((res, challenges))
    }

    /// Generate challenges by rejection sampling
    pub fn challenges<D: Default + Update + ExtendableOutput, W: Write + AsRef<[u8]>>(
        x: &Uint<MODULUS_LIMBS>,
        modulus: &Modulus<MODULUS_LIMBS>,
        nonce: &[u8],
        transcript: &mut W,
    ) -> [Uint<MODULUS_LIMBS>; NUM_CHALLENGES] {
        let mut chals = [Uint::<MODULUS_LIMBS>::ZERO; NUM_CHALLENGES];
        let mut ctr = 0;
        transcript.write_all(b"proof-of-gcd-is-1").unwrap();
        transcript.write_all(&uint_le_bytes(x)).unwrap();
        transcript.write_all(&uint_le_bytes(&modulus.0)).unwrap();
        // Not hashing in MODULUS_LIMBS since it can change if proof generation and verification happen on different arch. (64 bit vs 32 bit)
        transcript
            .write_all(modulus.size_for_hashing().as_slice())
            .unwrap();
        transcript.write_all(nonce).unwrap();
        let mut c_bytes = vec![0; Uint::<MODULUS_LIMBS>::BYTES];
        let mut i = 0_u32;
        while ctr < NUM_CHALLENGES as u32 {
            transcript.write_all(i.to_le_bytes().as_slice()).unwrap();
            D::digest_xof(transcript.as_ref(), &mut c_bytes);
            let c = Uint::<MODULUS_LIMBS>::from_le_slice(&c_bytes);
            if modulus.is_greater_than(&c) {
                chals[ctr as usize] = c;
                ctr += 1;
            }
            i += 1;
        }
        chals
    }
}

pub const fn num_challenges<const ALPHA: usize, const KAPPA: usize>() -> u32 {
    let l = log_base_2::<ALPHA>();
    ((l + KAPPA - 1) / l) as u32
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        safegcd_nlimbs,
        util::{get_1024_bit_primes, get_2048_bit_primes, get_coprime, timing_info},
    };
    use core::ops::Sub;
    use crypto_bigint::{U1024, U128, U2048, U256, U64};
    use rand_core::OsRng;
    use sha3::Shake256;
    use std::{ops::Add, time::Instant};

    macro_rules! check_given_primes {
        ( $num_iters: ident, $prime_type:ident, $primes: ident ) => {
            const KAPPA: usize = 128;
            const ALPHA: usize = 65537;
            const NUM_CHALLENGES: usize = num_challenges::<ALPHA, KAPPA>() as usize;

            const PRIME_LIMBS: usize = $prime_type::LIMBS;
            const PRIME_UNSAT_LIMBS: usize = safegcd_nlimbs!(Uint::<PRIME_LIMBS>::BITS as usize);
            const MODULUS_LIMBS: usize = PRIME_LIMBS * 2;
            const MODULUS_UNSAT_LIMBS: usize =
                safegcd_nlimbs!(Uint::<MODULUS_LIMBS>::BITS as usize);

            let mut rng = OsRng::default();
            let modulus = Modulus::<MODULUS_LIMBS>::new(&$primes);
            let primes_with_crt = PrimesWithPrecomp::from($primes.clone());
            let phi = euler_totient(&$primes.p, &$primes.q);
            let nonce = b"123";

            let mut prove_times = vec![];
            let mut prove_crt_times = vec![];
            let mut ver_times = vec![];

            println!(
                "Running {} iterations for {} bits prime",
                $num_iters,
                $prime_type::BITS
            );
            for _ in 0..$num_iters {
                let x = get_coprime(&mut rng, &phi, modulus.0.as_nz_ref());

                let mut transcript = vec![];
                let start = Instant::now();
                let proof = ProofGcdIsOne::<MODULUS_LIMBS, NUM_CHALLENGES, ALPHA, KAPPA>::new::<
                    Shake256,
                    _,
                    PRIME_LIMBS,
                    MODULUS_UNSAT_LIMBS,
                >(&x, $primes.clone(), &modulus, nonce, &mut transcript)
                .unwrap();
                prove_times.push(start.elapsed());

                assert_eq!(proof.num_responses(), NUM_CHALLENGES);

                let mut transcript = vec![];
                let start = Instant::now();
                proof
                    .verify::<Shake256, _>(&x, &modulus, nonce, &mut transcript)
                    .unwrap();
                ver_times.push(start.elapsed());

                let mut transcript = vec![];
                let start = Instant::now();
                let proof =
                    ProofGcdIsOne::<MODULUS_LIMBS, NUM_CHALLENGES, ALPHA, KAPPA>::new_given_precomputation::<
                        Shake256,
                        _,
                        PRIME_LIMBS,
                        PRIME_UNSAT_LIMBS,
                    >(
                        &x,
                        primes_with_crt.clone(),
                        &modulus,
                        nonce,
                        &mut transcript,
                    )
                    .unwrap();
                prove_crt_times.push(start.elapsed());

                assert_eq!(proof.num_responses(), NUM_CHALLENGES);

                let mut transcript = vec![];
                proof
                    .verify::<Shake256, _>(&x, &modulus, nonce, &mut transcript)
                    .unwrap();
            }

            println!("Proving time: {:?}", timing_info(prove_times));
            println!("Proving time with CRT: {:?}", timing_info(prove_crt_times));
            println!("Verification time: {:?}", timing_info(ver_times));
        };
    }

    #[test]
    fn bad_proofs() {
        const KAPPA: usize = 128;
        const ALPHA: usize = 65537;
        const NUM_CHALLENGES: usize = num_challenges::<ALPHA, KAPPA>() as usize;

        const PRIME_LIMBS: usize = U64::LIMBS;
        const MODULUS_LIMBS: usize = PRIME_LIMBS * 2;
        const MODULUS_UNSAT_LIMBS: usize = safegcd_nlimbs!(Uint::<MODULUS_LIMBS>::BITS as usize);

        let mut rng = OsRng::default();
        let primes = Primes::<PRIME_LIMBS>::new(&mut rng);
        let modulus = Modulus::<MODULUS_LIMBS>::new(&primes);
        let phi = euler_totient(&primes.p, &primes.q);

        let nonce = b"123";
        // Trying to invert p-1 should fail since its a factor of the totient
        let mut transcript = vec![];
        let x = primes.p.sub(&Uint::ONE).resize();
        assert!(matches!(
            ProofGcdIsOne::<MODULUS_LIMBS, NUM_CHALLENGES, ALPHA, KAPPA>::new::<
                Shake256,
                _,
                PRIME_LIMBS,
                MODULUS_UNSAT_LIMBS,
            >(&x, primes.clone(), &modulus, nonce, &mut transcript)
            .err()
            .unwrap(),
            Error::NotInvertible
        ));

        let x = get_coprime(&mut rng, &phi, modulus.0.as_nz_ref());
        let mut transcript = vec![];
        let mut proof = ProofGcdIsOne::<MODULUS_LIMBS, NUM_CHALLENGES, ALPHA, KAPPA>::new::<
            Shake256,
            _,
            PRIME_LIMBS,
            MODULUS_UNSAT_LIMBS,
        >(&x, primes.clone(), &modulus, nonce, &mut transcript)
        .unwrap();

        // Modulus is a multiple of a small prime
        let mut bad_modulus = modulus.clone();
        bad_modulus.0 = primes.p.widening_mul(&U64::from(5_u32)).to_odd().unwrap();
        let mut transcript = vec![];
        assert!(matches!(
            proof
                .verify::<Shake256, _>(&x, &bad_modulus, nonce, &mut transcript)
                .err()
                .unwrap(),
            Error::DivisibleByPrime(5)
        ));

        // Responses shouldn't be bigger than modulus
        proof.0[0] = modulus.0.as_nz_ref().add(&U128::ONE);
        let mut transcript = vec![];
        assert!(matches!(
            proof
                .verify::<Shake256, _>(&x, &modulus, nonce, &mut transcript)
                .err()
                .unwrap(),
            Error::GreaterThanModulus(0)
        ));

        // Responses should be valid
        proof.0[0] = proof.0[1];
        let mut transcript = vec![];
        assert!(matches!(
            proof
                .verify::<Shake256, _>(&x, &modulus, nonce, &mut transcript)
                .err()
                .unwrap(),
            Error::InvalidProofForIndex(0)
        ));
    }

    #[test]
    fn proof() {
        let mut rng = OsRng::default();
        let primes = Primes::<{ U256::LIMBS }>::new(&mut rng);
        let num_iters = 100;
        check_given_primes!(num_iters, U256, primes);
    }

    #[test]
    fn proof_with_1024_bit_primes() {
        let (p, q) = get_1024_bit_primes();
        let primes = Primes::from_primes(p, q);
        let num_iters = 10;
        check_given_primes!(num_iters, U1024, primes);
    }

    #[test]
    fn proof_with_2048_bit_primes() {
        let (p, q) = get_2048_bit_primes();
        let primes = Primes::from_primes(p, q);
        let num_iters = 10;
        check_given_primes!(num_iters, U2048, primes);
    }
}