fix(service-disco): service discovery key bucketing (#2592)

Author: SionoiS

libp2p/protocols/kademlia.nim

@@ -34,7 +34,8 @@ proc refreshTable*(
     if not (forceRefresh or bucket.isStale()):
       continue
 
-    let randomKey = randomKeyInBucket(rtable.selfId, i, kad.rng)
+    let randomKey =
+      randomKeyInBucket(rtable.selfId, i, kad.rng, rtable.config.maxBuckets)
     discard await kad.findNode(randomKey, rtable)
 
 proc bootstrap*(

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, bitops
 import chronos, chronicles, results
 import ./types
 import ./kademlia_metrics
@@ -40,17 +40,20 @@ proc new*(
     selfId: selfId, localNodeId: localNodeId.get(selfId), buckets: @[], config: config
   )
 
-proc bucketIndex*(
-    selfId, key: Key, hasher: Opt[XorDHasher], selfIdPreHashed = false
-): int =
+proc bucketIndex*(rtable: RoutingTable, key: Key): int =
   let
     selfHash =
-      if selfIdPreHashed:
-        selfId
+      if rtable.config.selfIdPreHashed:
+        rtable.selfId
       else:
-        selfId.hashFor(hasher)
-    keyHash = key.hashFor(hasher)
-  return xorDistance(selfHash, keyHash).leadingZeros
+        rtable.selfId.hashFor(rtable.config.hasher)
+    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)
 
 proc peerIndexInBucket(bucket: Bucket, nodeId: Key): Opt[int] =
   for i, p in bucket.peers:
@@ -97,13 +100,7 @@ proc insert*(rtable: RoutingTable, nodeId: Key): bool =
     debug "Cannot insert self in routing table", nodeId = nodeId
     return false # No self insertion
 
-  let idx = bucketIndex(
-    rtable.selfId, nodeId, rtable.config.hasher, rtable.config.selfIdPreHashed
-  )
-  if idx >= rtable.config.maxBuckets:
-    debug "Cannot insert node, max buckets have been reached",
-      nodeId = nodeId, bucketIdx = idx, maxBuckets = rtable.config.maxBuckets
-    return false
+  let idx = rtable.bucketIndex(nodeId)
 
   if idx >= rtable.buckets.len:
     # expand buckets lazily if needed
@@ -208,32 +205,34 @@ proc isStale*(bucket: Bucket): bool =
       return true
   return false
 
-proc randomKeyInBucket*(selfId: Key, bucketIndex: int, rng: Rng): Key =
-  var raw = selfId
-
-  # zero out higher bits
-  for i in 0 ..< bucketIndex:
-    let byteIdx = i div 8
-    let bitInByte = 7 - (i mod 8)
-    raw[byteIdx] = raw[byteIdx] and not (1'u8 shl bitInByte)
-
-  # flip the target bit
-  let tgtByte = bucketIndex div 8
-  let tgtBitInByte = 7 - (bucketIndex mod 8)
-  raw[tgtByte] = raw[tgtByte] xor (1'u8 shl tgtBitInByte)
-
-  # randomize lower bits of the boundary byte
-  let lsbMask = (1'u8 shl tgtBitInByte) - 1
-  if lsbMask != 0:
-    var rb: array[1, byte]
-    rng.generate(rb)
-    raw[tgtByte] = (raw[tgtByte] and not lsbMask) or (rb[0] and lsbMask)
+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
+    (byteIdx, rem) = divmod(leadingZeros, 8)
+    boundBitIdx = 7 - rem
+
+  var key = selfId
+
+  # For the boundary byte, from 0 to boundBitIdx the bits should be random.
+  for i in byteIdx ..< IdLength:
+    rng.generate(key[i])
+
+  # For the boundary byte, from boundBitIdx to 7 the bits should be the same as seflId.
+  for i in boundBitIdx .. 7:
+    if selfId[byteIdx].testBit(i):
+      key[byteIdx].setBit(i)
+    else:
+      key[byteIdx].clearBit(i)
 
-  # randomize remaining bytes
-  if tgtByte + 1 < raw.len:
-    rng.generate(raw.toOpenArray(tgtByte + 1, raw.len - 1))
+  # The bit at the boundary should be the opposite of selfId. So that the XOR is not 0.
+  if selfId[byteIdx].testBit(boundBitIdx) == key[byteIdx].testBit(boundBitIdx):
+    key[byteIdx].flipBit(boundBitIdx)
 
-  return raw
+  key
 
 proc allKeys*(bucket: Bucket): seq[Key] =
   return bucket.peers.mapIt(it.nodeId)

tests/libp2p/kademlia/test_find.nim

@@ -205,9 +205,7 @@ suite "KadDHT Find":
     check not kads[0].hasKey(kads[2].rtable.selfId)
 
     # Make kads[1]'s bucket stale to trigger refresh
-    let bucketIdx = bucketIndex(
-      kads[0].rtable.selfId, kads[1].rtable.selfId, kads[0].rtable.config.hasher
-    )
+    let bucketIdx = kads[0].rtable.bucketIndex(kads[1].rtable.selfId)
     makeBucketStale(kads[0].rtable.buckets[bucketIdx])
 
     check kads[0].rtable.buckets[bucketIdx].isStale()

tests/libp2p/kademlia/test_routing_table.nim

@@ -21,7 +21,7 @@ suite "KadDHT Routing Table":
     let other = testKey(0b10000000)
     discard rt.insert(other)
 
-    let idx = bucketIndex(selfId, other, Opt.none(XorDHasher))
+    let idx = rt.bucketIndex(other)
     check:
       rt.buckets.len > idx
       rt.buckets[idx].peers.len == 1
@@ -161,8 +161,10 @@ suite "KadDHT Routing Table":
 
   test "randomKeyInBucket returns id at correct distance":
     let selfId = testKey(0)
+    var rt =
+      RoutingTable.new(selfId, RoutingTableConfig.new(hasher = Opt.some(noOpHasher)))
     var rid = randomKeyInBucket(selfId, TargetBucket, rng())
-    let idx = bucketIndex(selfId, rid, Opt.some(noOpHasher))
+    let idx = rt.bucketIndex(rid)
     check:
       idx == TargetBucket
       rid != selfId

tests/libp2p/service_discovery/component/test_advertise_discover.nim

@@ -7,7 +7,6 @@ import
   ../../../../libp2p/[
     crypto/crypto,
     peerid,
-    protocols/kademlia/routing_table,
     protocols/kademlia/types,
     protocols/service_discovery,
     protocols/service_discovery/advertiser,
@@ -357,67 +356,6 @@ suite "Service Discovery Component - Advertise Discover":
         let foundB = await discovererNode.lookup(serviceBId)
         foundA.containsPeer(advertiserNode) and foundB.containsPeer(advertiserNode)
 
-  asyncTest "addProvidedService drops a known registrar at Kad bucketIndex 16":
-    # TODO: nim-libp2p#2499 service-disco: valid service-table peers are dropped when Kad bucketIndex is 16 or higher
-    # Precomputed private key to pin the registrar identity and service name.
-    # Their service table bucket index is 18 for this serviceId.
-    # A ServiceDiscovery service table accepts indexes 0 through 15.
-    const
-      serviceName = "service-bucket-indexing-component"
-      droppedRegistrarPrivateKey =
-        "080112404911f67066496ac09c4924a8ef86ec0d8532c9fb50046aa27fdb4132" &
-        "130a2c8ecf17e8057d0e5c6fcef5a3aeec98d70b43feffac4301c141a08071498ce89a1e"
-
-    let
-      conf = ServiceDiscoveryConfig.new(safetyParam = 0.0)
-      droppedRegistrarKey = PrivateKey.init(droppedRegistrarPrivateKey).get()
-      droppedRegistrarNode = setupServiceDiscoveryNode(
-        discoConfig = conf, privateKey = Opt.some(droppedRegistrarKey)
-      )
-      workingRegistrarNode = setupServiceDiscoveryNode(discoConfig = conf)
-      advertiserNode = setupServiceDiscoveryNode(discoConfig = conf)
-
-    startAndDeferStop(@[droppedRegistrarNode, workingRegistrarNode, advertiserNode])
-    await connect(droppedRegistrarNode, advertiserNode)
-    await connect(workingRegistrarNode, advertiserNode)
-
-    # The advertiser knows both registrars in the main Kad routing table.
-    # The droppedRegistrar sits just outside the service table bucket range.
-    # The workingRegistrar sits inside the service table bucket range.
-    let
-      service = makeServiceInfo(serviceName)
-      serviceId = service.id.hashServiceId()
-      droppedRegistrarPeerKey = droppedRegistrarNode.switch.peerInfo.peerId.toKey()
-      workingRegistrarPeerKey = workingRegistrarNode.switch.peerInfo.peerId.toKey()
-
-    check:
-      bucketIndex(
-        serviceId, droppedRegistrarPeerKey, Opt.none(XorDHasher), selfIdPreHashed = true
-      ) > conf.bucketsCount
-      bucketIndex(
-        serviceId, workingRegistrarPeerKey, Opt.none(XorDHasher), selfIdPreHashed = true
-      ) < conf.bucketsCount
-      advertiserNode.rtable.hasPeer(droppedRegistrarPeerKey)
-      advertiserNode.rtable.hasPeer(workingRegistrarPeerKey)
-
-    # addProvidedService seeds the service table from the main Kad routing table.
-    advertiserNode.addProvidedService(service)
-
-    # The in-range registrar is kept and receives the ad.
-    # The out-of-range registrar is absent from the service table.
-    # Only the remote task is tracked in running (local registration uses a dedicated loop).
-    let serviceTable = advertiserNode.rtManager.getTable(serviceId).get()
-    check:
-      serviceTable.hasPeer(workingRegistrarPeerKey)
-      not serviceTable.hasPeer(droppedRegistrarPeerKey)
-      advertiserNode.advertiser.running.len() == 1
-        # only the remote task (local registration is a separate dedicated loop)
-      droppedRegistrarNode.countAdsInCache(serviceId) == 0
-
-    checkUntilTimeout:
-      workingRegistrarNode.countAdsInCache(serviceId) == 1
-      droppedRegistrarNode.countAdsInCache(serviceId) == 0
-
   asyncTest "advertiser registers with peers discovered after addProvidedService":
     let conf = ServiceDiscoveryConfig.new(safetyParam = 0.0)
     let initialRegistrar = setupServiceDiscoveryNode(discoConfig = conf)

tests/libp2p/service_discovery/component/test_lookup_get_ads.nim

@@ -7,7 +7,6 @@ import
   ../../../../libp2p/[
     crypto/crypto,
     peerid,
-    protocols/kademlia/routing_table,
     protocols/kademlia/types,
     protocols/service_discovery/advertiser,
     protocols/service_discovery/discoverer,
@@ -19,56 +18,6 @@ import ../../../../libp2p/protocols/kademlia/protobuf as kad_protobuf
 import ../../../tools/[lifecycle, unittest]
 import ../utils
 
-const MaxRegistrarSetupAttempts = 1000
-
-proc bucketOf(r: ServiceDiscovery, serviceId: ServiceId): int =
-  bucketIndex(
-    serviceId,
-    r.switch.peerInfo.peerId.toKey(),
-    Opt.none(XorDHasher),
-    selfIdPreHashed = true,
-  )
-
-proc setupRegistrarsInDistinctBuckets(
-    conf: ServiceDiscoveryConfig, serviceId: ServiceId
-): tuple[queriedFirst, queriedSecond: ServiceDiscovery] =
-  ## Two registrars in distinct buckets of the routing table rooted at `serviceId`.
-  ## Returned in the order `lookup()` would visit them (lower bucket index first).
-  var a = setupServiceDiscoveryNode(discoConfig = conf)
-  var b = setupServiceDiscoveryNode(discoConfig = conf)
-  for _ in 0 ..< MaxRegistrarSetupAttempts:
-    if bucketOf(a, serviceId) != bucketOf(b, serviceId):
-      if bucketOf(a, serviceId) < bucketOf(b, serviceId):
-        return (a, b)
-      else:
-        return (b, a)
-
-    b = setupServiceDiscoveryNode(discoConfig = conf)
-
-  raiseAssert "could not find registrars in distinct buckets"
-
-proc setupRegistrarsInSameBucket(
-    conf: ServiceDiscoveryConfig, serviceId: ServiceId, count: int
-): seq[ServiceDiscovery] =
-  ## Registrars that land in one service routing table bucket.
-  doAssert count > 0, "count must be > 0"
-
-  var registrarsByBucket = initTable[int, seq[ServiceDiscovery]]()
-  for _ in 0 ..< MaxRegistrarSetupAttempts:
-    let registrar = setupServiceDiscoveryNode(discoConfig = conf)
-    let bucket = bucketOf(registrar, serviceId)
-    if bucket >= conf.bucketsCount:
-      continue
-
-    if not registrarsByBucket.hasKey(bucket):
-      registrarsByBucket[bucket] = @[]
-    registrarsByBucket[bucket].add(registrar)
-
-    if registrarsByBucket[bucket].len == count:
-      return registrarsByBucket[bucket]
-
-  raiseAssert "could not find enough registrars in one bucket"
-
 suite "Service Discovery Component - Lookup Get Ads":
   teardown:
     checkTrackers()
@@ -155,27 +104,25 @@ suite "Service Discovery Component - Lookup Get Ads":
     )
     let discovererNode = setupServiceDiscoveryNode(discoConfig = conf)
 
-    let serviceName = "service"
+    let (queriedFirst, queriedSecond, serviceName) =
+      setupRegistrarsInDistinctBuckets(conf)
     let serviceId = serviceName.hashServiceId()
-    let registrars = setupRegistrarsInDistinctBuckets(conf, serviceId)
 
-    startAndDeferStop(
-      @[discovererNode, registrars.queriedFirst, registrars.queriedSecond]
-    )
-    await connect(discovererNode, registrars.queriedFirst)
-    await connect(discovererNode, registrars.queriedSecond)
+    startAndDeferStop(@[discovererNode, queriedFirst, queriedSecond])
+    await connect(discovererNode, queriedFirst)
+    await connect(discovererNode, queriedSecond)
 
     # The first-queried registrar fills the result on its own.
     var firstBucketAds: seq[Advertisement]
     for _ in 0 ..< fLookup:
       firstBucketAds.add(makeAdvertisement(serviceName))
-    registrars.queriedFirst.registrar.cache[serviceId] = firstBucketAds
+    queriedFirst.registrar.cache[serviceId] = firstBucketAds
 
     # The other should never be queried, so should not appear in the result.
     let otherKey = randomKey()
     let otherAd = makeAdvertisement(serviceName, privateKey = otherKey)
     let otherPeerId = PeerId.init(otherKey).get()
-    registrars.queriedSecond.registrar.cache[serviceId] = @[otherAd]
+    queriedSecond.registrar.cache[serviceId] = @[otherAd]
 
     let found = await discovererNode.lookup(serviceId)
     check:
@@ -189,9 +136,8 @@ suite "Service Discovery Component - Lookup Get Ads":
     )
     let discovererNode = setupServiceDiscoveryNode(discoConfig = conf)
 
-    let serviceName = "service"
+    let (registrars, serviceName) = setupRegistrarsInSameBucket(conf, kLookup + 2)
     let serviceId = serviceName.hashServiceId()
-    let registrars = setupRegistrarsInSameBucket(conf, serviceId, kLookup + 2)
 
     startAndDeferStop(@[discovererNode] & registrars)
     for registrar in registrars:
@@ -208,21 +154,19 @@ suite "Service Discovery Component - Lookup Get Ads":
     let conf = ServiceDiscoveryConfig.new(safetyParam = 0.0, fLookup = 1)
     let discovererNode = setupServiceDiscoveryNode(discoConfig = conf)
 
-    let serviceName = "service"
+    let (queriedFirst, queriedSecond, serviceName) =
+      setupRegistrarsInDistinctBuckets(conf)
     let serviceId = serviceName.hashServiceId()
-    let registrars = setupRegistrarsInDistinctBuckets(conf, serviceId)
 
-    startAndDeferStop(
-      @[discovererNode, registrars.queriedFirst, registrars.queriedSecond]
-    )
-    await connect(discovererNode, registrars.queriedFirst)
-    await connect(discovererNode, registrars.queriedSecond)
+    startAndDeferStop(@[discovererNode, queriedFirst, queriedSecond])
+    await connect(discovererNode, queriedFirst)
+    await connect(discovererNode, queriedSecond)
 
     let firstBucketKey = randomKey()
     let firstBucketAd = makeAdvertisement(serviceName, privateKey = firstBucketKey)
     let secondBucketAd = makeAdvertisement(serviceName)
-    registrars.queriedFirst.registrar.cache[serviceId] = @[firstBucketAd]
-    registrars.queriedSecond.registrar.cache[serviceId] = @[secondBucketAd]
+    queriedFirst.registrar.cache[serviceId] = @[firstBucketAd]
+    queriedSecond.registrar.cache[serviceId] = @[secondBucketAd]
 
     let found = await discovererNode.lookup(serviceId)
     check: