diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim index a944f0b..fe538cc 100644 --- a/eth/p2p/discoveryv5/protocol.nim +++ b/eth/p2p/discoveryv5/protocol.nim @@ -73,7 +73,7 @@ ## This might be a concern for mobile devices. import - std/[tables, sets, options, math, sequtils], + std/[tables, sets, options, math, sequtils, algorithm], stew/shims/net as stewNet, json_serialization/std/net, stew/endians2, chronicles, chronos, stint, bearssl, eth/[rlp, keys, async_utils], @@ -633,6 +633,7 @@ proc lookupWorker(d: Protocol, destNode: Node, target: NodeId): let dists = lookupDistances(target, destNode.id) var i = 0 # TODO: We can make use of the multiple distances here now. + # Do findNode requests with different distances until we hit limits. while i < lookupRequestLimit and result.len < findNodeResultLimit: let r = await d.findNode(destNode, @[dists[i]]) # TODO: Handle failures better. E.g. stop on different failures than timeout @@ -648,21 +649,25 @@ proc lookup*(d: Protocol, target: NodeId): Future[seq[Node]] {.async, raises: [Exception, Defect].} = ## Perform a lookup for the given target, return the closest n nodes to the ## target. Maximum value for n is `BUCKET_SIZE`. - # TODO: Sort the returned nodes on distance - # Also use unseen nodes as a form of validation. - result = d.routingTable.neighbours(target, BUCKET_SIZE, seenOnly = false) - var asked = initHashSet[NodeId]() - asked.incl(d.localNode.id) - var seen = asked - for node in result: + # `closestNodes` holds the k closest nodes to target found, sorted by distance + # Unvalidated nodes are used for requests as a form of validation. + var closestNodes = d.routingTable.neighbours(target, BUCKET_SIZE, + seenOnly = false) + + var asked, seen = initHashSet[NodeId]() + asked.incl(d.localNode.id) # No need to ask our own node + seen.incl(d.localNode.id) # No need to discover our own node + for node in closestNodes: seen.incl(node.id) var pendingQueries = newSeqOfCap[Future[seq[Node]]](alpha) while true: var i = 0 - while i < result.len and pendingQueries.len < alpha: - let n = result[i] + # Doing `alpha` amount of requests at once as long as closer non queried + # nodes are discovered. + while i < closestNodes.len and pendingQueries.len < alpha: + let n = closestNodes[i] if not asked.containsOrIncl(n.id): pendingQueries.add(d.lookupWorker(n, target)) inc i @@ -682,10 +687,19 @@ proc lookup*(d: Protocol, target: NodeId): Future[seq[Node]] error "Resulting query should have beeen in the pending queries" let nodes = query.read + # TODO: Remove node on timed-out query? for n in nodes: if not seen.containsOrIncl(n.id): - if result.len < BUCKET_SIZE: - result.add(n) + # If it wasn't seen before, insert node while remaining sorted + closestNodes.insert(n, closestNodes.lowerBound(n, + proc(x: Node, n: Node): int = + cmp(distanceTo(x, target), distanceTo(n, target)) + )) + + if closestNodes.len > BUCKET_SIZE: + closestNodes.del(closestNodes.high()) + + return closestNodes proc lookupRandom*(d: Protocol): Future[seq[Node]] {.async, raises:[Exception, Defect].} = diff --git a/eth/p2p/discoveryv5/routing_table.nim b/eth/p2p/discoveryv5/routing_table.nim index 98ecebe..00cb68c 100644 --- a/eth/p2p/discoveryv5/routing_table.nim +++ b/eth/p2p/discoveryv5/routing_table.nim @@ -87,7 +87,7 @@ const DefaultTableIpLimits* = TableIpLimits(tableIpLimit: DefaultTableIpLimit, bucketIpLimit: DefaultBucketIpLimit) -proc distanceTo(n: Node, id: NodeId): UInt256 = +proc distanceTo*(n: Node, id: NodeId): UInt256 = ## Calculate the distance to a NodeId. n.id xor id diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim index 98cef69..782ba81 100644 --- a/tests/p2p/test_discoveryv5.nim +++ b/tests/p2p/test_discoveryv5.nim @@ -256,8 +256,7 @@ procSuite "Discovery v5 Tests": let target = nodes[i] let discovered = await nodes[nodeCount-1].lookup(target.localNode.id) debug "Lookup result", target = target.localNode, discovered - # if lookUp would return ordered on distance we could check discovered[0] - check discovered.contains(target.localNode) + check discovered[0] == target.localNode for node in nodes: await node.closeWait()