Merge branch 'master' into chore/kad/lookup-peers-before-insert

Author: gmelodie

libp2p/protocols/kademlia/find.nim

@@ -619,12 +619,30 @@ proc findPeer*(
 
   return ok(PeerInfo(peerId: target, addrs: kad.switch.peerStore[AddressBook][target]))
 
-proc findClosestPeers*(kad: KadDHT, target: Key): seq[Peer] =
-  let closestPeerKeys = kad.rtable.findClosest(target, kad.config.replication).filterIt(
-      it != kad.switch.peerInfo.peerId.toKey()
-    )
+proc findClosestPeers*(kad: KadDHT, target: Key, requester: PeerId): seq[Peer] =
+  ## Over-fetches by `excluded.len` so dropping self and `requester` still fills the reply.
+  let excluded = [kad.switch.peerInfo.peerId.toKey(), requester.toKey()]
+  let closestPeerKeys = kad.rtable
+    .findClosest(target, kad.config.replication + excluded.len)
+    .filterIt(it notin excluded)
+
+  return kad.switch.toPeers(
+    closestPeerKeys[0 ..< min(kad.config.replication, closestPeerKeys.len)]
+  )
+
+proc findNodeCloserPeers(kad: KadDHT, target: Key, requester: PeerId): seq[Peer] =
+  ## Also returns the target itself, which keeps client-mode peers resolvable.
+  let closest = kad.findClosestPeers(target, requester)
+  if target == requester.toKey():
+    return closest
+
+  let targetPeer = target.toPeer(kad.switch).valueOr:
+    return closest
+
+  if closest.len > 0 and closest[0].id == targetPeer.id:
+    return closest
 
-  return kad.switch.toPeers(closestPeerKeys)
+  return @[targetPeer] & closest
 
 method handleFindNode*(
     kad: KadDHT, stream: Stream, msg: Message
@@ -634,7 +652,8 @@ method handleFindNode*(
     return
 
   let response = Message(
-    msgType: Opt.some(MessageType.findNode), closerPeers: kad.findClosestPeers(msgKey)
+    msgType: Opt.some(MessageType.findNode),
+    closerPeers: kad.findNodeCloserPeers(msgKey, stream.peerId),
   )
   let encoded = response.encode(kad.config.hideConnectionStatus)
   kad_message_bytes_sent.inc(encoded.len.int64, labelValues = [$MessageType.findNode])

libp2p/protocols/kademlia/get.nim

@@ -177,7 +177,7 @@ method handleGetValue*(
     let response = Message(
       msgType: Opt.some(MessageType.getValue),
       key: Opt.some(key),
-      closerPeers: kad.findClosestPeers(key),
+      closerPeers: kad.findClosestPeers(key, stream.peerId),
     )
     let encoded = response.encode(kad.config.hideConnectionStatus)
     kad_message_bytes_sent.inc(encoded.len.int64, labelValues = [$MessageType.getValue])
@@ -197,7 +197,7 @@ method handleGetValue*(
         timeReceived: Opt.some(entryRecord.time),
       )
     ),
-    closerPeers: kad.findClosestPeers(key),
+    closerPeers: kad.findClosestPeers(key, stream.peerId),
   )
   let encoded = response.encode(kad.config.hideConnectionStatus)
   kad_message_bytes_sent.inc(encoded.len.int64, labelValues = [$MessageType.getValue])

libp2p/protocols/kademlia/provider.nim

@@ -391,7 +391,7 @@ proc handleGetProviders*(
   let response = Message(
     msgType: Opt.some(MessageType.getProviders),
     key: msg.key,
-    closerPeers: kad.findClosestPeers(msgKey),
+    closerPeers: kad.findClosestPeers(msgKey, stream.peerId),
     providerPeers: providers.toSeq(),
   )
   let encoded = response.encode(kad.config.hideConnectionStatus)

tests/libp2p/kademlia/test_find.nim

@@ -221,10 +221,11 @@ suite "KadDHT Find":
       kads[0].hasKey(kads[2].rtable.selfId)
 
   asyncTest "Find node with empty key returns closest peers":
-    let kads = setupKadSwitches(2)
+    let kads = setupKadSwitches(3)
     startAndDeferStop(kads)
 
-    await connect(kads[0], kads[1])
+    # Setup: kads[0] <-> kads[1], kads[0] <-> kads[2]
+    await connectHub(kads[0], kads[1 ..^ 1])
 
     # Send FIND_NODE with empty key directly
     let emptyKey: Key = @[]
@@ -235,9 +236,9 @@ suite "KadDHT Find":
     check:
       response.msgType == MessageType.findNode
       response.closerPeers.len == 1
-      response.closerPeers[0].id.get() == kads[1].rtable.selfId
+      response.closerPeers[0].id.get() == kads[2].rtable.selfId
 
-  asyncTest "Find node for own PeerID returns closest peers":
+  asyncTest "Find node for own PeerID excludes the requester":
     let kads = setupKadSwitches(3)
     startAndDeferStop(kads)
 
@@ -252,9 +253,78 @@ suite "KadDHT Find":
     let closerPeersIds = response.closerPeers.mapIt(it.id.get())
     check:
       response.msgType == MessageType.findNode
-      # kads[0] knows kads[1] and kads[2], should return both as closest peers
-      response.closerPeers.len == 2
-      kads[1].rtable.selfId in closerPeersIds
+      response.closerPeers.len == 1
+      kads[1].rtable.selfId notin closerPeersIds
+      kads[2].rtable.selfId in closerPeersIds
+
+  asyncTest "Find node for a known target returns it once":
+    let kads = setupKadSwitches(3)
+    startAndDeferStop(kads)
+
+    await connectHub(kads[0], kads[1 ..^ 1])
+
+    let response = (
+      await kads[1].dispatchFindNode(
+        kads[0].switch.peerInfo.peerId, kads[2].rtable.selfId
+      )
+    ).value()
+
+    check:
+      response.closerPeers.len == 1
+      response.closerPeers[0].id.get() == kads[2].rtable.selfId
+
+  asyncTest "Find node for an unknown target omits it":
+    let kads = setupKadSwitches(3)
+    startAndDeferStop(kads)
+
+    await connectHub(kads[0], kads[1 ..^ 1])
+
+    let unknownTarget = randomPeerId().toKey()
+    let response = (
+      await kads[1].dispatchFindNode(kads[0].switch.peerInfo.peerId, unknownTarget)
+    ).value()
+
+    let closerPeersIds = response.closerPeers.mapIt(it.id.get())
+    check:
+      unknownTarget notin closerPeersIds
+      kads[2].rtable.selfId in closerPeersIds
+
+  asyncTest "Find node returns a target known only from the address book":
+    let kads = setupKadSwitches(3)
+    startAndDeferStop(kads)
+
+    await connectHub(kads[0], kads[1 ..^ 1])
+
+    # A client-mode peer is in nobody's routing table, but its addresses are known.
+    let client = randomPeerId()
+    kads[0].switch.peerStore[AddressBook][client] =
+      @[MultiAddress.init("/ip4/127.0.0.1/tcp/9999").tryGet()]
+
+    let response = (
+      await kads[1].dispatchFindNode(kads[0].switch.peerInfo.peerId, client.toKey())
+    ).value()
+
+    check:
+      not kads[0].hasKey(client.toKey())
+      response.closerPeers[0].id.get() == client.toKey()
+
+  asyncTest "Find node excludes the requester for an unrelated target":
+    let kads = setupKadSwitches(3)
+    startAndDeferStop(kads)
+
+    await connectHub(kads[0], kads[1 ..^ 1])
+
+    # kads[0] knows kads[1] and kads[2], and both fit in a reply, so pre-filter
+    # kads[1] would get itself back even though the target is unrelated to it.
+    let response = (
+      await kads[1].dispatchFindNode(
+        kads[0].switch.peerInfo.peerId, randomPeerId().toKey()
+      )
+    ).value()
+
+    let closerPeersIds = response.closerPeers.mapIt(it.id.get())
+    check:
+      kads[1].rtable.selfId notin closerPeersIds
       kads[2].rtable.selfId in closerPeersIds
 
   asyncTest "Find node with empty routing table returns empty result":