Add gas meter to tokenfactory trackBeforeSend

CHANGELOG.md

@@ -64,6 +64,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#5855](https://github.com/osmosis-labs/osmosis/pull/5855) feat(x/cosmwasmpool): Sending token_in_max_amount to the contract before running contract msg
 * [#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
+* [#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{
@@ -121,8 +132,12 @@ func (k Keeper) callBeforeSendListener(ctx sdk.Context, from, to sdk.AccAddress,
 				return err
 			}
 
-			em := sdk.NewEventManager()
+			// if its track before send, apply gas meter to prevent infinite loop
+			if !blockBeforeSend {
+				ctx = ctx.WithGasMeter(sdk.NewGasMeter(types.TrackBeforeSendGasLimit))
+			}
 
+			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)

x/tokenfactory/keeper/before_send_test.go

@@ -124,3 +124,69 @@ 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()
+
+	// load wasm file
+	wasmFile := "./testdata/infinite_track_beforesend.wasm"
+	wasmCode, err := os.ReadFile(wasmFile)
+	s.Require().NoError(err)
+
+	for _, tc := range []struct {
+		name            string
+		useFactoryDenom bool
+		expectedError   bool
+	}{
+		{
+			name:            "sending tokenfactory denom from module to module with infinite contract should panic",
+			useFactoryDenom: true,
+		},
+		{
+			name:            "sending non-tokenfactory denom from module to module with infinite contract should not panic",
+			useFactoryDenom: false,
+		},
+	} {
+		s.Run(fmt.Sprintf("Case %s", tc.name), func() {
+			// setup test
+			s.SetupTest()
+
+			// 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, 1000000))
+			} else {
+				tokenToSend = sdk.NewCoins(sdk.NewInt64Coin("foo", 1000000))
+			}
+
+			// 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")
 )