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    fn generate_step_constraints(
 52        &self,
 53        cs: ConstraintSystemRef<F>,
 54        _i: usize,
 55        z_i: Vec<FpVar<F>>,
 56        external_inputs: Self::ExternalInputsVar,
 57    ) -> Result<Vec<FpVar<F>>, SynthesisError> {
 58        let pk_0_x = z_i[0].clone();
 59        let pk_0_y = z_i[1].clone();
 60        let pk_i_x = z_i[2].clone();
 61        let pk_i_y = z_i[3].clone();
 62        let mut degree = z_i[4].clone();
 63
 64        // get the 'msg' that has been signed, which is the hash of the previous-signer public key
 65        let mut poseidon = PoseidonSpongeVar::new(cs.clone(), &self.config);
 66        poseidon.absorb(&vec![pk_i_x, pk_i_y])?;
 67        let h = poseidon.squeeze_field_elements(1)?;
 68        let msg = h
 69            .first()
 70            .ok_or(ark_relations::r1cs::SynthesisError::Unsatisfiable)?;
 71
 72        // check that the last signer is signed by the new signer
 73        let ei: SigPkVar<C, GC> = external_inputs.into();
 74        let res = verify::<C, GC>(
 75            cs.clone(),
 76            self.config.clone(),
 77            ei.pk.clone(),
 78            (ei.sig_r, ei.sig_s),
 79            msg.clone(),
 80        )?;
 81        res.enforce_equal(&Boolean::<F>::TRUE)?;
 82
 83        // increment the degree
 84        degree = degree.clone() + FpVar::<F>::one();
 85
 86        let pk_i1_xy = ei.pk.to_constraint_field()?;
 87        Ok(vec![vec![pk_0_x, pk_0_y], pk_i1_xy, vec![degree]].concat())
 88    }
 89}
 90#[cfg(test)]
 91pub mod tests {
 92    use super::*;
 93    use ark_bn254::Fr;
 94    use ark_ec::AffineRepr;
 95    use ark_r1cs_std::{alloc::AllocVar, R1CSVar};
 96    use ark_relations::r1cs::ConstraintSystem;
 97    use ark_std::Zero;
 98    use rand::rngs::OsRng;
 99
100    use crate::signature::{gen_signatures, hash_pk};
101    use arkeddsa::ed_on_bn254_twist::{constraints::EdwardsVar, EdwardsProjective};
102    use folding_schemes::transcript::poseidon::poseidon_canonical_config;
103
104    #[test]
105    fn test_sig() {
106        let mut rng = OsRng;
107        let poseidon_config = poseidon_canonical_config::<Fr>();
108
109        const N: usize = 1;
110        let ext_inps = gen_signatures::<OsRng, EdwardsProjective>(&mut rng, &poseidon_config, 1);
111        let e = ext_inps[0].clone();
112
113        let msg = hash_pk(&poseidon_config, e.pk);
114
115        e.pk.verify(&poseidon_config, &msg, &e.sig).unwrap();
116    }
117
118    fn ensure_fcircuit_trait<FC: FCircuit<Fr>>(params: FC::Params) {
119        let _ = FC::new(params);
120    }
121
122    // test to check that the Sha256FCircuit computes the same values inside and outside the circuit
123    #[test]
124    fn test_fcircuit() {
125        let mut rng = rand::rngs::OsRng;
126        let poseidon_config = poseidon_canonical_config::<Fr>();
127
128        let pks_sigs = gen_signatures::<OsRng, EdwardsProjective>(&mut rng, &poseidon_config, 1);
129
130        // here `Fr` is the BN254::G1::Fr = ed_on_bn254_twist::EdwardsProjective::Fq
131        let cs = ConstraintSystem::<Fr>::new_ref();
132
133        type FC = EthDosCircuit<Fr, EdwardsProjective, EdwardsVar>;
134        ensure_fcircuit_trait::<FC>(poseidon_config.clone());
135
136        let circuit = FC::new(poseidon_config).unwrap();
137        let xy: (Fr, Fr) = pks_sigs[0].pk.0.xy().unwrap();
138        let pk0 = vec![xy.0, xy.1];
139        let z_i: Vec<Fr> = vec![pk0.clone(), pk0, vec![Fr::zero()]].concat();
140
141        let external_inputs_var =
142            SigPkVar::<EdwardsProjective, EdwardsVar>::new_witness(cs.clone(), || Ok(pks_sigs[0]))
143                .unwrap();
144
145        let z_iVar = Vec::<FpVar<Fr>>::new_witness(cs.clone(), || Ok(z_i)).unwrap();
146        let computed_z_i1Var = circuit
147            .generate_step_constraints(cs.clone(), 0, z_iVar.clone(), external_inputs_var)
148            .unwrap();
149        // check that the degree (in the last state) is 1, the amount of signatures verified
150        assert_eq!(computed_z_i1Var.value().unwrap()[4], Fr::from(1_u32));
151        assert!(cs.is_satisfied().unwrap());
152        dbg!(cs.num_constraints());
153        dbg!(&computed_z_i1Var.value().unwrap());
154    }
155}