fix: Burn and Recover msgs public key checks and cli commands (#13)

monerium · Oct 9, 2024 · f2109c8 · f2109c8
1 parent cebdf44
commit f2109c8
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 4 deletions.
diff --git a/client/cli/tx.go b/client/cli/tx.go
@@ -23,6 +23,8 @@ import (
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/client/tx"
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	"github.com/monerium/module-noble/v2/types"
 	"github.com/spf13/cobra"
 )
@@ -160,9 +162,9 @@ func TxAllowDenom() *cobra.Command {
 
 func TxBurn() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "burn [denom] [from] [amount] [signature]",
+		Use:   "burn [denom] [from] [amount] [signature] [pub_key]",
 		Short: "Transaction that burns a specific denom",
-		Args:  cobra.ExactArgs(4),
+		Args:  cobra.ExactArgs(5),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			clientCtx, err := client.GetClientTxContext(cmd)
 			if err != nil {
@@ -179,12 +181,22 @@ func TxBurn() *cobra.Command {
 				return err
 			}
 
+			var pubKey cryptotypes.PubKey
+			if err = clientCtx.Codec.UnmarshalInterfaceJSON([]byte(args[4]), &pubKey); err != nil {
+				return err
+			}
+			anyPubKey, err := codectypes.NewAnyWithValue(pubKey)
+			if err != nil {
+				return err
+			}
+
 			msg := &types.MsgBurn{
 				Denom:     args[0],
 				Signer:    clientCtx.GetFromAddress().String(),
 				From:      args[1],
 				Amount:    amount,
 				Signature: signature,
+				PubKey:    anyPubKey,
 			}
 
 			return tx.GenerateOrBroadcastTxCLI(clientCtx, cmd.Flags(), msg)
@@ -230,9 +242,9 @@ func TxMint() *cobra.Command {
 
 func TxRecover() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "recover [denom] [from] [to] [signature]",
+		Use:   "recover [denom] [from] [to] [signature] [pub_key]",
 		Short: "Recover balance of a specific denom from an account",
-		Args:  cobra.ExactArgs(4),
+		Args:  cobra.ExactArgs(5),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			clientCtx, err := client.GetClientTxContext(cmd)
 			if err != nil {
@@ -244,12 +256,22 @@ func TxRecover() *cobra.Command {
 				return err
 			}
 
+			var pubKey cryptotypes.PubKey
+			if err = clientCtx.Codec.UnmarshalInterfaceJSON([]byte(args[4]), &pubKey); err != nil {
+				return err
+			}
+			anyPubKey, err := codectypes.NewAnyWithValue(pubKey)
+			if err != nil {
+				return err
+			}
+
 			msg := &types.MsgRecover{
 				Denom:     args[0],
 				Signer:    clientCtx.GetFromAddress().String(),
 				From:      args[1],
 				To:        args[2],
 				Signature: signature,
+				PubKey:    anyPubKey,
 			}
 
 			return tx.GenerateOrBroadcastTxCLI(clientCtx, cmd.Flags(), msg)

diff --git a/keeper/msg_server.go b/keeper/msg_server.go
@@ -131,6 +131,9 @@ func (k msgServer) Burn(ctx context.Context, msg *types.MsgBurn) (*types.MsgBurn
 	if !k.IsSystem(ctx, msg.Denom, msg.Signer) {
 		return nil, types.ErrInvalidSystem
 	}
+	if msg.PubKey == nil {
+		return nil, types.ErrInvalidPubKey
+	}
 
 	var pubKey cryptotypes.PubKey
 	if err := k.cdc.UnpackAny(msg.PubKey, &pubKey); err != nil {
@@ -211,6 +214,9 @@ func (k msgServer) Recover(ctx context.Context, msg *types.MsgRecover) (*types.M
 	if !k.IsSystem(ctx, msg.Denom, msg.Signer) {
 		return nil, types.ErrInvalidSystem
 	}
+	if msg.PubKey == nil {
+		return nil, types.ErrInvalidPubKey
+	}
 
 	var pubKey cryptotypes.PubKey
 	if err := k.cdc.UnpackAny(msg.PubKey, &pubKey); err != nil {

diff --git a/keeper/msg_server_test.go b/keeper/msg_server_test.go
@@ -373,6 +373,15 @@ func TestBurn(t *testing.T) {
 	// ASSERT: The action should've failed due to invalid user address.
 	require.ErrorContains(t, err, "unable to decode user address")
 
+	// ACT: Attempt to burn with missing public key.
+	_, err = server.Burn(ctx, &types.MsgBurn{
+		Denom:  "ueure",
+		Signer: system.Address,
+		From:   utils.TestAccount().Address,
+	})
+	// ASSERT: The action should've failed due to invalid public key.
+	require.ErrorIs(t, err, types.ErrInvalidPubKey)
+
 	// ACT: Attempt to burn with invalid public key.
 	invalidPubKey, _ = codectypes.NewAnyWithValue(secp256k1.GenPrivKey().PubKey())
 	_, err = server.Burn(ctx, &types.MsgBurn{
@@ -597,6 +606,15 @@ func TestRecover(t *testing.T) {
 	// ASSERT: The action should've failed due to invalid public key.
 	require.ErrorIs(t, err, types.ErrInvalidPubKey)
 
+	// ACT: Attempt to recover with missing public key.
+	_, err = server.Recover(ctx, &types.MsgRecover{
+		Denom:  "ueure",
+		Signer: system.Address,
+		From:   "noble1rwvjzk28l38js7xx6mq23nrpghd8qqvxmj6ep2",
+	})
+	// ASSERT: The action should've failed due to invalid public key.
+	require.ErrorIs(t, err, types.ErrInvalidPubKey)
+
 	// ACT: Attempt to recover with invalid signature.
 	signature, _ := base64.StdEncoding.DecodeString("QBrRfIqjdBvXx9zaBcuiE9P5SVesxFO/He3deyx2OE0NoSNqwmSb7b5iP2UhZRI1duiOeho3+NETUkCBv14zjQ==")
 	_, err = server.Recover(ctx, &types.MsgRecover{