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