From 51d8cd4ade64545777814bdd1c110909c5855df4 Mon Sep 17 00:00:00 2001
From: Giovanni Petrantoni <7008900+sinkingsugar@users.noreply.github.com>
Date: Sat, 13 Feb 2021 13:39:32 +0900
Subject: [PATCH] [SEC] gossipsub - rebalanceMesh may prune up to D_lo on
 oversubscription (#531)

Fixes #403
---
 libp2p/protocols/pubsub/peertable.nim |  3 --
 tests/pubsub/testgossipinternal.nim   | 56 +++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/libp2p/protocols/pubsub/peertable.nim b/libp2p/protocols/pubsub/peertable.nim
index 15a1bf5..944fa2a 100644
--- a/libp2p/protocols/pubsub/peertable.nim
+++ b/libp2p/protocols/pubsub/peertable.nim
@@ -41,6 +41,3 @@ func peers*(table: PeerTable, topic: string): int =
     table[topic].len
   else:
     0
-
-func getPeers*(table: Table[string, HashSet[string]], topic: string): HashSet[string] =
-  table.getOrDefault(topic, initHashSet[string]())
diff --git a/tests/pubsub/testgossipinternal.nim b/tests/pubsub/testgossipinternal.nim
index b53ead7..7d259e0 100644
--- a/tests/pubsub/testgossipinternal.nim
+++ b/tests/pubsub/testgossipinternal.nim
@@ -509,3 +509,59 @@ suite "GossipSub internal":
 
     await conn.close()
     await gossipSub.switch.stop()
+
+  asyncTest "rebalanceMesh Degree Hi - audit scenario":
+    let gossipSub = TestGossipSub.init(newStandardSwitch())
+    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:
+      let conn = newBufferStream(noop)
+      conn.dir = Direction.In
+      conns &= conn
+      let peerInfo = PeerInfo.init(PrivateKey.random(ECDSA, rng[]).get())
+      conn.peerInfo = peerInfo
+      let peer = gossipSub.getPubSubPeer(peerInfo.peerId)
+      peer.score = 40.0
+      peer.sendConn = conn
+      gossipSub.onNewPeer(peer)
+      gossipSub.grafted(peer, topic)
+      gossipSub.peers[peerInfo.peerId] = peer
+      gossipSub.mesh[topic].incl(peer)
+
+    for i in 0..<7:
+      let conn = newBufferStream(noop)
+      conn.dir = Direction.Out
+      conns &= conn
+      let peerInfo = PeerInfo.init(PrivateKey.random(ECDSA, rng[]).get())
+      conn.peerInfo = peerInfo
+      let peer = gossipSub.getPubSubPeer(peerInfo.peerId)
+      peer.score = 10.0
+      peer.sendConn = conn
+      gossipSub.onNewPeer(peer)
+      gossipSub.grafted(peer, topic)
+      gossipSub.peers[peerInfo.peerId] = peer
+      gossipSub.mesh[topic].incl(peer)
+
+    check gossipSub.mesh[topic].len == 13
+    gossipSub.rebalanceMesh(topic)
+    # ensure we are above dlow
+    check gossipSub.mesh[topic].len > gossipSub.parameters.dLow
+    var outbound = 0
+    for peer in gossipSub.mesh[topic]:
+      if peer.sendConn.dir == Direction.Out:
+        inc outbound
+    # ensure we give priority and keep at least dOut outbound peers
+    check outbound >= gossipSub.parameters.dOut
+
+    await allFuturesThrowing(conns.mapIt(it.close()))
+    await gossipSub.switch.stop()