Add gas meter to tokenfactory trackBeforeSend (#5927)

* Finish :)

* Add comment

* Add changelog

* Catch panic

* Update comment

* Try changing to use child ctx

---------

Co-authored-by: Dev Ojha <ValarDragon@users.noreply.github.com>

CHANGELOG.md

@@ -66,6 +66,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#5893] (https://github.com/osmosis-labs/osmosis/pull/5893) Export createPosition method in CL so other modules can use it in testing
 * [#5870] (https://github.com/osmosis-labs/osmosis/pull/5870) Remove v14/ separator in protorev rest endpoints
 * [#5923] (https://github.com/osmosis-labs/osmosis/pull/5923) CL: Lower gas for initializing ticks
+* [#5927] (https://github.com/osmosis-labs/osmosis/pull/5927) Add gas metering to x/tokenfactory trackBeforeSend hook
 
 ### Minor improvements & Bug Fixes
 

x/tokenfactory/keeper/before_send.go

@@ -86,7 +86,15 @@ func (h Hooks) BlockBeforeSend(ctx sdk.Context, from, to sdk.AccAddress, amount
 
 // callBeforeSendListener iterates over each coin and sends corresponding sudo msg to the contract address stored in state.
 // If blockBeforeSend is true, sudoMsg wraps BlockBeforeSendMsg, otherwise sudoMsg wraps TrackBeforeSendMsg.
-func (k Keeper) callBeforeSendListener(ctx sdk.Context, from, to sdk.AccAddress, amount sdk.Coins, blockBeforeSend bool) error {
+// Note that we gas meter trackBeforeSend to prevent infinite contract calls.
+// CONTRACT: this should not be called in beginBlock or endBlock since out of gas will cause this method to panic.
+func (k Keeper) callBeforeSendListener(ctx sdk.Context, from, to sdk.AccAddress, amount sdk.Coins, blockBeforeSend bool) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = types.ErrTrackBeforeSendOutOfGas
+		}
+	}()
+
 	for _, coin := range amount {
 		cosmwasmAddress := k.GetBeforeSendHook(ctx, coin.Denom)
 		if cosmwasmAddress != "" {
@@ -98,6 +106,9 @@ func (k Keeper) callBeforeSendListener(ctx sdk.Context, from, to sdk.AccAddress,
 			var msgBz []byte
 
 			// get msgBz, either BlockBeforeSend or TrackBeforeSend
+			// Note that for trackBeforeSend, we need to gas meter computations to prevent infinite loop
+			// specifically because module to module sends are not gas metered.
+			// We don't need to do this for blockBeforeSend since blockBeforeSend is not called during module to module sends.
 			if blockBeforeSend {
 				msg := types.BlockBeforeSendSudoMsg{
 					BlockBeforeSend: types.BlockBeforeSendMsg{
@@ -120,12 +131,23 @@ func (k Keeper) callBeforeSendListener(ctx sdk.Context, from, to sdk.AccAddress,
 			if err != nil {
 				return err
 			}
-
 			em := sdk.NewEventManager()
 
-			_, err = k.contractKeeper.Sudo(ctx.WithEventManager(em), cwAddr, msgBz)
-			if err != nil {
-				return errorsmod.Wrapf(err, "failed to call before send hook for denom %s", coin.Denom)
+			// if its track before send, apply gas meter to prevent infinite loop
+			if blockBeforeSend {
+				_, err = k.contractKeeper.Sudo(ctx.WithEventManager(em), cwAddr, msgBz)
+				if err != nil {
+					return errorsmod.Wrapf(err, "failed to call before send hook for denom %s", coin.Denom)
+				}
+			} else {
+				childCtx := ctx.WithGasMeter(sdk.NewGasMeter(types.TrackBeforeSendGasLimit))
+				_, err = k.contractKeeper.Sudo(childCtx.WithEventManager(em), cwAddr, msgBz)
+				if err != nil {
+					return errorsmod.Wrapf(err, "failed to call before send hook for denom %s", coin.Denom)
+				}
+
+				// consume gas used for calling contract to the parent ctx
+				ctx.GasMeter().ConsumeGas(childCtx.GasMeter().GasConsumed(), "track before send gas")
 			}
 		}
 	}

x/tokenfactory/keeper/before_send_test.go

@@ -124,3 +124,78 @@ func (s *KeeperTestSuite) TestBeforeSendHook() {
 		})
 	}
 }
+
+// TestInfiniteTrackBeforeSend tests gas metering with infinite loop contract
+// to properly test if we are gas metering trackBeforeSend properly.
+func (s *KeeperTestSuite) TestInfiniteTrackBeforeSend() {
+	s.SkipIfWSL()
+
+	for _, tc := range []struct {
+		name            string
+		wasmFile        string
+		tokenToSend     sdk.Coins
+		useFactoryDenom bool
+		expectedError   bool
+	}{
+		{
+			name:            "sending tokenfactory denom from module to module with infinite contract should panic",
+			wasmFile:        "./testdata/infinite_track_beforesend.wasm",
+			useFactoryDenom: true,
+		},
+		{
+			name:            "sending non-tokenfactory denom from module to module with infinite contract should not panic",
+			wasmFile:        "./testdata/infinite_track_beforesend.wasm",
+			tokenToSend:     sdk.NewCoins(sdk.NewInt64Coin("foo", 1000000)),
+			useFactoryDenom: false,
+		},
+		{
+			name:            "Try using no 100 ",
+			wasmFile:        "./testdata/no100.wasm",
+			useFactoryDenom: true,
+		},
+	} {
+		s.Run(fmt.Sprintf("Case %s", tc.name), func() {
+			// setup test
+			s.SetupTest()
+
+			// load wasm file
+			wasmCode, err := os.ReadFile(tc.wasmFile)
+			s.Require().NoError(err)
+
+			// instantiate wasm code
+			codeID, _, err := s.contractKeeper.Create(s.Ctx, s.TestAccs[0], wasmCode, nil)
+			s.Require().NoError(err, "test: %v", tc.name)
+			cosmwasmAddress, _, err := s.contractKeeper.Instantiate(s.Ctx, codeID, s.TestAccs[0], s.TestAccs[0], []byte("{}"), "", sdk.NewCoins())
+			s.Require().NoError(err, "test: %v", tc.name)
+
+			// create new denom
+			res, err := s.msgServer.CreateDenom(sdk.WrapSDKContext(s.Ctx), types.NewMsgCreateDenom(s.TestAccs[0].String(), "bitcoin"))
+			s.Require().NoError(err, "test: %v", tc.name)
+			factoryDenom := res.GetNewTokenDenom()
+
+			var tokenToSend sdk.Coins
+			if tc.useFactoryDenom {
+				tokenToSend = sdk.NewCoins(sdk.NewInt64Coin(factoryDenom, 100))
+			} else {
+				tokenToSend = tc.tokenToSend
+			}
+
+			// send the mint module tokenToSend
+			s.FundModuleAcc("mint", tokenToSend)
+
+			// set beforesend hook to the new denom
+			// we register infinite loop contract here to test if we are gas metering properly
+			_, err = s.msgServer.SetBeforeSendHook(sdk.WrapSDKContext(s.Ctx), types.NewMsgSetBeforeSendHook(s.TestAccs[0].String(), factoryDenom, cosmwasmAddress.String()))
+			s.Require().NoError(err, "test: %v", tc.name)
+
+			// track before send suppresses in any case, thus we expect no error
+			err = s.App.BankKeeper.SendCoinsFromModuleToModule(s.Ctx, "mint", "distribution", tokenToSend)
+			s.Require().NoError(err)
+
+			// send should happen regardless of trackBeforeSend results
+			distributionModuleAddress := s.App.AccountKeeper.GetModuleAddress("distribution")
+			distributionModuleBalances := s.App.BankKeeper.GetAllBalances(s.Ctx, distributionModuleAddress)
+			s.Require().True(distributionModuleBalances.IsEqual(tokenToSend))
+		})
+	}
+}

x/tokenfactory/types/constants.go

@@ -0,0 +1,5 @@
+package types
+
+var (
+	TrackBeforeSendGasLimit = uint64(100_000)
+)

x/tokenfactory/types/errors.go

@@ -20,4 +20,5 @@ var (
 	ErrCreatorTooLong           = errorsmod.Register(ModuleName, 9, fmt.Sprintf("creator too long, max length is %d bytes", MaxCreatorLength))
 	ErrDenomDoesNotExist        = errorsmod.Register(ModuleName, 10, "denom does not exist")
 	ErrBurnFromModuleAccount    = errorsmod.Register(ModuleName, 11, "burning from Module Account is not allowed")
+	ErrTrackBeforeSendOutOfGas  = errorsmod.Register(ModuleName, 12, "gas meter hit maximum limit")
 )