feat: graft/prune events and mesh peer tagging (#383)

* reset to master

* feat: default to disabled tagging for graft / prune

* chore: linting fixes + restore package-lock.json

* refactor: place the tag function in sendGraft

* fix: ensure tags are pruned / grafted per topic

* test: added tagging tests to ensure value is properly set

* test: added test to ensure peer store is cleared of tags

* test: fix issue with flakey tagging test

* chore: remove vulnerabilities from pacakge-lock.json

* fix: handle error for peerStore merge

* fix: PR feedback adjustments

* refactor: make taggingMesh flag required

* chore: clean up feature

* chore: pr review

---------

Co-authored-by: Cayman <[email protected]>

Author: maschad

src/index.ts

@@ -63,7 +63,7 @@ import { removeFirstNItemsFromSet, removeItemsFromSet } from './utils/set.js'
 import { SimpleTimeCache } from './utils/time-cache.js'
 import type { GossipsubOptsSpec } from './config.js'
 import type {
-  Connection, Stream, PeerId, Peer, PeerStore,
+  Connection, Direction, Stream, PeerId, Peer, PeerStore,
   Message,
   PublishResult,
   PubSub,
@@ -189,6 +189,11 @@ export interface GossipsubOpts extends GossipsubOptsSpec, PubSubInit {
    * Limits to bound protobuf decoding
    */
   decodeRpcLimits?: DecodeRPCLimits
+
+  /**
+   * If true, will utilize the libp2p connection manager tagging system to prune/graft connections to peers, defaults to true
+   */
+  tagMeshPeers: boolean
 }
 
 export interface GossipsubMessage {
@@ -197,9 +202,17 @@ export interface GossipsubMessage {
   msg: Message
 }
 
+export interface MeshPeer {
+  peerId: string
+  topic: string
+  direction: Direction
+}
+
 export interface GossipsubEvents extends PubSubEvents {
   'gossipsub:heartbeat': CustomEvent
   'gossipsub:message': CustomEvent<GossipsubMessage>
+  'gossipsub:graft': CustomEvent<MeshPeer>
+  'gossipsub:prune': CustomEvent<MeshPeer>
 }
 
 enum GossipStatusCode {
@@ -408,6 +421,7 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
       fallbackToFloodsub: true,
       floodPublish: true,
       batchPublish: false,
+      tagMeshPeers: true,
       doPX: false,
       directPeers: [],
       D: constants.GossipsubD,
@@ -635,6 +649,11 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
         })
     }, constants.GossipsubDirectConnectInitialDelay)
 
+    if (this.opts.tagMeshPeers) {
+      this.addEventListener('gossipsub:graft', this.tagMeshPeer)
+      this.addEventListener('gossipsub:prune', this.untagMeshPeer)
+    }
+
     this.log('started')
   }
 
@@ -652,6 +671,11 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
     const { registrarTopologyIds } = this.status
     this.status = { code: GossipStatusCode.stopped }
 
+    if (this.opts.tagMeshPeers) {
+      this.removeEventListener('gossipsub:graft', this.tagMeshPeer)
+      this.removeEventListener('gossipsub:prune', this.untagMeshPeer)
+    }
+
     // unregister protocol and handlers
     const registrar = this.components.registrar
     await Promise.all(this.multicodecs.map(async (multicodec) => registrar.unhandle(multicodec)))
@@ -1507,6 +1531,7 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
       if (topicID == null) {
         return
       }
+
       const peersInMesh = this.mesh.get(topicID)
       if (peersInMesh == null) {
         // don't do PX when there is an unknown topic to avoid leaking our peers
@@ -1520,38 +1545,38 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
         return
       }
 
+      const backoffExpiry = this.backoff.get(topicID)?.get(id)
+
+      // This if/else chain contains the various cases of valid (and semi-valid) GRAFTs
+      // Most of these cases result in a PRUNE immediately being sent in response
+
       // we don't GRAFT to/from direct peers; complain loudly if this happens
       if (this.direct.has(id)) {
         this.log('GRAFT: ignoring request from direct peer %s', id)
         // this is possibly a bug from a non-reciprical configuration; send a PRUNE
         prune.push(topicID)
         // but don't px
         doPX = false
-        return
-      }
 
