From a645afd1b3a8adcb6a1f9e8e44aaf0b7d9f31bdd Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Mon, 20 Apr 2020 16:03:44 +0300
Subject: [PATCH] add backoff penalty for GRAFT floods

---
 gossipsub.go           | 22 +++++++++++++++++++---
 gossipsub_spam_test.go | 20 ++++++++++++++++++--
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/gossipsub.go b/gossipsub.go
index 4f83389..a0387bd 100644
--- a/gossipsub.go
+++ b/gossipsub.go
@@ -76,6 +76,9 @@ var (
 	// less than GossipSubPruneBackoff.
 	GossipSubGraftFloodThreshold = 10 * time.Second
 
+	// backoff penalty for GRAFT floods
+	GossipSubPruneBackoffPenalty = time.Hour
+
 	// Maximum number of messages to include in an IHAVE message. Also controls the maximum
 	// number of IHAVE ids we will accept and request with IWANT from a peer within a heartbeat,
 	// to protect from IHAVE floods. You should adjust this value from the default if your
@@ -503,15 +506,17 @@ func (gs *GossipSubRouter) handleGraft(p peer.ID, ctl *pb.ControlMessage) []*pb.
 		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
+				// and a penalty so that we don't GRAFT on this peer ourselves for a while
+				gs.addBackoffPenalty(p, topic)
 			} else {
 				prune = append(prune, topic)
+				// refresh the backoff
+				gs.addBackoff(p, topic)
 			}
 			continue
 		}
@@ -577,12 +582,23 @@ func (gs *GossipSubRouter) handlePrune(p peer.ID, ctl *pb.ControlMessage) {
 }
 
 func (gs *GossipSubRouter) addBackoff(p peer.ID, topic string) {
+	gs.doAddBackoff(p, topic, GossipSubPruneBackoff)
+}
+
+func (gs *GossipSubRouter) addBackoffPenalty(p peer.ID, topic string) {
+	gs.doAddBackoff(p, topic, GossipSubPruneBackoffPenalty)
+}
+
+func (gs *GossipSubRouter) doAddBackoff(p peer.ID, topic string, interval time.Duration) {
 	backoff, ok := gs.backoff[topic]
 	if !ok {
 		backoff = make(map[peer.ID]time.Time)
 		gs.backoff[topic] = backoff
 	}
-	backoff[p] = time.Now().Add(GossipSubPruneBackoff)
+	expire := time.Now().Add(interval)
+	if backoff[p].Before(expire) {
+		backoff[p] = expire
+	}
 }
 
 func (gs *GossipSubRouter) pxConnect(peers []*pb.PeerInfo) {
diff --git a/gossipsub_spam_test.go b/gossipsub_spam_test.go
index 73b46ed..18b21a3 100644
--- a/gossipsub_spam_test.go
+++ b/gossipsub_spam_test.go
@@ -318,8 +318,11 @@ func TestGossipsubAttackGRAFTDuringBackoff(t *testing.T) {
 	GossipSubPruneBackoff = 200 * time.Millisecond
 	originalGossipSubGraftFloodThreshold := GossipSubGraftFloodThreshold
 	GossipSubGraftFloodThreshold = 100 * time.Millisecond
+	originalGossipSubPruneBackoffPenalty := GossipSubPruneBackoffPenalty
+	GossipSubPruneBackoffPenalty = 500 * time.Millisecond
 	defer func() {
 		GossipSubPruneBackoff = originalGossipSubPruneBackoff
+		GossipSubPruneBackoffPenalty = originalGossipSubPruneBackoffPenalty
 		GossipSubGraftFloodThreshold = originalGossipSubGraftFloodThreshold
 	}()
 
@@ -414,8 +417,8 @@ func TestGossipsubAttackGRAFTDuringBackoff(t *testing.T) {
 						t.Fatalf("Expected %d PRUNE messages but got %d", 1, pc)
 					}
 
-					// Wait until after the prune backoff period
-					time.Sleep(GossipSubPruneBackoff * 2)
+					// Wait until after the prune backoff penalty period
+					time.Sleep(GossipSubPruneBackoffPenalty + time.Second)
 
 					// Send a GRAFT again to attempt to rejoin the mesh
 					writeMsg(&pb.RPC{
@@ -430,6 +433,19 @@ func TestGossipsubAttackGRAFTDuringBackoff(t *testing.T) {
 					if pc != 1 {
 						t.Fatalf("Expected %d PRUNE messages but got %d", 1, pc)
 					}
+
+					// make sure we are in the mesh of the legit host now
+					res := make(chan bool)
+					ps.eval <- func() {
+						mesh := ps.rt.(*GossipSubRouter).mesh[mytopic]
+						_, inMesh := mesh[attacker.ID()]
+						res <- inMesh
+					}
+
+					inMesh := <-res
+					if !inMesh {
+						t.Fatal("Expected to be in the mesh of the legitimate host")
+					}
 				}()
 			}
 		}