[p2p] feat: add maximum peers limit (#4272)

cmd/harmony/config_migrations.go

@@ -3,6 +3,7 @@ package main
 import (
 	"errors"
 	"fmt"
+	"sort"
 
 	goversion "github.com/hashicorp/go-version"
 	"github.com/pelletier/go-toml"
@@ -273,13 +274,42 @@ func init() {
 		confTree.Set("Version", "2.5.5")
 		return confTree
 	}
-
 	migrations["2.5.5"] = func(confTree *toml.Tree) *toml.Tree {
 		if confTree.Get("Log.Console") == nil {
 			confTree.Set("Log.Console", defaultConfig.Log.Console)
 		}
-
 		confTree.Set("Version", "2.5.6")
 		return confTree
 	}
+	migrations["2.5.6"] = func(confTree *toml.Tree) *toml.Tree {
+		if confTree.Get("P2P.MaxPeers") == nil {
+			confTree.Set("P2P.MaxPeers", defaultConfig.P2P.MaxPeers)
+		}
+		confTree.Set("Version", "2.5.7")
+		return confTree
+	}
+
+	// check that the latest version here is the same as in default.go
+	largestKey := getNextVersion(migrations)
+	if largestKey != tomlConfigVersion {
+		panic(fmt.Sprintf("next migration value: %s, toml version: %s", largestKey, tomlConfigVersion))
+	}
+}
+
+func getNextVersion(x map[string]configMigrationFunc) string {
+	versionMap := make(map[string]interface{}, 1)
+	versionMap["Version"] = "FakeVersion"
+	tree, _ := toml.TreeFromMap(versionMap)
+
+	// needs to be sorted in case the order is incorrect
+	keys := make([]string, len(x))
+	i := 0
+	for k := range x {
+		keys[i] = k
+		i++
+	}
+	sort.Strings(keys)
+	requiredFunc := x[keys[len(keys)-1]]
+	tree = requiredFunc(tree)
+	return tree.Get("Version").(string)
 }

cmd/harmony/default.go

@@ -5,7 +5,7 @@ import (
 	nodeconfig "github.com/harmony-one/harmony/internal/configs/node"
 )
 
-const tomlConfigVersion = "2.5.5"
+const tomlConfigVersion = "2.5.7"
 
 const (
 	defNetworkType = nodeconfig.Mainnet
@@ -31,6 +31,7 @@ var defaultConfig = harmonyconfig.HarmonyConfig{
 		DiscConcurrency:      nodeconfig.DefaultP2PConcurrency,
 		MaxConnsPerIP:        nodeconfig.DefaultMaxConnPerIP,
 		DisablePrivateIPScan: false,
+		MaxPeers:             nodeconfig.DefaultMaxPeers,
 	},
 	HTTP: harmonyconfig.HttpConfig{
 		Enabled:        true,

cmd/harmony/flags.go

@@ -62,6 +62,7 @@ var (
 		legacyKeyFileFlag,
 		p2pDisablePrivateIPScanFlag,
 		maxConnPerIPFlag,
+		maxPeersFlag,
 	}
 
 	httpFlags = []cli.Flag{
@@ -573,7 +574,12 @@ var (
 	}
 	maxConnPerIPFlag = cli.IntFlag{
 		Name:     "p2p.security.max-conn-per-ip",
-		Usage:    "maximum number of connections allowed per node",
+		Usage:    "maximum number of connections allowed per remote node, 0 means no limit",
+		DefValue: defaultConfig.P2P.MaxConnsPerIP,
+	}
+	maxPeersFlag = cli.IntFlag{
+		Name:     "p2p.security.max-peers",
+		Usage:    "maximum number of peers allowed, 0 means no limit",
 		DefValue: defaultConfig.P2P.MaxConnsPerIP,
 	}
 )
@@ -608,6 +614,10 @@ func applyP2PFlags(cmd *cobra.Command, config *harmonyconfig.HarmonyConfig) {
 		config.P2P.MaxConnsPerIP = cli.GetIntFlagValue(cmd, maxConnPerIPFlag)
 	}
 
+	if cli.IsFlagChanged(cmd, maxPeersFlag) {
+		config.P2P.MaxPeers = int64(cli.GetIntFlagValue(cmd, maxPeersFlag))
+	}
+
 	if cli.IsFlagChanged(cmd, p2pDisablePrivateIPScanFlag) {
 		config.P2P.DisablePrivateIPScan = cli.GetBoolFlagValue(cmd, p2pDisablePrivateIPScanFlag)
 	}

cmd/harmony/flags_test.go

@@ -64,6 +64,7 @@ func TestHarmonyFlags(t *testing.T) {
 					DiscConcurrency:      5,
 					MaxConnsPerIP:        5,
 					DisablePrivateIPScan: false,
+					MaxPeers:             defaultConfig.P2P.MaxPeers,
 				},
 				HTTP: harmonyconfig.HttpConfig{
 					Enabled:        true,
@@ -390,6 +391,7 @@ func TestP2PFlags(t *testing.T) {
 				DHTDataStore:         &defDataStore,
 				MaxConnsPerIP:        10,
 				DisablePrivateIPScan: false,
+				MaxPeers:             defaultConfig.P2P.MaxPeers,
 			},
 		},
 		{
@@ -400,6 +402,7 @@ func TestP2PFlags(t *testing.T) {
 				KeyFile:              "./key.file",
 				MaxConnsPerIP:        10,
 				DisablePrivateIPScan: false,
+				MaxPeers:             defaultConfig.P2P.MaxPeers,
 			},
 		},
 		{
@@ -411,6 +414,7 @@ func TestP2PFlags(t *testing.T) {
 				DiscConcurrency:      5,
 				MaxConnsPerIP:        5,
 				DisablePrivateIPScan: false,
+				MaxPeers:             defaultConfig.P2P.MaxPeers,
 			},
 		},
 		{
@@ -422,6 +426,19 @@ func TestP2PFlags(t *testing.T) {
 				DiscConcurrency:      nodeconfig.DefaultP2PConcurrency,
 				MaxConnsPerIP:        nodeconfig.DefaultMaxConnPerIP,
 				DisablePrivateIPScan: true,
+				MaxPeers:             defaultConfig.P2P.MaxPeers,
+			},
+		},
+		{
+			args: []string{"--p2p.security.max-peers", "100"},
+			expConfig: harmonyconfig.P2pConfig{
+				Port:                 nodeconfig.DefaultP2PPort,
+				IP:                   nodeconfig.DefaultPublicListenIP,
+				KeyFile:              "./.hmykey",
+				DiscConcurrency:      nodeconfig.DefaultP2PConcurrency,
+				MaxConnsPerIP:        nodeconfig.DefaultMaxConnPerIP,
+				DisablePrivateIPScan: defaultConfig.P2P.DisablePrivateIPScan,
+				MaxPeers:             100,
 			},
 		},
 	}

cmd/harmony/main.go

@@ -619,6 +619,7 @@ func createGlobalConfig(hc harmonyconfig.HarmonyConfig) (*nodeconfig.ConfigType,
 		DiscConcurrency:      hc.P2P.DiscConcurrency,
 		MaxConnPerIP:         hc.P2P.MaxConnsPerIP,
 		DisablePrivateIPScan: hc.P2P.DisablePrivateIPScan,
+		MaxPeers:             hc.P2P.MaxPeers,
 	})
 	if err != nil {
 		return nil, errors.Wrap(err, "cannot create P2P network host")

internal/configs/harmony/harmony.go

@@ -54,6 +54,7 @@ type P2pConfig struct {
 	DiscConcurrency      int     // Discovery Concurrency value
 	MaxConnsPerIP        int
 	DisablePrivateIPScan bool
+	MaxPeers             int64
 }
 
 type GeneralConfig struct {

internal/configs/node/network.go

@@ -59,7 +59,10 @@ const (
 	DefaultPrometheusPort = 9900
 	// DefaultP2PConcurrency is the default P2P concurrency, 0 means is set the default value of P2P Discovery, the actual value is 10
 	DefaultP2PConcurrency = 0
+	// DefaultMaxConnPerIP is the maximum number of connections to/from a remote IP
 	DefaultMaxConnPerIP = 10
+	// DefaultMaxPeers is the maximum number of remote peers, with 0 representing no limit
+	DefaultMaxPeers = 0
 )
 
 const (

p2p/host.go

@@ -87,6 +87,7 @@ type HostConfig struct {
 	DiscConcurrency      int
 	MaxConnPerIP         int
 	DisablePrivateIPScan bool
+	MaxPeers             int64
 }
 
 func init() {
@@ -182,7 +183,7 @@ func NewHost(cfg HostConfig) (Host, error) {
 	self.PeerID = p2pHost.ID()
 	subLogger := utils.Logger().With().Str("hostID", p2pHost.ID().Pretty()).Logger()
 
-	security := security.NewManager(cfg.MaxConnPerIP)
+	security := security.NewManager(cfg.MaxConnPerIP, cfg.MaxPeers)
 	// has to save the private key for host
 	h := &HostV2{
 		h:             p2pHost,

p2p/security/security.go

@@ -3,6 +3,7 @@ package security
 import (
 	"fmt"
 	"sync"
+	"sync/atomic"
 
 	"github.com/harmony-one/harmony/internal/utils"
 	libp2p_network "github.com/libp2p/go-libp2p-core/network"
@@ -17,14 +18,65 @@ type Security interface {
 
 type Manager struct {
 	maxConnPerIP int
+	maxPeers     int64
 
 	mutex sync.Mutex
-	peers sync.Map // All the connected nodes, key is the Peer's IP, value is the peer's ID array
+	peers peerMap // All the connected nodes, key is the Peer's IP, value is the peer's ID array
 }
 
-func NewManager(maxConnPerIP int) *Manager {
+type peerMap struct {
+	count int64
+	peers sync.Map
+}
+
+func (peerMap *peerMap) Len() int64 {
+	return atomic.LoadInt64(&peerMap.count)
+}
+
+func (peerMap *peerMap) Store(key, value interface{}) {
+	// only increment if you didn't have this key
+	hasKey := peerMap.HasKey(key)
+	peerMap.peers.Store(key, value)
+	if !hasKey {
+		atomic.AddInt64(&peerMap.count, 1)
+	}
+}
+
+func (peerMap *peerMap) HasKey(key interface{}) bool {
+	hasKey := false
+	peerMap.peers.Range(func(k, v interface{}) bool {
+		if k == key {
+			hasKey = true
+			return false
+		}
+		return true
+	})
+	return hasKey
+}
+
+func (peerMap *peerMap) Delete(key interface{}) {
+	peerMap.peers.Delete(key)
+	atomic.AddInt64(&peerMap.count, -1)
+}
+
+func (peerMap *peerMap) Load(key interface{}) (value interface{}, ok bool) {
+	return peerMap.peers.Load(key)
+}
+
+func (peerMap *peerMap) Range(f func(key, value any) bool) {
+	peerMap.peers.Range(f)
+}
+
+func NewManager(maxConnPerIP int, maxPeers int64) *Manager {
+	if maxConnPerIP < 0 {
+		panic("maximum connections per IP must not be negative")
+	}
+	if maxPeers < 0 {
+		panic("maximum peers must not be negative")
+	}
 	return &Manager{
 		maxConnPerIP: maxConnPerIP,
+		maxPeers:     maxPeers,
 	}
 }
 
@@ -32,12 +84,12 @@ func (m *Manager) OnConnectCheck(net libp2p_network.Network, conn libp2p_network
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	ip, err := getIP(conn)
+	remoteIp, err := getRemoteIP(conn)
 	if err != nil {
-		return errors.Wrap(err, "failed on get ip")
+		return errors.Wrap(err, "failed on get remote ip")
 	}
 
-	value, ok := m.peers.Load(ip)
+	value, ok := m.peers.Load(remoteIp)
 	if !ok {
 		value = []string{}
 	}
@@ -54,13 +106,24 @@ func (m *Manager) OnConnectCheck(net libp2p_network.Network, conn libp2p_network
 		peers = append(peers, peerID)
 	}
 
-	if len(peers) > m.maxConnPerIP {
-		utils.Logger().Warn().Int("len(peers)", len(peers)).Int("maxConnPerIP", m.maxConnPerIP).
-			Msg("Too much peers, closing")
+	if m.maxConnPerIP > 0 && len(peers) > m.maxConnPerIP {
+		utils.Logger().Warn().
+			Int("len(peers)", len(peers)).
+			Int("maxConnPerIP", m.maxConnPerIP).
+			Msgf("too many connections from %s, closing", remoteIp)
 		return net.ClosePeer(conn.RemotePeer())
 	}
 
-	m.peers.Store(ip, peers)
+	currentPeerCount := m.peers.Len()
+	// only limit addition if it's a new peer and not an existing peer with new connection
+	if m.maxPeers > 0 && currentPeerCount >= m.maxPeers && !m.peers.HasKey(remoteIp) {
+		utils.Logger().Warn().
+			Int64("connected peers", currentPeerCount).
+			Str("new peer", remoteIp).
+			Msg("too many peers, closing")
+		return net.ClosePeer(conn.RemotePeer())
+	}
+	m.peers.Store(remoteIp, peers)
 	return nil
 }
 
@@ -68,7 +131,7 @@ func (m *Manager) OnDisconnectCheck(conn libp2p_network.Conn) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	ip, err := getIP(conn)
+	ip, err := getRemoteIP(conn)
 	if err != nil {
 		return errors.Wrap(err, "failed on get ip")
 	}
@@ -87,9 +150,10 @@ func (m *Manager) OnDisconnectCheck(conn libp2p_network.Conn) error {
 	index, ok := find(peers, peerID)
 	if ok {
 		peers = append(peers[:index], peers[index+1:]...)
-		m.peers.Store(ip, peers)
 		if len(peers) == 0 {
 			m.peers.Delete(ip)
+		} else {
+			m.peers.Store(ip, peers)
 		}
 	}
 
@@ -106,7 +170,7 @@ func find(slice []string, val string) (int, bool) {
 	return -1, false
 }
 
-func getIP(conn libp2p_network.Conn) (string, error) {
+func getRemoteIP(conn libp2p_network.Conn) (string, error) {
 	for _, protocol := range conn.RemoteMultiaddr().Protocols() {
 		switch protocol.Code {
 		case ma.P_IP4:

p2p/security/security_test.go

@@ -7,7 +7,6 @@ import (
 	"time"
 
 	"github.com/libp2p/go-libp2p"
-	"github.com/libp2p/go-libp2p-core/crypto"
 	ic "github.com/libp2p/go-libp2p-core/crypto"
 	"github.com/libp2p/go-libp2p-core/host"
 	"github.com/libp2p/go-libp2p-core/network"
@@ -58,7 +57,7 @@ func TestManager_OnConnectCheck(t *testing.T) {
 	defer h1.Close()
 
 	fakeHost := &fakeHost{}
-	security := NewManager(2)
+	security := NewManager(2, 1)
 	h1.Network().Notify(fakeHost)
 	fakeHost.SetConnectCallback(security.OnConnectCheck)
 	fakeHost.SetDisconnectCallback(security.OnDisconnectCheck)
@@ -103,7 +102,7 @@ func TestManager_OnDisconnectCheck(t *testing.T) {
 	defer h1.Close()
 
 	fakeHost := &fakeHost{}
-	security := NewManager(2)
+	security := NewManager(2, 0)
 	h1.Network().Notify(fakeHost)
 	fakeHost.SetConnectCallback(security.OnConnectCheck)
 	fakeHost.SetDisconnectCallback(security.OnDisconnectCheck)
@@ -130,7 +129,7 @@ func TestManager_OnDisconnectCheck(t *testing.T) {
 }
 
 func newPeer(port int) (host.Host, error) {
-	priv, _, err := crypto.GenerateKeyPair(crypto.RSA, 2048)
+	priv, _, err := ic.GenerateKeyPair(ic.RSA, 2048)
 	if err != nil {
 		return nil, err
 	}
@@ -160,8 +159,8 @@ func (conn *fakeConn) ID() string                                        { retur
 func (conn *fakeConn) NewStream(context.Context) (network.Stream, error) { return nil, nil }
 func (conn *fakeConn) GetStreams() []network.Stream                      { return nil }
 func (conn *fakeConn) Stat() network.Stat                                { return network.Stat{} }
-func TestGetIP(t *testing.T) {
-	ip, err := getIP(&fakeConn{})
+func TestGetRemoteIP(t *testing.T) {
+	ip, err := getRemoteIP(&fakeConn{})
 	assert.Nil(t, err)
 	assert.Equal(t, "fe80::7802:31ff:fee9:c093", ip)
 }

rosetta/infra/harmony-mainnet.conf

@@ -1,4 +1,4 @@
-Version = "2.5.5"
+Version = "2.5.7"
 
 [BLSKeys]
   KMSConfigFile = ""
@@ -33,6 +33,7 @@ Version = "2.5.5"
   IsOffline = false
   NoStaking = true
   NodeType = "explorer"
+  RunElasticMode = false
   ShardID = 0
   TraceEnable = false
 
@@ -66,6 +67,7 @@ Version = "2.5.5"
   IP = "0.0.0.0"
   KeyFile = "./.hmykey"
   MaxConnsPerIP = 10
+  MaxPeers = 0
   Port = 9000
 
 [Pprof]
@@ -107,9 +109,9 @@ Version = "2.5.5"
   AccountSlots = 16
   AllowedTxsFile = "./.hmy/allowedtxs.txt"
   BlacklistFile = "./.hmy/blacklist.txt"
+  GlobalSlots = 5120
   LocalAccountsFile = "./.hmy/locals.txt"
   RosettaFixFile = "./rosetta_local_fix.csv"
-  GlobalSlots = 5120
 
 [WS]
   AuthPort = 9801

rosetta/infra/harmony-pstn.conf

@@ -1,4 +1,4 @@
-Version = "2.5.4"
+Version = "2.5.7"
 
 [BLSKeys]
   KMSConfigFile = ""
@@ -33,6 +33,7 @@ Version = "2.5.4"
   IsOffline = false
   NoStaking = true
   NodeType = "explorer"
+  RunElasticMode = false
   ShardID = 0
   TraceEnable = false
 
@@ -66,6 +67,7 @@ Version = "2.5.4"
   IP = "0.0.0.0"
   KeyFile = "./.hmykey"
   MaxConnsPerIP = 10
+  MaxPeers = 0
   Port = 9000
 
 [Pprof]
@@ -107,9 +109,9 @@ Version = "2.5.4"
   AccountSlots = 16
   AllowedTxsFile = "./.hmy/allowedtxs.txt"
   BlacklistFile = "./.hmy/blacklist.txt"
+  GlobalSlots = 5120
   LocalAccountsFile = "./.hmy/locals.txt"
   RosettaFixFile = ""
-  GlobalSlots = 5120
 
 [WS]
   AuthPort = 9801