refactor

cosmos · Jun 3, 2024 · bdb5fc1 · bdb5fc1
1 parent e0edef9
commit bdb5fc1
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 14 deletions.
diff --git a/x/feegrant/basic_fee.go b/x/feegrant/basic_fee.go
@@ -7,6 +7,7 @@ import (
 	errorsmod "cosmossdk.io/errors"
 
 	"cosmossdk.io/core/appmodule"
+	corecontext "cosmossdk.io/core/context"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 )
@@ -24,8 +25,10 @@ var _ FeeAllowanceI = (*BasicAllowance)(nil)
 // If remove is true (regardless of the error), the FeeAllowance will be deleted from storage
 // (eg. when it is used up). (See call to RevokeAllowance in Keeper.UseGrantedFees)
 func (a *BasicAllowance) Accept(ctx context.Context, fee sdk.Coins, _ []sdk.Msg) (bool, error) {
-	e := ctx.Value(ContextKey)
-	environment := e.(appmodule.Environment)
+	environment, ok := ctx.Value(corecontext.EnvironmentContextKey).(appmodule.Environment)
+	if !ok {
+		return true, errorsmod.Wrap(ErrFeeLimitExpired, "environment not set")
+	}
 	headerInfo := environment.HeaderService.HeaderInfo(ctx)
 	if a.Expiration != nil && a.Expiration.Before(headerInfo.Time) {
 		return true, errorsmod.Wrap(ErrFeeLimitExpired, "basic allowance")

diff --git a/x/feegrant/filtered_fee.go b/x/feegrant/filtered_fee.go
@@ -9,6 +9,7 @@ import (
 	errorsmod "cosmossdk.io/errors"
 
 	"cosmossdk.io/core/appmodule"
+	corecontext "cosmossdk.io/core/context"
 	"github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
@@ -91,8 +92,10 @@ func (a *AllowedMsgAllowance) Accept(ctx context.Context, fee sdk.Coins, msgs []
 
 func (a *AllowedMsgAllowance) allowedMsgsToMap(ctx context.Context) map[string]bool {
 	msgsMap := make(map[string]bool, len(a.AllowedMessages))
-	e := ctx.Value(ContextKey)
-	environment := e.(appmodule.Environment)
+	environment, ok := ctx.Value(corecontext.EnvironmentContextKey).(appmodule.Environment)
+	if !ok {
+		return nil
+	}
 	gasMeter := environment.GasService.GasMeter(ctx)
 	for _, msg := range a.AllowedMessages {
 		gasMeter.Consume(gasCostPerIteration, "check msg")
@@ -104,8 +107,10 @@ func (a *AllowedMsgAllowance) allowedMsgsToMap(ctx context.Context) map[string]b
 
 func (a *AllowedMsgAllowance) allMsgTypesAllowed(ctx context.Context, msgs []sdk.Msg) bool {
 	msgsMap := a.allowedMsgsToMap(ctx)
-	e := ctx.Value(ContextKey)
-	environment := e.(appmodule.Environment)
+	environment, ok := ctx.Value(corecontext.EnvironmentContextKey).(appmodule.Environment)
+	if !ok {
+		return false
+	}
 	gasMeter := environment.GasService.GasMeter(ctx)
 	for _, msg := range msgs {
 		gasMeter.Consume(gasCostPerIteration, "check msg")

diff --git a/x/feegrant/keeper/keeper.go b/x/feegrant/keeper/keeper.go
@@ -13,6 +13,7 @@ import (
 	"cosmossdk.io/x/auth/ante"
 	"cosmossdk.io/x/feegrant"
 
+	corecontext "cosmossdk.io/core/context"
 	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
@@ -232,9 +233,7 @@ func (k Keeper) UseGrantedFees(ctx context.Context, granter, grantee sdk.AccAddr
 		return err
 	}
 
-	ctx = context.WithValue(ctx, feegrant.ContextKey, k.Environment)
-
-	remove, err := grant.Accept(ctx, fee, msgs)
+	remove, err := grant.Accept(context.WithValue(ctx, corecontext.EnvironmentContextKey, k.Environment), fee, msgs)
 	if remove && err == nil {
 		// Ignoring the `revokeFeeAllowance` error, because the user has enough grants to perform this transaction.
 		_ = k.revokeAllowance(ctx, granter, grantee)

diff --git a/x/feegrant/key.go b/x/feegrant/key.go
@@ -2,7 +2,6 @@ package feegrant
 
 import (
 	"cosmossdk.io/collections"
-	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
 const (
@@ -11,8 +10,6 @@ const (
 
 	// StoreKey is the store key string for supply
 	StoreKey = ModuleName
-
-	ContextKey = sdk.ContextKey("client.context")
 )
 
 var (

diff --git a/x/feegrant/periodic_fee.go b/x/feegrant/periodic_fee.go
@@ -7,6 +7,7 @@ import (
 	errorsmod "cosmossdk.io/errors"
 
 	"cosmossdk.io/core/appmodule"
+	corecontext "cosmossdk.io/core/context"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 )
@@ -24,8 +25,10 @@ var _ FeeAllowanceI = (*PeriodicAllowance)(nil)
 // If remove is true (regardless of the error), the FeeAllowance will be deleted from storage
 // (eg. when it is used up). (See call to RevokeAllowance in Keeper.UseGrantedFees)
 func (a *PeriodicAllowance) Accept(ctx context.Context, fee sdk.Coins, _ []sdk.Msg) (bool, error) {
-	e := ctx.Value(ContextKey)
-	environment := e.(appmodule.Environment)
+	environment, ok := ctx.Value(corecontext.EnvironmentContextKey).(appmodule.Environment)
+	if !ok {
+		return true, errorsmod.Wrap(ErrFeeLimitExpired, "environment not set")
+	}
 	blockTime := environment.HeaderService.HeaderInfo(ctx).Time
 	if a.Basic.Expiration != nil && blockTime.After(*a.Basic.Expiration) {
 		return true, errorsmod.Wrap(ErrFeeLimitExpired, "absolute limit")