feat: DOS protection of non relay protocols - rate limit phase3

tests/common/test_all.nim

@@ -5,4 +5,7 @@ import
   ./test_envvar_serialization,
   ./test_protobuf_validation,
   ./test_sqlite_migrations,
-  ./test_parse_size
+  ./test_parse_size,
+  ./test_tokenbucket,
+  ./test_requestratelimiter,
+  ./test_timed_map

tests/common/test_requestratelimiter.nim

@@ -0,0 +1,84 @@
+#                Chronos Test Suite
+#            (c) Copyright 2022-Present
+#         Status Research & Development GmbH
+#
+#              Licensed under either of
+#  Apache License, version 2.0, (LICENSE-APACHEv2)
+#              MIT license (LICENSE-MIT)
+
+{.used.}
+
+import testutils/unittests
+import chronos, libp2p/stream/connection
+import std/[sequtils, options]
+
+import ../../waku/common/rate_limit/request_limiter
+import ../../waku/common/rate_limit/timed_map
+
+let proto = "ProtocolDescriptor"
+
+let conn1 = Connection(peerId: PeerId.random().tryGet())
+let conn2 = Connection(peerId: PeerId.random().tryGet())
+let conn3 = Connection(peerId: PeerId.random().tryGet())
+
+suite "RequestRateLimiter":
+  test "RequestRateLimiter Allow up to main bucket":
+    # keep limits low for easier calculation of ratios
+    let rateLimit: RateLimitSetting = (4, 2.minutes)
+    var limiter = newRequestRateLimiter(some(rateLimit))
+    # per peer tokens will be 6 / 4min
+    # as ratio is 2 in this case but max tokens are main tokens*ratio . 0.75
+    # notice meanwhile we have 8 global tokens over 2 period (4 mins) in sum
+    # See: waku/common/rate_limit/request_limiter.nim #func calcPeriodRatio
+
+    let now = Moment.now()
+    # with first use we register the peer also and start its timer
+    check limiter.checkUsage(proto, conn2, now) == true
+    for i in 0 ..< 3:
+      check limiter.checkUsage(proto, conn1, now) == true
+
+    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
+    for i in 0 ..< 3:
+      check limiter.checkUsage(proto, conn1, now + 3.minutes) == true
+
+    # conn1 reached the 75% of the main bucket over 2 periods of time
+    check limiter.checkUsage(proto, conn1, now + 3.minutes) == false
+
+    # conn2 has not used its tokens while we have 1 more tokens left in the main bucket
+    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
+
+  test "RequestRateLimiter Restrict overusing peer":
+    # keep limits low for easier calculation of ratios
+    let rateLimit: RateLimitSetting = (10, 2.minutes)
+    var limiter = newRequestRateLimiter(some(rateLimit))
+    # per peer tokens will be 15 / 4min
+    # as ratio is 2 in this case but max tokens are main tokens*ratio . 0.75
+    # notice meanwhile we have 20 tokens over 2 period (4 mins) in sum
+    # See: waku/common/rate_limit/request_limiter.nim #func calcPeriodRatio
+
+    let now = Moment.now()
+    # with first use we register the peer also and start its timer
+    for i in 0 ..< 10:
+      check limiter.checkUsage(proto, conn1, now) == true
+
+    # run out of main tokens but still used one more token from the peer's bucket
+    check limiter.checkUsage(proto, conn1, now) == false
+
+    for i in 0 ..< 4:
+      check limiter.checkUsage(proto, conn1, now + 3.minutes) == true
+
+    # conn1 reached the 75% of the main bucket over 2 periods of time
+    check limiter.checkUsage(proto, conn1, now + 3.minutes) == false
+
+    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
+    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
+    check limiter.checkUsage(proto, conn3, now + 3.minutes) == true
+    check limiter.checkUsage(proto, conn2, now + 3.minutes) == true
+    check limiter.checkUsage(proto, conn3, now + 3.minutes) == true
+
+    # conn1 gets replenished as the ratio was 2 giving twice as long replenish period than the main bucket
+    # see waku/common/rate_limit/request_limiter.nim #func calcPeriodRatio and calcPeerTokenSetting
+    check limiter.checkUsage(proto, conn1, now + 4.minutes) == true
+    # requests of other peers can also go
+    check limiter.checkUsage(proto, conn2, now + 4100.milliseconds) == true
+    check limiter.checkUsage(proto, conn3, now + 5.minutes) == true

tests/common/test_timed_map.nim

