fix(service-disco): extend IP-similarity scoring to IPv6 (#2805)

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Gabriel Cruz <8129788+gmelodie@users.noreply.github.com>

Author: 3 people

libp2p/protocols/service_discovery/registrar.nim

@@ -21,34 +21,35 @@ logScope:
 proc updateRegistrarMetrics(registrar: Registrar) {.raises: [].} =
   cd_registrar_cache_ads.set(registrar.cacheTimestamps.len.float64)
   cd_registrar_cache_services.set(registrar.cache.len.float64)
-  cd_iptree_unique_ips.set(registrar.ipTree.root.counter.float64)
+  cd_iptree_total_ips.set(
+    registrar.ipTree.root.counter.float64 + registrar.ipTree.root6.counter.float64
+  )
 
-proc filterIPv4(addrsInfos: seq[AddressInfo]): seq[IpAddress] {.raises: [].} =
+proc getIPs(addrsInfos: seq[AddressInfo]): seq[IpAddress] {.raises: [].} =
   var ips: seq[IpAddress]
   for addrInfo in addrsInfos:
     let multiAddr = addrInfo.address
     multiAddr.getIp().withValue(ip):
-      if ip.family == IpAddressFamily.IPv4:
-        ips.add(ip)
+      ips.add(ip)
   return ips
 
 proc adScore*(ipTree: IpTree, ad: Advertisement): float64 {.raises: [].} =
   ## Return the max score for this advertisement
 
   var maxScore = 0.0
-  for ip in ad.data.addresses.filterIPv4():
+  for ip in ad.data.addresses.getIPs():
     let score = ipTree.ipScore(ip)
     if score > maxScore:
       maxScore = score
 
   return maxScore
 
 proc insertAd*(ipTree: IpTree, ad: Advertisement) {.raises: [].} =
-  for ip in ad.data.addresses.filterIPv4():
+  for ip in ad.data.addresses.getIPs():
     ipTree.insertIp(ip)
 
 proc removeAd*(ipTree: IpTree, ad: Advertisement) {.raises: [].} =
-  for ip in ad.data.addresses.filterIPv4():
+  for ip in ad.data.addresses.getIPs():
     ipTree.removeIp(ip)
 
 proc isExpired(now, ts: Moment, expiry: Duration): bool =

libp2p/protocols/service_discovery/service_discovery_metrics.nim

@@ -42,7 +42,8 @@ declarePublicCounter cd_service_table_insertions,
   "peer insertions into service routing tables"
 
 # IpTree metrics
-declarePublicGauge cd_iptree_unique_ips, "unique IPs tracked in the IP tree"
+declarePublicGauge cd_iptree_total_ips,
+  "total IP entries currently tracked in the IP tree (counts repeat inserts of the same address, not distinct addresses)"
 
 # Advertiser metrics
 declarePublicGauge cd_advertiser_pending_actions,

libp2p/utils/iptree.nim

@@ -9,100 +9,131 @@ type
     left*, right*: IpTreeNode
 
   IpTree* = ref object
-    root*: IpTreeNode
+    root*: IpTreeNode ## IPv4 tree root: a 32-level binary tree.
+    root6*: IpTreeNode ## IPv6 tree root: a 128-level binary tree.
 
 proc new*(T: typedesc[IpTree]): T =
-  T(root: IpTreeNode(counter: 0))
+  T(root: IpTreeNode(counter: 0), root6: IpTreeNode(counter: 0))
+
+proc ipBytes(ip: IpAddress): seq[uint8] {.raises: [].} =
+  case ip.family
+  of IpAddressFamily.IPv4:
+    @(ip.address_v4)
+  of IpAddressFamily.IPv6:
+    @(ip.address_v6)
+
+proc treeRoot(ipTree: IpTree, ip: IpAddress): IpTreeNode =
+  case ip.family
+  of IpAddressFamily.IPv4: ipTree.root
+  of IpAddressFamily.IPv6: ipTree.root6
+
+proc step(v: IpTreeNode, goRight: bool, create: bool): IpTreeNode =
+  ## Returns the child of `v` in the direction given by `goRight`. With
+  ## `create`, a missing child is allocated and linked in, so the result is
+  ## never nil; otherwise a missing child yields nil.
+  var nxt = if goRight: v.right else: v.left
+  if nxt.isNil and create:
+    nxt = IpTreeNode(counter: 0)
+    if goRight:
+      v.right = nxt
+    else:
+      v.left = nxt
+  nxt
 
 proc insertIp*(ipTree: IpTree, ip: IpAddress) {.raises: [].} =
-  doAssert ip.family == IpAddressFamily.IPv4
-
-  var v = ipTree.root
+  ## Inserts an IP address into the IP tree, following its binary
+  ## representation and incrementing counters at each visited node.
+  ## IPv4 addresses walk the 32-level tree,
+  ## IPv6 addresses walk the 128-level tree.
+  var v = ipTree.treeRoot(ip)
   v.counter += 1
 
-  let bytes = ip.address_v4
+  let bytes = ip.ipBytes()
 
-  for i in 0 ..< 4:
-    let b = bytes[i]
+  for b in bytes:
     for bit in countdown(7, 0):
       let goRight = (b and (1'u8 shl bit)) != 0
-
-      if goRight:
-        if v.right.isNil:
-          v.right = IpTreeNode(counter: 0)
-        v = v.right
-      else:
-        if v.left.isNil:
-          v.left = IpTreeNode(counter: 0)
-        v = v.left
-
+      v = step(v, goRight, create = true)
       v.counter += 1
 
 proc removeIp*(ipTree: IpTree, ip: IpAddress) {.raises: [].} =
-  ## Removes an IPv4 address from the IP tree by decrementing counters along
-  ## the 32-bit path. Counters never go below zero. Only IPv4 is supported.
-  doAssert ip.family == IpAddressFamily.IPv4
-
-  if ipTree.root.counter == 0:
+  ## Removes an IP address from the IP tree by decrementing counters along
+  ## its binary-representation path, from the root through the leaf.
+  ## Counters never go below zero. Any trailing run of nodes left at counter
+  ## 0 with no children is unlinked, so a fully-removed address doesn't leave
+  ## dead nodes behind.
+  let root = ipTree.treeRoot(ip)
+  if root.counter == 0:
     return
 
-  var v = ipTree.root
-  let bytes = ip.address_v4
+  let bytes = ip.ipBytes()
 
-  var path: array[32, IpTreeNode]
-  var pathLen = 0
+  # path[0] is the root; path[i] for i > 0 is the node reached after the
+  # i-th bit. Sized for the deepest case (IPv6: root + 128 levels).
+  var path: array[129, IpTreeNode]
+  path[0] = root
+  var pathLen = 1
 
-  for i in 0 ..< 4:
-    let b = bytes[i]
+  var v = root
+  for b in bytes:
     for bit in countdown(7, 0):
+      let goRight = (b and (1'u8 shl bit)) != 0
+      v = step(v, goRight, create = false)
       if v.isNil or v.counter == 0:
         return
 
       path[pathLen] = v
       inc pathLen
 
-      let goLeft = (b and (1'u8 shl bit)) == 0
-      let nxt = if goLeft: v.left else: v.right
-      if nxt.isNil:
-        return
-      v = nxt
-
+  # Every node visited above has counter > 0 (checked as it was added to
+  # path), so this can't underflow.
   for j in 0 ..< pathLen:
-    let n = path[j]
-    if n.counter > 0:
-      dec n.counter
+    dec path[j].counter
+
+  # Prune the trailing run of now-empty leaves: walk from the leaf back
+  # toward the root, unlinking any node left with counter == 0 and no
+  # children, stopping at the first node that still has either.
+  for i in countdown(pathLen - 1, 1):
+    let n = path[i]
+    if n.counter > 0 or not n.left.isNil or not n.right.isNil:
+      break
+    let parent = path[i - 1]
+    if parent.left == n:
+      parent.left = nil
+    else:
+      parent.right = nil
 
 proc ipScore*(ipTree: IpTree, ip: IpAddress): float64 {.raises: [].} =
-  ## Returns an IP similarity score in [0.0, 1.0] for the given IPv4 address.
-  ## Asserts that `ip` is an IPv4 address.
+  ## Returns an IP similarity score in [0.0, 1.0] for the given IP address.
+  ## Supports both IPv4 (32-level tree) and IPv6 (128-level tree).
   ##
-  ## The score counts how many of the 32 prefix nodes along the IP's path have
+  ## The score counts how many of the prefix nodes along the IP's path have
   ## a counter exceeding the expected threshold (root.counter / 2^(depth+1)),
   ## where depth+1 is the tree level of the child node being evaluated.
   ## A high score means many existing IPs share the same subnet — a signal of
   ## Sybil-style clustering.
-  doAssert ip.family == IpAddressFamily.IPv4
-
-  if ipTree.root.counter == 0:
+  let root = ipTree.treeRoot(ip)
+  if root.counter == 0:
     return 0.0
 
-  var v = ipTree.root
+  var v = root
   var score = 0
-  let total = float64(ipTree.root.counter)
-  let bytes = ip.address_v4
+  let total = float64(root.counter)
+  let bytes = ip.ipBytes()
+  let nBits = bytes.len * 8
 
-  for i in 0 ..< 4:
-    let b = bytes[i]
+  var threshold = total * 0.5
+  for b in bytes:
     for bit in countdown(7, 0):
-      let depth = i * 8 + (7 - bit) # 0 .. 31; child node sits at tree level depth+1
-      let threshold = total / float64(1'u64 shl (depth + 1))
-
-      v = if (b and (1'u8 shl bit)) == 0: v.left else: v.right
+      let goRight = (b and (1'u8 shl bit)) != 0
+      v = step(v, goRight, create = false)
 
       if v.isNil:
-        return (float64(score) / 32.0)
+        return (float64(score) / float64(nBits))
 
       if float64(v.counter) > threshold:
         score += 1
 
-  (float64(score) / 32.0)
+      threshold *= 0.5
+
+  (float64(score) / float64(nBits))

tests/libp2p/service_discovery/component/test_advertise_discover.nim

@@ -391,22 +391,24 @@ suite "Service Discovery Component - Advertise Discover":
     checkUntilTimeout:
       lateRegistrar.countAdsInCache(serviceId) == 1
 
-  asyncTest "an IPv6 node registers three services one after another":
-    # TODO: vacp2p/nim-libp2p#2756
-    # A node listening only on IPv6 advertises only IPv6 addresses, which never
-    # enter the IPv4-only IP tree, so its advertisements draw no IP-similarity
-    # waiting time: each is Confirmed and cached, even under default settings.
+  asyncTest "an IPv6-only node is subject to IP-similarity throttling like an IPv4 node":
     let registrarNode = setupServiceDiscoveryNode()
     let advertiserNode = setupServiceDiscoveryNode(addresses = @[TcpAutoAddressIP6()])
     startAndDeferStop(@[registrarNode, advertiserNode])
     await connect(registrarNode, advertiserNode)
 
-    for i in 1 .. 3:
-      let service = makeServiceInfo("ipv6-service-" & $i)
-      let serviceId = service.id.hashServiceId()
-      advertiserNode.addProvidedService(service)
-      checkUntilTimeout:
-        registrarNode.countAdsInCache(serviceId) == 1
+    let firstService = makeServiceInfo("ipv6-service-1")
+    let firstServiceId = firstService.id.hashServiceId()
+    advertiserNode.addProvidedService(firstService)
+    checkUntilTimeout:
+      registrarNode.countAdsInCache(firstServiceId) == 1
+
+    let secondService = makeServiceInfo("ipv6-service-2")
+    let secondServiceId = secondService.id.hashServiceId()
+    advertiserNode.addProvidedService(secondService)
+    checkUntilTimeout:
+      registrarNode.registrar.boundService.hasKey(secondServiceId)
+    check registrarNode.countAdsInCache(secondServiceId) == 0
 
   asyncTest "a dual-stack node follows the normal path and is told to Wait":
     let registrarNode = setupServiceDiscoveryNode()

tests/libp2p/service_discovery/test_registrar.nim

@@ -637,7 +637,7 @@ suite "Service Discovery Registrar - Edge Cases":
 
     check w >= ZeroDuration
 
-  test "waitingTime with IPv6 addresses only (ignored in IP tree)":
+  test "waitingTime with IPv6 addresses only, tree empty":
     let registrar = Registrar.new()
     let discoConfig = ServiceDiscoveryConfig.new()
     let serviceId = makeServiceId()
@@ -649,6 +649,26 @@ suite "Service Discovery Registrar - Edge Cases":
 
     check w >= ZeroDuration
 
+  test "waitingTime with IPv6 addresses contributes IP similarity":
+    let registrar = Registrar.new()
+    let discoConfig = ServiceDiscoveryConfig.new()
+    let serviceId = makeServiceId()
+
+    registrar.ipTree.insertIp(
+      IpAddress(
+        family: IpAddressFamily.IPv6,
+        address_v6: [0'u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
+      )
+    )
+
+    let ipv6Addr = MultiAddress.init("/ip6/::1/tcp/9000").get()
+    let ad = makeAdvertisement(addrs = @[ipv6Addr])
+    let now = Moment.now()
+
+    let w = registrar.waitingTime(discoConfig, ad, 1000, serviceId, now)
+
+    check w > ZeroDuration
+
   test "waitingTime with mixed IPv4 and IPv6 addresses":
     let registrar = Registrar.new()
     let discoConfig = ServiceDiscoveryConfig.new()