From 7466f353451eae68183f86e042034b98a6755001 Mon Sep 17 00:00:00 2001
From: Callum Michael Waters <cmwaters19@gmail.com>
Date: Tue, 10 Mar 2020 18:13:45 +0100
Subject: [PATCH] mvp blacklist alg

---
 p2p/pex/addrbook.go      | 44 +++++++++++++++++++++++++++++++++++-----
 p2p/pex/known_address.go | 12 +++++++++++
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/p2p/pex/addrbook.go b/p2p/pex/addrbook.go
index dbba71345..3bf331172 100644
--- a/p2p/pex/addrbook.go
+++ b/p2p/pex/addrbook.go
@@ -22,8 +22,9 @@ import (
 )
 
 const (
-	bucketTypeNew = 0x01
-	bucketTypeOld = 0x02
+	bucketTypeNew  = 0x01
+	bucketTypeOld  = 0x02
+	defaultBanTime = 24 * time.Hour
 )
 
 // AddrBook is an address book used for tracking peers
@@ -87,6 +88,7 @@ type addrBook struct {
 	ourAddrs   map[string]struct{}
 	privateIDs map[p2p.ID]struct{}
 	addrLookup map[p2p.ID]*knownAddress // new & old
+	badPeers   map[p2p.ID]*knownAddress // blacklisted peers
 	bucketsOld []map[string]*knownAddress
 	bucketsNew []map[string]*knownAddress
 	nOld       int
@@ -108,6 +110,7 @@ func NewAddrBook(filePath string, routabilityStrict bool) AddrBook {
 		ourAddrs:          make(map[string]struct{}),
 		privateIDs:        make(map[p2p.ID]struct{}),
 		addrLookup:        make(map[p2p.ID]*knownAddress),
+		badPeers:          make(map[p2p.ID]*knownAddress),
 		filePath:          filePath,
 		routabilityStrict: routabilityStrict,
 	}
@@ -233,6 +236,7 @@ func (a *addrBook) HasAddress(addr *p2p.NetAddress) bool {
 
 // NeedMoreAddrs implements AddrBook - returns true if there are not have enough addresses in the book.
 func (a *addrBook) NeedMoreAddrs() bool {
+	a.ReinstateBadPeers()
 	return a.Size() < needAddressThreshold
 }
 
@@ -324,10 +328,24 @@ func (a *addrBook) MarkAttempt(addr *p2p.NetAddress) {
 	ka.markAttempt()
 }
 
-// MarkBad implements AddrBook. Currently it just ejects the address.
-// TODO: black list for some amount of time
+// MarkBad implements AddrBook. Kicks address out from book, places
+// the address in the badPeers pool.
 func (a *addrBook) MarkBad(addr *p2p.NetAddress) {
-	a.RemoveAddress(addr)
+	if a.addBadPeer(addr) {
+		a.RemoveAddress(addr)
+	}
+}
+
+func (a *addrBook) ReinstateBadPeers() {
+	for _, ka := range a.badPeers {
+		if !ka.isBanned(defaultBanTime) {
+			a.mtx.Lock()
+			bucket := a.calcNewBucket(ka.Addr, ka.Src)
+			a.addToNewBucket(ka, bucket)
+			delete(a.badPeers, ka.ID())
+			a.mtx.Unlock()
+		}
+	}
 }
 
 // GetSelection implements AddrBook.
@@ -725,6 +743,22 @@ func (a *addrBook) moveToOld(ka *knownAddress) {
 	}
 }
 
+func (a *addrBook) addBadPeer(addr *p2p.NetAddress) bool {
+	a.mtx.Lock()
+	defer a.mtx.Unlock()
+
+	// check it exists in addrbook
+	ka := a.addrLookup[addr.ID]
+	// check address is not already there
+	if _, alreadyBadPeer := a.badPeers[addr.ID]; ka != nil && !alreadyBadPeer {
+		// add to bad peer list
+		ka.ban()
+		a.badPeers[addr.ID] = ka
+		return true
+	}
+	return false
+}
+
 //---------------------------------------------------------------------
 // calculate bucket placements
 
diff --git a/p2p/pex/known_address.go b/p2p/pex/known_address.go
index af40d6ff0..72729a5d5 100644
--- a/p2p/pex/known_address.go
+++ b/p2p/pex/known_address.go
@@ -16,6 +16,7 @@ type knownAddress struct {
 	BucketType  byte            `json:"bucket_type"`
 	LastAttempt time.Time       `json:"last_attempt"`
 	LastSuccess time.Time       `json:"last_success"`
+	LastBanTime time.Time       `json:"last_ban_time"`
 }
 
 func newKnownAddress(addr *p2p.NetAddress, src *p2p.NetAddress) *knownAddress {
@@ -54,6 +55,17 @@ func (ka *knownAddress) markGood() {
 	ka.LastSuccess = now
 }
 
+func (ka *knownAddress) ban() {
+	ka.LastBanTime = time.Now()
+}
+
+func (ka *knownAddress) isBanned(banTime time.Duration) bool {
+	if ka.LastBanTime.Add(banTime).Before(time.Now()) {
+		return false
+	}
+	return true
+}
+
 func (ka *knownAddress) addBucketRef(bucketIdx int) int {
 	for _, bucket := range ka.Buckets {
 		if bucket == bucketIdx {