feat(rln-relay): multiple acceptable roots (#1177)

* feat(rln-relay): multiple acceptable roots

* fix(rln-relay): make sure onchain handler uses correct proc

* fix(rln-relay): typo

* style(rln-relay): convert const to pascalcase

* chore(rln-relay): address pr review

* fix(rln-relay): add require to test

* fix(rln-relay): add proc desc

* fix(rln-relay): ensure that group id commitments were added correctly

* fix(rln-relay): potential mem leak

* style(rln-relay): comments

* fix(rln-relay): magic number ambiguity

* chore(rln-relay): comment

Co-authored-by: Sanaz Taheri Boshrooyeh <35961250+staheri14@users.noreply.github.com>

Co-authored-by: Sanaz Taheri Boshrooyeh <35961250+staheri14@users.noreply.github.com>

tests/v2/test_waku_rln_relay.nim

@@ -2,7 +2,7 @@
 {.used.}
 
 import
-  std/options, sequtils, times,
+  std/options, sequtils, times, deques,
   testutils/unittests, chronos, chronicles, stint,
   stew/byteutils, stew/shims/net as stewNet,
   libp2p/crypto/crypto,
@@ -639,34 +639,39 @@ suite "Waku rln relay":
     check:
       verified.value() == false
 
-  test "invalidate messages with a valid, but stale root":
+  test "validate roots which are part of the acceptable window":
     # Setup: 
-    # This step consists of creating the rln instance,
+    # This step consists of creating the rln instance and waku-rln-relay,
     # Inserting members, and creating a valid proof with the merkle root
+    # create an RLN instance
     var rlnInstance = createRLNInstance()
     require:
-      rlnInstance.isOk() == true
+      rlnInstance.isOk()
     var rln = rlnInstance.value
 
+    let rlnRelay = WakuRLNRelay(rlnInstance:rln)
+
     let
       # create a membership key pair
-      memKeys = membershipKeyGen(rln).get()
+      memKeys = membershipKeyGen(rlnRelay.rlnInstance).get()
-      # peer's index in the Merkle Tree
+      # peer's index in the Merkle Tree. 
       index = 5
 
+    let membershipCount = AcceptableRootWindowSize + 5
+
     # Create a Merkle tree with random members
-    for i in 0..10:
+    for i in 0..membershipCount:
-      var memberIsAdded: bool = false
+      var memberIsAdded: RlnRelayResult[void]
       if (i == index):
         # insert the current peer's pk
-        memberIsAdded = rln.insertMember(memKeys.idCommitment)
+        memberIsAdded = rlnRelay.insertMember(memKeys.idCommitment)
       else:
         # create a new key pair
-        let memberKeys = rln.membershipKeyGen()
+        let memberKeys = rlnRelay.rlnInstance.membershipKeyGen()
-        memberIsAdded = rln.insertMember(memberKeys.get().idCommitment)
+        memberIsAdded = rlnRelay.insertMember(memberKeys.get().idCommitment)
-      # check the member is added
+      # require that the member is added
-      check:
+      require:
-        memberIsAdded
+        memberIsAdded.isOk()
 
     # Given: 
     # This step includes constructing a valid message with the latest merkle root
@@ -678,7 +683,7 @@ suite "Waku rln relay":
     debug "epoch in bytes", epochHex = epoch.toHex()
 
     # generate proof
-    let validProofRes = rln.proofGen(data = messageBytes,
+    let validProofRes = rlnRelay.rlnInstance.proofGen(data = messageBytes,
                                     memKeys = memKeys,
                                     memIndex = MembershipIndex(index),
                                     epoch = epoch)
@@ -687,38 +692,117 @@ suite "Waku rln relay":
     let validProof = validProofRes.value
 
     # validate the root (should be true)
-    let verified = rln.validateRoot(validProof.merkleRoot)
+    let verified = rlnRelay.validateRoot(validProof.merkleRoot)
 
     require:
-      verified.isOk()
+      verified == true
-      verified.value() == true
 
     # When: 
-    # This test depends on the local merkle tree root being different than a
+    # This test depends on the local merkle tree root being part of a
-    # new message with an older/different root
+    # acceptable set of roots, which is denoted by AcceptableRootWindowSize
-    # This can be simulated by removing a member, which changes the root of the tree
+    # The following action is equivalent to a member being removed upon listening to the events emitted by the contract
-    # Which is equivalent to a member being removed upon listening to the events emitted by the contract
-    # Progress the local tree by removing a member
-    discard rln.removeMember(MembershipIndex(0))
 
-    # Ensure the local tree root has changed
+    # Progress the local tree by removing members
-    let currentMerkleRoot = rln.getMerkleRoot()
+    for i in 0..AcceptableRootWindowSize - 2:
+      discard rlnRelay.removeMember(MembershipIndex(i))
+      # Ensure the local tree root has changed
+      let currentMerkleRoot = rlnRelay.rlnInstance.getMerkleRoot()
 
-    require:
+      require:
-      currentMerkleRoot.isOk()
+        currentMerkleRoot.isOk()
-      currentMerkleRoot.value() != validProof.merkleRoot
+        currentMerkleRoot.value() != validProof.merkleRoot
 
     # Then: 
-    # we try to verify a proof against this new merkle tree,
+    # we try to verify a root against this window,
-    # which should return false
+    # which should return true
-    # Try to send a message constructed with an older root
+    let olderRootVerified = rlnRelay.validateRoot(validProof.merkleRoot)
-    let olderRootVerified = rln.validateRoot(validProof.merkleRoot)
-
-    require:
-      olderRootVerified.isOk()
 
     check:
-      olderRootVerified.value() == false
+      olderRootVerified == true
+
+  test "invalidate roots which are not part of the acceptable window":
+    # Setup: 
+    # This step consists of creating the rln instance and waku-rln-relay,
+    # Inserting members, and creating a valid proof with the merkle root
+    
+    require:
+      AcceptableRootWindowSize < 10
+    
+    # create an RLN instance
+    var rlnInstance = createRLNInstance()
+    require:
+      rlnInstance.isOk()
+    var rln = rlnInstance.value
+
+    let rlnRelay = WakuRLNRelay(rlnInstance:rln)
+
+    let
+      # create a membership key pair
+      memKeys = membershipKeyGen(rlnRelay.rlnInstance).get()
+      # peer's index in the Merkle Tree. 
+      index = 6
+
+    let membershipCount = AcceptableRootWindowSize + 5 
+
+    # Create a Merkle tree with random members
+    for i in 0..membershipCount:
+      var memberIsAdded: RlnRelayResult[void]
+      if (i == index):
+        # insert the current peer's pk
+        memberIsAdded = rlnRelay.insertMember(memKeys.idCommitment)
+      else:
+        # create a new key pair
+        let memberKeys = rlnRelay.rlnInstance.membershipKeyGen()
+        memberIsAdded = rlnRelay.insertMember(memberKeys.get().idCommitment)
+      # require that the member is added
+      require:
+        memberIsAdded.isOk()
+
+    # Given: 
+    # This step includes constructing a valid message with the latest merkle root
+    # prepare the message
+    let messageBytes = "Hello".toBytes()
+
+    # prepare the epoch
+    var epoch: Epoch
+    debug "epoch in bytes", epochHex = epoch.toHex()
+
+    # generate proof
+    let validProofRes = rlnRelay.rlnInstance.proofGen(data = messageBytes,
+                                    memKeys = memKeys,
+                                    memIndex = MembershipIndex(index),
+                                    epoch = epoch)
+    require:
+      validProofRes.isOk()
+    let validProof = validProofRes.value
+
+    # validate the root (should be true)
+    let verified = rlnRelay.validateRoot(validProof.merkleRoot)
+
+    require:
+      verified == true
+
+    # When: 
+    # This test depends on the local merkle tree root being part of a
+    # acceptable set of roots, which is denoted by AcceptableRootWindowSize
+    # The following action is equivalent to a member being removed upon listening to the events emitted by the contract
+
+    # Progress the local tree by removing members
+    for i in 0..AcceptableRootWindowSize:
+      discard rlnRelay.removeMember(MembershipIndex(i))
+      # Ensure the local tree root has changed
+      let currentMerkleRoot = rlnRelay.rlnInstance.getMerkleRoot()
+      require:
+        currentMerkleRoot.isOk()
+        currentMerkleRoot.value() != validProof.merkleRoot
+
+    # Then: 
+    # we try to verify a proof against this window,
+    # which should return false
+    let olderRootVerified = rlnRelay.validateRoot(validProof.merkleRoot)
+
+    check:
+      olderRootVerified == false
 
   test "toEpoch and fromEpoch consistency check":
     # check edge cases
@@ -823,13 +907,15 @@ suite "Waku rln relay":
     doAssert(rlnInstance.isOk)
     var rln = rlnInstance.value
 
-    # add members
-    discard rln.addAll(groupIDCommitments)
-
     let
       wakuRlnRelay = WakuRLNRelay(membershipIndex: index,
           membershipKeyPair: groupKeyPairs[index], rlnInstance: rln)
 
+    # add members
+    let commitmentAddRes =  wakuRlnRelay.addAll(groupIDCommitments)
+    require:
+      commitmentAddRes.isOk()
+
     # get the current epoch time
     let time = epochTime()
 

waku/v2/protocol/waku_rln_relay/waku_rln_relay_types.nim

@@ -1,7 +1,7 @@
 {.push raises: [Defect].}
 
 import
-  std/tables,
+  std/[tables, deques],
   options, chronos, stint,
   web3,
   eth/keys,
@@ -9,23 +9,27 @@ import
   stew/arrayops,
   ../../utils/protobuf
 
+const AcceptableRootWindowSize* = 5
+
 when defined(rln) or (not defined(rln) and not defined(rlnzerokit)):
   ## Bn256 and RLN are Nim wrappers for the data types used in
   ## the rln library https://github.com/kilic/rln/blob/3bbec368a4adc68cd5f9bfae80b17e1bbb4ef373/src/ffi.rs
   type Bn256* = pointer
   type RLN*[E] = pointer
+  type RLNResult* = Result[RLN[Bn256], string]
 
 when defined(rlnzerokit):
   ## RLN is a Nim wrapper for the data types used in zerokit RLN
   type RLN* {.incompleteStruct.} = object
+  type RLNResult* = Result[ptr RLN, string]
+
+type RlnRelayResult*[T] = Result[T, string]
 
 type
   # identity key as defined in https://hackmd.io/tMTLMYmTR5eynw2lwK9n1w?view#Membership
   IDKey* = array[32, byte]
   # hash of identity key as defined ed in https://hackmd.io/tMTLMYmTR5eynw2lwK9n1w?view#Membership
   IDCommitment* = array[32, byte]
-
-type
   MerkleNode* = array[32, byte] # Each node of the Merkle tee is a Poseidon hash which is a 32 byte value
   Nullifier* = array[32, byte]
   Epoch* = array[32, byte]
@@ -121,6 +125,7 @@ when defined(rln) or (not defined(rln) and not defined(rlnzerokit)):
     # the log of nullifiers and Shamir shares of the past messages grouped per epoch
     nullifierLog*: Table[Epoch, seq[ProofMetadata]]
     lastEpoch*: Epoch # the epoch of the last published rln message
+    validMerkleRoots*: Deque[MerkleNode] # An array of valid merkle roots, which are updated in a FIFO fashion
 
 when defined(rlnzerokit):
   type WakuRLNRelay* = ref object
@@ -143,6 +148,8 @@ when defined(rlnzerokit):
     contentTopic*: string
     # the log of nullifiers and Shamir shares of the past messages grouped per epoch
     nullifierLog*: Table[Epoch, seq[ProofMetadata]]
+    validMerkleRoots*: Deque[MerkleNode] # An array of valid merkle roots, which are updated in a FIFO fashion
+
 
 type MessageValidationResult* {.pure.} = enum
   Valid, Invalid, Spam

waku/v2/protocol/waku_rln_relay/waku_rln_relay_utils.nim

@@ -1,9 +1,7 @@
 {.push raises: [Defect].}
 
 import
-  std/sequtils, tables, times,
+  std/[sequtils, tables, times, streams, os, deques],
-  std/streams,
-  std/os,
   chronicles, options, chronos, stint,
   confutils,
   web3, json,
@@ -24,13 +22,6 @@ import
 logScope:
   topics = "wakurlnrelayutils"
 
-when defined(rln) or (not defined(rln) and not defined(rlnzerokit)):
-  type RLNResult* = Result[RLN[Bn256], string]
-
-when defined(rlnzerokit):
-  type RLNResult* = Result[ptr RLN, string]
-
-type RlnRelayResult*[T] = Result[T, string]
 type MerkleNodeResult* = RlnRelayResult[MerkleNode]
 type RateLimitProofResult* = RlnRelayResult[RateLimitProof]
 type SpamHandler* = proc(wakuMessage: WakuMessage): void {.gcsafe, closure, raises: [Defect].}
@@ -400,16 +391,6 @@ when defined(rln) or (not defined(rln) and not defined(rlnzerokit)):
     var rootValue = cast[ptr MerkleNode] (root.`ptr`)[]
     return ok(rootValue)
 
-  proc validateRoot*(rlnInstance: RLN[Bn256], merkleRoot: MerkleNode): RlnRelayResult[bool] =
-    # Validate against the local merkle tree
-    let localTreeRoot = rlnInstance.getMerkleRoot()
-    if not localTreeRoot.isOk():
-      return err(localTreeRoot.error())
-    if localTreeRoot.value() == merkleRoot:
-      return ok(true)
-    else:
-      return ok(false)
-
   proc proofVerify*(rlnInstance: RLN[Bn256], data: openArray[byte], proof: RateLimitProof): RlnRelayResult[bool] =
     var
       proofBytes = serialize(proof, data)
@@ -518,16 +499,6 @@ when defined(rlnzerokit):
 
     return proofBytes
 
-  proc validateRoot*(rlnInstance: ptr RLN, proof: MerkleNode): RlnRelayResult[bool] =
-    # Validate against the local merkle tree
-    let localTreeRoot = rln.getMerkleRoot()
-    if not localTreeRoot.isOk():
-      return err(localTreeRoot.error())
-    if localTreeRoot.value() == merkleRoot:
-      return ok(true)
-    else:
-      return ok(false)
-
   proc proofVerify*(rlnInstance: ptr RLN, data: openArray[byte], proof: RateLimitProof): RlnRelayResult[bool] =
     var
       proofBytes = serialize(proof, data)
@@ -572,6 +543,44 @@ when defined(rlnzerokit):
     var rootValue = cast[ptr MerkleNode] (root.`ptr`)[]
     return ok(rootValue)
 
+
+proc updateValidRootQueue*(wakuRlnRelay: WakuRLNRelay, root: MerkleNode): void =
+  ## updates the valid Merkle root queue with the latest root and pops the oldest one when the capacity of `AcceptableRootWindowSize` is reached 
+  let overflowCount = wakuRlnRelay.validMerkleRoots.len() - AcceptableRootWindowSize
+  if overflowCount >= 0:
+    # Delete the oldest `overflowCount` elements in the deque (index 0..`overflowCount`)
+    for i in 0..overflowCount:
+      wakuRlnRelay.validMerkleRoots.popFirst() 
+  # Push the next root into the queue
+  wakuRlnRelay.validMerkleRoots.addLast(root)
+
+proc insertMember*(wakuRlnRelay: WakuRLNRelay, idComm: IDCommitment): RlnRelayResult[void] =
+  ## inserts a new id commitment into the local merkle tree, and adds the changed root to the 
+  ## queue of valid roots
+  let actionSucceeded = wakuRlnRelay.rlnInstance.insertMember(idComm)
+  if not actionSucceeded:
+    return err("could not insert id commitment into the merkle tree")
+
+  let rootAfterUpdate = ?wakuRlnRelay.rlnInstance.getMerkleRoot()
+  wakuRlnRelay.updateValidRootQueue(rootAfterUpdate)
+  return ok()
+  
+
+proc removeMember*(wakuRlnRelay: WakuRLNRelay, index: MembershipIndex): RlnRelayResult[void] =
+  ## removes a commitment from the local merkle tree at `index`, and adds the changed root to the
+  ## queue of valid roots
+  let actionSucceeded = wakuRlnRelay.rlnInstance.removeMember(index)
+  if not actionSucceeded:
+    return err("could not remove id commitment from the merkle tree")
+
+  let rootAfterUpdate = ?wakuRlnRelay.rlnInstance.getMerkleRoot()
+  wakuRlnRelay.updateValidRootQueue(rootAfterUpdate)
+  return ok()
+
+proc validateRoot*(wakuRlnRelay: WakuRLNRelay, root: MerkleNode): bool =
+  ## Validate against the window of roots stored in wakuRlnRelay.validMerkleRoots
+  return root in wakuRlnRelay.validMerkleRoots
+
 proc toMembershipKeyPairs*(groupKeys: seq[(string, string)]): seq[
     MembershipKeyPair] {.raises: [Defect, ValueError].} =
   ## groupKeys is  sequence of membership key tuples in the form of (identity key, identity commitment) all in the hexadecimal format
@@ -788,15 +797,9 @@ proc validateMessage*(rlnPeer: WakuRLNRelay, msg: WakuMessage,
         payload = string.fromBytes(msg.payload)
     return MessageValidationResult.Invalid
 
-  let merkleRootIsValidRes = rlnPeer.rlnInstance.validateRoot(msg.proof.merkleRoot)
+  if not rlnPeer.validateRoot(msg.proof.merkleRoot):
-
+    debug "invalid message: provided root does not belong to acceptable window of roots", provided=msg.proof.merkleRoot, validRoots=rlnPeer.validMerkleRoots
-  if merkleRootIsValidRes.isErr():
+    return MessageValidationResult.Invalid
-      debug "invalid message: could not validate the root"
-      return MessageValidationResult.Invalid
-
-  if not merkleRootIsValidRes.value():
-      debug "invalid message: received root does not match local root", payload = string.fromBytes(msg.payload)
-      return MessageValidationResult.Invalid
 
   # verify the proof
   let
@@ -854,25 +857,14 @@ proc appendRLNProof*(rlnPeer: WakuRLNRelay, msg: var WakuMessage,
   msg.proof = proof.value
   return true
 
-when defined(rln) or (not defined(rln) and not defined(rlnzerokit)):
+proc addAll*(wakuRlnRelay: WakuRLNRelay, list: seq[IDCommitment]): RlnRelayResult[void] =
-  proc addAll*(rlnInstance: RLN[Bn256], list: seq[IDCommitment]): bool =
+  # add members to the Merkle tree of the  `rlnInstance`
-    # add members to the Merkle tree of the  `rlnInstance`
+  for i in 0..list.len-1:
-    for i in 0..list.len-1:
+    let member = list[i]
-      let member = list[i]
+    let memberAdded = wakuRlnRelay.insertMember(member)
-      let member_is_added = rlnInstance.insertMember(member)
+    if not memberAdded.isOk():
-      if not member_is_added:
+      return err(memberAdded.error())
-        return false
+  return ok()
-    return true
-
-when defined(rlnzerokit):
-  proc addAll*(rlnInstance: ptr RLN, list: seq[IDCommitment]): bool =
-    # add members to the Merkle tree of the  `rlnInstance`
-    for i in 0..list.len-1:
-      let member = list[i]
-      let member_is_added = rlnInstance.insertMember(member)
-      if not member_is_added:
-        return false
-    return true
 
 # the types of inputs to this handler matches the MemberRegistered event/proc defined in the MembershipContract interface
 type RegistrationEventHandler  = proc(pubkey: Uint256, index: Uint256): void {.gcsafe, closure, raises: [Defect].}
@@ -986,12 +978,6 @@ proc mountRlnRelayStatic*(node: WakuNode,
   doAssert(rlnInstance.isOk)
   var rln = rlnInstance.value
 
-  # add members to the Merkle tree
-  for index in 0..group.len-1:
-    let member = group[index]
-    let member_is_added = rln.insertMember(member)
-    doAssert(member_is_added)
-
   # create the WakuRLNRelay
   var rlnPeer = WakuRLNRelay(membershipKeyPair: memKeyPair,
     membershipIndex: memIndex,
@@ -999,6 +985,12 @@ proc mountRlnRelayStatic*(node: WakuNode,
     pubsubTopic: pubsubTopic,
     contentTopic: contentTopic)
 
+    # add members to the Merkle tree
+  for index in 0..group.len-1:
+    let member = group[index]
+    let memberAdded = rlnPeer.insertMember(member)
+    doAssert(memberAdded.isOk())
+
   # adds a topic validator for the supplied pubsub topic at the relay protocol
   # messages published on this pubsub topic will be relayed upon a successful validation, otherwise they will be dropped
   # the topic validator checks for the correct non-spamming proof of the message
@@ -1077,9 +1069,9 @@ proc mountRlnRelayDynamic*(node: WakuNode,
     debug "a new key is added", pubkey=pubkey
     # assuming all the members arrive in order
     let pk = pubkey.toIDCommitment()
-    let isSuccessful = rlnPeer.rlnInstance.insertMember(pk)
+    let isSuccessful = rlnPeer.insertMember(pk)
     debug "received pk", pk=pk.toHex, index =index
-    doAssert(isSuccessful)
+    doAssert(isSuccessful.isOk())
 
   asyncSpawn rlnPeer.handleGroupUpdates(handler)
   debug "dynamic group management is started"