From 8236319a91b12297eca9cd8a4c2d1b81e7443e2f Mon Sep 17 00:00:00 2001
From: Giovanni Petrantoni <7008900+sinkingsugar@users.noreply.github.com>
Date: Mon, 22 Feb 2021 12:04:20 +0900
Subject: [PATCH] [SEC] gossipsub - handleGraft/handlePrune recommendations &
 notes (#530)

Fixes #401
---
 tests/pubsub/testgossipinternal.nim | 40 ++++++++++++++++++++++++++---
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/tests/pubsub/testgossipinternal.nim b/tests/pubsub/testgossipinternal.nim
index c109914..07fd6b4 100644
--- a/tests/pubsub/testgossipinternal.nim
+++ b/tests/pubsub/testgossipinternal.nim
@@ -543,18 +543,50 @@ suite "GossipSub internal":
     await conn.close()
     await gossipSub.switch.stop()
 
+  asyncTest "rebalanceMesh fail due to backoff":
+    let gossipSub = TestGossipSub.init(newStandardSwitch())
+    let topic = "foobar"
+    gossipSub.mesh[topic] = initHashSet[PubSubPeer]()
+    gossipSub.topicParams[topic] = TopicParams.init()
+
+    var conns = newSeq[Connection]()
+    gossipSub.gossipsub[topic] = initHashSet[PubSubPeer]()
+    for i in 0..<15:
+      let conn = newBufferStream(noop)
+      conns &= conn
+      let peerInfo = randomPeerInfo()
+      conn.peerInfo = peerInfo
+      let peer = gossipSub.getPubSubPeer(peerInfo.peerId)
+      peer.sendConn = conn
+      gossipSub.onNewPeer(peer)
+      gossipSub.peers[peerInfo.peerId] = peer
+      gossipSub.gossipsub[topic].incl(peer)
+      gossipSub.backingOff
+        .mgetOrPut(topic, initTable[PeerID, Moment]())
+        .add(peerInfo.peerId, Moment.now() + 1.hours)
+      let prunes = gossipSub.handleGraft(peer, @[ControlGraft(topicID: topic)])
+      # there must be a control prune due to violation of backoff
+      check prunes.len != 0
+
+    check gossipSub.peers.len == 15
+    gossipSub.rebalanceMesh(topic)
+    # expect 0 since they are all backing off
+    check gossipSub.mesh[topic].len == 0
+
+    await allFuturesThrowing(conns.mapIt(it.close()))
+    await gossipSub.switch.stop()
+
   asyncTest "rebalanceMesh Degree Hi - audit scenario":
     let gossipSub = TestGossipSub.init(newStandardSwitch())
+    let topic = "foobar"
+    gossipSub.mesh[topic] = initHashSet[PubSubPeer]()
+    gossipSub.topicParams[topic] = TopicParams.init()
     gossipSub.parameters.dScore = 4
     gossipSub.parameters.d = 6
     gossipSub.parameters.dOut = 3
     gossipSub.parameters.dHigh = 12
     gossipSub.parameters.dLow = 4
 
-    let topic = "foobar"
-    gossipSub.mesh[topic] = initHashSet[PubSubPeer]()
-    gossipSub.topicParams[topic] = TopicParams.init()
-
     var conns = newSeq[Connection]()
     gossipSub.gossipsub[topic] = initHashSet[PubSubPeer]()
     for i in 0..<6: