From 32ca1898d9e14bacefd4eb4c667f8d94462958b1 Mon Sep 17 00:00:00 2001
From: lchenut <ludovic@status.im>
Date: Tue, 10 May 2022 10:39:43 +0200
Subject: [PATCH] Gossipsub: Put Peer Exchange behind a flag (#715)

Add a flag to enable Peer Exchange in Gossipsub (disabled by default)
---
 libp2p/protocols/pubsub/gossipsub.nim          | 3 ++-
 libp2p/protocols/pubsub/gossipsub/behavior.nim | 2 ++
 libp2p/protocols/pubsub/gossipsub/types.nim    | 1 +
 tests/pubsub/testgossipsub.nim                 | 5 +++--
 tests/pubsub/utils.nim                         | 5 +++--
 5 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/libp2p/protocols/pubsub/gossipsub.nim b/libp2p/protocols/pubsub/gossipsub.nim
index 8354247ad..175e6956f 100644
--- a/libp2p/protocols/pubsub/gossipsub.nim
+++ b/libp2p/protocols/pubsub/gossipsub.nim
@@ -71,7 +71,8 @@ proc init*(_: type[GossipSubParams]): GossipSubParams =
       ipColocationFactorThreshold: 1.0,
       behaviourPenaltyWeight: -1.0,
       behaviourPenaltyDecay: 0.999,
-      disconnectBadPeers: false
+      disconnectBadPeers: false,
+      enablePX: false
     )
 
 proc validateParameters*(parameters: GossipSubParams): Result[void, cstring] =
diff --git a/libp2p/protocols/pubsub/gossipsub/behavior.nim b/libp2p/protocols/pubsub/gossipsub/behavior.nim
index 39418b399..66bd69ddb 100644
--- a/libp2p/protocols/pubsub/gossipsub/behavior.nim
+++ b/libp2p/protocols/pubsub/gossipsub/behavior.nim
@@ -79,6 +79,8 @@ proc handleBackingOff*(t: var BackoffTable, topic: string) {.raises: [Defect].}
       v[].del(peer)
 
 proc peerExchangeList*(g: GossipSub, topic: string): seq[PeerInfoMsg] {.raises: [Defect].} =
+  if not g.parameters.enablePX:
+    return @[]
   var peers = g.gossipsub.getOrDefault(topic, initHashSet[PubSubPeer]()).toSeq()
   peers.keepIf do (x: PubSubPeer) -> bool:
       x.score >= 0.0
diff --git a/libp2p/protocols/pubsub/gossipsub/types.nim b/libp2p/protocols/pubsub/gossipsub/types.nim
index a91c7b4ea..5736c01fd 100644
--- a/libp2p/protocols/pubsub/gossipsub/types.nim
+++ b/libp2p/protocols/pubsub/gossipsub/types.nim
@@ -138,6 +138,7 @@ type
     directPeers*: Table[PeerId, seq[MultiAddress]]
 
     disconnectBadPeers*: bool
+    enablePX*: bool
 
   BackoffTable* = Table[string, Table[PeerId, Moment]]
   ValidationSeenTable* = Table[MessageID, HashSet[PubSubPeer]]
diff --git a/tests/pubsub/testgossipsub.nim b/tests/pubsub/testgossipsub.nim
index f12a82952..416368dbf 100644
--- a/tests/pubsub/testgossipsub.nim
+++ b/tests/pubsub/testgossipsub.nim
@@ -945,7 +945,8 @@ suite "GossipSub":
     let
       nodes = generateNodes(
         2,
-        gossip = true) &
+        gossip = true,
+        enablePX = true) &
         generateNodes(1, gossip = true, sendSignedPeerRecord = true)
 
       # start switches
@@ -988,7 +989,7 @@ suite "GossipSub":
     )
     nodes[1].unsubscribe("foobar", handler)
 
-    await passed
+    await passed.wait(5.seconds)
 
     await allFuturesThrowing(
       nodes[0].switch.stop(),
diff --git a/tests/pubsub/utils.nim b/tests/pubsub/utils.nim
index 05f0299e5..993efdd01 100644
--- a/tests/pubsub/utils.nim
+++ b/tests/pubsub/utils.nim
@@ -42,7 +42,8 @@ proc generateNodes*(
   sign: bool = libp2p_pubsub_sign,
   sendSignedPeerRecord = false,
   unsubscribeBackoff = 1.seconds,
-  maxMessageSize: int = 1024 * 1024): seq[PubSub] =
+  maxMessageSize: int = 1024 * 1024,
+  enablePX: bool = false): seq[PubSub] =
 
   for i in 0..<num:
     let switch = newStandardSwitch(secureManagers = secureManagers, sendSignedPeerRecord = sendSignedPeerRecord)
@@ -55,7 +56,7 @@ proc generateNodes*(
         msgIdProvider = msgIdProvider,
         anonymize = anonymize,
         maxMessageSize = maxMessageSize,
-        parameters = (var p = GossipSubParams.init(); p.floodPublish = false; p.historyLength = 20; p.historyGossip = 20; p.unsubscribeBackoff = unsubscribeBackoff; p))
+        parameters = (var p = GossipSubParams.init(); p.floodPublish = false; p.historyLength = 20; p.historyGossip = 20; p.unsubscribeBackoff = unsubscribeBackoff; p.enablePX = enablePX; p))
       # set some testing params, to enable scores
       g.topicParams.mgetOrPut("foobar", TopicParams.init()).topicWeight = 1.0
       g.topicParams.mgetOrPut("foo", TopicParams.init()).topicWeight = 1.0