fix(service-disco): don't rescale bucket index by maxBuckets (#2841)

Author: gmelodie

libp2p/protocols/kademlia/routing_table.nim

@@ -20,6 +20,8 @@ proc new*(
     maxBuckets: int = DefaultMaxBuckets,
     selfIdPreHashed = false,
 ): T =
+  doAssert maxBuckets > 0 and maxBuckets <= MaxBucketsLimit,
+    "maxBuckets must be in 1 .. " & $MaxBucketsLimit
   RoutingTableConfig(
     replication: replication,
     hasher: hasher,
@@ -40,6 +42,9 @@ proc new*(
     selfId: selfId, localNodeId: localNodeId.get(selfId), buckets: @[], config: config
   )
 
+func bucketCount(maxBuckets: int): int =
+  clamp(maxBuckets, 1, MaxBucketsLimit)
+
 proc bucketIndex*(rtable: RoutingTable, key: Key): int =
   let
     selfHash =
@@ -50,10 +55,7 @@ proc bucketIndex*(rtable: RoutingTable, key: Key): int =
     keyHash = key.hashFor(rtable.config.hasher)
     lz = xorDistance(selfHash, keyHash).leadingZeros
 
-  if rtable.config.maxBuckets <= 1:
-    return 0
-
-  return min(((lz * rtable.config.maxBuckets) div 256), rtable.config.maxBuckets - 1)
+  min(lz, bucketCount(rtable.config.maxBuckets) - 1)
 
 proc peerIndexInBucket(bucket: Bucket, nodeId: Key): Opt[int] =
   for i, p in bucket.peers:
@@ -209,9 +211,7 @@ proc randomKeyInBucket*(
     selfId: Key, bucketIndex: int, rng: Rng, maxBuckets: int = DefaultMaxBuckets
 ): Key =
   let
-    index = clamp(bucketIndex, 0, (maxBuckets - 1))
-    buckets = max(1, maxBuckets)
-    leadingZeros = index * 256 div buckets
+    leadingZeros = clamp(bucketIndex, 0, bucketCount(maxBuckets) - 1)
     (byteIdx, rem) = divmod(leadingZeros, 8)
     boundBitIdx = 7 - rem
 

libp2p/protocols/kademlia/types.nim

@@ -11,7 +11,9 @@ import ./protobuf
 const
   IdLength* = 32 # 256-bit IDs
 
-  DefaultMaxBuckets* = 256
+  MaxBucketsLimit* = IdLength * 8
+    ## a bucket per shared prefix bit; deeper prefixes than the key is long cannot exist
+  DefaultMaxBuckets* = MaxBucketsLimit
   DefaultTimeout* = 5.seconds
   DefaultBucketRefreshTime* = 10.minutes
   DefaultBucketStaleTime* = 30.minutes

tests/libp2p/kademlia/test_routing_table.nim

@@ -169,6 +169,33 @@ suite "KadDHT Routing Table":
       idx == TargetBucket
       rid != selfId
 
+  test "bucketIndex is the prefix length, not a rescale of it":
+    let selfId = testKey(0)
+    var rt =
+      RoutingTable.new(selfId, RoutingTableConfig.new(hasher = Opt.some(noOpHasher)))
+
+    for lz in 0 .. 8:
+      let key = randomKeyInBucket(selfId, lz, rng())
+      check rt.bucketIndex(key) == lz
+
+  test "small maxBuckets keeps prefix resolution and saturates the last bucket":
+    let selfId = testKey(0)
+    const MaxBuckets = 16
+    var rt = RoutingTable.new(
+      selfId,
+      RoutingTableConfig.new(hasher = Opt.some(noOpHasher), maxBuckets = MaxBuckets),
+    )
+
+    # Prefix lengths below the cap each get their own bucket.
+    for lz in 0 ..< MaxBuckets - 1:
+      let key = randomKeyInBucket(selfId, lz, rng(), MaxBuckets)
+      check rt.bucketIndex(key) == lz
+
+    # Deeper prefixes all fall into the last bucket.
+    for lz in MaxBuckets - 1 .. MaxBuckets + 8:
+      let key = randomKeyInBucket(selfId, lz, rng(), DefaultMaxBuckets)
+      check rt.bucketIndex(key) == MaxBuckets - 1
+
   test "randomKey returns none for empty bucket":
     var bucket: Bucket
     check randomKey(bucket, rng()).isNone()