From 82f5ed5c3c38fd21c5a20d5ef6be9cb538f17b88 Mon Sep 17 00:00:00 2001
From: drklee3 <derrick@dlee.dev>
Date: Thu, 6 Apr 2023 13:21:56 -0700
Subject: [PATCH] feat: block new vesting create messages (#1539)

* Block new vesting create messages

* doc: Update VestingAccountDecorator doc comment
---
 app/ante/ante.go         |  2 +
 app/ante/vesting.go      | 37 ++++++++++++----
 app/ante/vesting_test.go | 92 ++++++++++++++++++++++++++++++++--------
 3 files changed, 105 insertions(+), 26 deletions(-)

diff --git a/app/ante/ante.go b/app/ante/ante.go
index a2bac8a9..2b6b81c8 100644
--- a/app/ante/ante.go
+++ b/app/ante/ante.go
@@ -143,6 +143,8 @@ func newCosmosAnteHandler(options cosmosHandlerOptions) sdk.AnteHandler {
 		NewAuthzLimiterDecorator(
 			sdk.MsgTypeURL(&evmtypes.MsgEthereumTx{}),
 			sdk.MsgTypeURL(&vesting.MsgCreateVestingAccount{}),
+			sdk.MsgTypeURL(&vesting.MsgCreatePermanentLockedAccount{}),
+			sdk.MsgTypeURL(&vesting.MsgCreatePeriodicVestingAccount{}),
 		),
 		authante.NewValidateBasicDecorator(),
 		authante.NewTxTimeoutHeightDecorator(),
diff --git a/app/ante/vesting.go b/app/ante/vesting.go
index 537a1528..b4b2267c 100644
--- a/app/ante/vesting.go
+++ b/app/ante/vesting.go
@@ -9,17 +9,38 @@ import (
 
 var _ sdk.AnteDecorator = VestingAccountDecorator{}
 
-// VestingAccountDecorator blocks MsgCreateVestingAccount from reaching the mempool
-type VestingAccountDecorator struct{}
-
-func NewVestingAccountDecorator() VestingAccountDecorator {
-	return VestingAccountDecorator{}
+// VestingAccountDecorator blocks vesting messages from reaching the mempool
+type VestingAccountDecorator struct {
+	disabledMsgTypeUrls []string
 }
 
-func (vad VestingAccountDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate bool, next sdk.AnteHandler) (newCtx sdk.Context, err error) {
+func NewVestingAccountDecorator() VestingAccountDecorator {
+	return VestingAccountDecorator{
+		disabledMsgTypeUrls: []string{
+			sdk.MsgTypeURL(&vesting.MsgCreateVestingAccount{}),
+			sdk.MsgTypeURL(&vesting.MsgCreatePermanentLockedAccount{}),
+			sdk.MsgTypeURL(&vesting.MsgCreatePeriodicVestingAccount{}),
+		},
+	}
+}
+
+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() {
-		if _, ok := msg.(*vesting.MsgCreateVestingAccount); ok {
-			return ctx, errorsmod.Wrap(sdkerrors.ErrUnauthorized, "MsgCreateVestingAccount not supported")
+		typeUrl := sdk.MsgTypeURL(msg)
+
+		for _, disabledTypeUrl := range vad.disabledMsgTypeUrls {
+			if typeUrl == disabledTypeUrl {
+				return ctx, errorsmod.Wrapf(
+					sdkerrors.ErrUnauthorized,
+					"MsgTypeURL %s not supported",
+					typeUrl,
+				)
+			}
 		}
 	}
 
diff --git a/app/ante/vesting_test.go b/app/ante/vesting_test.go
index 8c176bfa..12a98315 100644
--- a/app/ante/vesting_test.go
+++ b/app/ante/vesting_test.go
@@ -10,6 +10,7 @@ import (
 	"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/kava-labs/kava/app"
 	"github.com/kava-labs/kava/app/ante"
@@ -22,26 +23,81 @@ func TestVestingMempoolDecorator_MsgCreateVestingAccount_Unauthorized(t *testing
 
 	decorator := ante.NewVestingAccountDecorator()
 
-	tx, err := helpers.GenSignedMockTx(
-		rand.New(rand.NewSource(time.Now().UnixNano())),
-		txConfig,
-		[]sdk.Msg{
+	tests := []struct {
+		name       string
+		msg        sdk.Msg
+		wantHasErr bool
+		wantErr    string
+	}{
+		{
+			"MsgCreateVestingAccount",
 			vesting.NewMsgCreateVestingAccount(
 				testAddresses[0], testAddresses[1],
 				sdk.NewCoins(sdk.NewInt64Coin("ukava", 100_000_000)),
-				time.Date(1998, 1, 1, 0, 0, 0, 0, time.UTC).Unix(), false),
+				time.Date(1998, 1, 1, 0, 0, 0, 0, time.UTC).Unix(),
+				false,
+			),
+			true,
+			"MsgTypeURL /cosmos.vesting.v1beta1.MsgCreateVestingAccount not supported",
 		},
-		sdk.NewCoins(),
-		helpers.DefaultGenTxGas,
-		"testing-chain-id",
-		[]uint64{0},
-		[]uint64{0},
-		testPrivKeys[0],
-	)
-	require.NoError(t, err)
-	mmd := MockAnteHandler{}
-	ctx := sdk.Context{}.WithIsCheckTx(true)
-	_, err = decorator.AnteHandle(ctx, tx, false, mmd.AnteHandle)
-	require.Error(t, err)
-	require.Contains(t, err.Error(), "MsgCreateVestingAccount not supported")
+		{
+			"MsgCreateVestingAccount",
+			vesting.NewMsgCreatePermanentLockedAccount(
+				testAddresses[0], testAddresses[1],
+				sdk.NewCoins(sdk.NewInt64Coin("ukava", 100_000_000)),
+			),
+			true,
+			"MsgTypeURL /cosmos.vesting.v1beta1.MsgCreatePermanentLockedAccount not supported",
+		},
+		{
+			"MsgCreateVestingAccount",
+			vesting.NewMsgCreatePeriodicVestingAccount(
+				testAddresses[0], testAddresses[1],
+				time.Date(1998, 1, 1, 0, 0, 0, 0, time.UTC).Unix(),
+				nil,
+			),
+			true,
+			"MsgTypeURL /cosmos.vesting.v1beta1.MsgCreatePeriodicVestingAccount not supported",
+		},
+		{
+			"other messages not affected",
+			banktypes.NewMsgSend(
+				testAddresses[0], testAddresses[1],
+				sdk.NewCoins(sdk.NewInt64Coin("ukava", 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},
+				testPrivKeys[0],
+			)
+			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)
+			}
+		})
+	}
 }