From a26f8979cdb3bba16b532665cf481701cabb94e3 Mon Sep 17 00:00:00 2001
From: poma <Semenov.Roman@mail.ru>
Date: Fri, 12 Jul 2019 23:52:10 +0300
Subject: [PATCH] add verification test, fix bugs

---
 test/test_snark.js | 15 ++++++++--
 test/utils.js      | 70 +++++++++++++++++++++++++++++++++-------------
 2 files changed, 63 insertions(+), 22 deletions(-)

diff --git a/test/test_snark.js b/test/test_snark.js
index 084301d..5d6dcf2 100644
--- a/test/test_snark.js
+++ b/test/test_snark.js
@@ -1,4 +1,5 @@
 const fs = require('fs');
+const assert = require('assert');
 const circom = require("circom");
 const snarkjs = require("snarkjs");
 const circomlib = require('circomlib');
@@ -34,7 +35,7 @@ function generateDeposit() {
   const {root, path_elements, path_index} = await tree.path(1);
 
   // Circuit input
-  const input = stringifyBigInts({
+  const input = {
     // public
     root: root,
     nullifier: dep2.nullifier,
@@ -45,11 +46,21 @@ function generateDeposit() {
     secret: dep2.secret,
     pathElements: path_elements,
     pathIndex: path_index,
-  });
+  };
 
   console.log("Input:\n", input);
   console.time("Time");
   const proof = await utils.snarkProof(input);
   console.log("Proof:\n", proof);
   console.timeEnd("Time");
+
+  const verify = await utils.snarkVerify(proof);
+  assert(verify);
+
+  // try to cheat with recipient
+  proof.publicSignals[2] = '0x000000000000000000000000000000000000000000000000000000000000beef';
+  const verifyScam = await utils.snarkVerify(proof);
+  assert(!verifyScam);
+
+  console.log("Done.");
 })();
diff --git a/test/utils.js b/test/utils.js
index f53474d..9c5350d 100644
--- a/test/utils.js
+++ b/test/utils.js
@@ -1,3 +1,6 @@
+// This file is a bit of a mess because of different bigInt formats in websnark and snarkjs
+// It will be rewritten during browser integration
+
 const fs = require('fs');
 const circom = require("circom");
 const snarkjs = require("snarkjs");
@@ -8,12 +11,30 @@ const pedersen = circomlib.pedersenHash;
 const babyjub = circomlib.babyJub;
 const mimcsponge = circomlib.mimcsponge;
 const bigInt = snarkjs.bigInt;
-const buildGroth16 = require('../node_modules/websnark/src/groth16.js');
-const stringifyBigInts = require("../node_modules/websnark/tools/stringifybigint.js").stringifyBigInts;
-const unstringifyBigInts = require("../node_modules/websnark/tools/stringifybigint.js").unstringifyBigInts;
+const buildGroth16 = require('websnark/src/groth16');
+const stringifyBigInts = require("websnark/tools/stringifybigint").stringifyBigInts;
+const unstringifyBigInts = require("websnark/tools/stringifybigint").unstringifyBigInts;
+const stringifyBigInts2 = require("snarkjs/src/stringifybigint").stringifyBigInts;
+const unstringifyBigInts2 = require("snarkjs/src/stringifybigint").unstringifyBigInts;
 
 const rbigint = (nbytes) => snarkjs.bigInt.leBuff2int(crypto.randomBytes(nbytes));
 
+function unhexBigInts(o) {
+  if ((typeof(o) == "string") && (/^0x[0-9a-fA-F]+$/.test(o)))  {
+    return bigInt(o);
+  } else if (Array.isArray(o)) {
+    return o.map(unhexBigInts);
+  } else if (typeof o == "object") {
+    const res = {};
+    for (let k in o) {
+      res[k] = unhexBigInts(o[k]);
+    }
+    return res;
+  } else {
+    return o;
+  }
+}
+
 function pedersenHash(data) {
   return babyjub.unpackPoint(pedersen.hash(data))[0];
 }
@@ -46,6 +67,7 @@ function p256(o) {
 }
 
 function convertWitness(witness) {
+  witness = unstringifyBigInts(witness);
   const buffLen = witness.length * 32;
   const buff = new ArrayBuffer(buffLen);
   const h = {
@@ -53,9 +75,8 @@ function convertWitness(witness) {
     offset: 0
   };
   for (let i=0; i<witness.length; i++) {
-    for (let i=0; i<8; i++) {
-      //const v = witness[i].shiftRight(i*32).and(0xFFFFFFFF).toJSNumber();
-      const v = Number(witness[i].shr(i * 32).and(BigInt(0xFFFFFFFF)));
+    for (let j=0; j<8; j++) {
+      const v = witness[i].shiftRight(j*32).and(0xFFFFFFFF).toJSNumber();
       h.dataView.setUint32(h.offset, v, true);
       h.offset += 4;
     }
@@ -63,25 +84,34 @@ function convertWitness(witness) {
   return buff;
 }
 
-function toArrayBuffer(b) {
-  return b.buffer.slice(b.byteOffset, b.byteOffset + b.byteLength);
-}
-
 async function snarkProof(input) {
-  const circuit = new snarkjs.Circuit(unstringifyBigInts(require("../build/circuits/withdraw.json")));
-  const witnessArray = circuit.calculateWitness(input);
-  const witness = convertWitness(witnessArray);
-  const publicSignals = witnessArray.slice(1, circuit.nPubInputs + circuit.nOutputs + 1);
-  const key = toArrayBuffer(fs.readFileSync("../build/circuits/withdraw_proving_key.bin"));
+  input = unstringifyBigInts2(input);
+  const circuit = new snarkjs.Circuit(unstringifyBigInts2(require("../build/circuits/withdraw.json")));
+  const proving_key = fs.readFileSync("../build/circuits/withdraw_proving_key.bin");
+
+  const witness = circuit.calculateWitness(input);
+  const witnessBin = convertWitness(stringifyBigInts2(witness));
+  const publicSignals = witness.slice(1, circuit.nPubInputs + circuit.nOutputs + 1);
+
   const groth16 = await buildGroth16();
-  let proof = await groth16.proof(witness, key);
-  proof = unstringifyBigInts(proof);
-  return p256({
+  let proof = await groth16.proof(witnessBin, proving_key.buffer);
+  return p256(unstringifyBigInts2(stringifyBigInts({
     pi_a: [proof.pi_a[0], proof.pi_a[1]],
     pi_b: [[proof.pi_b[0][1], proof.pi_b[0][0]], [proof.pi_b[1][1], proof.pi_b[1][0]]],
     pi_c: [proof.pi_c[0], proof.pi_c[1]],
     publicSignals: publicSignals,
-  });
+  })));
 }
 
-module.exports = {rbigint, pedersenHash, snarkProof, mimcHash};
+async function snarkVerify(proof) {
+  proof = unhexBigInts(proof);
+  const verification_key = unstringifyBigInts2(require('../build/circuits/withdraw_verification_key.json'));
+  const data = {
+    pi_a: [proof.pi_a[0], proof.pi_a[1], bigInt(1)],
+    pi_b: [[proof.pi_b[0][1], proof.pi_b[0][0]], [proof.pi_b[1][1], proof.pi_b[1][0]], [bigInt(1), bigInt(0)]],
+    pi_c: [proof.pi_c[0], proof.pi_c[1], bigInt(1)]
+  };
+  return groth.isValid(verification_key, data, proof.publicSignals);
+}
+
+module.exports = {rbigint, pedersenHash, snarkProof, mimcHash, snarkVerify, unhexBigInts};