From 97f626df77425feb22bc1e1e0f5c206605a9d246 Mon Sep 17 00:00:00 2001
From: MSalopek <matija.salopek994@gmail.com>
Date: Thu, 1 Feb 2024 11:59:04 +0100
Subject: [PATCH] fix(x/feegrant)!: limit expired feegrant pruning to 200 per
 block (#19314)

---
 x/feegrant/keeper/keeper.go      | 10 ++++++-
 x/feegrant/keeper/keeper_test.go | 46 +++++++++++++++++++++++++++-----
 x/feegrant/module/abci.go        |  2 +-
 3 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/x/feegrant/keeper/keeper.go b/x/feegrant/keeper/keeper.go
index bb36f9c7964f..4510218cc6bb 100644
--- a/x/feegrant/keeper/keeper.go
+++ b/x/feegrant/keeper/keeper.go
@@ -323,17 +323,25 @@ func (k Keeper) addToFeeAllowanceQueue(ctx sdk.Context, grantKey []byte, exp *ti
 	store.Set(feegrant.FeeAllowancePrefixQueue(exp, grantKey), []byte{})
 }
 
+// NOTE: backport from v50
 // RemoveExpiredAllowances iterates grantsByExpiryQueue and deletes the expired grants.
-func (k Keeper) RemoveExpiredAllowances(ctx sdk.Context) {
+func (k Keeper) RemoveExpiredAllowances(ctx sdk.Context, limit int) {
 	exp := ctx.BlockTime()
 	store := ctx.KVStore(k.storeKey)
 	iterator := store.Iterator(feegrant.FeeAllowanceQueueKeyPrefix, sdk.InclusiveEndBytes(feegrant.AllowanceByExpTimeKey(&exp)))
 	defer iterator.Close()
 
+	count := 0
 	for ; iterator.Valid(); iterator.Next() {
 		store.Delete(iterator.Key())
 
 		granter, grantee := feegrant.ParseAddressesFromFeeAllowanceQueueKey(iterator.Key())
 		store.Delete(feegrant.FeeAllowanceKey(granter, grantee))
+
+		// limit the amount of iterations to avoid taking too much time
+		count++
+		if count == limit {
+			break
+		}
 	}
 }
diff --git a/x/feegrant/keeper/keeper_test.go b/x/feegrant/keeper/keeper_test.go
index 16d10a532cd6..348c6326bdf5 100644
--- a/x/feegrant/keeper/keeper_test.go
+++ b/x/feegrant/keeper/keeper_test.go
@@ -33,7 +33,7 @@ func TestKeeperTestSuite(t *testing.T) {
 }
 
 func (suite *KeeperTestSuite) SetupTest() {
-	suite.addrs = simtestutil.CreateIncrementalAccounts(4)
+	suite.addrs = simtestutil.CreateIncrementalAccounts(20)
 	key := sdk.NewKVStoreKey(feegrant.StoreKey)
 	testCtx := testutil.DefaultContextWithDB(suite.T(), key, sdk.NewTransientStoreKey("transient_test"))
 	encCfg := moduletestutil.MakeTestEncodingConfig(module.AppModuleBasic{})
@@ -41,10 +41,9 @@ func (suite *KeeperTestSuite) SetupTest() {
 	// setup gomock and initialize some globally expected executions
 	ctrl := gomock.NewController(suite.T())
 	suite.accountKeeper = feegranttestutil.NewMockAccountKeeper(ctrl)
-	suite.accountKeeper.EXPECT().GetAccount(gomock.Any(), suite.addrs[0]).Return(authtypes.NewBaseAccountWithAddress(suite.addrs[0])).AnyTimes()
-	suite.accountKeeper.EXPECT().GetAccount(gomock.Any(), suite.addrs[1]).Return(authtypes.NewBaseAccountWithAddress(suite.addrs[1])).AnyTimes()
-	suite.accountKeeper.EXPECT().GetAccount(gomock.Any(), suite.addrs[2]).Return(authtypes.NewBaseAccountWithAddress(suite.addrs[2])).AnyTimes()
-	suite.accountKeeper.EXPECT().GetAccount(gomock.Any(), suite.addrs[3]).Return(authtypes.NewBaseAccountWithAddress(suite.addrs[3])).AnyTimes()
+	for i := 0; i < len(suite.addrs); i++ {
+		suite.accountKeeper.EXPECT().GetAccount(gomock.Any(), suite.addrs[i]).Return(authtypes.NewBaseAccountWithAddress(suite.addrs[i])).AnyTimes()
+	}
 
 	suite.feegrantKeeper = keeper.NewKeeper(encCfg.Codec, key, suite.accountKeeper)
 	suite.ctx = testCtx.Ctx
@@ -447,7 +446,7 @@ func (suite *KeeperTestSuite) TestPruneGrants() {
 			}
 			err := suite.feegrantKeeper.GrantAllowance(suite.ctx, tc.granter, tc.grantee, tc.allowance)
 			suite.NoError(err)
-			suite.feegrantKeeper.RemoveExpiredAllowances(tc.ctx)
+			suite.feegrantKeeper.RemoveExpiredAllowances(tc.ctx, 5)
 			grant, err := suite.feegrantKeeper.GetAllowance(tc.ctx, tc.granter, tc.grantee)
 			if tc.expErrMsg != "" {
 				suite.Error(err)
@@ -461,3 +460,38 @@ func (suite *KeeperTestSuite) TestPruneGrants() {
 		})
 	}
 }
+
+func (suite *KeeperTestSuite) TestPruneGrantsEdgecases() {
+	eth := sdk.NewCoins(sdk.NewInt64Coin("eth", 123))
+	now := suite.ctx.BlockTime()
+	oneYearExpiry := now.AddDate(1, 0, 0)
+
+	granter := suite.addrs[1]
+	grantee1 := suite.addrs[2]
+	err := suite.feegrantKeeper.GrantAllowance(suite.ctx, granter, grantee1, &feegrant.BasicAllowance{
+		SpendLimit: eth,
+		Expiration: &oneYearExpiry,
+	})
+	suite.NoError(err)
+
+	grantee2 := suite.addrs[3]
+	err = suite.feegrantKeeper.GrantAllowance(suite.ctx, granter, grantee2, &feegrant.BasicAllowance{
+		SpendLimit: eth,
+		Expiration: &now,
+	})
+	suite.NoError(err)
+
+	// expect 2 active grants
+	grantsBeforePrune, err := suite.feegrantKeeper.AllowancesByGranter(suite.ctx, &feegrant.QueryAllowancesByGranterRequest{Granter: granter.String()})
+	suite.NoError(err)
+	suite.Len(grantsBeforePrune.Allowances, 2)
+
+	// use blocktime that would result in both grants being expired
+	expireCtx := suite.ctx.WithBlockTime(now.AddDate(1, 0, 1))
+
+	// expect 1 grant to be removed due to the imposed limit
+	suite.feegrantKeeper.RemoveExpiredAllowances(expireCtx, 1)
+	grantsAfterPrune, err := suite.feegrantKeeper.AllowancesByGranter(suite.ctx, &feegrant.QueryAllowancesByGranterRequest{Granter: granter.String()})
+	suite.NoError(err)
+	suite.Len(grantsAfterPrune.Allowances, 1)
+}
diff --git a/x/feegrant/module/abci.go b/x/feegrant/module/abci.go
index 34c19c528d6a..e6577806080e 100644
--- a/x/feegrant/module/abci.go
+++ b/x/feegrant/module/abci.go
@@ -6,5 +6,5 @@ import (
 )
 
 func EndBlocker(ctx sdk.Context, k keeper.Keeper) {
-	k.RemoveExpiredAllowances(ctx)
+	k.RemoveExpiredAllowances(ctx, 200)
 }