feat(rln-relay-v2): nonce/messageId manager (#2413)

* feat(rln-relay-v2): nonce/messageId manager

* fix: simplify

apps/chat2/chat2.nim

@@ -187,11 +187,11 @@ proc publish(c: Chat, line: string) =
   if not isNil(c.node.wakuRlnRelay):
     # for future version when we support more than one rln protected content topic,
     # we should check the message content topic as well
-    let success = c.node.wakuRlnRelay.appendRLNProof(message, float64(time))
-    if not success:
-      debug "could not append rate limit proof to the message", success=success
+    let appendRes = c.node.wakuRlnRelay.appendRLNProof(message, float64(time))
+    if appendRes.isErr():
+      debug "could not append rate limit proof to the message"
     else:
-      debug "rate limit proof is appended to the message", success=success
+      debug "rate limit proof is appended to the message"
       let decodeRes = RateLimitProof.init(message.proof)
       if decodeRes.isErr():
         error "could not decode the RLN proof"
@@ -514,14 +514,25 @@ proc processInput(rfd: AsyncFD, rng: ref HmacDrbgContext) {.async.} =
 
       echo "rln-relay preparation is in progress..."
 
-      let rlnConf = WakuRlnConfig(
-        rlnRelayDynamic: conf.rlnRelayDynamic,
-        rlnRelayCredIndex: conf.rlnRelayCredIndex,
-        rlnRelayEthContractAddress: conf.rlnRelayEthContractAddress,
-        rlnRelayEthClientAddress: conf.rlnRelayEthClientAddress,
-        rlnRelayCredPath: conf.rlnRelayCredPath,
-        rlnRelayCredPassword: conf.rlnRelayCredPassword
-      )
+      when defined(rln_v2):
+        let rlnConf = WakuRlnConfig(
+          rlnRelayDynamic: conf.rlnRelayDynamic,
+          rlnRelayCredIndex: conf.rlnRelayCredIndex,
+          rlnRelayEthContractAddress: conf.rlnRelayEthContractAddress,
+          rlnRelayEthClientAddress: conf.rlnRelayEthClientAddress,
+          rlnRelayCredPath: conf.rlnRelayCredPath,
+          rlnRelayCredPassword: conf.rlnRelayCredPassword,
+          rlnRelayUserMessageLimit: conf.rlnRelayUserMessageLimit,
+        )
+      else:
+        let rlnConf = WakuRlnConfig(
+          rlnRelayDynamic: conf.rlnRelayDynamic,
+          rlnRelayCredIndex: conf.rlnRelayCredIndex,
+          rlnRelayEthContractAddress: conf.rlnRelayEthContractAddress,
+          rlnRelayEthClientAddress: conf.rlnRelayEthClientAddress,
+          rlnRelayCredPath: conf.rlnRelayCredPath,
+          rlnRelayCredPassword: conf.rlnRelayCredPassword,
+        )
 
       waitFor node.mountRlnRelay(rlnConf,
                                 spamHandler=some(spamHandler))

apps/chat2/config_chat2.nim

@@ -266,6 +266,11 @@ type
       defaultValue: ""
       name: "rln-relay-cred-password" }: string
 
+    rlnRelayUserMessageLimit* {.
+      desc: "Set a user message limit for the rln membership registration. Must be a positive integer. Default is 1.",
+      defaultValue: 1,
+      name: "rln-relay-user-message-limit" .}: uint64
+
 # NOTE: Keys are different in nim-libp2p
 proc parseCmdArg*(T: type crypto.PrivateKey, p: string): T =
   try:

apps/wakunode2/app.nim

