From 32f02def3fab76bbbfb0f5f68fc90de1771b1fc6 Mon Sep 17 00:00:00 2001
From: Chen Chen <chen1233216@hotmail.com>
Date: Sat, 1 Feb 2025 02:02:23 +0800
Subject: [PATCH] [p2p] introduce account rate limit  (#4545)

---
 chainservice/builder.go       |  8 ++++++++
 chainservice/chainservice.go  | 25 +++++++++++++++++++++++++
 e2etest/local_actpool_test.go |  1 +
 p2p/agent.go                  |  3 +++
 server/itx/server.go          |  1 +
 5 files changed, 38 insertions(+)

diff --git a/chainservice/builder.go b/chainservice/builder.go
index 69664e37f9..dad8f4bd16 100644
--- a/chainservice/builder.go
+++ b/chainservice/builder.go
@@ -11,6 +11,7 @@ import (
 	"net/url"
 	"time"
 
+	"github.com/iotexproject/go-pkgs/cache"
 	"github.com/iotexproject/iotex-address/address"
 	"github.com/iotexproject/iotex-election/committee"
 	"github.com/iotexproject/iotex-proto/golang/iotextypes"
@@ -102,6 +103,13 @@ func (builder *Builder) SetP2PAgent(agent p2p.Agent) *Builder {
 	return builder
 }
 
+func (builder *Builder) SetAccountRateLimit(r int) *Builder {
+	builder.createInstance()
+	builder.cs.accRateLimitCfg = r
+	builder.cs.rateLimiters = cache.NewThreadSafeLruCache(10000)
+	return builder
+}
+
 // SetRPCStats sets the RPCStats instance
 func (builder *Builder) SetRPCStats(stats *nodestats.APILocalStats) *Builder {
 	builder.createInstance()
diff --git a/chainservice/chainservice.go b/chainservice/chainservice.go
index 0a4e234a10..56ab8bc800 100644
--- a/chainservice/chainservice.go
+++ b/chainservice/chainservice.go
@@ -11,8 +11,11 @@ import (
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
+	"go.uber.org/zap"
+	"golang.org/x/time/rate"
 	"google.golang.org/protobuf/proto"
 
+	"github.com/iotexproject/go-pkgs/cache"
 	"github.com/iotexproject/go-pkgs/hash"
 	"github.com/iotexproject/iotex-election/committee"
 	"github.com/iotexproject/iotex-proto/golang/iotexrpc"
@@ -79,6 +82,8 @@ type ChainService struct {
 	apiStats                 *nodestats.APILocalStats
 	blockTimeCalculator      *blockutil.BlockTimeCalculator
 	actionsync               *actsync.ActionSync
+	rateLimiters             cache.LRUCache
+	accRateLimitCfg          int
 }
 
 // Start starts the server
@@ -102,6 +107,17 @@ func (cs *ChainService) HandleAction(ctx context.Context, actPb *iotextypes.Acti
 	if err != nil {
 		return err
 	}
+	if cs.accRateLimitCfg > 0 {
+		sender := ""
+		if act.SenderAddress() != nil {
+			sender = act.SenderAddress().String()
+		}
+		limiter := cs.getRateLimiter(sender)
+		if !limiter.Allow() {
+			log.L().Debug("rate limit exceeded", zap.String("sender", act.SenderAddress().String()))
+			return nil
+		}
+	}
 	ctx = protocol.WithRegistry(ctx, cs.registry)
 	err = cs.actpool.Add(ctx, act)
 	if err != nil {
@@ -248,3 +264,12 @@ func (cs *ChainService) NewAPIServer(cfg api.Config, plugins map[int]interface{}
 
 	return svr, nil
 }
+
+func (cs *ChainService) getRateLimiter(sender string) *rate.Limiter {
+	limiter, exists := cs.rateLimiters.Get(sender)
+	if !exists {
+		limiter = rate.NewLimiter(rate.Limit(cs.accRateLimitCfg), 2*cs.accRateLimitCfg) // account limit request per second with a burst of *2
+		cs.rateLimiters.Add(sender, limiter)
+	}
+	return limiter.(*rate.Limiter)
+}
diff --git a/e2etest/local_actpool_test.go b/e2etest/local_actpool_test.go
index 56410381e4..b72406bd48 100644
--- a/e2etest/local_actpool_test.go
+++ b/e2etest/local_actpool_test.go
@@ -183,6 +183,7 @@ func newActPoolConfig(t *testing.T) (config.Config, error) {
 	cfg.ActPool.MinGasPriceStr = "0"
 	cfg.Consensus.Scheme = config.NOOPScheme
 	cfg.Network.Port = testutil.RandomPort()
+	cfg.Network.AccountRateLimit = 0
 
 	sk, err := crypto.GenerateKey()
 	if err != nil {
diff --git a/p2p/agent.go b/p2p/agent.go
index aa17a40fb4..c5038e97ca 100644
--- a/p2p/agent.go
+++ b/p2p/agent.go
@@ -95,6 +95,8 @@ type (
 		PrivateNetworkPSK string              `yaml:"privateNetworkPSK"`
 		MaxPeers          int                 `yaml:"maxPeers"`
 		MaxMessageSize    int                 `yaml:"maxMessageSize"`
+		// AccountRateLimit is the maximum number of requests per second per account.
+		AccountRateLimit int `yaml:"accountRateLimit"`
 	}
 
 	// Agent is the agent to help the blockchain node connect into the P2P networks and send/receive messages
@@ -145,6 +147,7 @@ var DefaultConfig = Config{
 	PrivateNetworkPSK: "",
 	MaxPeers:          30,
 	MaxMessageSize:    p2p.DefaultConfig.MaxMessageSize,
+	AccountRateLimit:  100,
 }
 
 // NewDummyAgent creates a dummy p2p agent
diff --git a/server/itx/server.go b/server/itx/server.go
index 8796795036..959b66da5c 100644
--- a/server/itx/server.go
+++ b/server/itx/server.go
@@ -73,6 +73,7 @@ func newServer(cfg config.Config, testing bool) (*Server, error) {
 	var cs *chainservice.ChainService
 	builder := chainservice.NewBuilder(cfg)
 	builder.SetP2PAgent(p2pAgent)
+	builder.SetAccountRateLimit(cfg.Network.AccountRateLimit)
 	rpcStats := nodestats.NewAPILocalStats()
 	builder.SetRPCStats(rpcStats)
 	if testing {