Merge pull request #1226 from privacy-scaling-explorations/feat/batch…

…-messages feat(poll): add message batch submission
privacy-scaling-explorations · Feb 22, 2024 · 03b685f · 03b685f
2 parents 6e0a67d + 005258b
commit 03b685f
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 6 deletions.
diff --git a/contracts/contracts/Poll.sol b/contracts/contracts/Poll.sol
@@ -71,6 +71,7 @@ contract Poll is Params, Utilities, SnarkCommon, Ownable, EmptyBallotRoots, IPol
   error MaciPubKeyLargerThanSnarkFieldSize();
   error StateAqAlreadyMerged();
   error StateAqSubtreesNeedMerge();
+  error InvalidBatchLength();
 
   event PublishMessage(Message _message, PubKey _encPubKey);
   event TopupMessage(Message _message);
@@ -197,6 +198,26 @@ contract Poll is Params, Utilities, SnarkCommon, Ownable, EmptyBallotRoots, IPol
     emit PublishMessage(_message, _encPubKey);
   }
 
+  /// @notice submit a message batch
+  /// @dev Can only be submitted before the voting deadline
+  /// @param _messages the messages
+  /// @param _encPubKeys the encrypted public keys
+  function publishMessageBatch(Message[] calldata _messages, PubKey[] calldata _encPubKeys) external {
+    if (_messages.length != _encPubKeys.length) {
+      revert InvalidBatchLength();
+    }
+
+    uint256 len = _messages.length;
+    for (uint256 i = 0; i < len; ) {
+      // an event will be published by this function already
+      publishMessage(_messages[i], _encPubKeys[i]);
+
+      unchecked {
+        i++;
+      }
+    }
+  }
+
   /// @inheritdoc IPoll
   function mergeMaciStateAqSubRoots(uint256 _numSrQueueOps, uint256 _pollId) public onlyOwner isAfterVotingDeadline {
     // This function cannot be called after the stateAq was merged

diff --git a/contracts/tests/Poll.test.ts b/contracts/tests/Poll.test.ts
@@ -36,6 +36,8 @@ describe("Poll", () => {
 
   const maciState = new MaciState(STATE_TREE_DEPTH);
 
+  const keypair = new Keypair();
+
   describe("deployment", () => {
     before(async () => {
       signer = await getDefaultSigner();
@@ -175,8 +177,6 @@ describe("Poll", () => {
 
   describe("publishMessage", () => {
     it("should publish a message to the Poll contract", async () => {
-      const keypair = new Keypair();
-
       const command = new PCommand(1n, keypair.pubKey, 0n, 9n, 1n, pollId, 0n);
 
       const signature = command.sign(keypair.privKey);
@@ -190,8 +190,6 @@ describe("Poll", () => {
     });
 
     it("should emit an event when publishing a message", async () => {
-      const keypair = new Keypair();
-
       const command = new PCommand(1n, keypair.pubKey, 0n, 9n, 1n, pollId, 0n);
 
       const signature = command.sign(keypair.privKey);
@@ -204,11 +202,45 @@ describe("Poll", () => {
       maciState.polls.get(pollId)?.publishMessage(message, keypair.pubKey);
     });
 
-    it("shold not allow to publish a message after the voting period ends", async () => {
+    it("should allow to publish a message batch", async () => {
+      const messages: [Message, PubKey][] = [];
+      for (let i = 0; i < 2; i += 1) {
+        const command = new PCommand(1n, keypair.pubKey, 0n, 9n, 1n, pollId, BigInt(i));
+        const signature = command.sign(keypair.privKey);
+        const sharedKey = Keypair.genEcdhSharedKey(keypair.privKey, coordinator.pubKey);
+        const message = command.encrypt(signature, sharedKey);
+        messages.push([message, keypair.pubKey]);
+      }
+
+      const tx = await pollContract.publishMessageBatch(
+        messages.map(([m]) => m.asContractParam()),
+        messages.map(([, k]) => k.asContractParam()),
+      );
+      const receipt = await tx.wait();
+      expect(receipt?.status).to.eq(1);
+
+      messages.forEach(([message, key]) => {
+        maciState.polls.get(pollId)?.publishMessage(message, key);
+      });
+    });
+
+    it("should throw when the message batch has messages length != encPubKeys length", async () => {
+      const command = new PCommand(1n, keypair.pubKey, 0n, 9n, 1n, pollId, 0n);
+      const signature = command.sign(keypair.privKey);
+      const sharedKey = Keypair.genEcdhSharedKey(keypair.privKey, coordinator.pubKey);
+      const message = command.encrypt(signature, sharedKey);
+      await expect(
+        pollContract.publishMessageBatch(
+          [message.asContractParam(), message.asContractParam()],
+          [keypair.pubKey.asContractParam()],
+        ),
+      ).to.be.revertedWithCustomError(pollContract, "InvalidBatchLength");
+    });
+
+    it("should not allow to publish a message after the voting period ends", async () => {
       const dd = await pollContract.getDeployTimeAndDuration();
       await timeTravel(signer.provider as unknown as EthereumProvider, Number(dd[0]) + 1);
 
-      const keypair = new Keypair();
       const command = new PCommand(1n, keypair.pubKey, 0n, 9n, 1n, pollId, 0n);
 
       const signature = command.sign(keypair.privKey);
@@ -219,6 +251,18 @@ describe("Poll", () => {
         pollContract.publishMessage(message.asContractParam(), keypair.pubKey.asContractParam(), { gasLimit: 300000 }),
       ).to.be.revertedWithCustomError(pollContract, "VotingPeriodOver");
     });
+
+    it("should not allow to publish a message batch after the voting period ends", async () => {
+      const command = new PCommand(1n, keypair.pubKey, 0n, 9n, 1n, pollId, 0n);
+      const signature = command.sign(keypair.privKey);
+      const sharedKey = Keypair.genEcdhSharedKey(keypair.privKey, coordinator.pubKey);
+      const message = command.encrypt(signature, sharedKey);
+      await expect(
+        pollContract.publishMessageBatch([message.asContractParam()], [keypair.pubKey.asContractParam()], {
+          gasLimit: 300000,
+        }),
+      ).to.be.revertedWithCustomError(pollContract, "VotingPeriodOver");
+    });
   });
 
   describe("Merge messages", () => {