@@ -462,15 +462,27 @@ proc setupProtocols(node: WakuNode,
 
   if conf.rlnRelay:
 
-    let rlnConf = WakuRlnConfig(
-      rlnRelayDynamic: conf.rlnRelayDynamic,
-      rlnRelayCredIndex: conf.rlnRelayCredIndex,
-      rlnRelayEthContractAddress: conf.rlnRelayEthContractAddress,
-      rlnRelayEthClientAddress: conf.rlnRelayEthClientAddress,
-      rlnRelayCredPath: conf.rlnRelayCredPath,
-      rlnRelayCredPassword: conf.rlnRelayCredPassword,
-      rlnRelayTreePath: conf.rlnRelayTreePath,
-    )
+    when defined(rln_v2):
+      let rlnConf = WakuRlnConfig(
+        rlnRelayDynamic: conf.rlnRelayDynamic,
+        rlnRelayCredIndex: conf.rlnRelayCredIndex,
+        rlnRelayEthContractAddress: conf.rlnRelayEthContractAddress,
+        rlnRelayEthClientAddress: conf.rlnRelayEthClientAddress,
+        rlnRelayCredPath: conf.rlnRelayCredPath,
+        rlnRelayCredPassword: conf.rlnRelayCredPassword,
+        rlnRelayTreePath: conf.rlnRelayTreePath,
+        rlnRelayUserMessageLimit: conf.rlnRelayUserMessageLimit,
+      )
+    else:
+      let rlnConf = WakuRlnConfig(
+        rlnRelayDynamic: conf.rlnRelayDynamic,
+        rlnRelayCredIndex: conf.rlnRelayCredIndex,
+        rlnRelayEthContractAddress: conf.rlnRelayEthContractAddress,
+        rlnRelayEthClientAddress: conf.rlnRelayEthClientAddress,
+        rlnRelayCredPath: conf.rlnRelayCredPath,
+        rlnRelayCredPassword: conf.rlnRelayCredPassword,
+        rlnRelayTreePath: conf.rlnRelayTreePath,
+      )
 
     try:
       waitFor node.mountRlnRelay(rlnConf)

tests/node/test_wakunode_relay_rln.nim

@@ -58,7 +58,7 @@ proc sendRlnMessage(
     payload: seq[byte] = "Hello".toBytes(),
 ): Future[bool] {.async.} =
   var message = WakuMessage(payload: payload, contentTopic: contentTopic)
-  doAssert(client.wakuRlnRelay.appendRLNProof(message, epochTime()))
+  doAssert(client.wakuRlnRelay.appendRLNProof(message, epochTime()).isOk())
   discard await client.publish(some(pubsubTopic), message)
   let isCompleted = await completionFuture.withTimeout(FUTURE_TIMEOUT)
   return isCompleted
@@ -249,18 +249,18 @@ suite "Waku RlnRelay - End to End":
           WakuMessage(payload: @payload150kibPlus, contentTopic: contentTopic)
 
       doAssert(
-        client.wakuRlnRelay.appendRLNProof(message1b, epoch + EpochUnitSeconds * 0)
+        client.wakuRlnRelay.appendRLNProof(message1b, epoch + EpochUnitSeconds * 0).isOk()
       )
       doAssert(
-        client.wakuRlnRelay.appendRLNProof(message1kib, epoch + EpochUnitSeconds * 1)
+        client.wakuRlnRelay.appendRLNProof(message1kib, epoch + EpochUnitSeconds * 1).isOk()
       )
       doAssert(
-        client.wakuRlnRelay.appendRLNProof(message150kib, epoch + EpochUnitSeconds * 2)
+        client.wakuRlnRelay.appendRLNProof(message150kib, epoch + EpochUnitSeconds * 2).isOk()
       )
       doAssert(
         client.wakuRlnRelay.appendRLNProof(
           message151kibPlus, epoch + EpochUnitSeconds * 3
-        )
+        ).isOk()
       )
 
       # When sending the 1B message

tests/waku_rln_relay/test_all.nim

@@ -4,4 +4,5 @@ import
   ./test_rln_group_manager_onchain,
   ./test_rln_group_manager_static,
   ./test_waku_rln_relay,
-  ./test_wakunode_rln_relay
+  ./test_wakunode_rln_relay,
+  ./test_rln_nonce_manager

tests/waku_rln_relay/test_rln_nonce_manager.nim

@@ -0,0 +1,51 @@
+{.used.}
+
+import 
+  testutils/unittests,
+  chronos,
+  os
+import 
+  ../../../waku/waku_rln_relay/nonce_manager
+
+
+suite "Nonce manager":
+  test "should initialize successfully":
+    let nm = NonceManager.init(nonceLimit = 100.uint)
+
+    check:
+      nm.nonceLimit == 100.uint
+      nm.nextNonce == 0.uint
+
+  test "should generate a new nonce":
+    let nm = NonceManager.init(nonceLimit = 100.uint)
+    let nonceRes = nm.get()
+
+    assert nonceRes.isOk(), $nonceRes.error
+
+    check:
+      nonceRes.get() == 0.uint
+      nm.nextNonce == 1.uint
+
+  test "should fail to generate a new nonce if limit is reached":
+    let nm = NonceManager.init(nonceLimit = 1.uint)
+    let nonceRes = nm.get()
+    let nonceRes2 = nm.get()
+
+    assert nonceRes.isOk(), $nonceRes.error
+    assert nonceRes2.isErr(), "Expected error, got: " & $nonceRes2.value
+
+    check:
+      nonceRes2.error.kind == NonceManagerErrorKind.NonceLimitReached
+
+  test "should generate a new nonce if epoch is crossed":
+    let nm = NonceManager.init(nonceLimit = 1.uint, epoch = float(0.000001))
+    let nonceRes = nm.get()
+    sleep(1)
+    let nonceRes2 = nm.get()
+
+    assert nonceRes.isOk(), $nonceRes.error
+    assert nonceRes2.isOk(), $nonceRes2.error
+
+    check:
+      nonceRes.value == 0.uint
+      nonceRes2.value == 0.uint

