test(gossipsub): control message (#1191)

Co-authored-by: Radoslaw Kaminski <radoslaw@status.im>
2026-01-09 02:38:19 -05:00 · 2025-05-13 16:54:07 +02:00
parent 4d0b4ecc22
commit 9997f3e3d3
3 changed files with 339 additions and 10 deletions
--- a/tests/pubsub/testgossipsubgossip.nim
+++ b/tests/pubsub/testgossipsubgossip.nim
@@ -715,3 +715,130 @@ suite "GossipSub Gossip Protocol":
      gossip0.mesh.getOrDefault("foobar").toSeq[0].codec == GossipSubCodec_10
    checkUntilTimeout:
      gossip1.mesh.getOrDefault("foobar").toSeq[0].codec == GossipSubCodec_10
+
+  asyncTest "IHAVE messages correctly advertise message ID to peers":
+    # Given 2 nodes
+    let
+      topic = "foo"
+      messageID = @[0'u8, 1, 2, 3]
+      ihaveMessage =
+        ControlMessage(ihave: @[ControlIHave(topicID: topic, messageIDs: @[messageID])])
+      numberOfNodes = 2
+      nodes = generateNodes(numberOfNodes, gossip = true, verifySignature = false)
+        .toGossipSub()
+      n0 = nodes[0]
+      n1 = nodes[1]
+
+    startNodesAndDeferStop(nodes)
+
+    # Given node1 has an IHAVE observer
+    var receivedIHave = newFuture[(string, seq[MessageId])]()
+    let checkForIhaves = proc(peer: PubSubPeer, msgs: var RPCMsg) =
+      if msgs.control.isSome:
+        for msg in msgs.control.get.ihave:
+          receivedIHave.complete((msg.topicID, msg.messageIDs))
+    n1.addObserver(PubSubObserver(onRecv: checkForIhaves))
+
+    # And the nodes are connected
+    await connectNodesStar(nodes)
+
+    # And both subscribe to the topic
+    subscribeAllNodes(nodes, topic, voidTopicHandler)
+    await waitForHeartbeat()
+
+    check:
+      n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+
+    # When an IHAVE message is sent
+    let p1 = n0.getOrCreatePeer(n1.peerInfo.peerId, @[GossipSubCodec_12])
+    n0.broadcast(@[p1], RPCMsg(control: some(ihaveMessage)), isHighPriority = false)
+    await waitForHeartbeat()
+
+    # Then the peer has the message ID
+    let r = await receivedIHave.waitForState(HEARTBEAT_TIMEOUT)
+    check:
+      r.isCompleted((topic, @[messageID]))
+
+  asyncTest "IWANT messages correctly request messages by their IDs":
+    # Given 2 nodes
+    let
+      topic = "foo"
+      messageID = @[0'u8, 1, 2, 3]
+      iwantMessage = ControlMessage(iwant: @[ControlIWant(messageIDs: @[messageID])])
+      numberOfNodes = 2
+      nodes = generateNodes(numberOfNodes, gossip = true, verifySignature = false)
+        .toGossipSub()
+      n0 = nodes[0]
+      n1 = nodes[1]
+
+    startNodesAndDeferStop(nodes)
+
+    # Given node1 has an IWANT observer
+    var receivedIWant = newFuture[seq[MessageId]]()
+    let checkForIwants = proc(peer: PubSubPeer, msgs: var RPCMsg) =
+      if msgs.control.isSome:
+        for msg in msgs.control.get.iwant:
+          receivedIWant.complete(msg.messageIDs)
+    n1.addObserver(PubSubObserver(onRecv: checkForIwants))
+
+    # And the nodes are connected
+    await connectNodesStar(nodes)
+
+    # And both subscribe to the topic
+    subscribeAllNodes(nodes, topic, voidTopicHandler)
+    await waitForHeartbeat()
+
+    check:
+      n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+
+    # When an IWANT message is sent
+    let p1 = n0.getOrCreatePeer(n1.peerInfo.peerId, @[GossipSubCodec_12])
+    n0.broadcast(@[p1], RPCMsg(control: some(iwantMessage)), isHighPriority = false)
+    await waitForHeartbeat()
+
+    # Then the peer has the message ID
+    let r = await receivedIWant.waitForState(HEARTBEAT_TIMEOUT)
+    check:
+      r.isCompleted(@[messageID])
+
+  asyncTest "IHAVE for message not held by peer triggers IWANT response to sender":
+    # Given 2 nodes
+    let
+      topic = "foo"
+      messageID = @[0'u8, 1, 2, 3]
+      ihaveMessage =
+        ControlMessage(ihave: @[ControlIHave(topicID: topic, messageIDs: @[messageID])])
+      numberOfNodes = 2
+      nodes = generateNodes(numberOfNodes, gossip = true, verifySignature = false)
+        .toGossipSub()
+      n0 = nodes[0]
+      n1 = nodes[1]
+
+    startNodesAndDeferStop(nodes)
+
+    # Given node1 has an IWANT observer
+    var receivedIWant = newFuture[seq[MessageId]]()
+    let checkForIwants = proc(peer: PubSubPeer, msgs: var RPCMsg) =
+      if msgs.control.isSome:
+        for msg in msgs.control.get.iwant:
+          receivedIWant.complete(msg.messageIDs)
+    n0.addObserver(PubSubObserver(onRecv: checkForIwants))
+
+    # And the nodes are connected
+    await connectNodesStar(nodes)
+
+    # And both nodes subscribe to the topic
+    subscribeAllNodes(nodes, topic, voidTopicHandler)
+    await waitForHeartbeat()
+
+    # When an IHAVE message is sent from node0
+    let p1 = n0.getOrCreatePeer(n1.peerInfo.peerId, @[GossipSubCodec_12])
+    n0.broadcast(@[p1], RPCMsg(control: some(ihaveMessage)), isHighPriority = false)
+    await waitForHeartbeat()
+
+    # Then node0 should receive an IWANT message from node1 (as node1 doesn't have the message)
+    let iWantResult = await receivedIWant.waitForState(HEARTBEAT_TIMEOUT)
+    check:
+      iWantResult.isCompleted(@[messageID])
--- a/tests/pubsub/testgossipsubmeshmanagement.nim
+++ b/tests/pubsub/testgossipsubmeshmanagement.nim
@@ -520,3 +520,200 @@ suite "GossipSub Mesh Management":
    nodes[1].subscribe("foobar", handler)

    await invalidDetected.wait(10.seconds)
+
+  asyncTest "GRAFT messages correctly add peers to mesh":
+    # Given 2 nodes
+    let
+      topic = "foobar"
+      graftMessage = ControlMessage(graft: @[ControlGraft(topicID: topic)])
+      numberOfNodes = 2
+      # First part of the hack: Weird dValues so peers are not GRAFTed automatically
+      dValues = DValues(dLow: some(0), dHigh: some(0), d: some(0), dOut: some(-1))
+      nodes = generateNodes(
+          numberOfNodes, gossip = true, verifySignature = false, dValues = some(dValues)
+        )
+        .toGossipSub()
+      n0 = nodes[0]
+      n1 = nodes[1]
+
+    startNodesAndDeferStop(nodes)
+
+    # And the nodes are connected
+    await connectNodesStar(nodes)
+
+    # And both subscribe to the topic
+    subscribeAllNodes(nodes, topic, voidTopicHandler)
+    await waitForHeartbeat()
+
+    # Because of the hack-ish dValues, the peers are added to gossipsub but not GRAFTed to mesh
+    check:
+      n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      not n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      not n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
+
+    # Stop both nodes in order to prevent GRAFT message to be sent by heartbeat 
+    await n0.stop()
+    await n1.stop()
+
+    # Second part of the hack
+    # Set values so peers can be GRAFTed
+    let newDValues =
+      some(DValues(dLow: some(1), dHigh: some(1), d: some(1), dOut: some(1)))
+    n0.parameters.applyDValues(newDValues)
+    n1.parameters.applyDValues(newDValues)
+
+    # When a GRAFT message is sent
+    let p0 = n1.getOrCreatePeer(n0.peerInfo.peerId, @[GossipSubCodec_12])
+    let p1 = n0.getOrCreatePeer(n1.peerInfo.peerId, @[GossipSubCodec_12])
+    n0.broadcast(@[p1], RPCMsg(control: some(graftMessage)), isHighPriority = false)
+    n1.broadcast(@[p0], RPCMsg(control: some(graftMessage)), isHighPriority = false)
+
+    await waitForPeersInTable(
+      nodes, topic, newSeqWith(numberOfNodes, 1), PeerTableType.Mesh
+    )
+
+    # Then the peers are GRAFTed
+    check:
+      n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
+
+  asyncTest "Received GRAFT for non-subscribed topic":
+    # Given 2 nodes
+    let
+      topic = "foo"
+      graftMessage = ControlMessage(graft: @[ControlGraft(topicID: topic)])
+      numberOfNodes = 2
+      nodes = generateNodes(numberOfNodes, gossip = true, verifySignature = false)
+        .toGossipSub()
+      n0 = nodes[0]
+      n1 = nodes[1]
+
+    startNodesAndDeferStop(nodes)
+
+    # And the nodes are connected
+    await connectNodesStar(nodes)
+
+    # And only node0 subscribes to the topic
+    nodes[0].subscribe(topic, voidTopicHandler)
+    await waitForHeartbeat()
+
+    check:
+      n0.topics.hasKey(topic)
+      not n1.topics.hasKey(topic)
+      not n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      not n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      not n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
+
+    # When a GRAFT message is sent
+    let p1 = n0.getOrCreatePeer(n1.peerInfo.peerId, @[GossipSubCodec_12])
+    n0.broadcast(@[p1], RPCMsg(control: some(graftMessage)), isHighPriority = false)
+    await waitForHeartbeat()
+
+    # Then the peer is not GRAFTed
+    check:
+      n0.topics.hasKey(topic)
+      not n1.topics.hasKey(topic)
+      not n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      not n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      not n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
+
+  asyncTest "PRUNE messages correctly removes peers from mesh":
+    # Given 2 nodes
+    let
+      topic = "foo"
+      backoff = 1
+      pruneMessage = ControlMessage(
+        prune: @[ControlPrune(topicID: topic, peers: @[], backoff: uint64(backoff))]
+      )
+      numberOfNodes = 2
+      nodes = generateNodes(numberOfNodes, gossip = true, verifySignature = false)
+        .toGossipSub()
+      n0 = nodes[0]
+      n1 = nodes[1]
+
+    startNodesAndDeferStop(nodes)
+
+    # And the nodes are connected
+    await connectNodesStar(nodes)
+
+    # And both subscribe to the topic
+    subscribeAllNodes(nodes, topic, voidTopicHandler)
+    await waitForHeartbeat()
+
+    check:
+      n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
+
+    # When a PRUNE message is sent
+    let p1 = n0.getOrCreatePeer(n1.peerInfo.peerId, @[GossipSubCodec_12])
+    n0.broadcast(@[p1], RPCMsg(control: some(pruneMessage)), isHighPriority = false)
+    await waitForHeartbeat()
+
+    # Then the peer is PRUNEd
+    check:
+      n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      not n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
+
+    # When another PRUNE message is sent
+    let p0 = n1.getOrCreatePeer(n0.peerInfo.peerId, @[GossipSubCodec_12])
+    n1.broadcast(@[p0], RPCMsg(control: some(pruneMessage)), isHighPriority = false)
+    await waitForHeartbeat()
+
+    # Then the peer is PRUNEd
+    check:
+      n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      not n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      not n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
+
+  asyncTest "Received PRUNE for non-subscribed topic":
+    # Given 2 nodes
+    let
+      topic = "foo"
+      pruneMessage =
+        ControlMessage(prune: @[ControlPrune(topicID: topic, peers: @[], backoff: 1)])
+      numberOfNodes = 2
+      nodes = generateNodes(numberOfNodes, gossip = true, verifySignature = false)
+        .toGossipSub()
+      n0 = nodes[0]
+      n1 = nodes[1]
+
+    startNodesAndDeferStop(nodes)
+
+    # And the nodes are connected
+    await connectNodesStar(nodes)
+
+    # And only node0 subscribes to the topic
+    n0.subscribe(topic, voidTopicHandler)
+    await waitForHeartbeat()
+
+    check:
+      n0.topics.hasKey(topic)
+      not n1.topics.hasKey(topic)
+      not n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      not n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      not n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
+
+    # When a PRUNE message is sent
+    let p1 = n0.getOrCreatePeer(n1.peerInfo.peerId, @[GossipSubCodec_12])
+    n0.broadcast(@[p1], RPCMsg(control: some(pruneMessage)), isHighPriority = false)
+    await waitForHeartbeat()
+
+    # Then the peer is not PRUNEd
+    check:
+      n0.topics.hasKey(topic)
+      not n1.topics.hasKey(topic)
+      not n0.gossipsub.hasPeerId(topic, n1.peerInfo.peerId)
+      n1.gossipsub.hasPeerId(topic, n0.peerInfo.peerId)
+      not n0.mesh.hasPeerId(topic, n1.peerInfo.peerId)
+      not n1.mesh.hasPeerId(topic, n0.peerInfo.peerId)
--- a/tests/pubsub/utils.nim
+++ b/tests/pubsub/utils.nim
@@ -78,7 +78,7 @@ func defaultMsgIdProvider*(m: Message): Result[MessageId, ValidationResult] =
      $m.data.hash & $m.topic.hash
  ok mid.toBytes()

-proc applyDValues(parameters: var GossipSubParams, dValues: Option[DValues]) =
+proc applyDValues*(parameters: var GossipSubParams, dValues: Option[DValues]) =
  if dValues.isNone:
    return
  let values = dValues.get
@@ -168,18 +168,21 @@ proc generateNodes*(
    switch.mount(pubsub)
    result.add(pubsub)

-proc connectNodes*(dialer: PubSub, target: PubSub) {.async.} =
+proc toGossipSub*(nodes: seq[PubSub]): seq[GossipSub] =
+  return nodes.mapIt(GossipSub(it))
+
+proc connectNodes*[T: PubSub](dialer: T, target: T) {.async.} =
  doAssert dialer.switch.peerInfo.peerId != target.switch.peerInfo.peerId,
    "Could not connect same peer"
  await dialer.switch.connect(target.peerInfo.peerId, target.peerInfo.addrs)

-proc connectNodesStar*(nodes: seq[PubSub]) {.async.} =
+proc connectNodesStar*[T: PubSub](nodes: seq[T]) {.async.} =
  for dialer in nodes:
    for node in nodes:
      if dialer.switch.peerInfo.peerId != node.switch.peerInfo.peerId:
        await connectNodes(dialer, node)

-proc connectNodesSparse*(nodes: seq[PubSub], degree: int = 2) {.async.} =
+proc connectNodesSparse*[T: PubSub](nodes: seq[T], degree: int = 2) {.async.} =
  if nodes.len < degree:
    raise
      (ref CatchableError)(msg: "nodes count needs to be greater or equal to degree!")
@@ -324,23 +327,25 @@ proc waitForPeersInTable*(
    )
    allSatisfied = checkState(nodes, topic, peerCounts, table, satisfied)

-proc startNodes*(nodes: seq[PubSub]) {.async.} =
+proc startNodes*[T: PubSub](nodes: seq[T]) {.async.} =
  await allFuturesThrowing(nodes.mapIt(it.switch.start()))

-proc stopNodes*(nodes: seq[PubSub]) {.async.} =
+proc stopNodes*[T: PubSub](nodes: seq[T]) {.async.} =
  await allFuturesThrowing(nodes.mapIt(it.switch.stop()))

-template startNodesAndDeferStop*(nodes: seq[PubSub]): untyped =
+template startNodesAndDeferStop*[T: PubSub](nodes: seq[T]): untyped =
  await startNodes(nodes)
  defer:
    await stopNodes(nodes)

-proc subscribeAllNodes*(nodes: seq[PubSub], topic: string, topicHandler: TopicHandler) =
+proc subscribeAllNodes*[T: PubSub](
+    nodes: seq[T], topic: string, topicHandler: TopicHandler
+) =
  for node in nodes:
    node.subscribe(topic, topicHandler)

-proc subscribeAllNodes*(
-    nodes: seq[PubSub], topic: string, topicHandlers: seq[TopicHandler]
+proc subscribeAllNodes*[T: PubSub](
+    nodes: seq[T], topic: string, topicHandlers: seq[TopicHandler]
 ) =
  if nodes.len != topicHandlers.len:
    raise (ref CatchableError)(msg: "nodes and topicHandlers count needs to match!")