Support findnode with multiple distances in discv5.1

eth/p2p/discoveryv5/dcli.nim

@@ -4,12 +4,14 @@ import
   stew/byteutils, confutils/std/net,
   eth/keys, eth/net/nat, enr, node
 
+### This is all just temporary to be compatible with both versions
 const UseDiscv51* {.booldefine.} = false
 
 when UseDiscv51:
   import protocolv1
 else:
   import protocol
+###
 
 type
   DiscoveryCmd* = enum
@@ -173,7 +175,10 @@ proc run(config: DiscoveryConf) =
     else:
       echo "No Pong message returned"
   of findnode:
-    let nodes = waitFor d.findNode(config.findNodeTarget, config.distance)
+    when UseDiscv51:
+      let nodes = waitFor d.findNode(config.findNodeTarget, @[config.distance])
+    else:
+      let nodes = waitFor d.findNode(config.findNodeTarget, config.distance)
     if nodes.isOk():
       echo "Received valid records:"
       for node in nodes[]:

eth/p2p/discoveryv5/protocolv1.nim

@@ -267,12 +267,22 @@ proc handlePing(d: Protocol, fromId: NodeId, fromAddr: Address,
 
 proc handleFindNode(d: Protocol, fromId: NodeId, fromAddr: Address,
     fn: FindNodeMessage, reqId: RequestId) =
-  if fn.distance == 0:
+  if fn.distances.len == 0:
+    d.sendNodes(fromId, fromAddr, reqId, [])
+  elif fn.distances.contains(0):
+    # A request for our own record.
+    # It would be a weird request if there are more distances next to 0
+    # requested, so in this case lets just pass only our own. TODO: OK?
     d.sendNodes(fromId, fromAddr, reqId, [d.localNode])
   else:
-    let distance = min(fn.distance, 256)
+    # TODO: Still deduplicate also?
-    d.sendNodes(fromId, fromAddr, reqId,
+    if fn.distances.all(proc (x: uint32): bool = return x <= 256):
-      d.routingTable.neighboursAtDistance(distance, seenOnly = true))
+      d.sendNodes(fromId, fromAddr, reqId,
+        d.routingTable.neighboursAtDistances(fn.distances, seenOnly = true))
+    else:
+      # At least one invalid distance, but the polite node we are, still respond
+      # with empty nodes.
+      d.sendNodes(fromId, fromAddr, reqId, [])
 
 proc handleMessage(d: Protocol, srcId: NodeId, fromAddr: Address,
     message: Message) {.raises:[Exception].} =
@@ -473,7 +483,7 @@ proc waitMessage(d: Protocol, fromNode: Node, reqId: RequestId):
   d.awaitedMessages[key] = result
 
 proc verifyNodesRecords*(enrs: openarray[Record], fromNode: Node,
-    distance: uint32): seq[Node] {.raises: [Defect].} =
+    distances: varargs[uint32]): seq[Node] {.raises: [Defect].} =
   ## Verify and convert ENRs to a sequence of nodes. Only ENRs that pass
   ## verification will be added. ENRs are verified for duplicates, invalid
   ## addresses and invalid distances.
@@ -500,11 +510,12 @@ proc verifyNodesRecords*(enrs: openarray[Record], fromNode: Node,
         trace "Nodes reply contained record with invalid ip-address",
           record = n.record.toURI, sender = fromNode.record.toURI, node = $n
         continue
-      # Check if returned node has the requested distance.
+      # Check if returned node has one of the requested distances.
-      if logDist(n.id, fromNode.id) != min(distance, 256):
+      if not distances.contains(logDist(n.id, fromNode.id)):
         warn "Nodes reply contained record with incorrect distance",
           record = n.record.toURI, sender = fromNode.record.toURI
         continue
+
       # No check on UDP port and thus any port is allowed, also the so called
       # "well-known" ports.
 
@@ -564,17 +575,17 @@ proc ping*(d: Protocol, toNode: Node):
     d.replaceNode(toNode)
     return err("Pong message not received in time")
 
-proc findNode*(d: Protocol, toNode: Node, distance: uint32):
+proc findNode*(d: Protocol, toNode: Node, distances: seq[uint32]):
     Future[DiscResult[seq[Node]]] {.async, raises: [Exception, Defect].} =
   ## Send a discovery findNode message.
   ##
   ## Returns the received nodes or an error.
   ## Received ENRs are already validated and converted to `Node`.
-  let reqId = d.sendMessage(toNode, FindNodeMessage(distance: distance))
+  let reqId = d.sendMessage(toNode, FindNodeMessage(distances: distances))
   let nodes = await d.waitNodes(toNode, reqId)
 
   if nodes.isOk:
-    let res = verifyNodesRecords(nodes.get(), toNode, distance)
+    let res = verifyNodesRecords(nodes.get(), toNode, distances)
     d.routingTable.setJustSeen(toNode)
     return ok(res)
   else:
@@ -596,8 +607,9 @@ proc lookupWorker(d: Protocol, destNode: Node, target: NodeId):
     Future[seq[Node]] {.async, raises: [Exception, Defect].} =
   let dists = lookupDistances(target, destNode.id)
   var i = 0
+  # TODO: We can make use of the multiple distances here now.
   while i < lookupRequestLimit and result.len < findNodeResultLimit:
-    let r = await d.findNode(destNode, dists[i])
+    let r = await d.findNode(destNode, @[dists[i]])
     # TODO: Handle failures better. E.g. stop on different failures than timeout
     if r.isOk:
       # TODO: I guess it makes sense to limit here also to `findNodeResultLimit`?
@@ -667,7 +679,7 @@ proc resolve*(d: Protocol, id: NodeId): Future[Option[Node]]
 
   let node = d.getNode(id)
   if node.isSome():
-    let request = await d.findNode(node.get(), 0)
+    let request = await d.findNode(node.get(), @[0'u32])
 
     # TODO: Handle failures better. E.g. stop on different failures than timeout
     if request.isOk() and request[].len > 0:
@@ -690,7 +702,7 @@ proc revalidateNode*(d: Protocol, n: Node)
   if pong.isOK():
     if pong.get().enrSeq > n.record.seqNum:
       # Request new ENR
-      let nodes = await d.findNode(n, 0)
+      let nodes = await d.findNode(n, @[0'u32])
       if nodes.isOk() and nodes[].len > 0:
         discard d.addNode(nodes[][0])
 

eth/p2p/discoveryv5/routing_table.nim

@@ -321,6 +321,20 @@ proc neighboursAtDistance*(r: RoutingTable, distance: uint32,
   # that are exactly the requested distance.
   keepIf(result, proc(n: Node): bool = logDist(n.id, r.thisNode.id) == distance)
 
+proc neighboursAtDistances*(r: RoutingTable, distances: seq[uint32],
+    k: int = BUCKET_SIZE, seenOnly = false): seq[Node] =
+  ## Return up to k neighbours at given logarithmic distances.
+  # TODO: This will currently return nodes with neighbouring distances on the
+  # first one prioritize. It might end up not including all the node distances
+  # requested. Need to rework the logic here and not use the neighbours call.
+  if distances.len > 0:
+    result = r.neighbours(idAtDistance(r.thisNode.id, distances[0]), k,
+      seenOnly)
+    # This is a bit silly, first getting closest nodes then to only keep the ones
+    # that are exactly the requested distances.
+    keepIf(result, proc(n: Node): bool =
+      distances.contains(logDist(n.id, r.thisNode.id)))
+
 proc len*(r: RoutingTable): int =
   for b in r.buckets: result += b.len
 

eth/p2p/discoveryv5/typesv1.nim

@@ -43,7 +43,7 @@ type
     port*: uint16
 
   FindNodeMessage* = object
-    distance*: uint32
+    distances*: seq[uint32]
 
   NodesMessage* = object
     total*: uint32

tests/p2p/test_discoveryv5.nim

@@ -5,6 +5,8 @@ import
   eth/p2p/discoveryv5/[enr, node, routing_table],
   ./discv5_test_helper
 
+
+### This is all just temporary to be compatible with both versions
 const UseDiscv51* {.booldefine.} = false
 
 when UseDiscv51:
@@ -16,6 +18,12 @@ else:
     eth/p2p/discoveryv5/[types, encoding],
     eth/p2p/discoveryv5/protocol as discv5_protocol
 
+  proc findNode*(d: discv5_protocol.Protocol, toNode: Node, distances: seq[uint32]):
+      Future[DiscResult[seq[Node]]] =
+    if distances.len > 0:
+      return d.findNode(toNode, distances[0])
+###
+
 procSuite "Discovery v5 Tests":
   let rng = newRng()
 
@@ -189,7 +197,7 @@ procSuite "Discovery v5 Tests":
       check idAtDistance(targetId, d) == parse(id, UInt256, 16)
 
   asyncTest "FindNode Test":
-    const dist = 253
+    const dist = 253'u32
     let
       mainNodeKey = PrivateKey.fromHex(
         "a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a617")[]
@@ -209,14 +217,14 @@ procSuite "Discovery v5 Tests":
 
     # Get ENR of the node itself
     var discovered =
-      await discv5_protocol.findNode(testNode, mainNode.localNode, 0)
+      await findNode(testNode, mainNode.localNode, @[0'u32])
     check:
       discovered.isOk
       discovered[].len == 1
       discovered[][0] == mainNode.localNode
     # Get ENRs of nodes added at provided logarithmic distance
     discovered =
-      await discv5_protocol.findNode(testNode, mainNode.localNode, dist)
+      await findNode(testNode, mainNode.localNode, @[dist])
     check discovered.isOk
     check discovered[].len == 10
     for n in nodes:
@@ -224,14 +232,14 @@ procSuite "Discovery v5 Tests":
 
     # Too high logarithmic distance, should return no nodes.
     discovered =
-      await discv5_protocol.findNode(testNode, mainNode.localNode, 4294967295'u32)
+      await findNode(testNode, mainNode.localNode, @[4294967295'u32])
     check:
       discovered.isOk
       discovered[].len == 0
 
     # Logarithmic distance of 256 should only return the testNode
     discovered =
-      await discv5_protocol.findNode(testNode, mainNode.localNode, 256)
+      await findNode(testNode, mainNode.localNode, @[256'u32])
     check:
       discovered.isOk
       discovered[].len == 1
@@ -239,7 +247,7 @@ procSuite "Discovery v5 Tests":
 
     # Empty bucket
     discovered =
-      await discv5_protocol.findNode(testNode, mainNode.localNode, 254)
+      await findNode(testNode, mainNode.localNode, @[254'u32])
     check discovered.isOk
     check discovered[].len == 0
 
@@ -249,7 +257,7 @@ procSuite "Discovery v5 Tests":
 
     # Full bucket
     discovered =
-      await discv5_protocol.findNode(testNode, mainNode.localNode, dist)
+      await findNode(testNode, mainNode.localNode, @[dist])
     check discovered.isOk
     check discovered[].len == 16
 
@@ -276,8 +284,8 @@ procSuite "Discovery v5 Tests":
       testNode = initDiscoveryNode(
         rng, PrivateKey.random(rng[]), localAddress(20302),
         @[mainNode.localNode.record])
-      discovered = await discv5_protocol.findNode(testNode, mainNode.localNode,
+      discovered = await findNode(testNode, mainNode.localNode,
-        closestDistance)
+        @[closestDistance])
 
     check discovered.isOk
     check closest in discovered[]
@@ -355,7 +363,7 @@ procSuite "Discovery v5 Tests":
       # Request the target ENR and manually add it to the routing table.
       # Ping for handshake based ENR passing will not work as our previous
       # session will still be in the LRU cache.
-      let nodes = await mainNode.findNode(targetNode.localNode, 0)
+      let nodes = await mainNode.findNode(targetNode.localNode, @[0'u32])
       check:
         nodes.isOk()
         nodes[].len == 1
@@ -575,5 +583,5 @@ procSuite "Discovery v5 Tests":
         recordInvalidDistance = enr.Record.init(
           1, pk, some(ValidIpAddress.init("12.13.14.15")), port, port)[]
         records = [recordInvalidDistance]
-        test = verifyNodesRecords(records, fromNode, 0)
+        test = verifyNodesRecords(records, fromNode, 0'u32)
       check test.len == 0

tests/p2p/test_discv51_encoding.nim

@@ -26,9 +26,9 @@ suite "Discovery v5 Protocol Message Encodings":
 
   test "FindNode Request":
     var p: FindNodeMessage
-    p.distance = 0x0100
+    p.distances = @[0x0100'u32]
     var reqId: RequestId = 1
-    check encodeMessage(p, reqId).toHex == "03c401820100"
+    check encodeMessage(p, reqId).toHex == "03c501c3820100"
 
   test "Nodes Response (empty)":
     var p: NodesMessage