fix: register codec for all messages (#1210)

* codec * add msgServer * fix tests * add new missing messages
zeta-chain · Oct 26, 2023 · ea9ba12 · ea9ba12
1 parent f9d9987
commit ea9ba12
Show file tree

Hide file tree

Showing 16 changed files with 123 additions and 57 deletions.
diff --git a/x/crosschain/keeper/gas_payment_test.go b/x/crosschain/keeper/gas_payment_test.go
@@ -226,12 +226,13 @@ func TestKeeper_PayGasNativeAndUpdateCctx(t *testing.T) {
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
+		fungibleMsgServer := fungiblekeeper.NewMsgServerImpl(*zk.FungibleKeeper)
 
 		// deploy gas coin and set fee params
 		chainID := getValidEthChainID(t)
 		deploySystemContracts(t, ctx, zk.FungibleKeeper, sdkk.EvmKeeper)
 		zrc20 := setupGasCoin(t, ctx, zk.FungibleKeeper, sdkk.EvmKeeper, chainID, "foobar", "foobar")
-		_, err := zk.FungibleKeeper.UpdateZRC20WithdrawFee(
+		_, err := fungibleMsgServer.UpdateZRC20WithdrawFee(
 			sdk.UnwrapSDKContext(ctx),
 			fungibletypes.NewMsgUpdateZRC20WithdrawFee(admin, zrc20.String(), sdk.NewUint(withdrawFee), math.Uint{}),
 		)
@@ -324,12 +325,13 @@ func TestKeeper_PayGasNativeAndUpdateCctx(t *testing.T) {
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
+		fungibleMsgServer := fungiblekeeper.NewMsgServerImpl(*zk.FungibleKeeper)
 
 		// deploy gas coin and set fee params
 		chainID := getValidEthChainID(t)
 		deploySystemContracts(t, ctx, zk.FungibleKeeper, sdkk.EvmKeeper)
 		zrc20 := setupGasCoin(t, ctx, zk.FungibleKeeper, sdkk.EvmKeeper, chainID, "foobar", "foobar")
-		_, err := zk.FungibleKeeper.UpdateZRC20WithdrawFee(
+		_, err := fungibleMsgServer.UpdateZRC20WithdrawFee(
 			sdk.UnwrapSDKContext(ctx),
 			fungibletypes.NewMsgUpdateZRC20WithdrawFee(admin, zrc20.String(), sdk.NewUint(withdrawFee), math.Uint{}),
 		)
@@ -367,6 +369,7 @@ func TestKeeper_PayGasInERC20AndUpdateCctx(t *testing.T) {
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
+		fungibleMsgServer := fungiblekeeper.NewMsgServerImpl(*zk.FungibleKeeper)
 
 		// deploy gas coin, erc20 and set fee params
 		chainID := getValidEthChainID(t)
@@ -383,7 +386,7 @@ func TestKeeper_PayGasInERC20AndUpdateCctx(t *testing.T) {
 			assetAddress,
 			"bar",
 		)
-		_, err := zk.FungibleKeeper.UpdateZRC20WithdrawFee(
+		_, err := fungibleMsgServer.UpdateZRC20WithdrawFee(
 			sdk.UnwrapSDKContext(ctx),
 			fungibletypes.NewMsgUpdateZRC20WithdrawFee(admin, gasZRC20.String(), sdk.NewUint(withdrawFee), math.Uint{}),
 		)
@@ -490,13 +493,14 @@ func TestKeeper_PayGasInERC20AndUpdateCctx(t *testing.T) {
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
+		fungibleMsgServer := fungiblekeeper.NewMsgServerImpl(*zk.FungibleKeeper)
 
 		// deploy gas coin, erc20 and set fee params
 		chainID := getValidEthChainID(t)
 		assetAddress := sample.EthAddress().String()
 		deploySystemContracts(t, ctx, zk.FungibleKeeper, sdkk.EvmKeeper)
 		gasZRC20 := setupGasCoin(t, ctx, zk.FungibleKeeper, sdkk.EvmKeeper, chainID, "foo", "foo")
-		_, err := zk.FungibleKeeper.UpdateZRC20WithdrawFee(
+		_, err := fungibleMsgServer.UpdateZRC20WithdrawFee(
 			sdk.UnwrapSDKContext(ctx),
 			fungibletypes.NewMsgUpdateZRC20WithdrawFee(admin, gasZRC20.String(), sdk.NewUint(withdrawFee), math.Uint{}),
 		)
@@ -534,6 +538,7 @@ func TestKeeper_PayGasInERC20AndUpdateCctx(t *testing.T) {
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
+		fungibleMsgServer := fungiblekeeper.NewMsgServerImpl(*zk.FungibleKeeper)
 
 		// deploy gas coin, erc20 and set fee params
 		chainID := getValidEthChainID(t)
@@ -550,7 +555,7 @@ func TestKeeper_PayGasInERC20AndUpdateCctx(t *testing.T) {
 			assetAddress,
 			"bar",
 		)
-		_, err := zk.FungibleKeeper.UpdateZRC20WithdrawFee(
+		_, err := fungibleMsgServer.UpdateZRC20WithdrawFee(
 			sdk.UnwrapSDKContext(ctx),
 			fungibletypes.NewMsgUpdateZRC20WithdrawFee(admin, gasZRC20.String(), sdk.NewUint(withdrawFee), math.Uint{}),
 		)
@@ -588,6 +593,7 @@ func TestKeeper_PayGasInERC20AndUpdateCctx(t *testing.T) {
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
+		fungibleMsgServer := fungiblekeeper.NewMsgServerImpl(*zk.FungibleKeeper)
 
 		// deploy gas coin, erc20 and set fee params
 		chainID := getValidEthChainID(t)
@@ -604,7 +610,7 @@ func TestKeeper_PayGasInERC20AndUpdateCctx(t *testing.T) {
 			assetAddress,
 			"bar",
 		)
-		_, err := zk.FungibleKeeper.UpdateZRC20WithdrawFee(
+		_, err := fungibleMsgServer.UpdateZRC20WithdrawFee(
 			sdk.UnwrapSDKContext(ctx),
 			fungibletypes.NewMsgUpdateZRC20WithdrawFee(admin, gasZRC20.String(), sdk.NewUint(withdrawFee), math.Uint{}),
 		)

diff --git a/x/crosschain/keeper/msg_server_whitelist_erc20.go b/x/crosschain/keeper/msg_server_whitelist_erc20.go
@@ -21,7 +21,7 @@ import (
 
 // WhitelistERC20 deploys a new zrc20, create a foreign coin object for the ERC20
 // and emit a crosschain tx to whitelist the ERC20 on the external chain
-func (k Keeper) WhitelistERC20(goCtx context.Context, msg *types.MsgWhitelistERC20) (*types.MsgWhitelistERC20Response, error) {
+func (k msgServer) WhitelistERC20(goCtx context.Context, msg *types.MsgWhitelistERC20) (*types.MsgWhitelistERC20Response, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 	if msg.Creator != k.zetaObserverKeeper.GetParams(ctx).GetAdminPolicyAccount(zetaObserverTypes.Policy_Type_group1) {
 		return nil, errorsmod.Wrap(sdkerrors.ErrUnauthorized, "Deploy can only be executed by the correct policy account")

diff --git a/x/crosschain/keeper/msg_server_whitelist_erc20_test.go b/x/crosschain/keeper/msg_server_whitelist_erc20_test.go
@@ -10,13 +10,15 @@ import (
 	"github.com/zeta-chain/zetacore/common"
 	keepertest "github.com/zeta-chain/zetacore/testutil/keeper"
 	"github.com/zeta-chain/zetacore/testutil/sample"
+	crosschainkeeper "github.com/zeta-chain/zetacore/x/crosschain/keeper"
 	"github.com/zeta-chain/zetacore/x/crosschain/types"
 	fungibletypes "github.com/zeta-chain/zetacore/x/fungible/types"
 )
 
 func TestKeeper_WhitelistERC20(t *testing.T) {
 	t.Run("can deploy and whitelist an erc20", func(t *testing.T) {
 		k, ctx, sdkk, zk := keepertest.CrosschainKeeper(t)
+		msgServer := crosschainkeeper.NewMsgServerImpl(*k)
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 
 		chainID := getValidEthChainID(t)
@@ -33,7 +35,7 @@ func TestKeeper_WhitelistERC20(t *testing.T) {
 		})
 
 		erc20Address := sample.EthAddress().Hex()
-		res, err := k.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
+		res, err := msgServer.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
 			Creator:      admin,
 			Erc20Address: erc20Address,
 			ChainId:      chainID,
@@ -63,7 +65,7 @@ func TestKeeper_WhitelistERC20(t *testing.T) {
 		require.Equal(t, uint64(100000), gasLimit.Uint64())
 
 		// Ensure that whitelist a new erc20 create a cctx with a different index
-		res, err = k.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
+		res, err = msgServer.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
 			Creator:      admin,
 			Erc20Address: sample.EthAddress().Hex(),
 			ChainId:      chainID,
@@ -79,9 +81,10 @@ func TestKeeper_WhitelistERC20(t *testing.T) {
 
 	t.Run("should fail if not authorized", func(t *testing.T) {
 		k, ctx, _, _ := keepertest.CrosschainKeeper(t)
+		msgServer := crosschainkeeper.NewMsgServerImpl(*k)
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 
-		_, err := k.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
+		_, err := msgServer.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
 			Creator:      sample.AccAddress(),
 			Erc20Address: sample.EthAddress().Hex(),
 			ChainId:      getValidEthChainID(t),
@@ -95,12 +98,13 @@ func TestKeeper_WhitelistERC20(t *testing.T) {
 
 	t.Run("should fail if invalid erc20 address", func(t *testing.T) {
 		k, ctx, _, zk := keepertest.CrosschainKeeper(t)
+		msgServer := crosschainkeeper.NewMsgServerImpl(*k)
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
 
-		_, err := k.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
+		_, err := msgServer.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
 			Creator:      admin,
 			Erc20Address: "invalid",
 			ChainId:      getValidEthChainID(t),
@@ -114,6 +118,7 @@ func TestKeeper_WhitelistERC20(t *testing.T) {
 
 	t.Run("should fail if foreign coin already exists for the asset", func(t *testing.T) {
 		k, ctx, _, zk := keepertest.CrosschainKeeper(t)
+		msgServer := crosschainkeeper.NewMsgServerImpl(*k)
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 
 		admin := sample.AccAddress()
@@ -126,7 +131,7 @@ func TestKeeper_WhitelistERC20(t *testing.T) {
 		fc.ForeignChainId = chainID
 		zk.FungibleKeeper.SetForeignCoins(ctx, fc)
 
-		_, err := k.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
+		_, err := msgServer.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
 			Creator:      admin,
 			Erc20Address: asset,
 			ChainId:      chainID,
@@ -140,14 +145,15 @@ func TestKeeper_WhitelistERC20(t *testing.T) {
 
 	t.Run("should fail if no tss set", func(t *testing.T) {
 		k, ctx, _, zk := keepertest.CrosschainKeeper(t)
+		msgServer := crosschainkeeper.NewMsgServerImpl(*k)
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 
 		chainID := getValidEthChainID(t)
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
 
 		erc20Address := sample.EthAddress().Hex()
-		_, err := k.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
+		_, err := msgServer.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
 			Creator:      admin,
 			Erc20Address: erc20Address,
 			ChainId:      chainID,
@@ -161,14 +167,15 @@ func TestKeeper_WhitelistERC20(t *testing.T) {
 
 	t.Run("should fail if nox valid chain ID", func(t *testing.T) {
 		k, ctx, _, zk := keepertest.CrosschainKeeper(t)
+		msgServer := crosschainkeeper.NewMsgServerImpl(*k)
 		k.GetAuthKeeper().GetModuleAccount(ctx, fungibletypes.ModuleName)
 
 		admin := sample.AccAddress()
 		setAdminPolicies(ctx, zk, admin)
 		k.SetTssAndUpdateNonce(ctx, *sample.Tss())
 
 		erc20Address := sample.EthAddress().Hex()
-		_, err := k.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
+		_, err := msgServer.WhitelistERC20(ctx, &types.MsgWhitelistERC20{
 			Creator:      admin,
 			Erc20Address: erc20Address,
 			ChainId:      10000,

diff --git a/x/crosschain/keeper/test_privnet.go b/x/crosschain/keeper/test_privnet.go
@@ -15,7 +15,8 @@ func (k Keeper) TestWhitelistERC20(ctx sdk.Context) error {
 	creator := k.zetaObserverKeeper.GetParams(ctx).GetAdminPolicyAccount(zetaObserverTypes.Policy_Type_group1)
 	msg := types.NewMsgWhitelistERC20(creator, types.ModuleAddressEVM.Hex(), common.GoerliChain().ChainId, "test", "testerc20", 17, 90_000)
 
-	_, err := k.WhitelistERC20(goCtx, msg)
+	msgServer := NewMsgServerImpl(k)
+	_, err := msgServer.WhitelistERC20(goCtx, msg)
 	if err != nil {
 		panic(err)
 	}

diff --git a/x/crosschain/types/codec.go b/x/crosschain/types/codec.go
@@ -9,24 +9,32 @@ import (
 
 func RegisterCodec(cdc *codec.LegacyAmino) {
 	cdc.RegisterConcrete(&MsgAddToOutTxTracker{}, "crosschain/AddToOutTxTracker", nil)
+	cdc.RegisterConcrete(&MsgAddToInTxTracker{}, "crosschain/AddToInTxTracker", nil)
 	cdc.RegisterConcrete(&MsgRemoveFromOutTxTracker{}, "crosschain/RemoveFromOutTxTracker", nil)
 	cdc.RegisterConcrete(&MsgCreateTSSVoter{}, "crosschain/CreateTSSVoter", nil)
 	cdc.RegisterConcrete(&MsgGasPriceVoter{}, "crosschain/GasPriceVoter", nil)
 	cdc.RegisterConcrete(&MsgNonceVoter{}, "crosschain/NonceVoter", nil)
-	cdc.RegisterConcrete(&MsgVoteOnObservedOutboundTxResponse{}, "crosschain/ReceiveConfirmation", nil)
-	cdc.RegisterConcrete(&MsgVoteOnObservedInboundTx{}, "crosschain/SendVoter", nil)
+	cdc.RegisterConcrete(&MsgVoteOnObservedOutboundTx{}, "crosschain/VoteOnObservedOutboundTx", nil)
+	cdc.RegisterConcrete(&MsgVoteOnObservedInboundTx{}, "crosschain/VoteOnObservedInboundTx", nil)
+	cdc.RegisterConcrete(&MsgWhitelistERC20{}, "crosschain/WhitelistERC20", nil)
+	cdc.RegisterConcrete(&MsgMigrateTssFunds{}, "crosschain/MigrateTssFunds", nil)
+	cdc.RegisterConcrete(&MsgUpdateTssAddress{}, "crosschain/UpdateTssAddress", nil)
 	cdc.RegisterConcrete(&MsgSetNodeKeys{}, "crosschain/SetNodeKeys", nil)
 }
 
 func RegisterInterfaces(registry cdctypes.InterfaceRegistry) {
 	registry.RegisterImplementations((*sdk.Msg)(nil),
 		&MsgAddToOutTxTracker{},
+		&MsgAddToInTxTracker{},
 		&MsgRemoveFromOutTxTracker{},
 		&MsgCreateTSSVoter{},
 		&MsgGasPriceVoter{},
 		&MsgNonceVoter{},
 		&MsgVoteOnObservedOutboundTx{},
 		&MsgVoteOnObservedInboundTx{},
+		&MsgWhitelistERC20{},
+		&MsgMigrateTssFunds{},
+		&MsgUpdateTssAddress{},
 		&MsgSetNodeKeys{},
 	)
 

diff --git a/x/fungible/keeper/begin_blocker_deploy_system_contracts_privnet.go b/x/fungible/keeper/begin_blocker_deploy_system_contracts_privnet.go
@@ -116,6 +116,7 @@ func (k Keeper) BlockOneDeploySystemContracts(goCtx context.Context) error {
 
 func (k Keeper) TestUpdateSystemContractAddress(goCtx context.Context) error {
 	ctx := sdk.UnwrapSDKContext(goCtx)
+	msgServer := NewMsgServerImpl(k)
 
 	wzeta, err := k.GetWZetaContractAddress(ctx)
 	if err != nil {
@@ -136,13 +137,14 @@ func (k Keeper) TestUpdateSystemContractAddress(goCtx context.Context) error {
 	}
 	creator := k.observerKeeper.GetParams(ctx).GetAdminPolicyAccount(observertypes.Policy_Type_group1)
 	msg := types.NewMsgUpdateSystemContract(creator, SystemContractAddress.Hex())
-	_, err = k.UpdateSystemContract(ctx, msg)
+	_, err = msgServer.UpdateSystemContract(ctx, msg)
 	k.Logger(ctx).Info("System contract updated", "new address", SystemContractAddress.String())
 	return err
 }
 
 func (k Keeper) TestUpdateZRC20WithdrawFee(goCtx context.Context) error {
 	ctx := sdk.UnwrapSDKContext(goCtx)
+	msgServer := NewMsgServerImpl(k)
 
 	foreignCoins := k.GetAllForeignCoins(ctx)
 	creator := k.observerKeeper.GetParams(ctx).GetAdminPolicyAccount(observertypes.Policy_Type_group1)
@@ -154,7 +156,7 @@ func (k Keeper) TestUpdateZRC20WithdrawFee(goCtx context.Context) error {
 			sdk.NewUint(uint64(foreignCoin.ForeignChainId)),
 			math.Uint{},
 		)
-		_, err := k.UpdateZRC20WithdrawFee(ctx, msg)
+		_, err := msgServer.UpdateZRC20WithdrawFee(ctx, msg)
 		if err != nil {
 			return err
 		}

diff --git a/x/fungible/keeper/msg_server_update_contract_bytecode.go b/x/fungible/keeper/msg_server_update_contract_bytecode.go
@@ -16,7 +16,7 @@ import (
 // Only a ZRC20 contract or the WZeta connector contract can be updated
 // IMPORTANT: the new contract bytecode must have the same storage layout as the old contract bytecode
 // the new contract can add new variable but cannot remove any existing variable
-func (k Keeper) UpdateContractBytecode(goCtx context.Context, msg *types.MsgUpdateContractBytecode) (*types.MsgUpdateContractBytecodeResponse, error) {
+func (k msgServer) UpdateContractBytecode(goCtx context.Context, msg *types.MsgUpdateContractBytecode) (*types.MsgUpdateContractBytecodeResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
 	// check authorization