feat: adds pruning for feegrant

x/feegrant/basic_fee.go

@@ -1,6 +1,8 @@
 package feegrant
 
 import (
+	time "time"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 )
@@ -52,3 +54,7 @@ func (a BasicAllowance) ValidateBasic() error {
 
 	return nil
 }
+
+func (a BasicAllowance) ExpiresAt() (*time.Time, error) {
+	return a.Expiration, nil
+}

x/feegrant/events.go

@@ -5,6 +5,7 @@ const (
 	EventTypeUseFeeGrant    = "use_feegrant"
 	EventTypeRevokeFeeGrant = "revoke_feegrant"
 	EventTypeSetFeeGrant    = "set_feegrant"
+	EventTypeUpdateFeeGrant = "update_feegrant"
 
 	AttributeKeyGranter = "granter"
 	AttributeKeyGrantee = "grantee"

x/feegrant/fees.go

@@ -1,6 +1,8 @@
 package feegrant
 
 import (
+	"time"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -22,4 +24,7 @@ type FeeAllowanceI interface {
 	// ValidateBasic should evaluate this FeeAllowance for internal consistency.
 	// Don't allow negative amounts, or negative periods for example.
 	ValidateBasic() error
+
+	// ExpiresAt returns the expiry time of the allowance.
+	ExpiresAt() (*time.Time, error)
 }

x/feegrant/filtered_fee.go

@@ -1,6 +1,8 @@
 package feegrant
 
 import (
+	"time"
+
 	"github.com/gogo/protobuf/proto"
 
 	"github.com/cosmos/cosmos-sdk/codec/types"
@@ -103,3 +105,11 @@ func (a *AllowedMsgAllowance) ValidateBasic() error {
 
 	return allowance.ValidateBasic()
 }
+
+func (a *AllowedMsgAllowance) ExpiresAt() (*time.Time, error) {
+	allowance, err := a.GetAllowance()
+	if err != nil {
+		return nil, err
+	}
+	return allowance.ExpiresAt()
+}

x/feegrant/keeper/keeper.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"fmt"
+	"time"
 
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/tendermint/tendermint/libs/log"
@@ -49,6 +50,25 @@ func (k Keeper) GrantAllowance(ctx sdk.Context, granter, grantee sdk.AccAddress,
 
 	store := ctx.KVStore(k.storeKey)
 	key := feegrant.FeeAllowanceKey(granter, grantee)
+
+	var exp *time.Time
+	existingGrant, err := k.getGrant(ctx, grantee, granter)
+	if err != nil && existingGrant.GetAllowance() != nil {
+		grant1, err := existingGrant.GetGrant()
+		if err != nil {
+			return err
+		}
+
+		exp, err = grant1.ExpiresAt()
+		if err != nil {
+			return err
+		}
+
+		if exp != nil {
+			k.removeFromGrantQueue(ctx, feegrant.FeeAllowancePrefixQueue(exp, key))
+		}
+	}
+
 	grant, err := feegrant.NewGrant(granter, grantee, feeAllowance)
 	if err != nil {
 		return err
@@ -60,6 +80,13 @@ func (k Keeper) GrantAllowance(ctx sdk.Context, granter, grantee sdk.AccAddress,
 	}
 
 	store.Set(key, bz)
+	expiration, err := feeAllowance.ExpiresAt()
+
+	if err != nil {
+		return err
+	} else if expiration != nil {
+		k.insertAllowanceKey(ctx, key, expiration)
+	}
 
 	ctx.EventManager().EmitEvent(
 		sdk.NewEvent(
@@ -72,6 +99,39 @@ func (k Keeper) GrantAllowance(ctx sdk.Context, granter, grantee sdk.AccAddress,
 	return nil
 }
 
+// UpdateAllowance updates the existing grant.
+func (k Keeper) UpdateAllowance(ctx sdk.Context, granter, grantee sdk.AccAddress, feeAllowance feegrant.FeeAllowanceI) error {
+	store := ctx.KVStore(k.storeKey)
+	key := feegrant.FeeAllowanceKey(granter, grantee)
+
+	_, err := k.getGrant(ctx, granter, grantee)
+	if err != nil {
+		return err
+	}
+
+	grant, err := feegrant.NewGrant(granter, grantee, feeAllowance)
+	if err != nil {
+		return err
+	}
+
+	bz, err := k.cdc.Marshal(&grant)
+	if err != nil {
+		return err
+	}
+
+	store.Set(key, bz)
+
+	ctx.EventManager().EmitEvent(
+		sdk.NewEvent(
+			feegrant.EventTypeUpdateFeeGrant,
+			sdk.NewAttribute(feegrant.AttributeKeyGranter, grant.Granter),
+			sdk.NewAttribute(feegrant.AttributeKeyGrantee, grant.Grantee),
+		),
+	)
+
+	return nil
+}
+
 // revokeAllowance removes an existing grant
 func (k Keeper) revokeAllowance(ctx sdk.Context, granter, grantee sdk.AccAddress) error {
 	_, err := k.getGrant(ctx, granter, grantee)
@@ -177,7 +237,7 @@ func (k Keeper) UseGrantedFees(ctx sdk.Context, granter, grantee sdk.AccAddress,
 	emitUseGrantEvent(ctx, granter.String(), grantee.String())
 
 	// if fee allowance is accepted, store the updated state of the allowance
-	return k.GrantAllowance(ctx, granter, grantee, grant)
+	return k.UpdateAllowance(ctx, granter, grantee, grant)
 }
 
 func emitUseGrantEvent(ctx sdk.Context, granter, grantee string) {
@@ -228,3 +288,26 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) (*feegrant.GenesisState, error) {
 		Allowances: grants,
 	}, err
 }
+
+func (k Keeper) removeFromGrantQueue(ctx sdk.Context, key []byte) {
+	store := ctx.KVStore(k.storeKey)
+	store.Delete(key)
+}
+
+func (k Keeper) insertAllowanceKey(ctx sdk.Context, grantKey []byte, exp *time.Time) {
+	store := ctx.KVStore(k.storeKey)
+	store.Set(feegrant.FeeAllowancePrefixQueue(exp, grantKey), grantKey)
+}
+
+// RemoveExpiredAllowances iterates grantsByExpiryQueue and deletes the expired grants.
+func (k Keeper) RemoveExpiredAllowances(ctx sdk.Context) {
+	exp := ctx.BlockTime()
+	store := ctx.KVStore(k.storeKey)
+	iterator := store.Iterator(feegrant.FeeAllowanceQueueKeyPrefix, sdk.InclusiveEndBytes(feegrant.AllowanceByExpTimeKey(&exp)))
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		store.Delete(iterator.Key())
+		store.Delete(iterator.Value())
+	}
+}

x/feegrant/keeper/keeper_test.go

@@ -257,5 +257,87 @@ func (suite *KeeperTestSuite) TestIterateGrants() {
 		suite.Require().Contains([]string{suite.addrs[0].String(), suite.addrs[2].String()}, grant.Granter)
 		return true
 	})
+}
+
+func (suite *KeeperTestSuite) TestPruneGrants() {
+	eth := sdk.NewCoins(sdk.NewInt64Coin("eth", 123))
+	now := suite.sdkCtx.BlockTime()
+	oneYearExpiry := now.AddDate(1, 0, 0)
+
+	testCases := []struct {
+		name      string
+		ctx       sdk.Context
+		granter   sdk.AccAddress
+		grantee   sdk.AccAddress
+		allowance feegrant.FeeAllowanceI
+		expErrMsg string
+	}{
+		{
+			name:    "grant not pruned from state",
+			ctx:     suite.sdkCtx,
+			granter: suite.addrs[0],
+			grantee: suite.addrs[1],
+			allowance: &feegrant.BasicAllowance{
+				SpendLimit: suite.atom,
+				Expiration: &now,
+			},
+		},
+		{
+			name:      "grant pruned from state after a block: error",
+			ctx:       suite.sdkCtx.WithBlockTime(now.AddDate(0, 0, 1)),
+			granter:   suite.addrs[2],
+			grantee:   suite.addrs[1],
+			expErrMsg: "not found",
+			allowance: &feegrant.BasicAllowance{
+				SpendLimit: eth,
+				Expiration: &now,
+			},
+		},
+		{
+			name:    "grant not pruned from state after a day: no error",
+			ctx:     suite.sdkCtx.WithBlockTime(now.AddDate(0, 0, 1)),
+			granter: suite.addrs[1],
+			grantee: suite.addrs[0],
+			allowance: &feegrant.BasicAllowance{
+				SpendLimit: eth,
+				Expiration: &oneYearExpiry,
+			},
+		},
+		{
+			name:      "grant pruned from state after a year: error",
+			ctx:       suite.sdkCtx.WithBlockTime(now.AddDate(1, 0, 1)),
+			granter:   suite.addrs[1],
+			grantee:   suite.addrs[2],
+			expErrMsg: "not found",
+			allowance: &feegrant.BasicAllowance{
+				SpendLimit: eth,
+				Expiration: &oneYearExpiry,
+			},
+		},
+		{
+			name:    "no expiry: no error",
+			ctx:     suite.sdkCtx.WithBlockTime(now.AddDate(1, 0, 0)),
+			granter: suite.addrs[1],
+			grantee: suite.addrs[2],
+			allowance: &feegrant.BasicAllowance{
+				SpendLimit: eth,
+				Expiration: &oneYearExpiry,
+			},
+		},
+	}
 
+	for _, tc := range testCases {
+		tc := tc
+		suite.Run(tc.name, func() {
+			suite.keeper.GrantAllowance(suite.sdkCtx, tc.granter, tc.grantee, tc.allowance)
+			suite.app.FeeGrantKeeper.RemoveExpiredAllowances(tc.ctx)
+			grant, err := suite.keeper.GetAllowance(tc.ctx, tc.granter, tc.grantee)
+			if tc.expErrMsg != "" {
+				suite.Error(err)
+				suite.Contains(err.Error(), tc.expErrMsg)
+			} else {
+				suite.NotNil(grant)
+			}
+		})
+	}
 }

x/feegrant/keeper/migrations.go

@@ -0,0 +1,21 @@
+package keeper
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	v045 "github.com/cosmos/cosmos-sdk/x/feegrant/migrations/v045"
+)
+
+// Migrator is a struct for handling in-place store migrations.
+type Migrator struct {
+	keeper Keeper
+}
+
+// NewMigrator returns a new Migrator.
+func NewMigrator(keeper Keeper) Migrator {
+	return Migrator{keeper: keeper}
+}
+
+// Migrate1to2 migrates from version 1 to 2.
+func (m Migrator) Migrate1to2(ctx sdk.Context) error {
+	return v045.MigrateStore(ctx, m.keeper.storeKey, m.keeper.cdc)
+}

x/feegrant/key.go

@@ -1,6 +1,8 @@
 package feegrant
 
 import (
+	time "time"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
 )
@@ -22,6 +24,9 @@ const (
 var (
 	// FeeAllowanceKeyPrefix is the set of the kvstore for fee allowance data
 	FeeAllowanceKeyPrefix = []byte{0x00}
+
+	// FeeAllowanceQueueKeyPrefix is the set of the kvstore for fee allowance keys data
+	FeeAllowanceQueueKeyPrefix = []byte{0x01}
 )
 
 // FeeAllowanceKey is the canonical key to store a grant from granter to grantee
@@ -34,3 +39,13 @@ func FeeAllowanceKey(granter sdk.AccAddress, grantee sdk.AccAddress) []byte {
 func FeeAllowancePrefixByGrantee(grantee sdk.AccAddress) []byte {
 	return append(FeeAllowanceKeyPrefix, address.MustLengthPrefix(grantee.Bytes())...)
 }
+
+// FeeAllowancePrefixQueue is the canonical key to store grant key.
+func FeeAllowancePrefixQueue(exp *time.Time, allowanceKey []byte) []byte {
+	allowanceByExpTimeKey := AllowanceByExpTimeKey(exp)
+	return append(allowanceByExpTimeKey, allowanceKey[1:]...)
+}
+
+func AllowanceByExpTimeKey(exp *time.Time) []byte {
+	return append(FeeAllowanceQueueKeyPrefix, sdk.FormatTimeBytes(*exp)...)
+}

x/feegrant/migrations/v045/keys.go

@@ -0,0 +1,5 @@
+package v045
+
+var (
+	ModuleName = "feegrant"
+)

x/feegrant/migrations/v045/store.go

@@ -0,0 +1,50 @@
+package v045
+
+import (
+	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	"github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/feegrant"
+)
+
+func addAllowancesByExpTimeQueue(ctx types.Context, store storetypes.KVStore, cdc codec.BinaryCodec) error {
+	prefixStore := prefix.NewStore(store, feegrant.FeeAllowanceKeyPrefix)
+	iterator := prefixStore.Iterator(nil, nil)
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		var grant feegrant.Grant
+		bz := iterator.Value()
+		if err := cdc.Unmarshal(bz, &grant); err != nil {
+			return err
+		}
+
+		grantInfo, err := grant.GetGrant()
+		if err != nil {
+			return err
+		}
+
+		exp, err := grantInfo.ExpiresAt()
+		if err != nil {
+			return err
+		}
+
+		if exp != nil {
+			key := iterator.Key()
+			if exp.Before(ctx.BlockTime()) {
+				prefixStore.Delete(key)
+			} else {
+				grantByExpTimeQueueKey := feegrant.FeeAllowancePrefixQueue(exp, key[1:])
+				store.Set(grantByExpTimeQueueKey, key)
+			}
+		}
+	}
+
+	return nil
+}
+
+func MigrateStore(ctx types.Context, storeKey storetypes.StoreKey, cdc codec.BinaryCodec) error {
+	store := ctx.KVStore(storeKey)
+	return addAllowancesByExpTimeQueue(ctx, store, cdc)
+}