From 89778cc947aede81a7c058fad2f109dfa7d3630c Mon Sep 17 00:00:00 2001
From: rjl493456442 <garyrong0905@gmail.com>
Date: Thu, 12 Mar 2020 20:32:24 +0800
Subject: [PATCH] les: ban bad light client

---
 les/peer.go           | 36 ++++++++++++++++++++++++++++++++++++
 les/server_handler.go |  9 +++++++--
 p2p/dial.go           |  5 +++--
 3 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/les/peer.go b/les/peer.go
index 28ec201bc9..a9f04c3b55 100644
--- a/les/peer.go
+++ b/les/peer.go
@@ -1030,6 +1030,24 @@ func (ps *clientPeerSet) allPeers() []*clientPeer {
 	return list
 }
 
+// banPeer terminates the connection and bans re-connection
+// or dialing in the short future.
+func (ps *clientPeerSet) banPeer(id string) error {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+
+	peer := ps.peers[id]
+	if peer == nil {
+		return errNotRegistered
+	}
+	delete(ps.peers, id)
+	for _, sub := range ps.subscribers {
+		sub.unregisterPeer(peer)
+	}
+	peer.Ban()
+	return nil
+}
+
 // close disconnects all peers. No new peers can be registered
 // after close has returned.
 func (ps *clientPeerSet) close() {
@@ -1181,6 +1199,24 @@ func (ps *serverPeerSet) allPeers() []*serverPeer {
 	return list
 }
 
+// banPeer terminates the connection and bans re-connection
+// or dialing in the short future.
+func (ps *serverPeerSet) banPeer(id string) error {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+
+	peer := ps.peers[id]
+	if peer == nil {
+		return errNotRegistered
+	}
+	delete(ps.peers, id)
+	for _, sub := range ps.subscribers {
+		sub.unregisterPeer(peer)
+	}
+	peer.Ban()
+	return nil
+}
+
 // close disconnects all peers. No new peers can be registered
 // after close has returned.
 func (ps *serverPeerSet) close() {
diff --git a/les/server_handler.go b/les/server_handler.go
index 186bdcbb03..0643b35558 100644
--- a/les/server_handler.go
+++ b/les/server_handler.go
@@ -108,7 +108,7 @@ func (h *serverHandler) runPeer(version uint, p *p2p.Peer, rw p2p.MsgReadWriter)
 	return h.handle(peer)
 }
 
-func (h *serverHandler) handle(p *clientPeer) error {
+func (h *serverHandler) handle(p *clientPeer) (err error) {
 	p.Log().Debug("Light Ethereum peer connected", "name", p.Name())
 
 	// Execute the LES handshake
@@ -152,7 +152,12 @@ func (h *serverHandler) handle(p *clientPeer) error {
 	connectedAt := mclock.Now()
 	defer func() {
 		wg.Wait() // Ensure all background task routines have exited.
-		h.server.peers.unregister(p.id)
+
+		if err == errTooManyInvalidRequest {
+			h.server.peers.banPeer(p.id)
+		} else {
+			h.server.peers.unregister(p.id)
+		}
 		h.server.clientPool.disconnect(p)
 		clientConnectionGauge.Update(int64(h.server.peers.len()))
 		connectionTimer.Update(time.Duration(mclock.Now() - connectedAt))
diff --git a/p2p/dial.go b/p2p/dial.go
index 1bc82e9d8b..d7d6dee611 100644
--- a/p2p/dial.go
+++ b/p2p/dial.go
@@ -123,7 +123,8 @@ type dialScheduler struct {
 	historyTimer     mclock.Timer
 	historyTimerTime mclock.AbsTime
 
-	// The bannedHistory keeps all recently banned nodes which won't be dialed for a short time
+	// The bannedHistory keeps all recently banned nodes
+	// which won't be dialed for a short time
 	bannedHistory expHeap
 
 	// for logStats
@@ -288,7 +289,7 @@ loop:
 			d.updateStaticPool(c.node.ID())
 
 		case c := <-d.banPeerCh:
-			// If it's not a static node and been banned,
+			// If it's not a static node and has been banned,
 			// add to the blacklist for dialing.
 			if d.static[c.node.ID()] == nil {
 				remoteIP := netutil.AddrIP(c.fd.RemoteAddr())