Merge pull request #1170 from privacy-scaling-explorations/fix/circui…

…t-unmatched-state-leaf

fix(circuits): enforce use of stateIndex from message

circuits/circom/messageValidator.circom

@@ -20,6 +20,7 @@ template MessageValidator() {
     validStateLeafIndex.in[0] <== stateTreeIndex;
     validStateLeafIndex.in[1] <== numSignUps;
 
+    // @todo check if we need this if we do the check inside processOne 
     // b) Whether the max vote option tree index is correct
     signal input voteOptionIndex;
     signal input maxVoteOptions;

circuits/circom/processMessages.circom

@@ -48,6 +48,8 @@ template ProcessMessages(
     var STATE_LEAF_PUB_Y_IDX = 1;
     var STATE_LEAF_VOICE_CREDIT_BALANCE_IDX = 2;
     var STATE_LEAF_TIMESTAMP_IDX = 3;
+
+    var N_BITS = 252;
 
     // Note that we sha256 hash some values from the contract, pass in the hash
     // as a public input, and pass in said values as private inputs. This saves
@@ -277,6 +279,7 @@ template ProcessMessages(
     component processors[batchSize]; 
     // topup type processor
     component processors2[batchSize]; 
+
     for (var i = batchSize - 1; i >= 0; i --) {
         // process it as vote type message
         processors[i] = ProcessOne(stateTreeDepth, voteOptionTreeDepth);
@@ -349,6 +352,7 @@ template ProcessMessages(
                     <== currentStateLeavesPathElements[i][j][k];
             }
         }
+
         // pick the correct result by msg type
         tmpStateRoot1[i] <== processors[i].newStateRoot * (2 - msgs[i][0]); 
         tmpStateRoot2[i] <== processors2[i].newStateRoot * (msgs[i][0] - 1);
@@ -378,6 +382,8 @@ template ProcessTopup(stateTreeDepth) {
     var STATE_LEAF_VOICE_CREDIT_BALANCE_IDX = 2;
     var STATE_LEAF_TIMESTAMP_IDX = 3;
 
+    var N_BITS = 252;
+
     signal input msgType;
     signal input stateTreeIndex;
     signal input amount;
@@ -395,9 +401,9 @@ template ProcessTopup(stateTreeDepth) {
     // msgType of topup command is 2
     amt <== amount * (msgType - 1); 
     index <== stateTreeIndex * (msgType - 1);
-    component validCreditBalance = LessEqThan(252);
+    component validCreditBalance = LessEqThan(N_BITS);
     // check stateIndex, if invalid index, set index and amount to zero
-    component validStateLeafIndex = LessEqThan(252);
+    component validStateLeafIndex = LessEqThan(N_BITS);
     validStateLeafIndex.in[0] <== index;
     validStateLeafIndex.in[1] <== numSignUps;
 
@@ -462,6 +468,8 @@ template ProcessOne(stateTreeDepth, voteOptionTreeDepth) {
     var STATE_LEAF_VOICE_CREDIT_BALANCE_IDX = 2;
     var STATE_LEAF_TIMESTAMP_IDX = 3;
 
+    var N_BITS = 252;
+
     signal input msgType;
     signal input numSignUps;
     signal input maxVoteOptions;
@@ -526,22 +534,34 @@ template ProcessOne(stateTreeDepth, voteOptionTreeDepth) {
 
     //  ----------------------------------------------------------------------- 
     // 2. If msgType = 0 and isValid is 0, generate indices for leaf 0
-    //  Otherwise, generate indices for commmand.stateIndex or topupStateIndex depending on msgType
+    // Otherwise, generate indices for commmand.stateIndex or topupStateIndex depending on msgType
     signal indexByType;
     signal tmpIndex1;
     signal tmpIndex2;
     tmpIndex1 <== cmdStateIndex * (2 - msgType);
     tmpIndex2 <== topupStateIndex * (msgType - 1);
     indexByType <== tmpIndex1 + tmpIndex2;
 
-    component stateIndexMux = Mux1();
-    stateIndexMux.s <== transformer.isValid + msgType - 1;
-    stateIndexMux.c[0] <== 0;
-    stateIndexMux.c[1] <== indexByType;
+    // we can validate if the state index is within the numSignups
+    // if not, we use 0
+    // this is because decryption of an invalid message 
+    // might result in random packed vals
+    component validStateLeafIndex = SafeLessThan(N_BITS);
+    validStateLeafIndex.in[0] <== indexByType;
+    validStateLeafIndex.in[1] <== numSignUps;
 
-    component stateLeafPathIndices = QuinGeneratePathIndices(stateTreeDepth);
-    stateLeafPathIndices.in <== stateIndexMux.out;
+    // use a mux to pick the correct index
+    component indexMux = Mux1();
+    indexMux.s <== validStateLeafIndex.out;
+    indexMux.c[0] <== 0;
+    indexMux.c[1] <== indexByType;
 
+    // @note that we expect a coordinator to send the state leaf corresponding to a message
+    // which specifies a valid state index. If this is not the case, the 
+    // proof will fail to generate. 
+    component stateLeafPathIndices = QuinGeneratePathIndices(stateTreeDepth);
+    stateLeafPathIndices.in <== indexMux.out;
+
     //  ----------------------------------------------------------------------- 
     // 3. Verify that the original state leaf exists in the given state root
     component stateLeafQip = QuinTreeInclusionProof(stateTreeDepth);
@@ -572,6 +592,7 @@ template ProcessOne(stateTreeDepth, voteOptionTreeDepth) {
             ballotQip.path_elements[i][j] <== ballotPathElements[i][j];
         }
     }
+
     ballotQip.root === currentBallotRoot;
 
     //  ----------------------------------------------------------------------- 
@@ -583,7 +604,7 @@ template ProcessOne(stateTreeDepth, voteOptionTreeDepth) {
     b <== currentVoteWeight * currentVoteWeight;
     c <== cmdNewVoteWeight * cmdNewVoteWeight;
 
-    component enoughVoiceCredits = SafeGreaterEqThan(252);
+    component enoughVoiceCredits = SafeGreaterEqThan(N_BITS);
     enoughVoiceCredits.in[0] <== stateLeaf[STATE_LEAF_VOICE_CREDIT_BALANCE_IDX] + b;
     enoughVoiceCredits.in[1] <== c;
 
@@ -592,8 +613,16 @@ template ProcessOne(stateTreeDepth, voteOptionTreeDepth) {
     isMessageValid.in[0] <== bothValid;
     isMessageValid.in[1] <== transformer.isValid + enoughVoiceCredits.out;
 
+    // check that the vote option index is < maxVoteOptions (0-indexed)
+    component validVoteOptionIndex = SafeLessThan(N_BITS);
+    validVoteOptionIndex.in[0] <== cmdVoteOptionIndex;
+    validVoteOptionIndex.in[1] <== maxVoteOptions;
+
+    // @note pick the correct vote option index based on whether the index is < max vote options 
+    // @todo can probably add one output to messageValidator and take from there
+    // or maybe we can remove altogther from messageValidator so we don't double check this
     component cmdVoteOptionIndexMux = Mux1();
-    cmdVoteOptionIndexMux.s <== isMessageValid.out;
+    cmdVoteOptionIndexMux.s <== validVoteOptionIndex.out;
     cmdVoteOptionIndexMux.c[0] <== 0;
     cmdVoteOptionIndexMux.c[1] <== cmdVoteOptionIndex;
 

circuits/circom/processMessagesNonQv.circom

@@ -49,6 +49,8 @@ template ProcessMessagesNonQv(
     var STATE_LEAF_PUB_Y_IDX = 1;
     var STATE_LEAF_VOICE_CREDIT_BALANCE_IDX = 2;
     var STATE_LEAF_TIMESTAMP_IDX = 3;
+
+    var N_BITS = 252;
 
     // Note that we sha256 hash some values from the contract, pass in the hash
     // as a public input, and pass in said values as private inputs. This saves
@@ -389,6 +391,8 @@ template ProcessOneNonQv(stateTreeDepth, voteOptionTreeDepth) {
     var STATE_LEAF_VOICE_CREDIT_BALANCE_IDX = 2;
     var STATE_LEAF_TIMESTAMP_IDX = 3;
 
+    var N_BITS = 252;
+
     signal input msgType;
     signal input numSignUps;
     signal input maxVoteOptions;
@@ -461,13 +465,25 @@ template ProcessOneNonQv(stateTreeDepth, voteOptionTreeDepth) {
     tmpIndex2 <== topupStateIndex * (msgType - 1);
     indexByType <== tmpIndex1 + tmpIndex2;
 
-    component stateIndexMux = Mux1();
-    stateIndexMux.s <== transformer.isValid + msgType - 1;
-    stateIndexMux.c[0] <== 0;
-    stateIndexMux.c[1] <== indexByType;
-
+    // we can validate if the state index is within the numSignups
+    // if not, we use 0
+    // this is because decryption of an invalid message 
+    // might result in random packed vals
+    component validStateLeafIndex = SafeLessThan(N_BITS);
+    validStateLeafIndex.in[0] <== indexByType;
+    validStateLeafIndex.in[1] <== numSignUps;
+
+    // use a mux to pick the correct index
+    component indexMux = Mux1();
+    indexMux.s <== validStateLeafIndex.out;
+    indexMux.c[0] <== 0;
+    indexMux.c[1] <== indexByType;
+
+    // @note that we expect a coordinator to send the state leaf corresponding to a message
+    // which specifies a valid state index. If this is not the case, the 
+    // proof will fail to generate. 
     component stateLeafPathIndices = QuinGeneratePathIndices(stateTreeDepth);
-    stateLeafPathIndices.in <== stateIndexMux.out;
+    stateLeafPathIndices.in <== indexMux.out;
 
     //  ----------------------------------------------------------------------- 
     // 3. Verify that the original state leaf exists in the given state root
@@ -510,7 +526,7 @@ template ProcessOneNonQv(stateTreeDepth, voteOptionTreeDepth) {
     b <== currentVoteWeight;
     c <== cmdNewVoteWeight;
 
-    component enoughVoiceCredits = SafeGreaterEqThan(252);
+    component enoughVoiceCredits = SafeGreaterEqThan(N_BITS);
     enoughVoiceCredits.in[0] <== stateLeaf[STATE_LEAF_VOICE_CREDIT_BALANCE_IDX] + b;
     enoughVoiceCredits.in[1] <== c;
 
@@ -519,8 +535,16 @@ template ProcessOneNonQv(stateTreeDepth, voteOptionTreeDepth) {
     isMessageValid.in[0] <== bothValid;
     isMessageValid.in[1] <== transformer.isValid + enoughVoiceCredits.out;
 
+    // check that the vote option index is < maxVoteOptions (0-indexed)
+    component validVoteOptionIndex = SafeLessThan(N_BITS);
+    validVoteOptionIndex.in[0] <== cmdVoteOptionIndex;
+    validVoteOptionIndex.in[1] <== maxVoteOptions;
+
+    // @note pick the correct vote option index based on whether the index is < max vote options 
+    // @todo can probably add one output to messageValidator and take from there
+    // or maybe we can remove altogther from messageValidator so we don't double check this
     component cmdVoteOptionIndexMux = Mux1();
-    cmdVoteOptionIndexMux.s <== isMessageValid.out;
+    cmdVoteOptionIndexMux.s <== validVoteOptionIndex.out;
     cmdVoteOptionIndexMux.c[0] <== 0;
     cmdVoteOptionIndexMux.c[1] <== cmdVoteOptionIndex;
 

circuits/circom/stateLeafAndBallotTransformer.circom

@@ -86,7 +86,7 @@ template StateLeafAndBallotTransformer() {
     messageValidator.voteWeight <== cmdNewVoteWeight;
 
     // if the message is valid then we swap out the public key
-    // we have to do this in two Mux one for pucKey[0]
+    // we have to do this in two Mux one for pubKey[0]
     // and one for pubKey[1]
     component newSlPubKey0Mux = Mux1();
     newSlPubKey0Mux.s <== messageValidator.isValid;

circuits/circom/stateLeafAndBallotTransformerNonQv.circom

@@ -86,7 +86,7 @@ template StateLeafAndBallotTransformerNonQv() {
     messageValidator.voteWeight <== cmdNewVoteWeight;
 
     // if the message is valid then we swap out the public key
-    // we have to do this in two Mux one for pucKey[0]
+    // we have to do this in two Mux one for pubKey[0]
     // and one for pubKey[1]
     component newSlPubKey0Mux = Mux1();
     newSlPubKey0Mux.s <== messageValidator.isValid;