Fluffy: Use neighboursInRange in neighborhoodGossip and filter out src node in handleFindContent (#3273)

Author: bhartnett

fluffy/network/wire/portal_protocol.nim

@@ -360,11 +360,6 @@ func getNode*(p: PortalProtocol, id: NodeId): Opt[Node] =
 func localNode*(p: PortalProtocol): Node =
   p.baseProtocol.localNode
 
-template neighbours*(
-    p: PortalProtocol, id: NodeId, k: int = BUCKET_SIZE, seenOnly = false
-): seq[Node] =
-  p.routingTable.neighbours(id, k, seenOnly)
-
 func distance(p: PortalProtocol, a, b: NodeId): UInt256 =
   p.routingTable.distance(a, b)
 
@@ -380,6 +375,34 @@ func inRange(
 template inRange*(p: PortalProtocol, contentId: ContentId): bool =
   p.inRange(p.localNode.id, p.dataRadius(), contentId)
 
+func neighbours*(
+    p: PortalProtocol,
+    id: NodeId,
+    k: int = BUCKET_SIZE,
+    seenOnly = false,
+    excluding = initHashSet[NodeId](),
+): seq[Node] =
+  func nodeNotExcluded(nodeId: NodeId): bool =
+    not excluding.contains(nodeId)
+
+  p.routingTable.neighbours(id, k, seenOnly, nodeNotExcluded)
+
+func neighboursInRange*(
+    p: PortalProtocol,
+    id: ContentId,
+    k: int = BUCKET_SIZE,
+    seenOnly = false,
+    excluding = initHashSet[NodeId](),
+): seq[Node] =
+  func nodeNotExcludedAndInRange(nodeId: NodeId): bool =
+    if excluding.contains(nodeId):
+      return false
+    let radius = p.radiusCache.get(nodeId).valueOr:
+      return false
+    p.inRange(nodeId, radius, id)
+
+  p.routingTable.neighbours(id, k, seenOnly, nodeNotExcludedAndInRange)
+
 func truncateEnrs(
     nodes: seq[Node], maxSize: int, enrOverhead: int
 ): List[ByteList[2048], 32] =
@@ -542,8 +565,10 @@ proc handleFindContent(
 
   # Node does not have the content, or content is not even in radius,
   # send closest neighbours to the requested content id.
+
   let
-    closestNodes = p.neighbours(NodeId(contentId), seenOnly = true)
+    closestNodes =
+      p.neighbours(contentId, seenOnly = true, excluding = toHashSet([srcId]))
     enrs = truncateEnrs(closestNodes, maxPayloadSize, enrOverhead)
   portal_content_enrs_packed.observe(enrs.len().int64, labelValues = [$p.protocolId])
 
@@ -1751,28 +1776,24 @@ proc neighborhoodGossip*(
   # table, but at the same time avoid unnecessary node lookups.
   # It might still cause issues in data getting propagated in a wider id range.
 
+  var excluding: HashSet[NodeId]
+  if srcNodeId.isSome():
+    excluding.incl(srcNodeId.get())
+
   var closestLocalNodes =
-    p.routingTable.neighbours(NodeId(contentId), BUCKET_SIZE, seenOnly = true)
+    p.neighboursInRange(contentId, BUCKET_SIZE, seenOnly = true, excluding)
 
   # Shuffling the order of the nodes in order to not always hit the same node
   # first for the same request.
   p.baseProtocol.rng[].shuffle(closestLocalNodes)
 
-  var gossipNodes: seq[Node]
-  for node in closestLocalNodes:
-    let radius = p.radiusCache.get(node.id).valueOr:
-      continue
-    if p.inRange(node.id, radius, contentId):
-      if srcNodeId.isNone() or node.id != srcNodeId.get():
-        gossipNodes.add(node)
-
   var numberOfGossipedNodes = 0
 
-  if not enableNodeLookup or gossipNodes.len() >= p.config.maxGossipNodes:
+  if not enableNodeLookup or closestLocalNodes.len() >= p.config.maxGossipNodes:
     # use local nodes for gossip
     portal_gossip_without_lookup.inc(labelValues = [$p.protocolId])
 
-    for node in gossipNodes:
+    for node in closestLocalNodes:
       let req = OfferRequest(dst: node, kind: Direct, contentList: contentList)
       await p.offerQueue.addLast(req)
       inc numberOfGossipedNodes

fluffy/tests/wire_protocol_tests/test_portal_wire_protocol.nim

@@ -157,16 +157,24 @@ procSuite "Portal Wire Protocol Tests":
     await proto2.stopPortalProtocol()
 
   asyncTest "FindContent/Content - send enrs":
-    let (proto1, proto2) = defaultTestSetup(rng)
+    let
+      proto1 = initPortalProtocol(rng, PrivateKey.random(rng[]), localAddress(20402))
+      proto2 = initPortalProtocol(rng, PrivateKey.random(rng[]), localAddress(20403))
+      proto3 = initPortalProtocol(rng, PrivateKey.random(rng[]), localAddress(20404))
+
+    # Make node1 know about node2, and node2 about node3
+    check proto1.addNode(proto2.localNode) == Added
+    check proto2.addNode(proto3.localNode) == Added
 
-    # ping in one direction to add, ping in the other to update as seen.
-    check (await proto1.baseProtocol.ping(proto2.localNode)).isOk()
-    check (await proto2.baseProtocol.ping(proto1.localNode)).isOk()
+    # node1 needs to know the radius of the nodes to determine if they are
+    # interested in content, so a ping is done.
+    check (await proto1.ping(proto2.localNode)).isOk()
+    check (await proto2.ping(proto3.localNode)).isOk()
 
     let contentKey = ContentKeyByteList.init(@[1'u8])
 
     # content does not exist so this should provide us with the closest nodes
-    # to the content, which is the only node in the routing table.
+    # to the content, which should only be node 3 because node 1 should be excluded
     let content = await proto1.findContentImpl(proto2.localNode, contentKey)
 
     check: