Merge pull request #80 from Rate-Limiting-Nullifier/feat/check-proof-…

…when-create-proof

Check proof when create proof

README.md

@@ -278,19 +278,18 @@ 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
+- Proof mismatches epoch, message, or rlnIdentifier
 - The snark proof itself is invalid
 
 ### Saving a proof
 User should save all proofs they receive to detect spams. You can save a proof by calling `rln.saveProof()`. The return value is an object indicating the status of the proof.
 
 ```typescript
 const result = await rln.saveProof(proof)
-// status can be "added", "breach", or "invalid".
-// - "added" means the proof is successfully added to the cache
-// - "breach" means the added proof breaches the rate limit, the result will be `breach`, in which case the `secret` will be recovered and is accessible by accessing `result.secret`
-// - "invalid" means the proof is invalid, either it's received before or verification fails
+// status can be VALID, DUPLICATE, BREACH.
+// - VALID means the proof is successfully added to the cache
+// - DUPLICATE means the proof is already saved before
+// - BREACH means the added proof breaches the rate limit, in which case the `secret` is recovered and is accessible by `result.secret`
 const status = result.status
 // if status is "breach", you can get the secret by
 const secret = result.secret

src/cache.ts

@@ -29,9 +29,9 @@ type CacheMap = {
 
 
 export enum Status {
-  ADDED = 'added',
-  BREACH = 'breach',
-  INVALID = 'invalid',
+  VALID,
+  DUPLICATE,
+  BREACH,
 }
 
 export type EvaluatedProof = {
@@ -42,10 +42,22 @@ export type EvaluatedProof = {
 }
 
 export interface ICache {
+  /**
+   * Add a proof to the cache and automatically evaluate it for rate limit breaches.
+   * @param proof CachedProof
+   * @returns an object with the status of the proof and the nullifier and secret if the proof is a breach
+   */
   addProof(proof: CachedProof): EvaluatedProof
+  /**
+   * Check the proof if it is either valid, duplicate, or breaching.
+   * Does not add the proof to the cache to avoid side effects.
+   * @param proof CachedProof
+   * @returns an object with the status of the proof and the nullifier and secret if the proof is a breach
+   */
+  checkProof(proof: CachedProof): EvaluatedProof
 }
 
-const DEFAULT_CACHE_SIZE = 100
+export const DEFAULT_CACHE_SIZE = 100
 /**
  * Cache for storing proofs and automatically evaluating them for rate limit breaches
  * in the memory.
@@ -68,7 +80,7 @@ export class MemoryCache implements ICache {
   }
 
   /**
-   *  Adds a proof to the cache
+   * Add a proof to the cache and automatically evaluate it for rate limit breaches.
    * @param proof CachedProof
    * @returns an object with the status of the proof and the nullifier and secret if the proof is a breach
    */
@@ -77,10 +89,29 @@ 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.VALID || 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, duplicate, or breaching.
+   * Does not add the proof to the cache to avoid side effects.
+   * @param proof CachedProof
+   * @returns an object with the status of the proof and the nullifier and secret if the proof is a breach
+   */
+  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 +123,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.INVALID, 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' }
+    // OK
+    if (proofs.length === 0) {
+      return { status: Status.VALID, nullifier: nullifier, msg: 'Proof added to cache' }
+    // Exists proof with same epoch and nullifier. Possible breach or duplicate proof
     } else {
-      // If there is only 1 proof, return added
-      return { status: Status.ADDED, nullifier: nullifier, msg: 'Proof added to cache' }
+      const sameProofs = this.cache[epochString][nullifier].filter(p => isSameProof(p, proof))
+      if (sameProofs.length > 0) {
+        return { status: Status.DUPLICATE, 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

@@ -2,7 +2,7 @@ import { Identity } from '@semaphore-protocol/identity'
 import { VerificationKey } from './types'
 import { calculateIdentityCommitment, calculateSignalHash } from './common'
 import { IRLNRegistry, ContractRLNRegistry } from './registry'
-import { MemoryCache, EvaluatedProof, ICache } from './cache'
+import { MemoryCache, EvaluatedProof, ICache, Status } from './cache'
 import { IMessageIDCounter, MemoryMessageIDCounter } from './message-id-counter'
 import { RLNFullProof, RLNProver, RLNVerifier } from './circuit-wrapper'
 import { ethers } from 'ethers'
@@ -44,9 +44,10 @@ export interface IRLN {
    */
   verifyProof(epoch: bigint, message: string, proof: RLNFullProof): Promise<boolean>
   /**
-   * Verify a RLNFullProof, save it to a cache, and detect if the proof is a spam.
-   * @param proof the RLNFullProof to be verified
-   * @returns EvaluatedProof the result
+   * Save a proof to the cache and check if it's a spam.
+   * @param proof the RLNFullProof to save and detect spam
+   * @returns result of the check. It could be VALID if the proof hasn't been seen,
+   * or DUPLICATE if the proof has been seen before, else BREACH means it could be spam.
    */
   saveProof(proof: RLNFullProof): Promise<EvaluatedProof>
 }
@@ -380,17 +381,38 @@ 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)
-    return this.prover.generateProof({
+    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 using the cache.
+    // Even if messageIDCounter is used, it is possible that the user restart and the counter is reset.
+    const res = await this.checkProof(proof)
+    if (res.status === Status.DUPLICATE) {
+      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.VALID) {
+      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
+      throw new Error('Unknown status')
+    }
   }
 
   /**
@@ -432,20 +454,20 @@ export class RLN implements IRLN {
   }
 
   /**
-   * Save the proof to the cache. If the proof is invalid, an error is thrown.
-   * Else, the proof is saved to the cache and is checked if it's a spam.
-   * @param proof the RLNFullProof to verify and save
-   * @returns the EvaluatedProof if the proof is valid, an error is thrown otherwise
+   * Save a proof to the cache and check if it's a spam.
+   * @param proof the RLNFullProof to save and detect spam
+   * @returns result of the check. `status` could be status.VALID if the proof is not a spam or invalid.
+   * Otherwise, it will be status.DUPLICATE or status.BREACH.
    */
   async saveProof(proof: RLNFullProof): Promise<EvaluatedProof> {
-    if (this.verifier === undefined) {
-      throw new Error('Verifier is not initialized')
-    }
-    if (!await this.verifier.verifyProof(this.rlnIdentifier, proof)) {
-      throw new Error('Invalid proof')
-    }
     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

@@ -1,5 +1,5 @@
 import { MemoryCache } from "../src"
-import { CachedProof, Status } from '../src/cache'
+import { CachedProof, DEFAULT_CACHE_SIZE, Status } from '../src/cache'
 import { fieldFactory } from "./utils"
 
 describe("MemoryCache", () => {
@@ -32,37 +32,51 @@ describe("MemoryCache", () => {
   })
 
   test("should have a cache length of 100", () => {
-    expect(cache.cacheLength).toBe(100)
+    expect(cache.cacheLength).toBe(DEFAULT_CACHE_SIZE)
   })
 
   test("should successfully add proof", () => {
+    const resultCheckProof = cache.checkProof(proof1)
+    expect(resultCheckProof.status).toBe(Status.VALID)
     const result1 = cache.addProof(proof1)
-    expect(result1.status).toBe(Status.ADDED)
+    expect(result1.status).toBe(resultCheckProof.status)
     expect(Object.keys(cache.cache).length
     ).toBe(1)
   })
 
   test("should detect breach and return secret", () => {
+    const resultCheckProof = cache.checkProof(proof2)
+    expect(resultCheckProof.status).toBe(Status.BREACH)
+    expect(resultCheckProof.secret).toBeGreaterThan(0)
     const result2 = cache.addProof(proof2)
-    expect(result2.status).toBe(Status.BREACH)
+    expect(result2.status).toBe(resultCheckProof.status)
     expect(result2.secret).toBeGreaterThan(0)
     expect(Object.keys(cache.cache).length
     ).toBe(1)
   })
 
   test("should check proof 3", () => {
+    const resultCheckProof = cache.checkProof(proof3)
+    expect(resultCheckProof.status).toBe(Status.VALID)
     const result3 = cache.addProof(proof3)
-    expect(result3.status).toBe(Status.ADDED)
+    expect(result3.status).toBe(resultCheckProof.status)
   })
 
   test("should check proof 4", () => {
+    const resultCheckProof = cache.checkProof(proof4)
+    expect(resultCheckProof.status).toBe(Status.VALID)
     const result4 = cache.addProof(proof4)
-    expect(result4.status).toBe(Status.ADDED)
+    expect(result4.status).toBe(resultCheckProof.status)
+    // two epochs are added to the cache now
+    expect(Object.keys(cache.cache).length
+    ).toBe(2)
   })
 
   test("should fail for proof 1 (duplicate proof)", () => {
     // Proof 1 is already in the cache
+    const resultCheckProof = cache.checkProof(proof1)
+    expect(resultCheckProof.status).toBe(Status.DUPLICATE)
     const result1 = cache.addProof(proof1)
-    expect(result1.status).toBe(Status.INVALID)
+    expect(result1.status).toBe(Status.DUPLICATE)
   });
 })

tests/circuit-wrapper.test.ts

@@ -21,7 +21,7 @@ describe("RLN", function () {
     const epoch = fieldFactory()
     const treeDepth = DEFAULT_REGISTRY_TREE_DEPTH
 
-    it("should generate valid proof", async function () {
+    test("should generate valid proof", async function () {
         const merkleProof = generateMerkleProof(rlnIdentifier, leaves, treeDepth, 0)
         const proof = await rlnProver.generateProof({
             rlnIdentifier,
@@ -40,7 +40,7 @@ describe("Withdraw", function () {
     const withdrawProver = new WithdrawProver(withdrawParams.wasmFilePath, withdrawParams.finalZkeyPath)
     const withdrawVerifier = new WithdrawVerifier(withdrawParams.verificationKey)
 
-    it("should generate valid proof", async function () {
+    test("should generate valid proof", async function () {
         const identitySecret = fieldFactory()
         const address = fieldFactory()
         const proof = await withdrawProver.generateProof({identitySecret, address})