From 002868e9de4e64f80059c1ee0fa3841a0b936245 Mon Sep 17 00:00:00 2001
From: mattverse <mattpark1028@gmail.com>
Date: Tue, 1 Aug 2023 00:19:37 +0900
Subject: [PATCH] Catch panic

---
 x/tokenfactory/keeper/before_send.go      |  8 +++++++-
 x/tokenfactory/keeper/before_send_test.go | 20 ++++++++------------
 x/tokenfactory/types/errors.go            |  1 +
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/x/tokenfactory/keeper/before_send.go b/x/tokenfactory/keeper/before_send.go
index 6a6a4a9aba4..21b0639a4f5 100644
--- a/x/tokenfactory/keeper/before_send.go
+++ b/x/tokenfactory/keeper/before_send.go
@@ -88,7 +88,13 @@ func (h Hooks) BlockBeforeSend(ctx sdk.Context, from, to sdk.AccAddress, amount
 // If blockBeforeSend is true, sudoMsg wraps BlockBeforeSendMsg, otherwise sudoMsg wraps TrackBeforeSendMsg.
 // 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) error {
+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 != "" {
diff --git a/x/tokenfactory/keeper/before_send_test.go b/x/tokenfactory/keeper/before_send_test.go
index f2011f22fe7..7186e903005 100644
--- a/x/tokenfactory/keeper/before_send_test.go
+++ b/x/tokenfactory/keeper/before_send_test.go
@@ -138,17 +138,15 @@ func (s *KeeperTestSuite) TestInfiniteTrackBeforeSend() {
 	for _, tc := range []struct {
 		name            string
 		useFactoryDenom bool
-		expectedPanic   bool
+		expectedError   bool
 	}{
 		{
 			name:            "sending tokenfactory denom from module to module with infinite contract should panic",
 			useFactoryDenom: true,
-			expectedPanic:   true,
 		},
 		{
 			name:            "sending non-tokenfactory denom from module to module with infinite contract should not panic",
 			useFactoryDenom: false,
-			expectedPanic:   false,
 		},
 	} {
 		s.Run(fmt.Sprintf("Case %s", tc.name), func() {
@@ -181,15 +179,13 @@ func (s *KeeperTestSuite) TestInfiniteTrackBeforeSend() {
 			_, 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)
 
-			if tc.expectedPanic {
-				s.Require().Panics(func() {
-					s.App.BankKeeper.SendCoinsFromModuleToModule(s.Ctx, "mint", "distribution", tokenToSend)
-				})
-			} else {
-				s.Require().NotPanics(func() {
-					s.App.BankKeeper.SendCoinsFromModuleToModule(s.Ctx, "mint", "distribution", tokenToSend)
-				})
-			}
+			// 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)
+
+			distributionModuleAddress := s.App.AccountKeeper.GetModuleAddress("distribution")
+			distributionModuleBalances := s.App.BankKeeper.GetAllBalances(s.Ctx, distributionModuleAddress)
+			s.Require().True(distributionModuleBalances.IsEqual(tokenToSend))
 		})
 	}
 }
diff --git a/x/tokenfactory/types/errors.go b/x/tokenfactory/types/errors.go
index a66d4b6b417..aa78ca24ae0 100644
--- a/x/tokenfactory/types/errors.go
+++ b/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")
 )