Merge pull request #1845 from privacy-scaling-explorations/fix/tally-…

…add-results fix(contracts): prevent adding additional tally results
privacy-scaling-explorations · Oct 3, 2024 · 07552f4 · 07552f4
2 parents 0d8d60c + a7f9b9e
commit 07552f4
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 93 deletions.
diff --git a/packages/contracts/contracts/Tally.sol b/packages/contracts/contracts/Tally.sol
@@ -20,6 +20,14 @@ contract Tally is Ownable, SnarkCommon, CommonUtilities, Hasher, DomainObjs, ITa
   uint256 internal constant TREE_ARITY = 2;
   uint256 internal constant VOTE_OPTION_TREE_ARITY = 5;
 
+  /// @notice Tally results
+  struct TallyResult {
+    /// Tally results value from tally.json
+    uint256 value;
+    /// Flag that this value was set and initialized
+    bool flag;
+  }
+
   /// @notice The commitment to the tally results. Its initial value is 0, but after
   /// the tally of each batch is proven on-chain via a zk-SNARK, it should be
   /// updated to:
@@ -53,7 +61,7 @@ contract Tally is Ownable, SnarkCommon, CommonUtilities, Hasher, DomainObjs, ITa
   Mode public immutable mode;
 
   // The tally results
-  mapping(uint256 => uint256) public tallyResults;
+  mapping(uint256 => TallyResult) public tallyResults;
 
   // The total tally results number
   uint256 public totalTallyResults;
@@ -391,8 +399,6 @@ contract Tally is Ownable, SnarkCommon, CommonUtilities, Hasher, DomainObjs, ITa
         i++;
       }
     }
-
-    totalTallyResults += voteOptionsLength;
   }
 
   /**
@@ -428,6 +434,13 @@ contract Tally is Ownable, SnarkCommon, CommonUtilities, Hasher, DomainObjs, ITa
       revert InvalidTallyVotesProof();
     }
 
-    tallyResults[_voteOptionIndex] = _tallyResult;
+    TallyResult storage previous = tallyResults[_voteOptionIndex];
+
+    if (!previous.flag) {
+      totalTallyResults++;
+    }
+
+    previous.flag = true;
+    previous.value = _tallyResult;
   }
 }
diff --git a/packages/contracts/package.json b/packages/contracts/package.json
@@ -151,7 +151,7 @@
     "benchmark:polygon-amoy": "pnpm run benchmark --network polygon_amoy"
   },
   "dependencies": {
-    "@nomicfoundation/hardhat-ethers": "^3.0.6",
+    "@nomicfoundation/hardhat-ethers": "^3.0.8",
     "@nomicfoundation/hardhat-toolbox": "^5.0.0",
     "@openzeppelin/contracts": "^5.0.2",
     "@openzeppelin/merkle-tree": "^1.0.7",

diff --git a/packages/contracts/tests/Tally.test.ts b/packages/contracts/tests/Tally.test.ts
@@ -372,6 +372,23 @@ describe("TallyVotes", () => {
 
       const indices = tallyData.results.tally.map((_, index) => index);
 
+      await tallyContract
+        .addTallyResults(
+          indices,
+          tallyData.results.tally,
+          tallyResultProofs,
+          tallyData.results.salt,
+          tallyData.totalSpentVoiceCredits.commitment,
+          newPerVOSpentVoiceCreditsCommitment,
+        )
+        .then((tx) => tx.wait());
+
+      const initialResults = await Promise.all(indices.map((index) => tallyContract.tallyResults(index)));
+      const initialTotalResults = await tallyContract.totalTallyResults();
+
+      expect(initialTotalResults).to.equal(tallyData.results.tally.length);
+      expect(initialResults.map((result) => result.value)).to.deep.equal(tallyData.results.tally);
+
       await tallyContract
         .addTallyResults(
           indices,
@@ -386,8 +403,8 @@ describe("TallyVotes", () => {
       const results = await Promise.all(indices.map((index) => tallyContract.tallyResults(index)));
       const totalResults = await tallyContract.totalTallyResults();
 
-      expect(totalResults).to.equal(tallyData.results.tally.length);
-      expect(results).to.deep.equal(tallyData.results.tally);
+      expect(initialTotalResults).to.equal(totalResults);
+      expect(initialResults).to.deep.equal(results);
 
       const onChainNewTallyCommitment = await tallyContract.tallyCommitment();
       expect(tallyGeneratedInputs!.newTallyCommitment).to.eq(onChainNewTallyCommitment.toString());