From 1645516276134c8cc4529cc69d82e0fb390435fe Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Mon, 20 Apr 2020 10:26:32 +0300
Subject: [PATCH] add protections from GRAFT floods

---
 gossipsub.go | 32 ++++++++++++++++++++++----------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/gossipsub.go b/gossipsub.go
index 1f0212e..4ff1da3 100644
--- a/gossipsub.go
+++ b/gossipsub.go
@@ -70,6 +70,11 @@ var (
 
 	// Number of peers to opportunistically graft
 	GossipSubOpportunisticGraftPeers = 2
+
+	// If a GRAFT comes before GossipSubGraftFloodThreshold has ellapsed since the last PRUNE,
+	// then there is no PRUNE response emitted. This protects against GRAFT floods and should be
+	// less than GossipSubPruneBackoff.
+	GossipSubGraftFloodThreshold = 10 * time.Second
 )
 
 // NewGossipSub returns a new PubSub object using GossipSubRouter as the router.
@@ -459,6 +464,23 @@ func (gs *GossipSubRouter) handleGraft(p peer.ID, ctl *pb.ControlMessage) []*pb.
 			continue
 		}
 
+		// make sure we are not backing off that peer
+		expire, backoff := gs.backoff[topic][p]
+		if backoff && now.Before(expire) {
+			log.Debugf("GRAFT: ignoring backed off peer %s", p)
+			// refresh the backoff
+			gs.addBackoff(p, topic)
+			// check the flood cutoff -- is the GRAFT coming too fast?
+			floodCutoff := expire.Add(GossipSubGraftFloodThreshold - GossipSubPruneBackoff)
+			if now.Before(floodCutoff) {
+				// no prune, and no PX either
+				doPX = false
+			} else {
+				prune = append(prune, topic)
+			}
+			continue
+		}
+
 		// check the score
 		if score < 0 {
 			// we don't GRAFT peers with negative score
@@ -472,16 +494,6 @@ func (gs *GossipSubRouter) handleGraft(p peer.ID, ctl *pb.ControlMessage) []*pb.
 			continue
 		}
 
-		// make sure we are not backing off that peer
-		expire, backoff := gs.backoff[topic][p]
-		if backoff && now.Before(expire) {
-			log.Debugf("GRAFT: ignoring backed off peer %s", p)
-			// refresh the backoff
-			gs.addBackoff(p, topic)
-			prune = append(prune, topic)
-			continue
-		}
-
 		log.Debugf("GRAFT: add mesh link from %s in %s", p, topic)
 		gs.tracer.Graft(p, topic)
 		peers[p] = struct{}{}