-      // make sure we are not backing off that peer
-      const expire = this.backoff.get(topicID)?.get(id)
-      if (typeof expire === 'number' && now < expire) {
+        // make sure we are not backing off that peer
+      } else if (typeof backoffExpiry === 'number' && now < backoffExpiry) {
         this.log('GRAFT: ignoring backed off peer %s', id)
         // add behavioral penalty
         this.score.addPenalty(id, 1, ScorePenalty.GraftBackoff)
         // no PX
         doPX = false
         // check the flood cutoff -- is the GRAFT coming too fast?
-        const floodCutoff = expire + this.opts.graftFloodThreshold - this.opts.pruneBackoff
+        const floodCutoff = backoffExpiry + this.opts.graftFloodThreshold - this.opts.pruneBackoff
         if (now < floodCutoff) {
           // extra penalty
           this.score.addPenalty(id, 1, ScorePenalty.GraftBackoff)
         }
         // refresh the backoff
         this.addBackoff(id, topicID)
         prune.push(topicID)
-        return
-      }
 
-      // check the score
-      if (score < 0) {
+        // check the score
+      } else if (score < 0) {
         // we don't GRAFT peers with negative score
         this.log('GRAFT: ignoring peer %s with negative score: score=%d, topic=%s', id, score, topicID)
         // we do send them PRUNE however, because it's a matter of protocol correctness
@@ -1560,23 +1585,24 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
         doPX = false
         // add/refresh backoff so that we don't reGRAFT too early even if the score decays
         this.addBackoff(id, topicID)
-        return
-      }
 
-      // check the number of mesh peers; if it is at (or over) Dhi, we only accept grafts
-      // from peers with outbound connections; this is a defensive check to restrict potential
-      // mesh takeover attacks combined with love bombing
-      if (peersInMesh.size >= this.opts.Dhi && !(this.outbound.get(id) ?? false)) {
+        // check the number of mesh peers; if it is at (or over) Dhi, we only accept grafts
+        // from peers with outbound connections; this is a defensive check to restrict potential
+        // mesh takeover attacks combined with love bombing
+      } else if (peersInMesh.size >= this.opts.Dhi && !(this.outbound.get(id) ?? false)) {
         prune.push(topicID)
         this.addBackoff(id, topicID)
-        return
-      }
 
-      this.log('GRAFT: Add mesh link from %s in %s', id, topicID)
-      this.score.graft(id, topicID)
-      peersInMesh.add(id)
+        // valid graft
+      } else {
+        this.log('GRAFT: Add mesh link from %s in %s', id, topicID)
+        this.score.graft(id, topicID)
+        peersInMesh.add(id)
+
+        this.metrics?.onAddToMesh(topicID, InclusionReason.Subscribed, 1)
+      }
 
-      this.metrics?.onAddToMesh(topicID, InclusionReason.Subscribed, 1)
+      this.safeDispatchEvent<MeshPeer>('gossipsub:graft', { detail: { peerId: id, topic: topicID, direction: 'inbound' } })
     })
 
     if (prune.length === 0) {
@@ -1627,10 +1653,12 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
             score,
             topicID
           )
-          continue
+        } else {
+          await this.pxConnect(peers)
         }
-        await this.pxConnect(peers)
       }
+
+      this.safeDispatchEvent<MeshPeer>('gossipsub:prune', { detail: { peerId: id, topic: topicID, direction: 'inbound' } })
     }
   }
 
@@ -2325,6 +2353,21 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
 
     this.metrics?.onRpcSent(rpc, rpcBytes.length)
 
+    if (rpc.control?.graft != null) {
+      for (const topic of rpc.control?.graft) {
+        if (topic.topicID != null) {
+          this.safeDispatchEvent<MeshPeer>('gossipsub:graft', { detail: { peerId: id, topic: topic.topicID, direction: 'outbound' } })
+        }
+      }
+    }
+    if (rpc.control?.prune != null) {
+      for (const topic of rpc.control?.prune) {
+        if (topic.topicID != null) {
+          this.safeDispatchEvent<MeshPeer>('gossipsub:prune', { detail: { peerId: id, topic: topic.topicID, direction: 'outbound' } })
+        }
+      }
+    }
+
     return true
   }
 
