feat: check proof is valid when creating proof

README.md

@@ -278,8 +278,7 @@ const proofResult = await rln.verifyProof(epoch, message, proof) // true or fals
 ```
 
 A proof can be invalid in the following conditions:
-- The proof is not for you. You're using a different `rlnIdentifier`
-- The proof is not for the current epoch or message
+- Proof mismatches epoch, message, or rlnIdentifier
 - The snark proof itself is invalid
 
 ### Saving a proof

src/cache.ts

@@ -43,6 +43,12 @@ export type EvaluatedProof = {
 
 export interface ICache {
   addProof(proof: CachedProof): EvaluatedProof
+  /**
+   * Check the proof if it is either valid, seen, or breaching.
+   * Does not add the proof to the cache to avoid side effects.
+   * @param proof CachedProof
+   */
+  checkProof(proof: CachedProof): EvaluatedProof
 }
 
 const DEFAULT_CACHE_SIZE = 100
@@ -77,10 +83,28 @@ export class MemoryCache implements ICache {
     // Since `BigInt` can't be used as key, use String instead
     const epochString = String(proof.epoch)
     const nullifier = String(proof.nullifier)
+    // Check if the proof status
+    const resCheckProof = this.checkProof(proof)
+    // Only add the proof to the cache automatically if it's not seen before.
+    if (resCheckProof.status === Status.ADDED || resCheckProof.status === Status.BREACH) {
+      // Add proof to cache
+      this.cache[epochString][nullifier].push(proof)
+    }
+    return resCheckProof
+  }
 
-    this.evaluateEpoch(epochString)
+  /**
+   * Check the proof if it is either valid, seen, or breaching.
+   * Does not add the proof to the cache to avoid side effects.
+   * @param proof CachedProof
+   */
+  checkProof(proof: CachedProof): EvaluatedProof {
+    const epochString = String(proof.epoch)
+    const nullifier = String(proof.nullifier)
+    this.shiftEpochs(epochString)
     // If nullifier doesn't exist for this epoch, create an empty array
     this.cache[epochString][nullifier] = this.cache[epochString][nullifier] || []
+    const proofs = this.cache[epochString][nullifier]
 
     // Check if the proof already exists. It's O(n) but it's not a big deal since n is exactly the
     // rate limit and it's usually small.
@@ -92,33 +116,26 @@ export class MemoryCache implements ICache {
         BigInt(proof1.nullifier) === BigInt(proof2.nullifier)
       )
     }
-    const sameProofs = this.cache[epochString][nullifier].filter(p => isSameProof(p, proof))
-    if (sameProofs.length > 0) {
-      return { status: Status.SEEN, msg: 'Proof already exists' }
-    }
-
-    // Add proof to cache
-    this.cache[epochString][nullifier].push(proof)
-
-    // Check if there is more than 1 proof for this nullifier for this epoch
-    return this.evaluateNullifierAtEpoch(epochString, nullifier)
-  }
-
-  private evaluateNullifierAtEpoch(epoch: string, nullifier: string): EvaluatedProof {
-    const proofs = this.cache[epoch][nullifier]
-    if (proofs.length > 1) {
-      // If there is more than 1 proof, return breach and secret
-      const [x1, y1] = [BigInt(proofs[0].x), BigInt(proofs[0].y)]
-      const [x2, y2] = [BigInt(proofs[1].x), BigInt(proofs[1].y)]
-      const secret = shamirRecovery(x1, x2, y1, y2)
-      return { status: Status.BREACH, nullifier: nullifier, secret: secret, msg: 'Rate limit breach, secret attached' }
-    } else {
-      // If there is only 1 proof, return added
+    // OK
+    if (proofs.length === 0) {
       return { status: Status.ADDED, nullifier: nullifier, msg: 'Proof added to cache' }
+    // Exists proof with same epoch and nullifier. Possible breach or duplicate proof
+    } else {
+      const sameProofs = this.cache[epochString][nullifier].filter(p => isSameProof(p, proof))
+      if (sameProofs.length > 0) {
+        return { status: Status.SEEN, msg: 'Proof already exists' }
+      } else {
+        const otherProof = proofs[0]
+        // Breach. Return secret
+        const [x1, y1] = [BigInt(proof.x), BigInt(proof.y)]
+        const [x2, y2] = [BigInt(otherProof.x), BigInt(otherProof.y)]
+        const secret = shamirRecovery(x1, x2, y1, y2)
+        return { status: Status.BREACH, nullifier: nullifier, secret: secret, msg: 'Rate limit breach, secret attached' }
+      }
     }
   }
 
-  private evaluateEpoch(epoch: string) {
+  private shiftEpochs(epoch: string) {
     if (this.cache[epoch]) {
       // If epoch already exists, return
       return

src/message-id-counter.ts

@@ -1,7 +1,16 @@
+/**
+ * Always return **next** the current counter value and increment the counter.
+ * @param epoch epoch of the message
+ * @throws Error if the counter exceeds the message limit
+ */
 export interface IMessageIDCounter {
   messageLimit: bigint;
+  /**
+   * Return the current counter value and increment the counter.
+   *
+   * @param epoch
+   */
   getMessageIDAndIncrement(epoch: bigint): Promise<bigint>
-  peekNextMessageID(epoch: bigint): Promise<bigint>
 }
 
 type EpochMap = {
@@ -22,14 +31,6 @@ export class MemoryMessageIDCounter implements IMessageIDCounter {
     return this._messageLimit
   }
 
-  async peekNextMessageID(epoch: bigint): Promise<bigint> {
-    const epochStr = epoch.toString()
-    if (this.epochToMessageID[epochStr] === undefined) {
-      return BigInt(0)
-    }
-    return this.epochToMessageID[epochStr]
-  }
-
   async getMessageIDAndIncrement(epoch: bigint): Promise<bigint> {
     const epochStr = epoch.toString()
     // Initialize the message id counter if it doesn't exist

src/rln.ts

@@ -381,26 +381,33 @@ export class RLN implements IRLN {
       )
     }
     const merkleProof = await this.registry.generateMerkleProof(this.identityCommitment)
-    const messageID = await this.messageIDCounter.getMessageIDAndIncrement(epoch)
+    // NOTE: get the message id and increment the counter.
+    // Even if the message is not sent, the counter is still incremented.
+    // It's intended to avoid any possibly for user to reuse the same message id.
+    const messageId = await this.messageIDCounter.getMessageIDAndIncrement(epoch)
     const userMessageLimit = await this.registry.getMessageLimit(this.identityCommitment)
     const proof = await this.prover.generateProof({
       rlnIdentifier: this.rlnIdentifier,
       identitySecret: this.identitySecret,
       userMessageLimit: userMessageLimit,
-      messageId: messageID,
+      messageId,
       merkleProof,
       x: calculateSignalHash(message),
       epoch,
     })
-    // Double check if the proof will spam or not.
+    // Double check if the proof will spam or not using the cache.
     // Even if messageIDCounter is used, it is possible that the user restart and the counter is reset.
-    // In this case, the user can still spam.
-    const res = await this.saveProof(proof)
+    const res = await this.checkProof(proof)
     if (res.status === Status.SEEN) {
       throw new Error('Proof has been generated before')
     } else if (res.status === Status.BREACH) {
       throw new Error('Proof will spam')
     } else if (res.status === Status.ADDED) {
+      const resSaveProof = await this.saveProof(proof)
+      if (resSaveProof.status !== res.status) {
+        // Sanity check
+        throw new Error('Status of save proof and check proof mismatch')
+      }
       return proof
     } else {
       // Sanity check
@@ -453,11 +460,14 @@ export class RLN implements IRLN {
    * or 'seen' if the proof has been saved before, else 'breach' if the proof is a spam.
    */
   async saveProof(proof: RLNFullProof): Promise<EvaluatedProof> {
-    if (this.verifier === undefined) {
-      throw new Error('Verifier is not initialized')
-    }
     const { snarkProof, epoch } = proof
     const { x, y, nullifier } = snarkProof.publicSignals
     return this.cache.addProof({ x, y, nullifier, epoch })
   }
+
+  private async checkProof(proof: RLNFullProof): Promise<EvaluatedProof> {
+    const { snarkProof, epoch } = proof
+    const { x, y, nullifier } = snarkProof.publicSignals
+    return this.cache.checkProof({ x, y, nullifier, epoch })
+  }
 }

tests/cache.test.ts

@@ -58,6 +58,9 @@ describe("MemoryCache", () => {
   test("should check proof 4", () => {
     const result4 = cache.addProof(proof4)
     expect(result4.status).toBe(Status.ADDED)
+    // two epochs are added to the cache now
+    expect(Object.keys(cache.cache).length
+    ).toBe(2)
   })
 
   test("should fail for proof 1 (duplicate proof)", () => {

tests/message-id-counter.test.ts

@@ -1,15 +1,8 @@
-// export interface IMessageIDCounter {
-//     messageLimit: bigint;
-//     getMessageIDAndIncrement(epoch: bigint): Promise<bigint>
-//     peekNextMessageID(epoch: bigint): Promise<bigint>
-//   }
-
-import { MemoryMessageIDCounter } from "../src/message-id-counter";
-
+import { FakeMessageIDCounter } from "./utils";
 
 describe('MessageIDCounter', () => {
     const messageLimit = BigInt(1)
-    const messageIDCounter = new MemoryMessageIDCounter(messageLimit)
+    const messageIDCounter = new FakeMessageIDCounter(messageLimit)
 
     it('should return correct message limit', async () => {
         expect(messageIDCounter.messageLimit).toEqual(messageLimit)

tests/rln.test.ts

@@ -1,21 +1,11 @@
 import { RLN, RLNFullProof } from "../src";
 import { ICache, MemoryCache, Status } from "../src/cache";
 import { rlnParams, withdrawParams } from "./configs";
-import { MemoryMessageIDCounter } from "../src/message-id-counter";
 import { ethers } from "ethers";
 import { setupTestingContracts } from "./factories";
-import { fieldFactory } from "./utils";
+import { FakeMessageIDCounter, fieldFactory } from "./utils";
 import { MemoryRLNRegistry } from "../src/registry";
 
-class FakeMessageIDCounter extends MemoryMessageIDCounter {
-    reset(epoch: bigint) {
-        const epochStr = epoch.toString()
-        if (this.epochToMessageID[epochStr] === undefined) {
-            return;
-        }
-        this.epochToMessageID[epochStr] = BigInt(0)
-    }
-}
 
 describe("RLN", function () {
     describe("constructor params", function () {
@@ -58,9 +48,6 @@ describe("RLN", function () {
             await expect(async () => {
                 await rln.verifyProof(randomEpoch, randomMessage, mockProof)
             }).rejects.toThrow("Verifier is not initialized");
-            await expect(async () => {
-                await rln.saveProof(mockProof)
-            }).rejects.toThrow("Verifier is not initialized");
         });
     });
 
@@ -108,9 +95,6 @@ describe("RLN", function () {
             await expect(async () => {
                 await rln.verifyProof(randomEpoch, randomMessage, mockProof)
             }).rejects.toThrow("Verifier is not initialized");
-            await expect(async () => {
-                await rln.saveProof(mockProof)
-            }).rejects.toThrow("Verifier is not initialized");
         });
 
     });
@@ -123,6 +107,7 @@ describe("RLN", function () {
         const epoch1 = BigInt(1);
         const message0 = "abc";
         const message1 = "abcd";
+        const message2 = "abcde";
 
         let deployed;
         let waitUntilFreezePeriodPassed: () => Promise<void>
@@ -223,6 +208,19 @@ describe("RLN", function () {
             expect(allRateCommitments[0]).toBe(await rlnA0.getRateCommitment());
         });
 
+        it("should be able to set message id counter", async function () {
+            // Test: set a new message id counter with zero message limit
+            // I.e. no message can be sent
+            const zeroMessageLimit = BigInt(0)
+            const newMessageIDCounter = new FakeMessageIDCounter(zeroMessageLimit)
+            await rlnA0.setMessageIDCounter(newMessageIDCounter)
+            await expect(async () => {
+                await rlnA0.createProof(epoch0, message1)
+            }).rejects.toThrow(`Message ID counter exceeded message limit ${zeroMessageLimit}`);
+            // Change it back to make other tests work
+            await rlnA0.setMessageIDCounter(messageIDCounterA0)
+        })
+
         it("should be able to create proof", async function () {
             const messageIDBefore = await messageIDCounterA0.peekNextMessageID(epoch0);
             proofA00 = await rlnA0.createProof(epoch0, message0);
@@ -288,31 +286,34 @@ describe("RLN", function () {
             // messageLimitA1 is 1, so A1 can only create 1 proof per epoch
             // Test: can save the first proof
             proofA10 = await rlnA1.createProof(epoch0, message0);
-            // Test: fails when saving duplicate proof since it has been saved in createProof
+            // Test: status should be 'seen' when saving duplicate proof since it has been saved in createProof
             const resA10Again = await rlnA1.saveProof(proofA10);
             expect(resA10Again.status).toBe(Status.SEEN);
 
-            // Reset messageIDCounterA1 at epoch0 to force it create a proof
-            // when it already exceeds `messageLimitA1`
-            messageIDCounterA1.reset(epoch0);
+            // Reset messageIDCounterA1 at epoch0 to bypass the message id counter and
+            // let it create a proof when it already exceeds `messageLimitA1`.
+            await rlnA1.setMessageIDCounter(new FakeMessageIDCounter(BigInt(messageLimitA1)));
+
             // Test: even messageIDCounter is reset, there is another guard `cache`
             // to prevent creating more than `messageLimitA1` proofs
             await expect(async () => {
-                await rlnA1.createProof(epoch0, message0);
+                await rlnA1.createProof(epoch0, message1);
             }).rejects.toThrow("Proof will spam");
 
-            // Reset cache too
+            // Reset cache too, to allow rln createProof
             cacheA1.cache[epoch0.toString()] = {}
-
+            // Reset messageIDCounterA1 again to bypass the message id counter since it increments
+            // the message id counter when `createProof` even if it fails.
+            await rlnA1.setMessageIDCounter(new FakeMessageIDCounter(BigInt(messageLimitA1)));
             // Test: number of proofs per epoch exceeds `messageLimitA1`, breach/ slashed when `saveProof`
             proofA11 = await rlnA1.createProof(epoch0, message1);
-            const resA11 = await rlnA1.saveProof(proofA11);
-            expect(resA11.status).toBe(Status.BREACH);
-            if (resA11.secret === undefined) {
+            const resA10AgainAgain = await rlnA1.saveProof(proofA10);
+            expect(resA10AgainAgain.status).toBe(Status.BREACH);
+            if (resA10AgainAgain.secret === undefined) {
                 throw new Error("secret should not be undefined")
             }
             // Test: but A1 cannot slash itself
-            const secret = resA11.secret;
+            const secret = resA10AgainAgain.secret;
             await expect(async () => {
                 await rlnA1.slash(secret)
             }).rejects.toThrow('execution reverted: "RLN, slash: self-slashing is prohibited"');
@@ -321,7 +322,7 @@ describe("RLN", function () {
             await rlnA1.createProof(epoch1, message1);
         });
 
-        it("should reveal its secret by others and get slashed", async function () {
+        it("should reveal its secret and get slashed by others", async function () {
             // Test: A0 is up-to-date and receives more than `messageLimitA1` proofs,
             // so A1's secret is breached by A0
             const resA10 = await rlnA0.saveProof(proofA10);
@@ -344,6 +345,5 @@ describe("RLN", function () {
             // Test: verifyProof fails since proofA10.rlnIdentifier mismatches rlnB's rlnIdentifier
             expect(await rlnB.verifyProof(epoch0, message0, proofA10)).toBe(false);
         });
-        // TODO: Add tests to set messageIDCounter
     });
 });

tests/utils.ts

@@ -5,6 +5,7 @@ import { BigNumber } from '@ethersproject/bignumber'
 import { keccak256 } from '@ethersproject/keccak256'
 
 import { Fq } from "../src/common"
+import { MemoryMessageIDCounter } from "../src/message-id-counter";
 
 export function fieldFactory(excludes?: bigint[], trials: number = 100): bigint {
     if (excludes) {
@@ -20,6 +21,16 @@ export function fieldFactory(excludes?: bigint[], trials: number = 100): bigint
     }
 }
 
+export class FakeMessageIDCounter extends MemoryMessageIDCounter {
+    async peekNextMessageID(epoch: bigint): Promise<bigint> {
+        const epochStr = epoch.toString()
+        if (this.epochToMessageID[epochStr] === undefined) {
+          return BigInt(0)
+        }
+        return this.epochToMessageID[epochStr]
+    }
+}
+
 
 function calculateZeroValue(id: bigint): bigint {
     const hexStr = BigNumber.from(id).toTwos(256).toHexString()