From 81a41354b3caff15a827e4d68f9aba5e8301e14b Mon Sep 17 00:00:00 2001
From: YoshihitoAso <yoshihito.asoh@gmail.com>
Date: Thu, 8 Feb 2024 19:54:06 +0900
Subject: [PATCH] fix: memory leak from policy registration

---
 consensus/istanbul/config.go                  | 13 +++++++
 .../istanbul/validator/proposerpolicy_test.go | 34 ++++++++++++++-----
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/consensus/istanbul/config.go b/consensus/istanbul/config.go
index cd362c481..b14ce4c6a 100644
--- a/consensus/istanbul/config.go
+++ b/consensus/istanbul/config.go
@@ -24,6 +24,7 @@ import (
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/math"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/naoina/toml"
 )
@@ -35,6 +36,8 @@ const (
 	Sticky
 )
 
+const MaxValidatorSetInRegistry = 128 // Max number of ValidatorSet in the registry
+
 // ProposerPolicy represents the Validator Proposer Policy
 type ProposerPolicy struct {
 	Id         ProposerPolicyId    // Could be RoundRobin or Sticky
@@ -111,7 +114,13 @@ func (p *ProposerPolicy) RegisterValidatorSet(valSet ValidatorSet) {
 		p.registry = []ValidatorSet{valSet}
 	} else {
 		p.registry = append(p.registry, valSet)
+		// Non-validators never call ClearRegistry
+		// Limit the registry length to MaxValidatorSetInRegistry to prevent unexpected leaks
+		if len(p.registry) > MaxValidatorSetInRegistry {
+			p.registry = p.registry[1:]
+		}
 	}
+	log.Debug("Validator Policy Registry", "length", p.GetRegistrySize())
 }
 
 // ClearRegistry removes any ValidatorSet from the ProposerPolicy registry
@@ -122,6 +131,10 @@ func (p *ProposerPolicy) ClearRegistry() {
 	p.registry = nil
 }
 
+func (p *ProposerPolicy) GetRegistrySize() int {
+	return len(p.registry)
+}
+
 type Config struct {
 	RequestTimeout           uint64                `toml:",omitempty"` // The timeout for each Istanbul round in milliseconds.
 	BlockPeriod              uint64                `toml:",omitempty"` // Default minimum difference between two consecutive block's timestamps in second
diff --git a/consensus/istanbul/validator/proposerpolicy_test.go b/consensus/istanbul/validator/proposerpolicy_test.go
index 0159787bf..0465af9b4 100644
--- a/consensus/istanbul/validator/proposerpolicy_test.go
+++ b/consensus/istanbul/validator/proposerpolicy_test.go
@@ -24,15 +24,19 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+var (
+	addr1    = common.HexToAddress("0xc53f2189bf6d7bf56722731787127f90d319e112")
+	addr2    = common.HexToAddress("0xed2d479591fe2c5626ce09bca4ed2a62e00e5bc2")
+	addr3    = common.HexToAddress("0xc8417f834995aaeb35f342a67a4961e19cd4735c")
+	addr4    = common.HexToAddress("0x784ae51f5013b51c8360afdf91c6bc5a16f586ea")
+	addr5    = common.HexToAddress("0xecf0974e6f0630fd91ea4da8399cdb3f59e5220f")
+	addr6    = common.HexToAddress("0x411c4d11acd714b82a5242667e36de14b9e1d10b")
+	addr7    = common.HexToAddress("0x681381b3D0DaaC179d95aCc9e22E23da2DA670f6")
+	addrSet  = []common.Address{addr1, addr2, addr3, addr4, addr5, addr6}
+	addrSet2 = []common.Address{addr7, addr1, addr2, addr3, addr4, addr5}
+)
+
 func TestProposerPolicy(t *testing.T) {
-	addr1 := common.HexToAddress("0xc53f2189bf6d7bf56722731787127f90d319e112")
-	addr2 := common.HexToAddress("0xed2d479591fe2c5626ce09bca4ed2a62e00e5bc2")
-	addr3 := common.HexToAddress("0xc8417f834995aaeb35f342a67a4961e19cd4735c")
-	addr4 := common.HexToAddress("0x784ae51f5013b51c8360afdf91c6bc5a16f586ea")
-	addr5 := common.HexToAddress("0xecf0974e6f0630fd91ea4da8399cdb3f59e5220f")
-	addr6 := common.HexToAddress("0x411c4d11acd714b82a5242667e36de14b9e1d10b")
-
-	addrSet := []common.Address{addr1, addr2, addr3, addr4, addr5, addr6}
 	addressSortedByByte := []common.Address{addr6, addr4, addr1, addr3, addr5, addr2}
 	addressSortedByString := []common.Address{addr6, addr4, addr1, addr2, addr5, addr3}
 
@@ -51,3 +55,17 @@ func TestProposerPolicy(t *testing.T) {
 		assert.Equal(t, addressSortedByString[i].Hex(), valList[i].String(), "validatorSet not string sorted")
 	}
 }
+
+func TestProposerPolicyRegistration(t *testing.T) {
+	// Test that registration can't go beyond MaxValidatorSetInRegistry limit
+	pp := istanbul.NewRoundRobinProposerPolicy()
+	pp2 := istanbul.NewRoundRobinProposerPolicy()
+	valSet := NewSet(addrSet, pp)
+	valSet2 := NewSet(addrSet2, pp2)
+
+	for i := 0; i < istanbul.MaxValidatorSetInRegistry; i++ {
+		pp.RegisterValidatorSet(valSet)
+	}
+	pp.RegisterValidatorSet(valSet2)
+	assert.Equal(t, istanbul.MaxValidatorSetInRegistry, pp.GetRegistrySize(), "validator set not dropped")
+}