@@ -3023,6 +3066,26 @@ export class GossipSub extends TypedEventEmitter<GossipsubEvents> implements Pub
 
     metrics.registerScoreWeights(sw)
   }
+
+  private readonly tagMeshPeer = (evt: CustomEvent<MeshPeer>): void => {
+    const { peerId, topic } = evt.detail
+    this.components.peerStore.merge(peerIdFromString(peerId), {
+      tags: {
+        [topic]: {
+          value: 100
+        }
+      }
+    }).catch((err) => { this.log.error('Error tagging peer %s with topic %s', peerId, topic, err) })
+  }
+
+  private readonly untagMeshPeer = (evt: CustomEvent<MeshPeer>): void => {
+    const { peerId, topic } = evt.detail
+    this.components.peerStore.merge(peerIdFromString(peerId), {
+      tags: {
+        [topic]: undefined
+      }
+    }).catch((err) => { this.log.error('Error untagging peer %s with topic %s', peerId, topic, err) })
+  }
 }
 
 export function gossipsub (

test/gossip.spec.ts

@@ -99,6 +99,67 @@ describe('gossip', () => {
     expect(publishResult.recipients).to.deep.equal([])
   })
 
+  it('should tag peers', async function () {
+    this.timeout(10e4)
+    const nodeA = nodes[0]
+    const nodeB = nodes[1]
+    const topic = 'Z'
+
+    const twoNodes = [nodeA, nodeB]
+
+    const graftPromises = twoNodes.map(async (n) => pEvent(n.pubsub, 'gossipsub:graft'))
+
+    // add subscriptions to each node
+    twoNodes.forEach((n) => { n.pubsub.subscribe(topic) })
+
+    // every node connected to every other
+    await connectAllPubSubNodes(twoNodes)
+
+    // await grafts
+    await Promise.all(graftPromises)
+
+    // await mesh rebalancing
+    await Promise.all(twoNodes.map(async (n) => pEvent(n.pubsub, 'gossipsub:heartbeat')))
+
+    const peerInfoA = await nodeA.components.peerStore.get(nodeB.components.peerId).catch((e) => undefined)
+    const peerInfoB = await nodeB.components.peerStore.get(nodeA.components.peerId).catch((e) => undefined)
+    expect(peerInfoA?.tags.get(topic)?.value).to.equal(100)
+    expect(peerInfoB?.tags.get(topic)?.value).to.equal(100)
+  })
+
+  it('should remove the tags upon pruning', async function () {
+    this.timeout(10e4)
+    const nodeA = nodes[0]
+    const nodeB = nodes[1]
+    const topic = 'Z'
+
+    const twoNodes = [nodeA, nodeB]
+
+    const subscriptionPromises = nodes.map(async (n) => pEvent(n.pubsub, 'subscription-change'))
+    // add subscriptions to each node
+    twoNodes.forEach((n) => { n.pubsub.subscribe(topic) })
+
+    // every node connected to every other
+    await connectAllPubSubNodes(nodes)
+
+    // await for subscriptions to be transmitted
+    await Promise.all(subscriptionPromises)
+
+    // await mesh rebalancing
+    await Promise.all(twoNodes.map(async (n) => pEvent(n.pubsub, 'gossipsub:heartbeat')))
+
+    twoNodes.forEach((n) => { n.pubsub.unsubscribe(topic) })
+
+    // await for unsubscriptions to be transmitted
+    // await mesh rebalancing
+    await Promise.all(twoNodes.map(async (n) => pEvent(n.pubsub, 'gossipsub:heartbeat')))
+
+    const peerInfoA = await nodeA.components.peerStore.get(nodeB.components.peerId).catch((e) => undefined)
+    const peerInfoB = await nodeB.components.peerStore.get(nodeA.components.peerId).catch((e) => undefined)
+    expect(peerInfoA?.tags.get(topic)).to.be.undefined()
+    expect(peerInfoB?.tags.get(topic)).to.be.undefined()
+  })
+
   it('should reject incoming messages bigger than maxInboundDataLength limit', async function () {
     this.timeout(10e4)
     const nodeA = nodes[0]