tests/waku_rln_relay/test_waku_rln_relay.nim

@@ -687,9 +687,9 @@ suite "Waku rln relay":
 
     # ensure proofs are added
     require:
-      proofAdded1
-      proofAdded2
-      proofAdded3
+      proofAdded1.isOk()
+      proofAdded2.isOk()
+      proofAdded3.isOk()
 
     # validate messages
     # validateMessage proc checks the validity of the message fields and adds it to the log (if valid)

tests/waku_rln_relay/test_wakunode_rln_relay.nim

@@ -85,7 +85,7 @@ procSuite "WakuNode - RLN relay":
 
     # prepare the epoch
     var message = WakuMessage(payload: @payload, contentTopic: contentTopic)
-    doAssert(node1.wakuRlnRelay.appendRLNProof(message, epochTime()))
+    doAssert(node1.wakuRlnRelay.appendRLNProof(message, epochTime()).isOk())
 
 
     ## node1 publishes a message with a rate limit proof, the message is then relayed to node2 which in turn
@@ -155,12 +155,12 @@ procSuite "WakuNode - RLN relay":
 
     for i in 0..<3:
       var message = WakuMessage(payload: ("Payload_" & $i).toBytes(), contentTopic: contentTopics[0])
-      doAssert(nodes[0].wakuRlnRelay.appendRLNProof(message, epochTime))
+      doAssert(nodes[0].wakuRlnRelay.appendRLNProof(message, epochTime).isOk())
       messages1.add(message)
 
     for i in 0..<3:
       var message = WakuMessage(payload: ("Payload_" & $i).toBytes(), contentTopic: contentTopics[1])
-      doAssert(nodes[1].wakuRlnRelay.appendRLNProof(message, epochTime))
+      doAssert(nodes[1].wakuRlnRelay.appendRLNProof(message, epochTime).isOk())
       messages2.add(message)
 
     # publish 3 messages from node[0] (last 2 are spam, window is 10 secs)
@@ -346,9 +346,9 @@ procSuite "WakuNode - RLN relay":
 
     #  check proofs are added correctly
     check:
-      proofAdded1
-      proofAdded2
-      proofAdded3
+      proofAdded1.isOk()
+      proofAdded2.isOk()
+      proofAdded3.isOk()
 
     #  relay handler for node3
     var completionFut1 = newFuture[bool]()
@@ -452,9 +452,9 @@ procSuite "WakuNode - RLN relay":
 
     #  check proofs are added correctly
     check:
-      proofAdded1
-      proofAdded2
-      proofAdded3
+      proofAdded1.isOk()
+      proofAdded2.isOk()
+      proofAdded3.isOk()
 
     #  relay handler for node2
     var completionFut1 = newFuture[bool]()

waku/waku_api/jsonrpc/relay/handlers.nim

