fix: ensure refresh of nearby buckets

Author: gmelodie

libp2p/protocols/kademlia.nim

@@ -36,11 +36,10 @@ proc refreshTable*(
     if not (forceRefresh or bucket.isStale()):
       continue
 
-    # Buckets too near to target are covered by the findNode on selfId above.
-    let randomKey = randomKeyInBucket(rtable, i, kad.rng).valueOr:
+    let target = rtable.refreshTarget(i, kad.rng).valueOr:
       trace "No refresh target for bucket", bucket = i
       continue
-    discard await kad.findNode(randomKey, rtable)
+    discard await kad.findNode(target, rtable)
 
 proc bootstrap*(
     kad: KadDHT, forceRefresh = false

libp2p/protocols/kademlia/routing_table.nim

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0 OR MIT
 # Copyright (c) Status Research & Development GmbH
 
-import algorithm, sequtils
+import algorithm, sequtils, math
 import chronos, chronicles, results
 import ./types
 import ./kademlia_metrics
@@ -57,12 +57,14 @@ func bucketIndexFor(rtable: RoutingTable, selfHash: Key, key: Key): int =
 func bucketIndex*(rtable: RoutingTable, key: Key): int =
   rtable.bucketIndexFor(rtable.selfHash(), key)
 
-func minLeadingZeros(rtable: RoutingTable, bucketIndex: int): int =
+func minLeadingZeros*(rtable: RoutingTable, bucketIndex: int): int =
   ## Fewest leading zeros in the XOR distance that lands a key in `bucketIndex`.
+  ## Inverts `bucketIndexFor`, which floors; the inverse must round up, or it
+  ## names an lz that falls into the bucket below when maxBuckets ∤ 256.
   if rtable.config.maxBuckets <= 1:
     return 0
 
-  bucketIndex * 256 div rtable.config.maxBuckets
+  ceilDiv(bucketIndex * 256, rtable.config.maxBuckets)
 
 proc peerIndexInBucket(bucket: Bucket, nodeId: Key): Opt[int] =
   for i, p in bucket.peers:
@@ -243,3 +245,15 @@ proc randomKey*(bucket: Bucket, rng: Rng): Opt[Key] =
     proc(e: NodeEntry): Key =
       e.nodeId
   )
+
+proc refreshTarget*(rtable: RoutingTable, bucketIndex: int, rng: Rng): Opt[Key] =
+  ## Key to run a findNode against in order to refresh `bucketIndex`.
+  let random = randomKeyInBucket(rtable, bucketIndex, rng)
+  if random.isSome():
+    return random
+
+  # Buckets near self need a shared prefix too long to draw at random. A peer
+  # the bucket already holds has that prefix by construction.
+  if bucketIndex notin 0 ..< rtable.buckets.len:
+    return Opt.none(Key)
+  rtable.buckets[bucketIndex].randomKey(rng)

tests/libp2p/kademlia/test_routing_table.nim

@@ -12,6 +12,19 @@ proc testKey*(x: byte): Key =
   buf[31] = x
   return @buf
 
+proc randomTestKey(): Key =
+  var buf = newSeqUninit[byte](IdLength)
+  rng().generate(buf)
+  buf
+
+proc keyWithLeadingZeros(n: int): Key =
+  ## Key whose XOR distance to an all-zero selfId under `noOpHasher` has
+  ## exactly `n` leading zero bits.
+  doAssert n < IdLength * 8, "an all-zero key has no first set bit"
+  var buf: array[IdLength, byte]
+  buf[n div 8] = 0x80'u8 shr (n mod 8)
+  return @buf
+
 suite "KadDHT Routing Table":
   const TargetBucket = 6
 
@@ -162,6 +175,38 @@ suite "KadDHT Routing Table":
     bucket.peers = @[NodeEntry(nodeId: testKey(1), lastSeen: Moment.now())]
     check isStale(bucket) == false
 
+  test "minLeadingZeros names the first leading-zero count that lands in a bucket":
+    let selfId = testKey(0)
+    for maxBuckets in [256, 16, 100, 3]:
+      var rt = RoutingTable.new(
+        selfId,
+        RoutingTableConfig.new(hasher = Opt.some(noOpHasher), maxBuckets = maxBuckets),
+      )
+
+      for i in 0 ..< maxBuckets:
+        let lz = rt.minLeadingZeros(i)
+        check rt.bucketIndex(keyWithLeadingZeros(lz)) == i
+        # A key one bit short of the boundary belongs to a lower bucket.
+        if i > 0:
+          check rt.bucketIndex(keyWithLeadingZeros(lz - 1)) < i
+
+  test "minLeadingZeros rounds up when maxBuckets does not divide 256":
+    let selfId = testKey(0)
+
+    proc table(maxBuckets: int): RoutingTable =
+      RoutingTable.new(
+        selfId,
+        RoutingTableConfig.new(hasher = Opt.some(noOpHasher), maxBuckets = maxBuckets),
+      )
+
+    check:
+      table(256).minLeadingZeros(12) == 12
+      # A 16-way split gives each bucket 16 bits of prefix.
+      table(16).minLeadingZeros(1) == 16
+      table(16).minLeadingZeros(12) == 192
+      # 1 * 256 / 100 is 2.56, and lz 2 still lands in bucket 0.
+      table(100).minLeadingZeros(1) == 3
+
   test "randomKeyInBucket returns id at correct distance":
     let selfId = testKey(0)
     var rt =
@@ -199,6 +244,44 @@ suite "KadDHT Routing Table":
     # Bucket 1 of a 16-way split needs 16 shared prefix bits — beyond the cap.
     check randomKeyInBucket(rt, 1, rng()).isNone()
 
+  test "refreshTarget falls back to a peer of a bucket too near to search for":
+    # Service-discovery tables split the id space into few, very wide buckets,
+    # so every bucket but the first is out of reach of a random draw.
+    let selfId = testKey(0)
+    var rt = RoutingTable.new(selfId, RoutingTableConfig.new(maxBuckets = 16))
+
+    var inserted: seq[Key]
+    while inserted.len < 3:
+      let key = randomTestKey()
+      if rt.bucketIndex(key) != 1:
+        continue
+      discard rt.insert(key)
+      inserted.add(key)
+
+    check randomKeyInBucket(rt, 1, rng()).isNone()
+
+    let target = rt.refreshTarget(1, rng()).expect("bucket 1 holds peers")
+    check:
+      target in inserted
+      rt.bucketIndex(target) == 1
+
+  test "refreshTarget returns none for an unreachable empty bucket":
+    let selfId = testKey(0)
+    var rt = RoutingTable.new(selfId, RoutingTableConfig.new(maxBuckets = 16))
+
+    check rt.refreshTarget(1, rng()).isNone()
+
+  test "refreshTarget prefers a random key when the bucket is reachable":
+    let selfId = testKey(0)
+    var rt = RoutingTable.new(selfId, RoutingTableConfig.new())
+    let peer = rt.keyInBucket(TargetBucket)
+    discard rt.insert(peer)
+
+    let target = rt.refreshTarget(TargetBucket, rng()).expect("bucket is reachable")
+    check:
+      target != peer
+      rt.bucketIndex(target) == TargetBucket
+
   test "randomKey returns none for empty bucket":
     var bucket: Bucket
     check randomKey(bucket, rng()).isNone()