1use ark_crypto_primitives::sponge::{
  2    poseidon::{PoseidonConfig, PoseidonSponge},
  3    Absorb, CryptographicSponge,
  4};
  5use ark_ec::{AffineRepr, CurveGroup};
  6use ark_ff::{BigInteger, PrimeField};
  7use ark_r1cs_std::alloc::{AllocVar, AllocationMode};
  8use ark_r1cs_std::boolean::Boolean;
  9use ark_r1cs_std::prelude::CurveVar;
 10use ark_relations::r1cs::{Namespace, SynthesisError};
 11use ark_serialize::CanonicalSerialize;
 12use ark_std::{rand::Rng, Zero};
 13use core::borrow::Borrow;
 14use rand_core::CryptoRngCore;
 15use std::fmt::Debug;
 16
 17use arkeddsa::{signature::Signature, PublicKey, SigningKey};
 18
 19use crate::fcircuit::CF;
 20
 21// recall, here C = ed_on_bn254, so C::BaseField = BN254::ScalarField
 22#[derive(Clone, Copy, Debug, PartialEq)]
 23pub struct SigPk<C: CurveGroup> {
 24    pub pk: PublicKey<C>,
 25    pub sig: Signature<C>,
 26}
 27impl<C: CurveGroup> Default for SigPk<C> {
 28    fn default() -> Self {
 29        Self {
 30            pk: PublicKey(C::zero().into_affine()),
 31            sig: Signature::new(C::zero().into_affine(), C::ScalarField::zero()),
 32        }
 33    }
 34}
 35impl<C: CurveGroup> SigPk<C> {
 36    pub fn to_bytes(&self) -> Vec<u8> {
 37        let sig_bytes = self.sig.to_bytes();
 38        let pk_bytes = self.pk.to_bytes();
 39        vec![sig_bytes, pk_bytes].concat()
 40    }
 41    pub fn from_bytes(b: Vec<u8>) -> Self {
 42        let u_point_size = C::Affine::generator().serialized_size(ark_serialize::Compress::No);
 43
 44        let sig = Signature::<C>::from_bytes(&b[..32 + u_point_size]).unwrap();
 45        let pk = PublicKey::<C>::from_bytes(&b[32 + u_point_size..]).unwrap();
 46        Self { pk, sig }
 47    }
 48}
 49
 50#[derive(Clone, Debug)]
 51pub struct SigPkVar<C: CurveGroup, GC: CurveVar<C, CF<C>>> {
 52    pub pk: GC,
 53    pub sig_r: GC,
 54    pub sig_s: Vec<Boolean<CF<C>>>,
 55}
 56impl<C: CurveGroup, GC: CurveVar<C, CF<C>>> Default for SigPkVar<C, GC> {
 57    fn default() -> Self {
 58        Self {
 59            pk: GC::zero(),
 60            sig_r: GC::zero(),
 61            sig_s: vec![Boolean::<CF<C>>::FALSE; 253], // TODO 253-> fieldbitsize
 62        }
 63    }
 64}
 65
 66impl<C, GC> AllocVar<SigPk<C>, CF<C>> for SigPkVar<C, GC>
 67where
 68    C: CurveGroup,
 69    GC: CurveVar<C, CF<C>>,
 70{
 71    fn new_variable<T: Borrow<SigPk<C>>>(
 72        cs: impl Into<Namespace<CF<C>>>,
 73        f: impl FnOnce() -> Result<T, SynthesisError>,
 74        mode: AllocationMode,
 75    ) -> Result<Self, SynthesisError> {
 76        f().and_then(|val| {
 77            let cs = cs.into();
 78
 79            let e = val.borrow();
 80            let pk = GC::new_variable(cs.clone(), || Ok(e.pk.0), mode)?;
 81            let sig_r = GC::new_variable(cs.clone(), || Ok(e.sig.r), mode)?;
 82            let sig_s = Vec::<Boolean<CF<C>>>::new_variable(
 83                cs.clone(),
 84                || Ok(e.sig.s.into_bigint().to_bits_le()),
 85                mode,
 86            )?;
 87            let v = Self { pk, sig_r, sig_s };
 88
 89            Ok(v)
 90        })
 91    }
 92}
 93
 94pub fn hash_pk<C: CurveGroup>(
 95    poseidon_config: &PoseidonConfig<C::BaseField>,
 96    pk: PublicKey<C>,
 97) -> C::BaseField
 98where
 99    C::BaseField: PrimeField + Absorb,
100{
101    let mut poseidon = PoseidonSponge::new(poseidon_config);
102    let (pk_x, pk_y): (C::BaseField, C::BaseField) = pk.xy().unwrap();
103    poseidon.absorb(&vec![pk_x, pk_y]);
104    let k = poseidon.squeeze_field_elements::<C::BaseField>(1);
105    *k.first().unwrap()
106}
107
108// returns a vector of signatures & publickeys, where each signature is of the previous public key
109pub fn gen_signatures<R: Rng + CryptoRngCore, C: CurveGroup>(
110    rng: &mut R,
111    poseidon_config: &PoseidonConfig<C::BaseField>,
112    steps: usize,
113) -> Vec<SigPk<C>>
114where
115    C::BaseField: PrimeField + Absorb,
116{
117    let mut prev_pk = None;
118    let mut res: Vec<SigPk<C>> = Vec::new();
119    for _ in 0..steps {
120        let extinp = gen_sig(rng, poseidon_config, prev_pk);
121        res.push(extinp);
122        prev_pk = Some(extinp.pk);
123    }
124    res
125}
126
127// generates a new secret key, and signs the given `prev_pk` with it. If the `prev_pk==None`, it
128// will use the newly generated public key as the prev_pk.
129pub fn gen_sig<R: Rng + CryptoRngCore, C: CurveGroup>(
130    rng: &mut R,
131    poseidon_config: &PoseidonConfig<C::BaseField>,
132    prev_pk: Option<PublicKey<C>>,
133) -> SigPk<C>
134where
135    C::BaseField: PrimeField + Absorb,
136{
137    let sk = SigningKey::<C>::generate::<blake2::Blake2b512>(rng).unwrap();
138    let pk = sk.public_key();
139
140    // if prev_pk!=None, use it, else, set the new pk to it
141    let prev_pk = if prev_pk.is_some() {
142        prev_pk.unwrap()
143    } else {
144        *pk
145    };
146
147    let msg = hash_pk(poseidon_config, prev_pk);
148
149    let sig = sk
150        .sign::<blake2::Blake2b512>(&poseidon_config, &msg)
151        .unwrap();
152    pk.verify(&poseidon_config, &msg, &sig).unwrap();
153    SigPk {
154        pk: pk.clone(),
155        sig,
156    }
157}