use map

okx · Jul 9, 2024 · c73afcb · c73afcb
1 parent 8867016
commit c73afcb
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 35 deletions.
diff --git a/app/config/config.go b/app/config/config.go
@@ -139,7 +139,7 @@ type OecConfig struct {
 	maxTxLimitPerPeer uint64
 
 	enableConsensusIPWhitelist bool
-	consensusIPWhitelist       []string
+	consensusIPWhitelist       map[string]bool
 }
 
 const (
@@ -829,7 +829,7 @@ func (c *OecConfig) GetEnableConsensusIPWhitelist() bool {
 	return c.enableConsensusIPWhitelist
 }
 
-func (c *OecConfig) GetConsensusIPWhitelist() []string {
+func (c *OecConfig) GetConsensusIPWhitelist() map[string]bool {
 	return c.consensusIPWhitelist
 }
 
@@ -861,7 +861,7 @@ func (c *OecConfig) SetEnableConsensusIPWhitelist(value bool) {
 func (c *OecConfig) SetConsensusIPWhitelist(value string) {
 	ipList := resolveNodeKeyWhitelist(value)
 	for _, ip := range ipList {
-		c.consensusIPWhitelist = append(c.consensusIPWhitelist, strings.TrimSpace(ip))
+		c.consensusIPWhitelist[strings.TrimSpace(ip)] = true
 	}
 }
 

diff --git a/libs/tendermint/blockchain/v0/reactor.go b/libs/tendermint/blockchain/v0/reactor.go
@@ -3,13 +3,13 @@ package v0
 import (
 	"errors"
 	"fmt"
+	cfg "github.com/okex/exchain/libs/tendermint/config"
 	"reflect"
 	"sync"
 	"time"
 
 	amino "github.com/tendermint/go-amino"
 
-	cfg "github.com/okex/exchain/libs/tendermint/config"
 	"github.com/okex/exchain/libs/tendermint/libs/log"
 	"github.com/okex/exchain/libs/tendermint/p2p"
 	sm "github.com/okex/exchain/libs/tendermint/state"
@@ -203,20 +203,6 @@ func (bcR *BlockchainReactor) respondToPeer(msg *bcBlockRequestMessage,
 
 // Receive implements Reactor by handling 4 types of messages (look below).
 func (bcR *BlockchainReactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
-	if cfg.DynamicConfig.GetEnableConsensusIPWhitelist() {
-		okIP := false
-		for _, ip := range cfg.DynamicConfig.GetConsensusIPWhitelist() {
-			if src.RemoteIP().String() == ip {
-				okIP = true
-				break
-			}
-		}
-		if !okIP {
-			bcR.Logger.Error("consensus msg:IP not in whitelist", "IP", src.RemoteIP().String())
-			return
-		}
-	}
-
 	msg, err := decodeMsg(msgBytes)
 	if err != nil {
 		bcR.Logger.Error("Error decoding message", "src", src, "chId", chID, "msg", msg, "err", err, "bytes", msgBytes)
@@ -236,6 +222,13 @@ func (bcR *BlockchainReactor) Receive(chID byte, src p2p.Peer, msgBytes []byte)
 	case *bcBlockRequestMessage:
 		bcR.respondToPeer(msg, src)
 	case *bcBlockResponseMessage:
+		if cfg.DynamicConfig.GetEnableConsensusIPWhitelist() {
+			okIP := cfg.DynamicConfig.GetConsensusIPWhitelist()[src.RemoteIP().String()]
+			if !okIP {
+				bcR.Logger.Error("consensus msg:IP not in whitelist", "IP", src.RemoteIP().String())
+				return
+			}
+		}
 		bcR.Logger.Info("AddBlock.", "Height", msg.Block.Height, "Peer", src.ID())
 		bcR.pool.AddBlock(src.ID(), msg, len(msgBytes))
 	case *bcStatusRequestMessage:

diff --git a/libs/tendermint/config/dynamic_config_okchain.go b/libs/tendermint/config/dynamic_config_okchain.go
@@ -42,7 +42,7 @@ type IDynamicConfig interface {
 	GetPendingPoolBlacklist() string
 	GetMaxTxLimitPerPeer() uint64
 	GetEnableConsensusIPWhitelist() bool
-	GetConsensusIPWhitelist() []string
+	GetConsensusIPWhitelist() map[string]bool
 }
 
 var DynamicConfig IDynamicConfig = MockDynamicConfig{}
@@ -238,6 +238,6 @@ func (c MockDynamicConfig) GetMaxTxLimitPerPeer() uint64 {
 
 func (c MockDynamicConfig) GetEnableConsensusIPWhitelist() bool { return false }
 
-func (c MockDynamicConfig) GetConsensusIPWhitelist() []string {
-	return []string{}
+func (c MockDynamicConfig) GetConsensusIPWhitelist() map[string]bool {
+	return map[string]bool{}
 }
diff --git a/libs/tendermint/consensus/reactor.go b/libs/tendermint/consensus/reactor.go
@@ -344,13 +344,7 @@ func (conR *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 	}
 
 	if cfg.DynamicConfig.GetEnableConsensusIPWhitelist() {
-		okIP := false
-		for _, ip := range cfg.DynamicConfig.GetConsensusIPWhitelist() {
-			if src.RemoteIP().String() == ip {
-				okIP = true
-				break
-			}
-		}
+		okIP := cfg.DynamicConfig.GetConsensusIPWhitelist()[src.RemoteIP().String()]
 		if !okIP {
 			conR.Logger.Error("consensus msg:IP not in whitelist", "IP", src.RemoteIP().String())
 			return

diff --git a/libs/tendermint/evidence/reactor.go b/libs/tendermint/evidence/reactor.go
@@ -65,13 +65,7 @@ func (evR *Reactor) AddPeer(peer p2p.Peer) {
 // It adds any received evidence to the evpool.
 func (evR *Reactor) Receive(chID byte, src p2p.Peer, msgBytes []byte) {
 	if cfg.DynamicConfig.GetEnableConsensusIPWhitelist() {
-		okIP := false
-		for _, ip := range cfg.DynamicConfig.GetConsensusIPWhitelist() {
-			if src.RemoteIP().String() == ip {
-				okIP = true
-				break
-			}
-		}
+		okIP := cfg.DynamicConfig.GetConsensusIPWhitelist()[src.RemoteIP().String()]
 		if !okIP {
 			evR.Logger.Error("consensus msg:IP not in whitelist", "IP", src.RemoteIP().String())
 			return