1use anyhow::Result;
  2
  3use plonky2::iop::{
  4    target::Target,
  5    witness::{PartialWitness, PartitionWitness, Witness, WitnessWrite},
  6};
  7use plonky2::field::{
  8    goldilocks_field::GoldilocksField,
  9    types::{Field, PrimeField64},
 10};
 11use plonky2::hash::poseidon::PoseidonHash;
 12use plonky2::plonk::{
 13    circuit_builder::CircuitBuilder,
 14    config::GenericConfig,
 15};
 16
 17use crate::{
 18    mod65537::Mod65537Builder,
 19    schnorr::{SchnorrPublicKey, SchnorrSignature},
 20};
 21
 22type GoldF = GoldilocksField;
 23
 24pub struct MessageTarget {
 25    msg: Vec<Target>,
 26}
 27
 28impl MessageTarget {
 29    fn new_with_size(builder: &mut CircuitBuilder<GoldF, 2>, n: usize) -> Self {
 30        Self {
 31            msg: builder.add_virtual_targets(n),
 32        }
 33    }
 34
 35    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, msg: &Vec<GoldF>) -> Result<()> {
 36        assert!(msg.len() == self.msg.len());
 37        for (&t, &x) in self.msg.iter().zip(msg.iter()) {
 38            pw.set_target(t, x)?;
 39        }
 40
 41        Ok(())
 42    }
 43}
 44
 45pub struct SchnorrSignatureTarget {
 46    s: Target,
 47    e: Target,
 48}
 49
 50impl SchnorrSignatureTarget {
 51    fn new_virtual(builder: &mut CircuitBuilder<GoldF, 2>) -> Self {
 52        let s = builder.add_virtual_target();
 53        let e = builder.add_virtual_target();
 54        Self{ s, e }
 55    }
 56
 57    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, sig: &SchnorrSignature) -> Result<()> {
 58        pw.set_target(self.s, GoldilocksField::from_canonical_u64(sig.s))?;
 59        pw.set_target(self.e, GoldilocksField::from_canonical_u64(sig.e))?;
 60        Ok(())
 61    }
 62}
 63
 64pub struct SchnorrPublicKeyTarget {
 65    pk: Target,
 66}
 67
 68impl SchnorrPublicKeyTarget {
 69    fn new_virtual(builder: &mut CircuitBuilder<GoldF, 2>) -> Self {
 70        Self{ pk: builder.add_virtual_target() }
 71    }
 72
 73    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, pk: &SchnorrPublicKey) -> Result<()> {
 74        pw.set_target(self.pk, pk.pk)?;
 75        Ok(())
 76    }
 77}
 78
 79pub struct SchnorrBuilder {}
 80
 81impl SchnorrBuilder {
 82
 83
 84    pub fn constrain_sig <
 85        C: GenericConfig<2, F = GoldF>,
 86    > (
 87        &self,
 88        builder: &mut CircuitBuilder::<GoldF, 2>,
 89        sig: &SchnorrSignatureTarget,
 90        msg: &MessageTarget,
 91        pk: &SchnorrPublicKeyTarget,
 92    ) -> () {
 93        let PRIME_GROUP_GEN: Target = builder.constant(GoldF::from_canonical_u64(6612579038192137166));
 94        let PRIME_GROUP_ORDER: Target = builder.constant(GoldF::from_canonical_u64(65537));
 95        const num_bits_exp: usize = 32;
 96
 97        /* here's the direct verification calculation,
 98        which we verify in-circuit
 99        let r: GoldF = Self::pow(self.PRIME_GROUP_GEN, sig.s)
100            * Self::pow(pk.pk, sig.e);
101        let e_v: u64 = self.hash_insecure(&r, msg);
102        e_v == sig.e   */
103
104        let gs: Target = builder.exp(PRIME_GROUP_GEN, sig.s, num_bits_exp);
105        let pe: Target = builder.exp(pk.pk, sig.e, num_bits_exp);
106        let r: Target = builder.mul(gs, pe);
107
108        // compute hash
109        // note that it's safe to clone Targets since they just contain indices
110        let hash_input: Vec<Target> = std::iter::once(r)
111            .chain(msg.msg.iter().cloned()) 
112            .collect();
113        let hash_output: Target = builder.hash_n_to_hash_no_pad::<PoseidonHash>(
114            hash_input,
115        ).elements[0]; // whoops have to take mod group order;
116
117        let e: Target = Mod65537Builder::mod_65537(builder, hash_output);
118
119        // enforce equality
120        builder.connect(e, sig.e);
121    }
122}
123
124#[cfg(test)]
125mod tests{
126    use crate::schnorr::{SchnorrPublicKey, SchnorrSecretKey, SchnorrSigner, SchnorrSignature};
127    use crate::schnorr_prover::{MessageTarget, Mod65537Builder, SchnorrBuilder, SchnorrPublicKeyTarget, SchnorrSignatureTarget};
128    use plonky2::hash::poseidon::Poseidon;
129    use plonky2::iop::{
130        target::Target,
131        witness::{PartialWitness, PartitionWitness, Witness, WitnessWrite},
132    };
133    use plonky2::plonk::{
134        circuit_builder::CircuitBuilder,
135        circuit_data::{
136            CircuitConfig, 
137            CircuitData, 
138            CommonCircuitData, 
139            VerifierCircuitData, 
140            VerifierOnlyCircuitData
141        },
142        config::{GenericConfig, PoseidonGoldilocksConfig},
143    };
144    use plonky2::field::{
145        goldilocks_field::GoldilocksField,
146        types::Field,
147    };
148    use rand;
149
150    #[test]
151    fn test_schnorr() -> () {
152        const D: usize = 2;
153        type C = PoseidonGoldilocksConfig;
154        type F = <C as GenericConfig<D>>::F;
155
156        let mut rng: rand::rngs::ThreadRng = rand::thread_rng();
157
158        let config = CircuitConfig::standard_recursion_config();
159        let mut builder = CircuitBuilder::<F, D>::new(config);
160
161        let sb: SchnorrBuilder = SchnorrBuilder{};
162
163        // create keypair, message, signature
164        let sk: SchnorrSecretKey = SchnorrSecretKey{ sk: 133 };
165        let ss = SchnorrSigner::new();
166        let pk: SchnorrPublicKey = ss.keygen(&sk);
167        let msg: Vec<GoldilocksField> = ss.u64_into_goldilocks_vec(
168            vec![1500, 1600, 0, 0, 0]
169        );
170        let msg_size: usize = msg.len();
171        let sig: SchnorrSignature = ss.sign(&msg, &sk, &mut rng);
172
173        let pk_targ = SchnorrPublicKeyTarget::new_virtual(&mut builder);
174        let sig_targ = SchnorrSignatureTarget::new_virtual(&mut builder);
175        let msg_targ = MessageTarget::new_with_size(&mut builder, msg_size);
176
177        
178        // instead of verifying we're going to prove the verification
179        sb.constrain_sig::<PoseidonGoldilocksConfig> (
180            &mut builder,
181            &sig_targ,
182            &msg_targ,
183            &pk_targ
184        );
185
186        // assign witnesses for verification
187        let mut pw: PartialWitness<F> = PartialWitness::new();
188        pk_targ.set_witness(&mut pw, &pk);
189        sig_targ.set_witness(&mut pw, &sig);
190        msg_targ.set_witness(&mut pw, &msg);
191
192
193        let data = builder.build::<C>();
194        let proof = data.prove(pw).unwrap();
195    }
196}