refactor(x/feegrant): set environment in context (#20529)

cosmos · Jun 3, 2024 · d1aab15 · d1aab15
1 parent 74d89bf
commit d1aab15
Show file tree

Hide file tree

Showing 9 changed files with 106 additions and 21 deletions.
diff --git a/x/feegrant/CHANGELOG.md b/x/feegrant/CHANGELOG.md
@@ -31,6 +31,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* [#20529](https://github.com/cosmos/cosmos-sdk/pull/20529) `Accept` on the `FeeAllowanceI` interface now expects the feegrant environment in the `context.Context`.
 * [#19450](https://github.com/cosmos/cosmos-sdk/pull/19450) Migrate module to use `appmodule.Environment` instead of passing individual services.
 
 ### Consensus Breaking Changes

diff --git a/x/feegrant/basic_fee.go b/x/feegrant/basic_fee.go
@@ -2,8 +2,11 @@ package feegrant
 
 import (
 	"context"
+	"fmt"
 	"time"
 
+	"cosmossdk.io/core/appmodule"
+	corecontext "cosmossdk.io/core/context"
 	errorsmod "cosmossdk.io/errors"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -23,8 +26,12 @@ 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) {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	if a.Expiration != nil && a.Expiration.Before(sdkCtx.HeaderInfo().Time) {
+	environment, ok := ctx.Value(corecontext.EnvironmentContextKey).(appmodule.Environment)
+	if !ok {
+		return false, fmt.Errorf("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/basic_fee_test.go b/x/feegrant/basic_fee_test.go
@@ -1,12 +1,15 @@
 package feegrant_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"cosmossdk.io/core/appmodule/v2"
+	corecontext "cosmossdk.io/core/context"
 	"cosmossdk.io/core/header"
 	storetypes "cosmossdk.io/store/types"
 	"cosmossdk.io/x/feegrant"
@@ -137,9 +140,11 @@ func TestBasicFeeValidAllow(t *testing.T) {
 			require.NoError(t, err)
 
 			ctx := testCtx.Ctx.WithHeaderInfo(header.Info{Time: tc.blockTime})
-
 			// now try to deduct
-			removed, err := tc.allowance.Accept(ctx, tc.fee, []sdk.Msg{})
+			removed, err := tc.allowance.Accept(context.WithValue(ctx, corecontext.EnvironmentContextKey, appmodule.Environment{
+				HeaderService: mockHeaderService{},
+				GasService:    mockGasService{},
+			}), tc.fee, []sdk.Msg{})
 			if !tc.accept {
 				require.Error(t, err)
 				return

diff --git a/x/feegrant/filtered_fee.go b/x/feegrant/filtered_fee.go
@@ -2,10 +2,13 @@ package feegrant
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/cosmos/gogoproto/proto"
 
+	"cosmossdk.io/core/appmodule"
+	corecontext "cosmossdk.io/core/context"
 	errorsmod "cosmossdk.io/errors"
 
 	"github.com/cosmos/cosmos-sdk/codec/types"
@@ -70,7 +73,11 @@ func (a *AllowedMsgAllowance) SetAllowance(allowance FeeAllowanceI) error {
 
 // Accept method checks for the filtered messages has valid expiry
 func (a *AllowedMsgAllowance) Accept(ctx context.Context, fee sdk.Coins, msgs []sdk.Msg) (bool, error) {
-	if !a.allMsgTypesAllowed(ctx, msgs) {
+	allowed, err := a.allMsgTypesAllowed(ctx, msgs)
+	if err != nil {
+		return false, err
+	}
+	if !allowed {
 		return false, errorsmod.Wrap(ErrMessageNotAllowed, "message does not exist in allowed messages")
 	}
 
@@ -88,28 +95,43 @@ func (a *AllowedMsgAllowance) Accept(ctx context.Context, fee sdk.Coins, msgs []
 	return remove, err
 }
 
-func (a *AllowedMsgAllowance) allowedMsgsToMap(ctx context.Context) map[string]bool {
+func (a *AllowedMsgAllowance) allowedMsgsToMap(ctx context.Context) (map[string]bool, error) {
 	msgsMap := make(map[string]bool, len(a.AllowedMessages))
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	environment, ok := ctx.Value(corecontext.EnvironmentContextKey).(appmodule.Environment)
+	if !ok {
+		return nil, fmt.Errorf("environment not set")
+	}
+	gasMeter := environment.GasService.GasMeter(ctx)
 	for _, msg := range a.AllowedMessages {
-		sdkCtx.GasMeter().ConsumeGas(gasCostPerIteration, "check msg")
+		if err := gasMeter.Consume(gasCostPerIteration, "check msg"); err != nil {
+			return nil, err
+		}
 		msgsMap[msg] = true
 	}
 
-	return msgsMap
+	return msgsMap, nil
 }
 
-func (a *AllowedMsgAllowance) allMsgTypesAllowed(ctx context.Context, msgs []sdk.Msg) bool {
-	msgsMap := a.allowedMsgsToMap(ctx)
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
+func (a *AllowedMsgAllowance) allMsgTypesAllowed(ctx context.Context, msgs []sdk.Msg) (bool, error) {
+	msgsMap, err := a.allowedMsgsToMap(ctx)
+	if err != nil {
+		return false, err
+	}
+	environment, ok := ctx.Value(corecontext.EnvironmentContextKey).(appmodule.Environment)
+	if !ok {
+		return false, fmt.Errorf("environment not set")
+	}
+	gasMeter := environment.GasService.GasMeter(ctx)
 	for _, msg := range msgs {
-		sdkCtx.GasMeter().ConsumeGas(gasCostPerIteration, "check msg")
+		if err := gasMeter.Consume(gasCostPerIteration, "check msg"); err != nil {
+			return false, err
+		}
 		if !msgsMap[sdk.MsgTypeURL(msg)] {
-			return false
+			return false, nil
 		}
 	}
 
-	return true
+	return true, nil
 }
 
 // ValidateBasic implements FeeAllowance and enforces basic sanity checks

diff --git a/x/feegrant/filtered_fee_test.go b/x/feegrant/filtered_fee_test.go
@@ -1,12 +1,15 @@
 package feegrant_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"cosmossdk.io/core/appmodule/v2"
+	corecontext "cosmossdk.io/core/context"
 	"cosmossdk.io/core/header"
 	storetypes "cosmossdk.io/store/types"
 	banktypes "cosmossdk.io/x/bank/types"
@@ -156,7 +159,10 @@ func TestFilteredFeeValidAllow(t *testing.T) {
 			require.NoError(t, err)
 
 			// now try to deduct
-			removed, err := allowance.Accept(ctx, tc.fee, []sdk.Msg{&call})
+			removed, err := allowance.Accept(context.WithValue(ctx, corecontext.EnvironmentContextKey, appmodule.Environment{
+				HeaderService: mockHeaderService{},
+				GasService:    mockGasService{},
+			}), tc.fee, []sdk.Msg{&call})
 			if !tc.accept {
 				require.Error(t, err)
 				return

diff --git a/x/feegrant/keeper/keeper.go b/x/feegrant/keeper/keeper.go
@@ -7,6 +7,7 @@ import (
 
 	"cosmossdk.io/collections"
 	"cosmossdk.io/core/appmodule"
+	corecontext "cosmossdk.io/core/context"
 	"cosmossdk.io/core/event"
 	"cosmossdk.io/core/log"
 	errorsmod "cosmossdk.io/errors"
@@ -232,7 +233,7 @@ func (k Keeper) UseGrantedFees(ctx context.Context, granter, grantee sdk.AccAddr
 		return err
 	}
 
-	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/mock_test.go b/x/feegrant/mock_test.go
@@ -0,0 +1,32 @@
+package feegrant_test
+
+import (
+	"context"
+
+	coregas "cosmossdk.io/core/gas"
+	coreheader "cosmossdk.io/core/header"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+type mockHeaderService struct{}
+
+func (h mockHeaderService) HeaderInfo(ctx context.Context) coreheader.Info {
+	return sdk.UnwrapSDKContext(ctx).HeaderInfo()
+}
+
+type mockGasService struct {
+	coregas.Service
+}
+
+func (m mockGasService) GasMeter(ctx context.Context) coregas.Meter {
+	return mockGasMeter{}
+}
+
+type mockGasMeter struct {
+	coregas.Meter
+}
+
+func (m mockGasMeter) Consume(amount coregas.Gas, descriptor string) error {
+	return nil
+}
diff --git a/x/feegrant/periodic_fee.go b/x/feegrant/periodic_fee.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"time"
 
+	"cosmossdk.io/core/appmodule"
+	corecontext "cosmossdk.io/core/context"
 	errorsmod "cosmossdk.io/errors"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -23,9 +25,11 @@ 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) {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	blockTime := sdkCtx.HeaderInfo().Time
-
+	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")
 	}

diff --git a/x/feegrant/periodic_fee_test.go b/x/feegrant/periodic_fee_test.go
@@ -1,12 +1,15 @@
 package feegrant_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"cosmossdk.io/core/appmodule/v2"
+	corecontext "cosmossdk.io/core/context"
 	"cosmossdk.io/core/header"
 	storetypes "cosmossdk.io/store/types"
 	"cosmossdk.io/x/feegrant"
@@ -219,7 +222,11 @@ func TestPeriodicFeeValidAllow(t *testing.T) {
 
 			ctx := testCtx.Ctx.WithHeaderInfo(header.Info{Time: tc.blockTime})
 			// now try to deduct
-			remove, err := tc.allow.Accept(ctx, tc.fee, []sdk.Msg{})
+			// Set environment to ctx
+			remove, err := tc.allow.Accept(context.WithValue(ctx, corecontext.EnvironmentContextKey, appmodule.Environment{
+				HeaderService: mockHeaderService{},
+				GasService:    mockGasService{},
+			}), tc.fee, []sdk.Msg{})
 			if !tc.accept {
 				require.Error(t, err)
 				return