add statetransition circuit

* refactor crypto/elgamal package
* move some parts of state into statetransition circuit

api/process.go

@@ -65,12 +65,12 @@ func (a *API) newProcess(w http.ResponseWriter, r *http.Request) {
 		ErrGenericInternalServerError.Withf("could not marshal ballot mode: %v", err).Write(w)
 		return
 	}
-	state, err := st.Initialize(p.CensusRoot, ballotmode, publicKey.Marshal())
-	if err != nil {
+
+	if err := st.Initialize(p.CensusRoot, ballotmode, publicKey.Marshal()); err != nil {
 		ErrGenericInternalServerError.Withf("could not initialize state: %v", err).Write(w)
 		return
 	}
-	root, err := state.RootAsBigInt()
+	root, err := st.RootAsBigInt()
 	if err != nil {
 		ErrGenericInternalServerError.Withf("could not get state root: %v", err).Write(w)
 		return

circuits/statetransition/circuit.go

@@ -0,0 +1,144 @@
+package statetransition
+
+import (
+	"github.com/consensys/gnark/frontend"
+	"github.com/consensys/gnark/std/hash/mimc"
+	"github.com/vocdoni/gnark-crypto-primitives/elgamal"
+	"github.com/vocdoni/gnark-crypto-primitives/utils"
+)
+
+const (
+	// votes that were processed in AggregatedProof
+	VoteBatchSize = 10
+)
+
+type Circuit struct {
+	// ---------------------------------------------------------------------------------------------
+	// PUBLIC INPUTS
+
+	// list of root hashes
+	RootHashBefore frontend.Variable `gnark:",public"`
+	RootHashAfter  frontend.Variable `gnark:",public"`
+	NumNewVotes    frontend.Variable `gnark:",public"`
+	NumOverwrites  frontend.Variable `gnark:",public"`
+
+	// ---------------------------------------------------------------------------------------------
+	// SECRET INPUTS
+
+	AggregatedProof frontend.Variable // mock, this should be a zkProof
+
+	ProcessID     MerkleProof
+	CensusRoot    MerkleProof
+	BallotMode    MerkleProof
+	EncryptionKey MerkleProof
+	ResultsAdd    MerkleTransition
+	ResultsSub    MerkleTransition
+	Ballot        [VoteBatchSize]MerkleTransition
+	Commitment    [VoteBatchSize]MerkleTransition
+}
+
+// Define declares the circuit's constraints
+func (circuit Circuit) Define(api frontend.API) error {
+	hashFn := func(api frontend.API, data ...frontend.Variable) (frontend.Variable, error) {
+		h, err := mimc.NewMiMC(api)
+		if err != nil {
+			return 0, err
+		}
+		h.Write(data...)
+		return h.Sum(), nil
+	}
+
+	circuit.VerifyAggregatedZKProof(api)
+	circuit.VerifyMerkleProofs(api, hashFn)
+	circuit.VerifyMerkleTransitions(api, hashFn)
+	circuit.VerifyBallots(api)
+	return nil
+}
+
+func (circuit Circuit) VerifyAggregatedZKProof(api frontend.API) {
+	// all of the following values compose the preimage that is hashed
+	// to produce the public input needed to verify AggregatedProof.
+	// they are extracted from the MerkleProofs:
+	// ProcessID     := circuit.ProcessID.Value
+	// CensusRoot    := circuit.CensusRoot.Value
+	// BallotMode    := circuit.BallotMode.Value
+	// EncryptionKey := circuit.EncryptionKey.Value
+	// Nullifiers    := circuit.Ballot[i].NewKey
+	// Ballots       := circuit.Ballot[i].NewValue
+	// Addressess    := circuit.Commitment[i].NewKey
+	// Commitments   := circuit.Commitment[i].NewValue
+
+	api.Println("verify AggregatedZKProof mock:", circuit.AggregatedProof) // mock
+
+	packedInputs := func() frontend.Variable {
+		for i, p := range []MerkleProof{
+			circuit.ProcessID,
+			circuit.CensusRoot,
+			circuit.BallotMode,
+			circuit.EncryptionKey,
+		} {
+			api.Println("packInputs mock", i, p.Value) // mock
+		}
+		for i := range circuit.Ballot {
+			api.Println("packInputs mock nullifier", i, circuit.Ballot[i].NewKey) // mock
+			api.Println("packInputs mock ballot", i, circuit.Ballot[i].NewValue)  // mock
+		}
+		for i := range circuit.Commitment {
+			api.Println("packInputs mock address", i, circuit.Commitment[i].NewKey)      // mock
+			api.Println("packInputs mock commitment", i, circuit.Commitment[i].NewValue) // mock
+		}
+		return 1 // mock, should return hash of packed inputs
+	}
+
+	api.AssertIsEqual(packedInputs(), 1) // TODO: mock, should actually verify AggregatedZKProof
+}
+
+func (circuit Circuit) VerifyMerkleProofs(api frontend.API, hFn utils.Hasher) {
+	api.Println("verify ProcessID, CensusRoot, BallotMode and EncryptionKey belong to RootHashBefore")
+	circuit.ProcessID.VerifyProof(api, hFn, circuit.RootHashBefore)
+	circuit.CensusRoot.VerifyProof(api, hFn, circuit.RootHashBefore)
+	circuit.BallotMode.VerifyProof(api, hFn, circuit.RootHashBefore)
+	circuit.EncryptionKey.VerifyProof(api, hFn, circuit.RootHashBefore)
+}
+
+func (circuit Circuit) VerifyMerkleTransitions(api frontend.API, hFn utils.Hasher) {
+	// verify chain of tree transitions, order here is fundamental.
+	api.Println("tree transition starts with RootHashBefore:", prettyHex(circuit.RootHashBefore))
+	root := circuit.RootHashBefore
+	for i := range circuit.Ballot {
+		root = circuit.Ballot[i].Verify(api, hFn, root)
+	}
+	for i := range circuit.Commitment {
+		root = circuit.Commitment[i].Verify(api, hFn, root)
+	}
+	root = circuit.ResultsAdd.Verify(api, hFn, root)
+	root = circuit.ResultsSub.Verify(api, hFn, root)
+	api.Println("and final root is", prettyHex(root), "should be equal to RootHashAfter", prettyHex(circuit.RootHashAfter))
+	api.AssertIsEqual(root, circuit.RootHashAfter)
+}
+
+// VerifyBallots counts the ballots using homomorphic encrpytion
+func (circuit Circuit) VerifyBallots(api frontend.API) {
+	ballotSum, overwrittenSum, zero := elgamal.NewCiphertext(), elgamal.NewCiphertext(), elgamal.NewCiphertext()
+	var ballotCount, overwrittenCount frontend.Variable = 0, 0
+
+	for _, b := range circuit.Ballot {
+		// TODO: check that Hash(NewCiphertext) matches b.NewValue
+		// and Hash(OldCiphertext) matches b.OldValue
+		ballotSum.Add(api, ballotSum,
+			elgamal.NewCiphertext().Select(api, b.IsInsertOrUpdate(api), &b.NewCiphertext, zero))
+
+		overwrittenSum.Add(api, overwrittenSum,
+			elgamal.NewCiphertext().Select(api, b.IsUpdate(api), &b.OldCiphertext, zero))
+
+		ballotCount = api.Add(ballotCount, api.Select(b.IsInsertOrUpdate(api), 1, 0))
+		overwrittenCount = api.Add(overwrittenCount, api.Select(b.IsUpdate(api), 1, 0))
+	}
+
+	circuit.ResultsAdd.NewCiphertext.AssertIsEqual(api,
+		circuit.ResultsAdd.OldCiphertext.Add(api, &circuit.ResultsAdd.OldCiphertext, ballotSum))
+	circuit.ResultsSub.NewCiphertext.AssertIsEqual(api,
+		circuit.ResultsSub.OldCiphertext.Add(api, &circuit.ResultsSub.OldCiphertext, overwrittenSum))
+	api.AssertIsEqual(circuit.NumNewVotes, ballotCount)
+	api.AssertIsEqual(circuit.NumOverwrites, overwrittenCount)
+}