fix(app.go): disable vesting account creation to prevent contract a…

…ddress bricking (#923) * add disabled messages * add other vesting messages * option refactor * add vesting ante handler * implement test for vesting ante handler * lint issues * fix lint 2 * fix lint 3 --------- Co-authored-by: Tanmay <tanmay.bhattacharya.smit@gmail.com>
zeta-chain · Aug 11, 2023 · e82e86a · e82e86a
1 parent b0d8cc8
commit e82e86a
Show file tree

Hide file tree

Showing 6 changed files with 192 additions and 8 deletions.
diff --git a/app/ante/ante_test.go b/app/ante/ante_test.go
@@ -0,0 +1,18 @@
+package ante_test
+
+import sdk "github.com/cosmos/cosmos-sdk/types"
+
+var _ sdk.AnteHandler = (&MockAnteHandler{}).AnteHandle
+
+// MockAnteHandler mocks an AnteHandler
+type MockAnteHandler struct {
+	WasCalled bool
+	CalledCtx sdk.Context
+}
+
+// AnteHandle implements AnteHandler
+func (mah *MockAnteHandler) AnteHandle(ctx sdk.Context, _ sdk.Tx, _ bool) (sdk.Context, error) {
+	mah.WasCalled = true
+	mah.CalledCtx = ctx
+	return ctx, nil
+}
diff --git a/app/ante/handler_options.go b/app/ante/handler_options.go
@@ -23,18 +23,16 @@ import (
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/x/auth/ante"
 	"github.com/cosmos/cosmos-sdk/x/auth/migrations/legacytx"
-	"github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
 	ibcante "github.com/cosmos/ibc-go/v6/modules/core/ante"
 	ethante "github.com/evmos/ethermint/app/ante"
 	ethermint "github.com/evmos/ethermint/types"
-	evmtypes "github.com/evmos/ethermint/x/evm/types"
 )
 
 func NewLegacyCosmosAnteHandlerEip712(options ethante.HandlerOptions) sdk.AnteHandler {
 	return sdk.ChainAnteDecorators(
 		ethante.RejectMessagesDecorator{}, // reject MsgEthereumTxs
-		NewAuthzLimiterDecorator(sdk.MsgTypeURL(&evmtypes.MsgEthereumTx{}), // disable the Msg types that cannot be included on an authz.MsgExec msgs field
-			sdk.MsgTypeURL(&types.MsgCreateVestingAccount{})),
+		NewAuthzLimiterDecorator(options.DisabledAuthzMsgs...),
+		NewVestingAccountDecorator(),
 		ante.NewSetUpContextDecorator(),
 		ante.NewValidateBasicDecorator(),
 		ante.NewTxTimeoutHeightDecorator(),
@@ -73,8 +71,8 @@ func newEthAnteHandler(options ethante.HandlerOptions) sdk.AnteHandler {
 func newCosmosAnteHandler(options ethante.HandlerOptions) sdk.AnteHandler {
 	return sdk.ChainAnteDecorators(
 		ethante.RejectMessagesDecorator{}, // reject MsgEthereumTxs
-		NewAuthzLimiterDecorator(sdk.MsgTypeURL(&evmtypes.MsgEthereumTx{}), // disable the Msg types that cannot be included on an authz.MsgExec msgs field
-			sdk.MsgTypeURL(&types.MsgCreateVestingAccount{})),
+		NewAuthzLimiterDecorator(options.DisabledAuthzMsgs...),
+		NewVestingAccountDecorator(),
 		ante.NewSetUpContextDecorator(),
 		ante.NewExtensionOptionsDecorator(options.ExtensionOptionChecker),
 		ante.NewValidateBasicDecorator(),
@@ -97,8 +95,8 @@ func newCosmosAnteHandler(options ethante.HandlerOptions) sdk.AnteHandler {
 func newCosmosAnteHandlerNoGasLimit(options ethante.HandlerOptions) sdk.AnteHandler {
 	return sdk.ChainAnteDecorators(
 		ethante.RejectMessagesDecorator{}, // reject MsgEthereumTxs
-		NewAuthzLimiterDecorator(sdk.MsgTypeURL(&evmtypes.MsgEthereumTx{}), // disable the Msg types that cannot be included on an authz.MsgExec msgs field
-			sdk.MsgTypeURL(&types.MsgCreateVestingAccount{})),
+		NewAuthzLimiterDecorator(options.DisabledAuthzMsgs...),
+		NewVestingAccountDecorator(),
 		NewSetUpContextDecorator(),
 		ante.NewExtensionOptionsDecorator(options.ExtensionOptionChecker),
 		ante.NewValidateBasicDecorator(),

diff --git a/app/ante/vesting.go b/app/ante/vesting.go
@@ -0,0 +1,50 @@
+package ante
+
+import (
+	errorsmod "cosmossdk.io/errors"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	vesting "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
+)
+
+var _ sdk.AnteDecorator = VestingAccountDecorator{}
+
+// VestingAccountDecorator blocks vesting messages from reaching the mempool
+type VestingAccountDecorator struct {
+	disabledMsgTypeURLs []string
+}
+
+// NewVestingAccountDecorator creates a decorator to block vesting messages from reaching the mempool
+func NewVestingAccountDecorator() VestingAccountDecorator {
+	return VestingAccountDecorator{
+		disabledMsgTypeURLs: []string{
+			sdk.MsgTypeURL(&vesting.MsgCreateVestingAccount{}),
+			sdk.MsgTypeURL(&vesting.MsgCreatePermanentLockedAccount{}),
+			sdk.MsgTypeURL(&vesting.MsgCreatePeriodicVestingAccount{}),
+		},
+	}
+}
+
+// AnteHandle implements AnteDecorator
+func (vad VestingAccountDecorator) AnteHandle(
+	ctx sdk.Context,
+	tx sdk.Tx,
+	simulate bool,
+	next sdk.AnteHandler,
+) (newCtx sdk.Context, err error) {
+	for _, msg := range tx.GetMsgs() {
+		typeURL := sdk.MsgTypeURL(msg)
+
+		for _, disabledTypeURL := range vad.disabledMsgTypeURLs {
+			if typeURL == disabledTypeURL {
+				return ctx, errorsmod.Wrapf(
+					sdkerrors.ErrUnauthorized,
+					"MsgTypeURL %s not supported",
+					typeURL,
+				)
+			}
+		}
+	}
+
+	return next(ctx, tx, simulate)
+}
diff --git a/app/ante/vesting_test.go b/app/ante/vesting_test.go
@@ -0,0 +1,104 @@
+package ante_test
+
+import (
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/cosmos/cosmos-sdk/simapp/helpers"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	vesting "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	"github.com/stretchr/testify/require"
+
+	"github.com/zeta-chain/zetacore/app"
+	"github.com/zeta-chain/zetacore/app/ante"
+	"github.com/zeta-chain/zetacore/testutil/sample"
+)
+
+func TestVesting_AnteHandle(t *testing.T) {
+	txConfig := app.MakeEncodingConfig().TxConfig
+
+	testPrivKey, testAddress := sample.PrivKeyAddressPair()
+	_, testAddress2 := sample.PrivKeyAddressPair()
+
+	decorator := ante.NewVestingAccountDecorator()
+
+	tests := []struct {
+		name       string
+		msg        sdk.Msg
+		wantHasErr bool
+		wantErr    string
+	}{
+		{
+			"MsgCreateVestingAccount",
+			vesting.NewMsgCreateVestingAccount(
+				testAddress, testAddress2,
+				sdk.NewCoins(sdk.NewInt64Coin("azeta", 100_000_000)),
+				time.Date(1998, 1, 1, 0, 0, 0, 0, time.UTC).Unix(),
+				false,
+			),
+			true,
+			"MsgTypeURL /cosmos.vesting.v1beta1.MsgCreateVestingAccount not supported",
+		},
+		{
+			"MsgCreatePermanentLockedAccount",
+			vesting.NewMsgCreatePermanentLockedAccount(
+				testAddress, testAddress2,
+				sdk.NewCoins(sdk.NewInt64Coin("azeta", 100_000_000)),
+			),
+			true,
+			"MsgTypeURL /cosmos.vesting.v1beta1.MsgCreatePermanentLockedAccount not supported",
+		},
+		{
+			"MsgCreatePeriodicVestingAccount",
+			vesting.NewMsgCreatePeriodicVestingAccount(
+				testAddress, testAddress2,
+				time.Date(1998, 1, 1, 0, 0, 0, 0, time.UTC).Unix(),
+				nil,
+			),
+			true,
+			"MsgTypeURL /cosmos.vesting.v1beta1.MsgCreatePeriodicVestingAccount not supported",
+		},
+		{
+			"Non blocked message",
+			banktypes.NewMsgSend(
+				testAddress, testAddress2,
+				sdk.NewCoins(sdk.NewInt64Coin("azeta", 100_000_000)),
+			),
+			false,
+			"",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tx, err := helpers.GenSignedMockTx(
+				rand.New(rand.NewSource(time.Now().UnixNano())),
+				txConfig,
+				[]sdk.Msg{
+					tt.msg,
+				},
+				sdk.NewCoins(),
+				helpers.DefaultGenTxGas,
+				"testing-chain-id",
+				[]uint64{0},
+				[]uint64{0},
+				testPrivKey,
+			)
+			require.NoError(t, err)
+
+			mmd := MockAnteHandler{}
+			ctx := sdk.Context{}.WithIsCheckTx(true)
+
+			_, err = decorator.AnteHandle(ctx, tx, false, mmd.AnteHandle)
+
+			if tt.wantHasErr {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.wantErr)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}
diff --git a/app/app.go b/app/app.go
@@ -566,6 +566,12 @@ func New(
 		SignModeHandler: encodingConfig.TxConfig.SignModeHandler(),
 		SigGasConsumer:  evmante.DefaultSigVerificationGasConsumer,
 		MaxTxGasWanted:  maxGasWanted,
+		DisabledAuthzMsgs: []string{
+			sdk.MsgTypeURL(&evmtypes.MsgEthereumTx{}), // disable the Msg types that cannot be included on an authz.MsgExec msgs field
+			sdk.MsgTypeURL(&vestingtypes.MsgCreateVestingAccount{}),
+			sdk.MsgTypeURL(&vestingtypes.MsgCreatePermanentLockedAccount{}),
+			sdk.MsgTypeURL(&vestingtypes.MsgCreatePeriodicVestingAccount{}),
+		},
 	}
 
 	anteHandler, err := ante.NewAnteHandler(options)

diff --git a/testutil/sample/sample.go b/testutil/sample/sample.go
@@ -11,3 +11,11 @@ func AccAddress() string {
 	addr := pk.Address()
 	return sdk.AccAddress(addr).String()
 }
+
+// PrivKeyAddressPair returns a private key, address pair
+func PrivKeyAddressPair() (*ed25519.PrivKey, sdk.AccAddress) {
+	privKey := ed25519.GenPrivKey()
+	addr := privKey.PubKey().Address()
+
+	return privKey, sdk.AccAddress(addr)
+}