@@ -0,0 +1,60 @@
+{.used.}
+
+import unittest2
+import chronos/timer
+import ../../waku/common/rate_limit/timed_map
+
+suite "TimedMap":
+  test "put/get":
+    var cache = TimedMap[int, string].init(5.seconds)
+
+    let now = Moment.now()
+    check:
+      cache.mgetOrPut(1, "1", now) == "1"
+      cache.mgetOrPut(1, "1", now + 1.seconds) == "1"
+      cache.mgetOrPut(2, "2", now + 4.seconds) == "2"
+
+    check:
+      1 in cache
+      2 in cache
+
+    check:
+      cache.mgetOrPut(3, "3", now + 6.seconds) == "3"
+      # expires 1
+
+    check:
+      1 notin cache
+      2 in cache
+      3 in cache
+
+      cache.addedAt(2) == now + 4.seconds
+
+    check:
+      cache.mgetOrPut(2, "modified2", now + 8.seconds) == "2" # refreshes 2
+      cache.mgetOrPut(4, "4", now + 12.seconds) == "4" # expires 3
+
+    check:
+      2 in cache
+      3 notin cache
+      4 in cache
+
+    check:
+      cache.remove(4).isSome()
+      4 notin cache
+
+    check:
+      cache.mgetOrPut(100, "100", now + 100.seconds) == "100" # expires everything
+      100 in cache
+      2 notin cache
+
+  test "enough items to force cache heap storage growth":
+    var cache = TimedMap[int, string].init(5.seconds)
+
+    let now = Moment.now()
+    for i in 101 .. 100000:
+      check:
+        cache.mgetOrPut(i, $i, now) == $i
+
+    for i in 101 .. 100000:
+      check:
+        i in cache

tests/common/test_tokenbucket.nim

@@ -0,0 +1,69 @@
+#                Chronos Test Suite
+#            (c) Copyright 2022-Present
+#         Status Research & Development GmbH
+#
+#              Licensed under either of
+#  Apache License, version 2.0, (LICENSE-APACHEv2)
+#              MIT license (LICENSE-MIT)
+
+{.used.}
+
+import testutils/unittests
+import chronos
+import ../../waku/common/rate_limit/token_bucket
+
+suite "Token Bucket":
+  test "TokenBucket Sync test - strict":
+    var bucket = TokenBucket.newStrict(1000, 1.milliseconds)
+    let
+      start = Moment.now()
+      fullTime = start + 1.milliseconds
+    check:
+      bucket.tryConsume(800, start) == true
+      bucket.tryConsume(200, start) == true
+      # Out of budget
+      bucket.tryConsume(100, start) == false
+      bucket.tryConsume(800, fullTime) == true
+      bucket.tryConsume(200, fullTime) == true
+      # Out of budget
+      bucket.tryConsume(100, fullTime) == false
+
+  test "TokenBucket Sync test - compensating":
+    var bucket = TokenBucket.new(1000, 1.milliseconds)
+    let
+      start = Moment.now()
+      fullTime = start + 1.milliseconds
+    check:
+      bucket.tryConsume(800, start) == true
+      bucket.tryConsume(200, start) == true
+      # Out of budget
+      bucket.tryConsume(100, start) == false
+      bucket.tryConsume(800, fullTime) == true
+      bucket.tryConsume(200, fullTime) == true
+      # Due not using the bucket for a full period the compensation will satisfy this request
+      bucket.tryConsume(100, fullTime) == true
+
+  test "TokenBucket Max compensation":
+    var bucket = TokenBucket.new(1000, 1.minutes)
+    var reqTime = Moment.now()
+
+    check bucket.tryConsume(1000, reqTime)
+    check bucket.tryConsume(1, reqTime) == false
+    reqTime += 1.minutes
+    check bucket.tryConsume(500, reqTime) == true
+    reqTime += 1.minutes
+    check bucket.tryConsume(1000, reqTime) == true
+    reqTime += 10.seconds
+    # max compensation is 25% so try to consume 250 more
+    check bucket.tryConsume(250, reqTime) == true
+    reqTime += 49.seconds
+    # out of budget within the same period
+    check bucket.tryConsume(1, reqTime) == false
+
+  test "TokenBucket Short replenish":
+    var bucket = TokenBucket.new(15000, 1.milliseconds)
+    let start = Moment.now()
+    check bucket.tryConsume(15000, start)
+    check bucket.tryConsume(1, start) == false
+
+    check bucket.tryConsume(15000, start + 1.milliseconds) == true

tests/waku_filter_v2/test_all.nim

@@ -1,3 +1,4 @@
 {.used.}
 
-import ./test_waku_client, ./test_waku_filter_protocol
+import
+  ./test_waku_client, ./test_waku_filter_protocol, ./test_waku_filter_dos_protection