chore(rln-relay): Verify proofs based on bandwidth usage

waku-org · Jul 6, 2023 · 2bf7b48 · 2bf7b48
1 parent 4f0bdf9
commit 2bf7b48
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 6 deletions.
diff --git a/apps/wakunode2/app.nim b/apps/wakunode2/app.nim
@@ -383,7 +383,8 @@ proc setupProtocols(node: WakuNode,
         rlnRelayEthAccountAddress: conf.rlnRelayEthAccountAddress,
         rlnRelayCredPath: conf.rlnRelayCredPath,
         rlnRelayCredentialsPassword: conf.rlnRelayCredentialsPassword,
-        rlnRelayTreePath: conf.rlnRelayTreePath
+        rlnRelayTreePath: conf.rlnRelayTreePath,
+        rlnRelayBandwidthCutoff: conf.rlnRelayBandwidthCutoff
       )
 
       try:

diff --git a/apps/wakunode2/external_config.nim b/apps/wakunode2/external_config.nim
@@ -211,6 +211,11 @@ type
       defaultValue: ""
       name: "rln-relay-tree-path" }: string
 
+    rlnRelayBandwidthCutoff* {.
+      desc: "Message rate in bytes/sec after which verification of proofs should happen",
+      defaultValue: 1_000_000
+      name: "rln-relay-bandwidth-cutff" }: int
+
     staticnodes* {.
       desc: "Peer multiaddr to directly connect with. Argument may be repeated."
       name: "staticnode" }: seq[string]

diff --git a/tests/v2/waku_rln_relay/test_wakunode_rln_relay.nim b/tests/v2/waku_rln_relay/test_wakunode_rln_relay.nim
@@ -142,6 +142,7 @@ procSuite "WakuNode - RLN relay":
       rlnRelayContentTopic: contentTopic,
       rlnRelayCredIndex: 1.uint,
       rlnRelayTreePath: genTempPath("rln_tree", "wakunode_4"),
+      rlnRelayBandwidthCutoff: 1,
     ))
 
     await node1.start()
@@ -154,6 +155,7 @@ procSuite "WakuNode - RLN relay":
       rlnRelayContentTopic: contentTopic,
       rlnRelayCredIndex: 2.uint,
       rlnRelayTreePath: genTempPath("rln_tree", "wakunode_5"),
+      rlnRelayBandwidthCutoff: 1,
     ))
 
     await node2.start()
@@ -166,6 +168,7 @@ procSuite "WakuNode - RLN relay":
       rlnRelayContentTopic: contentTopic,
       rlnRelayCredIndex: 3.uint,
       rlnRelayTreePath: genTempPath("rln_tree", "wakunode_6"),
+      rlnRelayBandwidthCutoff: 1,
     ))
 
     await node3.start()
@@ -248,6 +251,7 @@ procSuite "WakuNode - RLN relay":
       rlnRelayContentTopic: contentTopic,
       rlnRelayCredIndex: 1.uint,
       rlnRelayTreePath: genTempPath("rln_tree", "wakunode_7"),
+      rlnRelayBandwidthCutoff: 1,
     ))
 
     await node1.start()
@@ -261,6 +265,7 @@ procSuite "WakuNode - RLN relay":
       rlnRelayContentTopic: contentTopic,
       rlnRelayCredIndex: 2.uint,
       rlnRelayTreePath: genTempPath("rln_tree", "wakunode_8"),
+      rlnRelayBandwidthCutoff: 1,
     ))
 
     await node2.start()
@@ -274,6 +279,7 @@ procSuite "WakuNode - RLN relay":
       rlnRelayContentTopic: contentTopic,
       rlnRelayCredIndex: 3.uint,
       rlnRelayTreePath: genTempPath("rln_tree", "wakunode_9"),
+      rlnRelayBandwidthCutoff: 1,
     ))
 
     await node3.start()

diff --git a/waku/v2/waku_rln_relay/rln_relay.nim b/waku/v2/waku_rln_relay/rln_relay.nim
@@ -5,7 +5,7 @@ else:
 
 import
   std/[algorithm, sequtils, strutils, tables, times, os, deques],
-  chronicles, options, chronos, stint,
+  chronicles, options, chronos, chronos/ratelimit, stint,
   confutils,
   web3, json,
   web3/ethtypes,
@@ -23,7 +23,8 @@ import
   ./protocol_metrics
 import
   ../waku_core,
-  ../waku_keystore
+  ../waku_keystore,
+  ../utils/collector
 
 logScope:
   topics = "waku rln_relay"
@@ -41,6 +42,7 @@ type WakuRlnConfig* = object
   rlnRelayCredPath*: string
   rlnRelayCredentialsPassword*: string
   rlnRelayTreePath*: string
+  rlnRelayBandwidthCutoff*: int
 
 proc createMembershipList*(rln: ptr RLN, n: int): RlnRelayResult[(
     seq[RawMembershipCredentials], string
@@ -86,6 +88,7 @@ type WakuRLNRelay* = ref object of RootObj
   nullifierLog*: Table[Epoch, seq[ProofMetadata]]
   lastEpoch*: Epoch # the epoch of the last published rln message
   groupManager*: GroupManager
+  messageBucket*: TokenBucket
 
 proc hasDuplicate*(rlnPeer: WakuRLNRelay,
                    proofMetadata: ProofMetadata): RlnRelayResult[bool] =
@@ -160,14 +163,16 @@ proc absDiff*(e1, e2: Epoch): uint64 =
   else:
     return epoch2 - epoch1
 
-proc validateMessage*(rlnPeer: WakuRLNRelay, msg: WakuMessage,
-    timeOption: Option[float64] = none(float64)): MessageValidationResult =
+proc validateMessage*(rlnPeer: WakuRLNRelay, 
+                      msg: WakuMessage,
+                      timeOption = none(float64)): MessageValidationResult =
   ## validate the supplied `msg` based on the waku-rln-relay routing protocol i.e.,
   ## the `msg`'s epoch is within MaxEpochGap of the current epoch
   ## the `msg` has valid rate limit proof
   ## the `msg` does not violate the rate limit
   ## `timeOption` indicates Unix epoch time (fractional part holds sub-seconds)
   ## if `timeOption` is supplied, then the current epoch is calculated based on that
+
   let decodeRes = RateLimitProof.init(msg.proof)
   if decodeRes.isErr():
     return MessageValidationResult.Invalid
@@ -286,6 +291,14 @@ proc generateRlnValidator*(wakuRlnRelay: WakuRLNRelay,
   let contentTopic = wakuRlnRelay.contentTopic
   proc validator(topic: string, message: messages.Message): Future[pubsub.ValidationResult] {.async.} =
     trace "rln-relay topic validator is called"
+
+    ## Check if enough tokens can be consumed from the message bucket
+    if wakuRlnRelay.messageBucket.tryConsume(message.data.len):
+      # message is within the bandwidth limit
+      return pubsub.ValidationResult.Accept
+    else:
+      info "message bandwidth limit exceeded, running rate limit proof validation"
+
     let decodeRes = WakuMessage.decode(message.data)
     if decodeRes.isOk():
       let
@@ -379,7 +392,8 @@ proc mount(conf: WakuRlnConfig,
 
   return WakuRLNRelay(pubsubTopic: conf.rlnRelayPubsubTopic,
                       contentTopic: conf.rlnRelayContentTopic,
-                      groupManager: groupManager)
+                      groupManager: groupManager,
+                      messageBucket: TokenBucket.new(conf.rlnRelayBandwidthCutoff))
 
 
 proc new*(T: type WakuRlnRelay,