elegant the code

[bytom/bytom.git] / netsync / block_keeper.go

diff --git a/netsync/block_keeper.go b/netsync/block_keeper.go
index 703e2bf..5badaf2 100644 (file)
--- a/netsync/block_keeper.go
+++ b/netsync/block_keeper.go
@@ -29,6 +29,7 @@ var (
        errGetBlockByHash  = errors.New("Get block by hash error")
        errBroadcastStatus = errors.New("Broadcast new status block error")
        errReqBlock        = errors.New("Request block error")
+       errPeerNotRegister = errors.New("peer is not registered")
 )
 
 //TODO: add retry mechanism
@@ -74,7 +75,13 @@ func (bk *blockKeeper) BlockRequestWorker(peerID string, maxPeerHeight uint64) e
        orphanNum := uint64(0)
        reqNum := uint64(0)
        isOrphan := false
-       for num <= maxPeerHeight && num > 0 {
+       bkPeer, ok := bk.peers.Peer(peerID)
+       if !ok {
+               log.Info("peer is not registered")
+               return errPeerNotRegister
+       }
+       swPeer := bkPeer.getPeer()
+       for 0 < num && num <= maxPeerHeight {
                if isOrphan {
                        reqNum = orphanNum
                } else {
@@ -83,18 +90,24 @@ func (bk *blockKeeper) BlockRequestWorker(peerID string, maxPeerHeight uint64) e
                block, err := bk.BlockRequest(peerID, reqNum)
                if errors.Root(err) == errPeerDropped || errors.Root(err) == errGetBlockTimeout || errors.Root(err) == errReqBlock {
                        log.WithField("Peer abnormality. PeerID: ", peerID).Info(err)
-                       peer := bk.peers.Peer(peerID)
-                       if peer == nil {
+                       if bkPeer == nil {
                                log.Info("peer is not registered")
                                break
                        }
-                       log.Info("Peer communication abnormality")
-                       bk.sw.StopPeerGracefully(peer.Peer)
+                       log.Info("Block keeper request block error. Stop peer.")
+                       bk.sw.StopPeerGracefully(swPeer)
                        break
                }
                isOrphan, err = bk.chain.ProcessBlock(block)
                if err != nil {
-                       bk.sw.AddScamPeer(bk.peers.Peer(peerID).getPeer())
+                       if bkPeer == nil {
+                               log.Info("peer is deleted")
+                               break
+                       }
+                       if ban := bkPeer.addBanScore(20, 0, "block process error"); ban {
+                               bk.sw.AddBannedPeer(swPeer)
+                               bk.sw.StopPeerGracefully(swPeer)
+                       }
                        log.WithField("hash:", block.Hash()).Errorf("blockKeeper fail process block %v ", err)
                        break
                }
@@ -115,10 +128,19 @@ func (bk *blockKeeper) BlockRequestWorker(peerID string, maxPeerHeight uint64) e
                        return errGetBlockByHash
                }
 
-               if err := bk.peers.BroadcastNewStatus(block); err != nil {
+               peers, err := bk.peers.BroadcastNewStatus(block)
+               if err != nil {
                        log.Errorf("Failed on broadcast new status block %v", err)
                        return errBroadcastStatus
                }
+               for _, peer := range peers {
+                       if peer == nil {
+                               return errPeerNotRegister
+                       }
+                       swPeer := peer.getPeer()
+                       log.Info("Block keeper broadcast block error. Stop peer.")
+                       bk.sw.StopPeerGracefully(swPeer)
+               }
        }
        return nil
 }
@@ -171,7 +193,13 @@ func (bk *blockKeeper) txsProcessWorker() {
                log.Info("Receive new tx from remote peer. TxID:", tx.ID.String())
                bk.peers.MarkTransaction(txsResponse.peerID, &tx.ID)
                if isOrphan, err := bk.chain.ValidateTx(tx); err != nil && isOrphan == false {
-                       bk.sw.AddScamPeer(bk.peers.Peer(txsResponse.peerID).getPeer())
+                       if bkPeer, ok := bk.peers.Peer(txsResponse.peerID); ok {
+                               swPeer := bkPeer.getPeer()
+                               if ban := bkPeer.addBanScore(10, 0, "tx error"); ban {
+                                       bk.sw.AddBannedPeer(swPeer)
+                                       bk.sw.StopPeerGracefully(swPeer)
+                               }
+                       }
                }
        }
 }