From 4110c252af6ba4202fe948d1a4fc7c35fa84e3e5 Mon Sep 17 00:00:00 2001
From: Callum Michael Waters <cmwaters19@gmail.com>
Date: Wed, 11 Mar 2020 15:44:22 +0100
Subject: [PATCH] make banTime an argument and set default in PEX reactor
 instead of AddrBook

---
 p2p/pex/addrbook.go      | 15 +++++++--------
 p2p/pex/addrbook_test.go | 28 ++++++++++++++++++++++++----
 p2p/pex/pex_reactor.go   |  7 +++++--
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/p2p/pex/addrbook.go b/p2p/pex/addrbook.go
index c3351a600..66a9118e3 100644
--- a/p2p/pex/addrbook.go
+++ b/p2p/pex/addrbook.go
@@ -22,9 +22,8 @@ import (
 )
 
 const (
-	bucketTypeNew  = 0x01
-	bucketTypeOld  = 0x02
-	defaultBanTime = 24 * time.Hour
+	bucketTypeNew = 0x01
+	bucketTypeOld = 0x02
 )
 
 // AddrBook is an address book used for tracking peers
@@ -60,7 +59,7 @@ type AddrBook interface {
 	// Mark address
 	MarkGood(p2p.ID)
 	MarkAttempt(*p2p.NetAddress)
-	MarkBad(*p2p.NetAddress) // Move peer to bad peers list
+	MarkBad(*p2p.NetAddress, time.Duration) // Move peer to bad peers list
 	// Add bad peers back to addrBook
 	ReinstateBadPeers()
 
@@ -326,11 +325,11 @@ func (a *addrBook) MarkAttempt(addr *p2p.NetAddress) {
 
 // MarkBad implements AddrBook. Kicks address out from book, places
 // the address in the badPeers pool.
-func (a *addrBook) MarkBad(addr *p2p.NetAddress) {
+func (a *addrBook) MarkBad(addr *p2p.NetAddress, banTime time.Duration) {
 	a.mtx.Lock()
 	defer a.mtx.Unlock()
 
-	if a.addBadPeer(addr) {
+	if a.addBadPeer(addr, banTime) {
 		a.removeAddress(addr)
 	}
 }
@@ -751,7 +750,7 @@ func (a *addrBook) removeAddress(addr *p2p.NetAddress) {
 	a.removeFromAllBuckets(ka)
 }
 
-func (a *addrBook) addBadPeer(addr *p2p.NetAddress) bool {
+func (a *addrBook) addBadPeer(addr *p2p.NetAddress, banTime time.Duration) bool {
 	a.mtx.Lock()
 	defer a.mtx.Unlock()
 
@@ -761,7 +760,7 @@ func (a *addrBook) addBadPeer(addr *p2p.NetAddress) bool {
 	if ka != nil {
 		if _, alreadyBadPeer := a.badPeers[addr.ID]; !alreadyBadPeer {
 			// add to bad peer list
-			ka.ban(defaultBanTime)
+			ka.ban(banTime)
 			a.badPeers[addr.ID] = ka
 		}
 		return true
diff --git a/p2p/pex/addrbook_test.go b/p2p/pex/addrbook_test.go
index 363958c44..eaa3280f2 100644
--- a/p2p/pex/addrbook_test.go
+++ b/p2p/pex/addrbook_test.go
@@ -3,14 +3,13 @@ package pex
 import (
 	"encoding/hex"
 	"fmt"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"io/ioutil"
 	"math"
 	"os"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
 	"github.com/tendermint/tendermint/libs/log"
 	tmmath "github.com/tendermint/tendermint/libs/math"
 	tmrand "github.com/tendermint/tendermint/libs/rand"
@@ -343,7 +342,7 @@ func TestAddrBookGetSelectionWithBias(t *testing.T) {
 		}
 	}
 
-	got, expected := int((float64(good)/float64(len(selection)))*100), (100 - biasTowardsNewAddrs)
+	got, expected := int((float64(good)/float64(len(selection)))*100), 100-biasTowardsNewAddrs
 
 	// compute some slack to protect against small differences due to rounding:
 	slack := int(math.Round(float64(100) / float64(len(selection))))
@@ -544,6 +543,27 @@ func TestMultipleAddrBookAddressSelection(t *testing.T) {
 	}
 }
 
+//func TestBanBadPeers(t *testing.T) {
+//	fname := createTempFileName("addrbook_test")
+//	defer deleteTempFile(fname)
+//
+//	book := NewAddrBook(fname, true)
+//	book.SetLogger(log.TestingLogger())
+//
+//	addr := randIPv4Address(t)
+//	book.AddAddress(addr, addr)
+//
+//	book.MarkBad(addr, 1 * time.Second)
+//
+//	assert.False(t, book.HasAddress(addr))
+//
+//	time.Sleep(1 * time.Second)
+//
+//	book.ReinstateBadPeers()
+//
+//	assert.True(t, book.HasAddress(addr))
+//}
+
 func assertMOldAndNNewAddrsInSelection(t *testing.T, m, n int, addrs []*p2p.NetAddress, book *addrBook) {
 	nOld, nNew := countOldAndNewAddrsInSelection(addrs, book)
 	assert.Equal(t, m, nOld, "old addresses")
diff --git a/p2p/pex/pex_reactor.go b/p2p/pex/pex_reactor.go
index bce9b40ff..afe29e16e 100644
--- a/p2p/pex/pex_reactor.go
+++ b/p2p/pex/pex_reactor.go
@@ -50,6 +50,9 @@ const (
 	// Especially in the beginning, node should have more trusted peers than
 	// untrusted.
 	biasToSelectNewPeers = 30 // 70 to select good peers
+
+	// if a peer is marked bad, it will be banned for at least this time period
+	defaultBanTime = 24 * time.Hour
 )
 
 type errMaxAttemptsToDial struct {
@@ -535,7 +538,7 @@ func (r *Reactor) dialPeer(addr *p2p.NetAddress) error {
 		// failed to connect to. Then we can clean up attemptsToDial, which acts as
 		// a blacklist currently.
 		// https://github.com/tendermint/tendermint/issues/3572
-		r.book.MarkBad(addr)
+		r.book.MarkBad(addr, defaultBanTime)
 		return errMaxAttemptsToDial{}
 	}
 
@@ -747,7 +750,7 @@ func markAddrInBookBasedOnErr(addr *p2p.NetAddress, book AddrBook, err error) {
 	// TODO: detect more "bad peer" scenarios
 	switch err.(type) {
 	case p2p.ErrSwitchAuthenticationFailure:
-		book.MarkBad(addr)
+		book.MarkBad(addr, defaultBanTime)
 	default:
 		book.MarkAttempt(addr)
 	}