1use ark_crypto_primitives::sponge::{
  2    constraints::CryptographicSpongeVar,
  3    poseidon::{constraints::PoseidonSpongeVar, PoseidonConfig},
  4    Absorb,
  5};
  6use ark_ec::CurveGroup;
  7use ark_ff::{Field, PrimeField};
  8use ark_r1cs_std::prelude::CurveVar;
  9use ark_r1cs_std::{
 10    boolean::Boolean,
 11    eq::EqGadget,
 12    fields::{fp::FpVar, FieldVar},
 13};
 14use ark_relations::r1cs::{ConstraintSystemRef, SynthesisError};
 15use ark_std::marker::PhantomData;
 16use std::fmt::Debug;
 17
 18use arkeddsa::constraints::verify;
 19use folding_schemes::{frontend::FCircuit, Error};
 20
 21use crate::signature::{SigPk, SigPkVar};
 22
 23pub type CF<C> = <<C as CurveGroup>::BaseField as Field>::BasePrimeField;
 24
 25/// Test circuit to be folded
 26#[derive(Clone, Debug)]
 27pub struct EthDosCircuit<F: PrimeField, C: CurveGroup, GC: CurveVar<C, F>> {
 28    _c: PhantomData<C>,
 29    _gc: PhantomData<GC>,
 30    config: PoseidonConfig<F>,
 31}
 32impl<F: PrimeField, C: CurveGroup, GC: CurveVar<C, F>> FCircuit<F> for EthDosCircuit<F, C, GC>
 33where
 34    F: Absorb,
 35    C: CurveGroup<BaseField = F>,
 36{
 37    type Params = PoseidonConfig<F>;
 38    type ExternalInputs = SigPk<C>;
 39    type ExternalInputsVar = SigPkVar<C, GC>;
 40
 41    fn new(config: Self::Params) -> Result<Self, Error> {
 42        Ok(Self {
 43            _c: PhantomData,
 44            _gc: PhantomData,
 45            config,
 46        })
 47    }
 48    fn state_len(&self) -> usize {
 49        5
 50    }
 51    // This method defines the logic that is done in-circuit at each folding step
 52    fn generate_step_constraints(
 53        &self,
 54        cs: ConstraintSystemRef<F>,
 55        _i: usize,
 56        z_i: Vec<FpVar<F>>,
 57        external_inputs: Self::ExternalInputsVar,
 58    ) -> Result<Vec<FpVar<F>>, SynthesisError> {
 59        // get the values from the state, where: state = [ pk_0, pk_i, i]
 60        let pk_0_x = z_i[0].clone();
 61        let pk_0_y = z_i[1].clone();
 62        let pk_i_x = z_i[2].clone();
 63        let pk_i_y = z_i[3].clone();
 64        let mut degree = z_i[4].clone();
 65
 66        // get the 'pk_i_hashed' value, which is the hash of the pk_i, and is the value that has
 67        // been signed by the new public key (pk_i+1)
 68        let mut poseidon = PoseidonSpongeVar::new(cs.clone(), &self.config);
 69        poseidon.absorb(&vec![pk_i_x, pk_i_y])?;
 70        let h = poseidon.squeeze_field_elements(1)?;
 71        let pk_i_hashed = h
 72            .first()
 73            .ok_or(ark_relations::r1cs::SynthesisError::Unsatisfiable)?;
 74
 75        // check that the last signer's public key (pk_i) hashed (=pk_i_hashed) is signed by the
 76        // new signer public key (pk_i+1)
 77        let res = verify::<C, GC>(
 78            cs.clone(),
 79            self.config.clone(),
 80            external_inputs.pk.clone(), // pk_{i+1}
 81            (external_inputs.sig_r, external_inputs.sig_s),
 82            pk_i_hashed.clone(),
 83        )?;
 84        res.enforce_equal(&Boolean::<F>::TRUE)?;
 85
 86        // increment the degree
 87        degree = degree.clone() + FpVar::<F>::one();
 88
 89        // return the new IVC state, where we place the pk_{i+1} at the place where previously had
 90        // the pk_i, together with the new updated degree of distance value
 91        let pk_i1_xy = external_inputs.pk.to_constraint_field()?;
 92        Ok([vec![pk_0_x, pk_0_y], pk_i1_xy, vec![degree]].concat())
 93    }
 94}
 95#[cfg(test)]
 96pub mod tests {
 97    use super::*;
 98    use ark_bn254::Fr;
 99    use ark_ec::AffineRepr;
100    use ark_r1cs_std::{alloc::AllocVar, R1CSVar};
101    use ark_relations::r1cs::ConstraintSystem;
102    use ark_std::Zero;
103    use rand::rngs::OsRng;
104
105    use crate::signature::{gen_signatures, hash_pk};
106    use arkeddsa::ed_on_bn254_twist::{constraints::EdwardsVar, EdwardsProjective};
107    use folding_schemes::transcript::poseidon::poseidon_canonical_config;
108
109    #[test]
110    fn test_sig() {
111        let mut rng = OsRng;
112        let poseidon_config = poseidon_canonical_config::<Fr>();
113
114        const N: usize = 1;
115        let ext_inps = gen_signatures::<OsRng, EdwardsProjective>(&mut rng, &poseidon_config, 1);
116        let e = ext_inps[0].clone();
117
118        let msg = hash_pk(&poseidon_config, e.pk);
119
120        e.pk.verify(&poseidon_config, &msg, &e.sig).unwrap();
121    }
122
123    fn ensure_fcircuit_trait<FC: FCircuit<Fr>>(params: FC::Params) {
124        let _ = FC::new(params);
125    }
126
127    // test to check that the Sha256FCircuit computes the same values inside and outside the circuit
128    #[test]
129    fn test_fcircuit() {
130        let mut rng = rand::rngs::OsRng;
131        let poseidon_config = poseidon_canonical_config::<Fr>();
132
133        let pks_sigs = gen_signatures::<OsRng, EdwardsProjective>(&mut rng, &poseidon_config, 1);
134
135        // here `Fr` is the BN254::G1::Fr = ed_on_bn254_twist::EdwardsProjective::Fq
136        let cs = ConstraintSystem::<Fr>::new_ref();
137
138        type FC = EthDosCircuit<Fr, EdwardsProjective, EdwardsVar>;
139        ensure_fcircuit_trait::<FC>(poseidon_config.clone());
140
141        let circuit = FC::new(poseidon_config).unwrap();
142        let xy: (Fr, Fr) = pks_sigs[0].pk.0.xy().unwrap();
143        let pk0 = vec![xy.0, xy.1];
144        let z_i: Vec<Fr> = vec![pk0.clone(), pk0, vec![Fr::zero()]].concat();
145
146        let external_inputs_var =
147            SigPkVar::<EdwardsProjective, EdwardsVar>::new_witness(cs.clone(), || Ok(pks_sigs[0]))
148                .unwrap();
149
150        let z_iVar = Vec::<FpVar<Fr>>::new_witness(cs.clone(), || Ok(z_i)).unwrap();
151        let computed_z_i1Var = circuit
152            .generate_step_constraints(cs.clone(), 0, z_iVar.clone(), external_inputs_var)
153            .unwrap();
154        // check that the degree (in the last state) is 1, the amount of signatures verified
155        assert_eq!(computed_z_i1Var.value().unwrap()[4], Fr::from(1_u32));
156        assert!(cs.is_satisfied().unwrap());
157        dbg!(cs.num_constraints());
158        dbg!(&computed_z_i1Var.value().unwrap());
159    }
160}