From 0c43425d2f4f9f172783898a2c520fa0f2789c12 Mon Sep 17 00:00:00 2001
From: Zhou Zhiyao <ZHOU0250@e.ntu.edu.sg>
Date: Fri, 29 Nov 2019 23:55:31 +0800
Subject: [PATCH] restrict removed raft node in p2p connection (#884)

---
 p2p/server.go   | 14 ++++++++++++++
 raft/handler.go | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/p2p/server.go b/p2p/server.go
index 38f73289d..437f9a96a 100644
--- a/p2p/server.go
+++ b/p2p/server.go
@@ -190,6 +190,9 @@ type Server struct {
 	loopWG        sync.WaitGroup // loop, listenLoop
 	peerFeed      event.Feed
 	log           log.Logger
+
+	// raft peers info
+	checkPeerInRaft func(*enode.Node) bool
 }
 
 type peerOpFunc func(map[enode.ID]*Peer)
@@ -927,6 +930,13 @@ func (srv *Server) setupConn(c *conn, flags connFlag, dialDest *enode.Node) erro
 	}
 	clog := srv.log.New("id", c.node.ID(), "addr", c.fd.RemoteAddr(), "conn", c.flags)
 
+	// If raft is running, check if the dialing node is in the raft cluster
+	// Node doesn't belong to raft cluster is not allowed to join the p2p network
+	if srv.checkPeerInRaft != nil && !srv.checkPeerInRaft(c.node) {
+		log.Trace("incoming connection peer is not in the raft cluster", "enode.id", c.node.ID())
+		return nil
+	}
+
 	//START - QUORUM Permissioning
 	currentNode := srv.NodeInfo().ID
 	cnodeName := srv.NodeInfo().Name
@@ -1116,3 +1126,7 @@ func (srv *Server) PeersInfo() []*PeerInfo {
 	}
 	return infos
 }
+
+func (srv *Server) SetCheckPeerInRaft(f func(*enode.Node) bool) {
+	srv.checkPeerInRaft = f
+}
diff --git a/raft/handler.go b/raft/handler.go
index 3d79c2014..0ea823374 100755
--- a/raft/handler.go
+++ b/raft/handler.go
@@ -143,6 +143,8 @@ func (pm *ProtocolManager) Start(p2pServer *p2p.Server) {
 	pm.p2pServer = p2pServer
 	pm.minedBlockSub = pm.eventMux.Subscribe(core.NewMinedBlockEvent{})
 	pm.startRaft()
+	// update raft peers info to p2p server
+	pm.p2pServer.SetCheckPeerInRaft(pm.peerExist)
 	go pm.minedBroadcastLoop()
 }
 
@@ -305,6 +307,18 @@ func (pm *ProtocolManager) isNodeAlreadyInCluster(node *enode.Node) error {
 	return nil
 }
 
+func (pm *ProtocolManager) peerExist(node *enode.Node) bool {
+	pm.mu.RLock()
+	defer pm.mu.RUnlock()
+
+	for _, p := range pm.peers {
+		if node.ID() == p.p2pNode.ID() {
+			return true
+		}
+	}
+	return false
+}
+
 func (pm *ProtocolManager) ProposeNewPeer(enodeId string) (uint16, error) {
 	parsedUrl, _ := url.Parse(enodeId)
 	node, err := enode.ParseV4(enodeId)