test(core): add test cases for processMessages()

core/ts/__tests__/MaciState.test.ts

@@ -109,7 +109,7 @@ describe("MaciState", function () {
       accumulatorQueue.mergeSubRoots(0);
       accumulatorQueue.merge(treeDepths.messageTreeDepth);
 
-      expect(accumulatorQueue.getRoot(treeDepths.messageTreeDepth).toString()).to.eq(msgTree.root.toString());
+      expect(accumulatorQueue.getRoot(treeDepths.messageTreeDepth).toString()).to.eq(poll.messageTree.root.toString());
     });
 
     it("packProcessMessageSmallVals and unpackProcessMessageSmallVals", () => {

core/ts/__tests__/ProcessMessage.test.ts → core/ts/__tests__/ProcessMessages.test.ts

@@ -1,13 +1,127 @@
-/* eslint-disable */
-describe("Message processing", () => {
-  describe("Process a batch of messages", () => {
+import { expect } from "chai";
+
+import { Signature } from "maci-crypto";
+import { PCommand, Message, Keypair, PubKey } from "maci-domainobjs";
+
+import { STATE_TREE_DEPTH, MaciState, Poll } from "..";
+import { AssertionError } from "assert";
+
+const voiceCreditBalance = BigInt(100);
+const duration = 30;
+const maxValues = {
+  maxUsers: 25,
+  maxMessages: 25,
+  maxVoteOptions: 25,
+};
+const treeDepths = {
+  intStateTreeDepth: 2,
+  messageTreeDepth: 3,
+  messageTreeSubDepth: 2,
+  voteOptionTreeDepth: 4,
+};
+const messageBatchSize = 25;
+
+class TestHarness {
+  maciState = new MaciState(STATE_TREE_DEPTH);
+  coordinatorKeypair = new Keypair();
+  poll: Poll;
+  pollId: number;
+  users: Keypair[] = [];
+  stateIndices = new Map<Keypair, number>();
+
+  constructor() {
+    this.pollId = this.maciState.deployPoll(
+      duration,
+      BigInt(Math.floor(Date.now() / 1000) + duration),
+      maxValues,
+      treeDepths,
+      messageBatchSize,
+      this.coordinatorKeypair,
+    );
+    this.poll = this.maciState.polls[this.pollId];
+  }
+
+  createUsers = (numUsers: number): Keypair[] => {
+    for (let i = 0; i < numUsers; i++) {
+      const user = new Keypair();
+      this.users.push(user);
+      const stateIndex = this.signup(user);
+      this.stateIndices.set(user, stateIndex);
+    }
+    return this.users;
+  };
+
+  signup = (user: Keypair): number => {
+    const timestamp = BigInt(Math.floor(Date.now() / 1000));
+    const stateIndex = this.maciState.signUp(user.pubKey, voiceCreditBalance, timestamp);
+    return stateIndex;
+  };
+
+  vote = (user: Keypair, stateIndex: number, voteOptionIndex: bigint, voteWeight: bigint, nonce: bigint): void => {
+    const { command, signature } = this.createCommand(user, stateIndex, voteOptionIndex, voteWeight, nonce);
+
+    const { message, encPubKey } = this.createMessage(command, signature, this.coordinatorKeypair);
+
+    this.poll.publishMessage(message, encPubKey);
+  };
+
+  createMessage = (
+    command: PCommand,
+    signature: Signature,
+    coordinatorKeypair: Keypair,
+  ): { message: Message; encPubKey: PubKey } => {
+    const ecdhKeypair = new Keypair();
+    const sharedKey = Keypair.genEcdhSharedKey(ecdhKeypair.privKey, coordinatorKeypair.pubKey);
+    const message = command.encrypt(signature, sharedKey);
+    return { message, encPubKey: ecdhKeypair.pubKey };
+  };
+
+  createCommand = (
+    user: Keypair,
+    stateIndex: number,
+    voteOptionIndex: bigint,
+    voteWeight: bigint,
+    nonce: bigint,
+  ): { command: PCommand; signature: Signature } => {
+    const command = new PCommand(
+      BigInt(stateIndex),
+      user.pubKey,
+      voteOptionIndex,
+      voteWeight,
+      nonce,
+      BigInt(this.pollId),
+    );
+
+    const signature = command.sign(user.privKey);
+
+    return { command, signature };
+  };
+
+  finalizePoll = (): void => {
+    this.poll.processMessages(this.pollId);
+    this.poll.tallyVotes();
+  };
+
+  getStateIndex = (user: Keypair): number => {
+    return this.stateIndices.get(user);
+  };
+}
+
+describe("Poll message processing and validation", function () {
+  // set timeout to 30 seconds
+  this.timeout(30000);
+
+  let testHarness: TestHarness;
+  let poll: Poll;
+
+  describe("Sanity checks", () => {
     beforeEach(async () => {
       testHarness = new TestHarness();
 
       poll = testHarness.poll;
     });
 
-    it("processMessage() should process a valid message", async () => {
+    it("processMessages() should process a valid message", async () => {
       const voteOptionIndex = BigInt(0);
       const voteWeight = BigInt(9);
       const nonce = BigInt(1);
@@ -16,11 +130,16 @@ describe("Message processing", () => {
       testHarness.vote(users[0], testHarness.getStateIndex(users[0]), voteOptionIndex, voteWeight, nonce);
       testHarness.finalizePoll();
 
-      expect(poll.messages.length).to.eq(1);
-      expect(poll.tallyResult[0]).to.eq(BigInt(9));
+      const messageLengthResult = poll.messages.length;
+      const expectedNumVotes = users.length;
+      expect(messageLengthResult).to.eq(expectedNumVotes);
+
+      const tallyResult = poll.tallyResult[0];
+      const expectedTallyResult = BigInt(9);
+      expect(tallyResult).to.eq(expectedTallyResult);
     });
 
-    it("processMessage() should not process messages twice", async () => {
+    it("processMessages() should not process messages twice", async () => {
       const voteOptionIndex = BigInt(0);
       const voteWeight = BigInt(9);
       const nonce = BigInt(1);
@@ -31,15 +150,20 @@ describe("Message processing", () => {
 
       expect(() => {
         poll.processMessages(testHarness.pollId);
-      }).to.throw();
+      }).to.throw(AssertionError, /No more messages to process/);
 
       poll.tallyVotes();
 
-      expect(poll.messages.length).to.eq(1);
-      expect(poll.tallyResult[0]).to.eq(BigInt(9));
+      const messageLengthResult = poll.messages.length;
+      const expectedNumVotes = users.length;
+      expect(messageLengthResult).to.eq(expectedNumVotes);
+
+      const tallyResult = poll.tallyResult[0];
+      const expectedTallyResult = BigInt(9);
+      expect(tallyResult).to.eq(expectedTallyResult);
     });
 
-    it("processMessage() should not process a message with an incorrect nonce", async () => {
+    it("processMessages() should not process a message with an incorrect nonce", async () => {
       const voteOptionIndex = BigInt(0);
       const voteWeight = BigInt(9);
 
@@ -69,7 +193,7 @@ describe("Message processing", () => {
     // the square of the vote weight. Since the maximum voice credit is 100 here,
     // the vote weight can only be a value between 1 and 10
     // (as these are the square roots of numbers up to 100).
-    it("processMessage() should not process a message with an incorrect vote weight", async () => {
+    it("processMessages() should not process a message with an incorrect vote weight", async () => {
       const voteOptionIndex = BigInt(0);
       const nonce = BigInt(1);
 
@@ -96,20 +220,17 @@ describe("Message processing", () => {
       expect(tallyResult).to.eq(expectedTallyResult);
     });
 
-    it("processMessage() should not process a message with an incorrect state tree index", async () => {
+    it("processMessages() should not process a message with an incorrect state tree index", async () => {
       const voteOptionIndex = BigInt(0);
       const nonce = BigInt(1);
+      const voteWeight = BigInt(9);
       const numVotes = 5;
 
       const users = testHarness.createUsers(5);
-      // generate a bunch of invalid votes with vote weights that are not between 1 and 10
-      let voteWeight: bigint;
-      for (let i = 0; i < users.length; i++) {
-        do {
-          voteWeight = BigInt(Math.floor(Math.random() * 100) - 50);
-        } while (BigInt(1) <= voteWeight && voteWeight <= BigInt(10));
 
-        testHarness.vote(users[i], testHarness.getStateIndex(users[i]), voteOptionIndex, voteWeight, nonce);
+      for (let i = 0; i < users.length; i++) {
+        // generate a bunch of invalid votes with incorrect state tree index
+        testHarness.vote(users[i], testHarness.getStateIndex(users[i]) + 1, voteOptionIndex, voteWeight, nonce);
       }
 
       testHarness.finalizePoll();
@@ -123,7 +244,7 @@ describe("Message processing", () => {
       expect(tallyResult).to.eq(expectedTallyResult);
     });
 
-    it("processMessage() should not process a message with an incorrect signature", async () => {
+    it("processMessages() should not process a message with an incorrect signature", async () => {
       const voteOptionIndex = BigInt(0);
       const voteWeight = BigInt(9);
       const nonce = BigInt(1);
@@ -166,7 +287,7 @@ describe("Message processing", () => {
       expect(tallyResult).to.eq(expectedTallyResult);
     });
 
-    it("processMessage() should not process a message with an invalid coordinator key", async () => {
+    it("processMessages() should not process a message with an invalid coordinator key", async () => {
       const voteOptionIndex = BigInt(0);
       const voteWeight = BigInt(9);
       const nonce = BigInt(1);