Merge pull request #50 from worldcoin/dcbuild3r/import-setup-v2

PRO-427: add trusted setup proving system importing

main.go

@@ -106,6 +106,57 @@ func main() {
 					return nil
 				},
 			},
+			{
+				Name: "import-setup",
+				Flags: []cli.Flag{
+					&cli.StringFlag{Name: "output", Usage: "Output file", Required: true},
+					&cli.StringFlag{Name: "pk", Usage: "Proving key", Required: true},
+					&cli.StringFlag{Name: "vk", Usage: "Verifying key", Required: true},
+					&cli.StringFlag{Name: "mode", Usage: "insertion/deletion", Required: true},
+					&cli.UintFlag{Name: "tree-depth", Usage: "Merkle tree depth", Required: true},
+					&cli.UintFlag{Name: "batch-size", Usage: "Batch size", Required: true},
+				},
+				Action: func(context *cli.Context) error {
+					path := context.String("output")
+					pk := context.String("pk")
+					vk := context.String("vk")
+					mode := context.String("mode")
+					treeDepth := uint32(context.Uint("tree-depth"))
+					batchSize := uint32(context.Uint("batch-size"))
+					var system *prover.ProvingSystem
+					var err error
+
+					logging.Logger().Info().Msg("Importing setup")
+
+					if mode == server.InsertionMode {
+						system, err = prover.ImportInsertionSetup(treeDepth, batchSize, pk, vk)
+
+						if err != nil {
+							return err
+						}
+
+					} else if mode == server.DeletionMode {
+						system, err = prover.ImportDeletionSetup(treeDepth, batchSize, pk, vk)
+
+						if err != nil {
+							return err
+						}
+					} else {
+						return fmt.Errorf("Invalid mode: %s", mode)
+					}
+					file, err := os.Create(path)
+					defer file.Close()
+					if err != nil {
+						return err
+					}
+					written, err := system.WriteTo(file)
+					if err != nil {
+						return err
+					}
+					logging.Logger().Info().Int64("bytesWritten", written).Msg("proving system written to file")
+					return nil
+				},
+			},
 			{
 				Name: "export-solidity",
 				Flags: []cli.Flag{

prover/circuit_utils.go

@@ -1,10 +1,14 @@
 package prover
 
 import (
+	"fmt"
 	"io"
+	"os"
 	"strconv"
+	"worldcoin/gnark-mbu/logging"
 	"worldcoin/gnark-mbu/prover/poseidon"
 
+	"github.com/consensys/gnark-crypto/ecc"
 	"github.com/consensys/gnark/backend/groth16"
 	"github.com/consensys/gnark/constraint"
 	"github.com/consensys/gnark/frontend"
@@ -182,6 +186,35 @@ func (gadget DeletionProof) DefineGadget(api frontend.API) interface{} {
 	return root
 }
 
+// Trusted setup utility functions
+// Taken from: https://github.com/bnb-chain/zkbnb/blob/master/common/prove/proof_keys.go#L19
+func LoadProvingKey(filepath string) (pk groth16.ProvingKey, err error) {
+	logging.Logger().Info().Msg("start reading proving key")
+	pk = groth16.NewProvingKey(ecc.BN254)
+	f, _ := os.Open(filepath)
+	_, err = pk.ReadFrom(f)
+	if err != nil {
+		return pk, fmt.Errorf("read file error")
+	}
+	f.Close()
+
+	return pk, nil
+}
+
+// Taken from: https://github.com/bnb-chain/zkbnb/blob/master/common/prove/proof_keys.go#L32
+func LoadVerifyingKey(filepath string) (verifyingKey groth16.VerifyingKey, err error) {
+	logging.Logger().Info().Msg("start reading verifying key")
+	verifyingKey = groth16.NewVerifyingKey(ecc.BN254)
+	f, _ := os.Open(filepath)
+	_, err = verifyingKey.ReadFrom(f)
+	if err != nil {
+		return verifyingKey, fmt.Errorf("read file error")
+	}
+	f.Close()
+
+	return verifyingKey, nil
+}
+
 // ReducedModRCheck Checks a little-endian array of bits asserting that it represents a number that
 // is less than the field modulus R.
 type ReducedModRCheck struct {

prover/deletion_proving_system.go

@@ -143,3 +143,34 @@ func (ps *ProvingSystem) VerifyDeletion(inputHash big.Int, proof *Proof) error {
 	}
 	return groth16.Verify(proof.Proof, ps.VerifyingKey, witness)
 }
+
+func ImportDeletionSetup(treeDepth uint32, batchSize uint32, pkPath string, vkPath string) (*ProvingSystem, error) {
+	proofs := make([][]frontend.Variable, batchSize)
+	for i := 0; i < int(batchSize); i++ {
+		proofs[i] = make([]frontend.Variable, treeDepth)
+	}
+	circuit := DeletionMbuCircuit{
+		Depth:        int(treeDepth),
+		BatchSize:    int(batchSize),
+		IdComms:      make([]frontend.Variable, batchSize),
+		MerkleProofs: proofs,
+	}
+	ccs, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, &circuit)
+	if err != nil {
+		return nil, err
+	}
+
+	pk, err := LoadProvingKey(pkPath)
+
+	if err != nil {
+		return nil, err
+	}
+
+	vk, err := LoadVerifyingKey(vkPath)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &ProvingSystem{treeDepth, batchSize, pk, vk, ccs}, nil
+}

prover/insertion_circuit.go

@@ -3,7 +3,9 @@ package prover
 import (
 	"worldcoin/gnark-mbu/prover/keccak"
 
+	"github.com/consensys/gnark-crypto/ecc"
 	"github.com/consensys/gnark/frontend"
+	"github.com/consensys/gnark/frontend/cs/r1cs"
 	"github.com/reilabs/gnark-lean-extractor/v2/abstractor"
 )
 
@@ -72,3 +74,34 @@ func (circuit *InsertionMbuCircuit) Define(api frontend.API) error {
 
 	return nil
 }
+
+func ImportInsertionSetup(treeDepth uint32, batchSize uint32, pkPath string, vkPath string) (*ProvingSystem, error) {
+	proofs := make([][]frontend.Variable, batchSize)
+	for i := 0; i < int(batchSize); i++ {
+		proofs[i] = make([]frontend.Variable, treeDepth)
+	}
+	circuit := InsertionMbuCircuit{
+		Depth:        int(treeDepth),
+		BatchSize:    int(batchSize),
+		IdComms:      make([]frontend.Variable, batchSize),
+		MerkleProofs: proofs,
+	}
+	ccs, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, &circuit)
+	if err != nil {
+		return nil, err
+	}
+
+	pk, err := LoadProvingKey(pkPath)
+
+	if err != nil {
+		return nil, err
+	}
+
+	vk, err := LoadVerifyingKey(vkPath)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &ProvingSystem{treeDepth, batchSize, pk, vk, ccs}, nil
+}