1use anyhow::Result;
  2
  3use plonky2::iop::{
  4    generator::{GeneratedValues, SimpleGenerator},
  5    target::{BoolTarget, Target},
  6    witness::{PartialWitness, PartitionWitness, Witness, WitnessWrite},
  7};
  8use plonky2::field::{
  9    extension::{Extendable, FieldExtension},
 10    goldilocks_field::GoldilocksField,
 11    types::{Field, PrimeField64},
 12};
 13use plonky2::hash::{
 14    hash_types::RichField,
 15    poseidon::PoseidonHash,
 16};
 17use plonky2::plonk::{
 18    circuit_builder::CircuitBuilder,
 19    circuit_data::{
 20        CircuitConfig, 
 21        CircuitData,
 22        CommonCircuitData, 
 23        VerifierCircuitData, 
 24        VerifierOnlyCircuitData
 25    },
 26    config::{GenericConfig, PoseidonGoldilocksConfig},
 27    proof::ProofWithPublicInputs,
 28};
 29use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
 30
 31use crate::schnorr::{SchnorrPublicKey, SchnorrSignature};
 32
 33type GoldF = GoldilocksField;
 34
 35#[derive(Debug, Default)]
 36pub struct Mod65537Generator {
 37    a: Target,
 38    q: Target,
 39    r: Target,
 40}
 41
 42impl SimpleGenerator<GoldF, 2> for Mod65537Generator {
 43    fn id(&self) -> String {
 44        "Mod65537Generator".to_string()
 45    }
 46    fn dependencies(&self) -> Vec<Target> {
 47        vec![self.a]
 48    }
 49
 50    fn run_once(
 51        &self,
 52        witness: &PartitionWitness<GoldF>,
 53        out_buffer: &mut GeneratedValues<GoldF>,
 54    ) -> Result<()> {
 55        let a = witness.get_target(self.a);
 56        let a64 = a.to_canonical_u64();
 57        let q64 = a64 / 65537;
 58        let r64 = a64 % 65537;
 59
 60        out_buffer.set_target(self.q, GoldF::from_canonical_u64(q64));
 61        out_buffer.set_target(self.r, GoldF::from_canonical_u64(r64));
 62
 63        Ok(())
 64    }
 65
 66    fn serialize(&self, dst: &mut Vec<u8>, common_data: &CommonCircuitData<GoldF, 2>) -> IoResult<()> {
 67        dst.write_target(self.a)?;
 68        dst.write_target(self.q)?;
 69        dst.write_target(self.r)?;
 70        Ok(())
 71    }
 72
 73    fn deserialize(src: &mut Buffer, common_data: &CommonCircuitData<GoldF, 2>) -> IoResult<Self>
 74    where
 75        Self: Sized 
 76    {
 77            let a = src.read_target()?;
 78            let q = src.read_target()?;
 79            let r = src.read_target()?;
 80            Ok(Self { a, q, r })
 81    }
 82}
 83
 84pub struct Mod65537Builder {}
 85
 86impl Mod65537Builder {
 87    // Reduce a modulo the constant 65537
 88    // where a is the canonical representative for an element of the field
 89    // (meaning: 0 \leq a < p)
 90
 91    // To prove this, write
 92    // a = 65537 * q + r, and do range checks to check that:
 93    // 0 <= q <= floor(p / 65537)
 94    // 0 <= r < 65537
 95    // (these first two checks guarantee that a lies in the range [0, p + 65536])
 96    // if q = floor(p / 65537) then r = 0
 97    // (note that p % 65537 == 1 so this is the only possibility)
 98    pub(crate) fn mod_65537 (
 99        builder: &mut CircuitBuilder::<GoldF, 2>,
100        a: Target,
101    ) -> Target {
102        let q = builder.add_virtual_target();
103        let r = builder.add_virtual_target();
104
105        // the Mod65537Generator will assign values to q and r later
106        builder.add_simple_generator( Mod65537Generator { a, q, r } );
107
108        // impose four constraints
109        // 1. a = 65537 * q + r
110        let t65537 = builder.constant(GoldF::from_canonical_u64(65537));
111        let a_copy = builder.mul_add(t65537, q, r);
112        builder.connect(a, a_copy);
113
114        // 2. 0 <= q <= floor(p / 65537)
115        // max_q is 281470681743360 = floor(p / 65537) = (p-1) / 65537 = 2^48 - 2^32
116        let max_q = builder.constant(GoldF::from_canonical_u64(281470681743360));
117        builder.range_check(q, 48);
118        let diff_q = builder.sub(max_q, q);
119        builder.range_check(diff_q, 48);
120
121        // 3. 0 <= r < 65537
122        let max_r = builder.constant(GoldF::from_canonical_u64(65537));
123        builder.range_check(r, 17);
124        let diff_r = builder.sub(max_r, r);
125        builder.range_check(diff_r, 17);
126
127        // 4. if q = floor(p / 65537) then r = 0
128        let q_equals_max = builder.is_equal(q, max_q);
129        let prod_temp = builder.mul(q_equals_max.target, r);
130        let zero_temp = builder.zero();
131        builder.connect(prod_temp, zero_temp);
132
133        // throw in the Generator to tell builder how to compute r
134        builder.add_simple_generator( Mod65537Generator {a, q, r} );
135
136        r
137    }
138}
139
140pub struct MessageTarget {
141    msg: Vec<Target>,
142}
143
144impl MessageTarget {
145    fn new_with_size(builder: &mut CircuitBuilder<GoldF, 2>, n: usize) -> Self {
146        Self {
147            msg: builder.add_virtual_targets(n),
148        }
149    }
150
151    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, msg: &Vec<GoldF>) -> Result<()> {
152        assert!(msg.len() == self.msg.len());
153        for (&t, &x) in self.msg.iter().zip(msg.iter()) {
154            pw.set_target(t, x)?;
155        }
156
157        Ok(())
158    }
159}
160
161pub struct SchnorrSignatureTarget {
162    s: Target,
163    e: Target,
164}
165
166impl SchnorrSignatureTarget {
167    fn new_virtual(builder: &mut CircuitBuilder<GoldF, 2>) -> Self {
168        let s = builder.add_virtual_target();
169        let e = builder.add_virtual_target();
170        Self{ s, e }
171    }
172
173    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, sig: &SchnorrSignature) -> Result<()> {
174        pw.set_target(self.s, GoldilocksField::from_canonical_u64(sig.s))?;
175        pw.set_target(self.e, GoldilocksField::from_canonical_u64(sig.e))?;
176        Ok(())
177    }
178}
179
180pub struct SchnorrPublicKeyTarget {
181    pk: Target,
182}
183
184impl SchnorrPublicKeyTarget {
185    fn new_virtual(builder: &mut CircuitBuilder<GoldF, 2>) -> Self {
186        Self{ pk: builder.add_virtual_target() }
187    }
188
189    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, pk: &SchnorrPublicKey) -> Result<()> {
190        pw.set_target(self.pk, pk.pk)?;
191        Ok(())
192    }
193}
194
195pub struct SchnorrBuilder {}
196
197impl SchnorrBuilder {
198
199
200    pub fn constrain_sig <
201        C: GenericConfig<2, F = GoldF>,
202    > (
203        &self,
204        builder: &mut CircuitBuilder::<GoldF, 2>,
205        sig: &SchnorrSignatureTarget,
206        msg: &MessageTarget,
207        pk: &SchnorrPublicKeyTarget,
208    ) -> () {
209        let PRIME_GROUP_GEN: Target = builder.constant(GoldF::from_canonical_u64(6612579038192137166));
210        let PRIME_GROUP_ORDER: Target = builder.constant(GoldF::from_canonical_u64(65537));
211        const num_bits_exp: usize = 32;
212
213        /* here's the direct verification calculation,
214        which we verify in-circuit
215        let r: GoldF = Self::pow(self.PRIME_GROUP_GEN, sig.s)
216            * Self::pow(pk.pk, sig.e);
217        let e_v: u64 = self.hash_insecure(&r, msg);
218        e_v == sig.e   */
219
220        let gs: Target = builder.exp(PRIME_GROUP_GEN, sig.s, num_bits_exp);
221        let pe: Target = builder.exp(pk.pk, sig.e, num_bits_exp);
222        let r: Target = builder.mul(gs, pe);
223
224        // compute hash
225        // note that it's safe to clone Targets since they just contain indices
226        let hash_input: Vec<Target> = std::iter::once(r)
227            .chain(msg.msg.iter().cloned()) 
228            .collect();
229        let hash_output: Target = builder.hash_n_to_hash_no_pad::<PoseidonHash>(
230            hash_input,
231        ).elements[0]; // whoops have to take mod group order;
232
233        let e: Target = Mod65537Builder::mod_65537(builder, hash_output);
234
235        // enforce equality
236        builder.connect(e, sig.e);
237    }
238}
239
240#[cfg(test)]
241mod tests{
242    use crate::schnorr::{SchnorrPublicKey, SchnorrSecretKey, SchnorrSigner, SchnorrSignature};
243    use crate::schnorr_prover::{MessageTarget, Mod65537Builder, SchnorrBuilder, SchnorrPublicKeyTarget, SchnorrSignatureTarget};
244    use plonky2::hash::poseidon::Poseidon;
245    use plonky2::iop::{
246        target::Target,
247        witness::{PartialWitness, PartitionWitness, Witness, WitnessWrite},
248    };
249    use plonky2::plonk::{
250        circuit_builder::CircuitBuilder,
251        circuit_data::{
252            CircuitConfig, 
253            CircuitData, 
254            CommonCircuitData, 
255            VerifierCircuitData, 
256            VerifierOnlyCircuitData
257        },
258        config::{GenericConfig, PoseidonGoldilocksConfig},
259    };
260    use plonky2::field::{
261        goldilocks_field::GoldilocksField,
262        types::Field,
263    };
264    use rand;
265
266    #[test]
267    fn test_mod65537() -> () {
268        const D: usize = 2;
269        const p: u64 = 18446744069414584321;  // the Goldilocks prime
270        type C = PoseidonGoldilocksConfig;
271        type F = <C as GenericConfig<D>>::F;
272
273        let config = CircuitConfig::standard_recursion_config();
274        let mut builder = CircuitBuilder::<F, D>::new(config);
275
276        let a64: Vec<u64> = vec![0, 1, 2, 65535, 65536, 65537, p - 4, p - 3, p - 2, p - 1];
277
278        let a: Vec<Target> = a64
279            .iter()
280            .map(|x| builder.constant(GoldilocksField::from_canonical_u64(*x)))
281            .collect();
282
283        let r: Vec<Target> = a.iter()
284            .map(|targ| Mod65537Builder::mod_65537(&mut builder, *targ))
285            .collect();
286
287        // check that the outputs are correct,
288        // obviously you don't need this in your own code
289        let r_expected64: Vec<u64> = a64.iter().map(|x| x % 65537).collect();
290        println!("Expected residues mod 64: {:?}", r_expected64);
291        let r_expected: Vec<Target> = r_expected64.iter()
292            .map(|x| builder.constant(GoldilocksField::from_canonical_u64(*x)))
293            .collect();
294        r.iter().zip(r_expected.iter())
295            .for_each(|(x, y)| builder.connect(*x, *y));
296        
297        let mut pw: PartialWitness<F> = PartialWitness::new();
298        let data = builder.build::<C>();
299        let proof = data.prove(pw).unwrap();
300
301        ()
302    }
303
304    #[test]
305    fn test_schnorr() -> () {
306        const D: usize = 2;
307        type C = PoseidonGoldilocksConfig;
308        type F = <C as GenericConfig<D>>::F;
309
310        let mut rng: rand::rngs::ThreadRng = rand::thread_rng();
311
312        let config = CircuitConfig::standard_recursion_config();
313        let mut builder = CircuitBuilder::<F, D>::new(config);
314
315        let sb: SchnorrBuilder = SchnorrBuilder{};
316
317        // create keypair, message, signature
318        let sk: SchnorrSecretKey = SchnorrSecretKey{ sk: 133 };
319        let ss = SchnorrSigner::new();
320        let pk: SchnorrPublicKey = ss.keygen(&sk);
321        let msg: Vec<GoldilocksField> = ss.u64_into_goldilocks_vec(
322            vec![1500, 1600, 0, 0, 0]
323        );
324        let msg_size: usize = msg.len();
325        let sig: SchnorrSignature = ss.sign(&msg, &sk, &mut rng);
326
327        let pk_targ = SchnorrPublicKeyTarget::new_virtual(&mut builder);
328        let sig_targ = SchnorrSignatureTarget::new_virtual(&mut builder);
329        let msg_targ = MessageTarget::new_with_size(&mut builder, msg_size);
330
331        
332        // instead of verifying we're going to prove the verification
333        sb.constrain_sig::<PoseidonGoldilocksConfig> (
334            &mut builder,
335            &sig_targ,
336            &msg_targ,
337            &pk_targ
338        );
339
340        // assign witnesses for verification
341        let mut pw: PartialWitness<F> = PartialWitness::new();
342        pk_targ.set_witness(&mut pw, &pk);
343        sig_targ.set_witness(&mut pw, &sig);
344        msg_targ.set_witness(&mut pw, &msg);
345
346
347        let data = builder.build::<C>();
348        let proof = data.prove(pw).unwrap();
349    }
350}