@@ -102,9 +102,8 @@ proc installRelayApiHandlers*(node: WakuNode, server: RpcServer, cache: MessageC
     # if RLN is mounted, append the proof to the message
     if not node.wakuRlnRelay.isNil():
       # append the proof to the message
-      let success = node.wakuRlnRelay.appendRLNProof(message,
-                                                    float64(getTime().toUnix()))
-      if not success:
+      node.wakuRlnRelay.appendRLNProof(message,
+                                       float64(getTime().toUnix())).isOkOr:
         raise newException(ValueError, "Failed to publish: error appending RLN proof to message")
       # validate the message before sending it
       let result = node.wakuRlnRelay.validateMessageAndUpdateLog(message)
@@ -201,9 +200,8 @@ proc installRelayApiHandlers*(node: WakuNode, server: RpcServer, cache: MessageC
     # if RLN is mounted, append the proof to the message
     if not node.wakuRlnRelay.isNil():
       # append the proof to the message
-      let success = node.wakuRlnRelay.appendRLNProof(message,
-                                                    float64(getTime().toUnix()))
-      if not success:
+      node.wakuRlnRelay.appendRLNProof(message,
+                                       float64(getTime().toUnix())).isOkOr:
         raise newException(ValueError, "Failed to publish: error appending RLN proof to message")
       # validate the message before sending it
       let result = node.wakuRlnRelay.validateMessageAndUpdateLog(message)

waku/waku_api/rest/relay/handlers.nim

@@ -130,9 +130,8 @@ proc installRelayApiHandlers*(router: var RestRouter, node: WakuNode, cache: Mes
     # if RLN is mounted, append the proof to the message
     if not node.wakuRlnRelay.isNil():
       # append the proof to the message
-      let success = node.wakuRlnRelay.appendRLNProof(message,
-                                                    float64(getTime().toUnix()))
-      if not success:
+      node.wakuRlnRelay.appendRLNProof(message,
+                                       float64(getTime().toUnix())).isOkOr:
         return RestApiResponse.internalServerError("Failed to publish: error appending RLN proof to message")
 
     (await node.wakuRelay.validateMessage(pubsubTopic, message)).isOkOr:
@@ -219,7 +218,7 @@ proc installRelayApiHandlers*(router: var RestRouter, node: WakuNode, cache: Mes
 
     # if RLN is mounted, append the proof to the message
     if not node.wakuRlnRelay.isNil():
-      if not node.wakuRlnRelay.appendRLNProof(message, float64(getTime().toUnix())):
+      node.wakuRlnRelay.appendRLNProof(message, float64(getTime().toUnix())).isOkOr:
         return RestApiResponse.internalServerError(
           "Failed to publish: error appending RLN proof to message")
 

waku/waku_rln_relay/group_manager/group_manager_base.nim

@@ -188,10 +188,12 @@ when defined(rln_v2):
       return err("user message limit is not set")
     waku_rln_proof_generation_duration_seconds.nanosecondTime:
       let proof = proofGen(rlnInstance = g.rlnInstance,
-                                data = data,
-                                memKeys = g.idCredentials.get(),
-                                memIndex = g.membershipIndex.get(),
-                                epoch = epoch).valueOr:
+                           data = data,
+                           membership = g.idCredentials.get(),
+                           index = g.membershipIndex.get(),
+                           epoch = epoch,
+                           userMessageLimit = g.userMessageLimit.get(),
+                           messageId = messageId).valueOr:
         return err("proof generation failed: " & $error)
     return ok(proof)
 else:

waku/waku_rln_relay/nonce_manager.nim

@@ -0,0 +1,67 @@
+when (NimMajor, NimMinor) < (1, 4):
+  {.push raises: [Defect].}
+else:
+  {.push raises: [].}
+
+import
+  chronos,
+  stew/results,
+  times
+import
+  ./constants
+
+export
+  chronos,
+  times,
+  results,
+  constants
+
+# This module contains the NonceManager interface
+# The NonceManager is responsible for managing the messageId used to generate RLN proofs
+# It should be used to fetch a new messageId every time a proof is generated
+# It refreshes the messageId every `epoch` seconds
+
+type
+  Nonce* = uint64
+  NonceManager* = ref object of RootObj
+    epoch*: float64
+    nextNonce*: Nonce
+    lastNonceTime*: float64
+    nonceLimit*: Nonce
+
+  NonceManagerErrorKind* = enum
+    NonceLimitReached
+
+  NonceManagerError* = object
+    kind*: NonceManagerErrorKind
+    error*: string
+
+  NonceManagerResult*[T] = Result[T, NonceManagerError]
+
+proc `$`*(ne: NonceManagerError): string =
+  case ne.kind
+  of NonceLimitReached:
+    return "NonceLimitReached: " & ne.error
+
+proc init*(T: type NonceManager, nonceLimit: Nonce, epoch = EpochUnitSeconds): T =
+  return NonceManager(
+    epoch: epoch,
+    nextNonce: 0,
+    lastNonceTime: 0,
+    nonceLimit: nonceLimit
+  )
+
+
+proc get*(n: NonceManager): NonceManagerResult[Nonce] =
+  let now = getTime().toUnixFloat()
+  var retNonce = n.nextNonce
+
+  if now - n.lastNonceTime >= n.epoch: retNonce = 0
+  n.nextNonce = retNonce + 1
+  n.lastNonceTime = now
+
+  if retNonce >= n.nonceLimit:
+    return err(NonceManagerError(kind: NonceLimitReached, 
+                                 error: "Nonce limit reached. Please wait for the next epoch"))
+
+  return ok(retNonce)