Add query decrypted tx method to confidential transfers client

proto/confidentialtransfers/query.proto

@@ -5,6 +5,7 @@ import "cosmos/base/query/v1beta1/pagination.proto";
 import "gogoproto/gogo.proto";
 import "google/api/annotations.proto";
 import "confidentialtransfers/confidential.proto";
+import "confidentialtransfers/zk.proto";
 
 option go_package = "github.com/sei-protocol/sei-chain/x/confidentialtransfers/types";
 
@@ -60,4 +61,61 @@ message DecryptedCtAccount {
   uint32   pending_balance_credit_counter = 5;
   string   available_balance = 6; // decrypted available balance
   uint64   decryptable_available_balance = 7; // decrypted aes encrypted available balance
-}
+}
+
+// Decrypted version of ApplyPendingBalance
+message ApplyPendingBalanceDecrypted {
+  option (gogoproto.equal) = false;
+  option (gogoproto.goproto_getters) = false;
+
+  string address = 1;
+  string denom = 2;
+  uint64 new_decryptable_available_balance = 3;
+  uint32 current_pending_balance_counter = 4;
+  string current_available_balance = 5;
+}
+
+// Decrypted version of InitializeAccount
+message InitializeAccountDecrypted {
+  option (gogoproto.equal) = false;
+  option (gogoproto.goproto_getters) = false;
+
+  string from_address = 1;
+  string denom = 2;
+  bytes pubkey = 3;
+  uint32 pending_balance_lo = 4;
+  uint64 pending_balance_hi = 5;
+  string available_balance = 6;
+  uint64 decryptable_balance = 7;
+  InitializeAccountMsgProofs proofs = 8;
+}
+
+// Decrypted version of Withdraw
+message WithdrawDecrypted {
+  option (gogoproto.equal) = false;
+  option (gogoproto.goproto_getters) = false;
+
+  string from_address = 1;
+  string denom = 2;
+  uint64 amount = 3;
+  uint64 decryptable_balance = 4;
+  string remaining_balance_commitment = 5;
+  WithdrawMsgProofs proofs = 6;
+}
+
+// Decrypted version of Transfer
+message TransferDecrypted {
+  option (gogoproto.equal) = false;
+  option (gogoproto.goproto_getters) = false;
+
+  string from_address = 1;
+  string to_address = 2;
+  string denom = 3;
+  uint32 transfer_amount_lo = 4;
+  uint32 transfer_amount_hi = 5;
+  uint64 total_transfer_amount = 6;
+  string remaining_balance_commitment = 7;
+  string decryptable_balance = 8;
+  TransferMsgProofs proofs = 9;
+  repeated string auditors = 10;
+}

x/confidentialtransfers/client/cli/query.go

@@ -1,11 +1,18 @@
 package cli
 
 import (
+	"context"
+	"crypto/ecdsa"
+	"encoding/hex"
 	"fmt"
 
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/flags"
+	"github.com/cosmos/cosmos-sdk/codec"
+	cdctypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/tx"
+	"github.com/gogo/protobuf/proto"
 	"github.com/sei-protocol/sei-chain/x/confidentialtransfers/types"
 	"github.com/sei-protocol/sei-cryptography/pkg/encryption"
 	"github.com/sei-protocol/sei-cryptography/pkg/encryption/elgamal"
@@ -17,7 +24,7 @@ const decryptAvailableBalanceFlag = "decrypt-available-balance"
 // GetQueryCmd returns the cli query commands for the minting module.
 func GetQueryCmd() *cobra.Command {
 	confidentialTransfersQueryCmd := &cobra.Command{
-		Use:                        types.ModuleName,
+		Use:                        types.ShortModuleName,
 		Short:                      "Querying commands for the confidential transfer module",
 		DisableFlagParsing:         true,
 		SuggestionsMinimumDistance: 2,
@@ -27,6 +34,7 @@ func GetQueryCmd() *cobra.Command {
 	confidentialTransfersQueryCmd.AddCommand(
 		GetCmdQueryAccount(),
 		GetCmdQueryAllAccount(),
+		GetCmdQueryTx(),
 	)
 
 	return confidentialTransfersQueryCmd
@@ -178,3 +186,126 @@ func queryAllAccounts(cmd *cobra.Command, args []string) error {
 
 	return nil
 }
+
+// GetCmdQueryTx implements a command to query a tx by it's transaction hash and return it in it's decrypted state.
+func GetCmdQueryTx() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "tx [hash]",
+		Short: "Query the confidential transaction and decrypts it",
+		Long: "Query the confidential transaction by it's tx hash and decrypts it using the private key of the account in --from" +
+			"Pass the --from flag to decrypt the account" +
+			"Pass the --decrypt-available-balance flag to attempt to decrypt the available balance. (This is an expensive operation)",
+		Args: cobra.ExactArgs(1),
+		RunE: queryDecryptedTx,
+	}
+
+	flags.AddQueryFlagsToCmd(cmd)
+	cmd.Flags().String(flags.FlagFrom, "", "Name or address of private key to decrypt the account")
+	cmd.Flags().Bool(decryptAvailableBalanceFlag, false, "Set this to attempt to decrypt the available balance")
+	return cmd
+}
+
+func queryDecryptedTx(cmd *cobra.Command, args []string) error {
+	clientCtx, err := client.GetClientQueryContext(cmd)
+	if err != nil {
+		return err
+	}
+
+	txHashHex := args[0]
+
+	from, err := cmd.Flags().GetString(flags.FlagFrom)
+	if err != nil {
+		return err
+	}
+	fromAddr, _, _, err := client.GetFromFields(clientCtx, clientCtx.Keyring, from)
+	if err != nil {
+		return err
+	}
+
+	decryptAvailableBalance, err := cmd.Flags().GetBool(decryptAvailableBalanceFlag)
+	if err != nil {
+		return err
+	}
+
+	// Decode the transaction hash from hex to bytes
+	txHash, err := hex.DecodeString(txHashHex)
+	if err != nil {
+		return fmt.Errorf("failed to decode transaction hash: %w", err)
+	}
+
+	// Connect to the gRPC client
+	node, err := clientCtx.GetNode()
+	if err != nil {
+		return fmt.Errorf("failed to connect to node: %w", err)
+	}
+
+	// Query the transaction using Tendermint's Tx endpoint
+	res, err := node.Tx(context.Background(), txHash, false)
+	if err != nil {
+		return fmt.Errorf("failed to fetch transaction: %w", err)
+	}
+
+	// Decode the transaction
+	var rawTx tx.Tx
+	if err := clientCtx.Codec.Unmarshal(res.Tx, &rawTx); err != nil {
+		return fmt.Errorf("failed to unmarshal transaction: %w", err)
+	}
+
+	decryptor := elgamal.NewTwistedElgamal()
+	privateKey, err := getPrivateKey(cmd)
+
+	if err != nil {
+		return err
+	}
+	msgPrinted := false
+	for _, msg := range rawTx.Body.Messages {
+		result, foundMsg, err := handleDecryptableMessage(clientCtx.Codec, msg, decryptor, privateKey, decryptAvailableBalance, fromAddr.String())
+		if !foundMsg {
+			continue
+		} else {
+			if err != nil {
+				return err
+			}
+			err = clientCtx.PrintProto(result)
+			msgPrinted = true
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	if !msgPrinted {
+		return fmt.Errorf("no decryptable message found in the transaction")
+	}
+
+	return nil
+}
+
+// Helper function to unmarshal a message and run its Decrypt() method
+func handleDecryptableMessage(
+	cdc codec.Codec,
+	msgAny *cdctypes.Any,
+	decryptor *elgamal.TwistedElGamal,
+	privKey *ecdsa.PrivateKey,
+	decryptAvailableBalance bool,
+	address string) (msg proto.Message, foundDecryptableMsg bool, error error) {
+	// Try to unmarshal the message as one of the known types
+	var sdkmsg sdk.Msg
+	err := cdc.UnpackAny(msgAny, &sdkmsg)
+	if err != nil {
+		return nil, false, nil
+	}
+
+	decryptable, ok := sdkmsg.(types.Decryptable)
+	if !ok {
+		return nil, false, nil
+	}
+
+	// Successfully unmarshaled as decryptable message, run the Decrypt() method
+	result, err := decryptable.Decrypt(decryptor, *privKey, decryptAvailableBalance, address)
+	if err != nil {
+		return nil, true, fmt.Errorf("failed to decrypt message: %w", err)
+	}
+
+	return result, true, nil
+}

x/confidentialtransfers/client/cli/tx.go

@@ -451,7 +451,7 @@ func makeApplyPendingBalanceCmd(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	msg, err := types.NewMsgApplyPendingBalance(
+	applyPendingBalance, err := types.NewApplyPendingBalance(
 		*privKey,
 		address,
 		denom,
@@ -464,6 +464,12 @@ func makeApplyPendingBalanceCmd(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+
+	msg := types.NewMsgApplyPendingBalanceProto(applyPendingBalance)
+
+	if err != nil {
+		return err
+	}
 
 	if err = msg.ValidateBasic(); err != nil {
 		return err

x/confidentialtransfers/keeper/msg_server_test.go

@@ -813,7 +813,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalance() {
 	initialState, _ := suite.SetupAccountState(testPk, DefaultTestDenom, 10, initialAvailableBalance, initialPendingBalance, 1000)
 
 	// Create an apply pending balance request
-	req, err := types.NewMsgApplyPendingBalance(
+	applyPendingBalance, _ := types.NewApplyPendingBalance(
 		*testPk,
 		testAddr.String(),
 		DefaultTestDenom,
@@ -823,8 +823,9 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalance() {
 		initialState.PendingBalanceLo,
 		initialState.PendingBalanceHi)
 
+	req := types.NewMsgApplyPendingBalanceProto(applyPendingBalance)
 	// Execute the apply pending balance
-	_, err = suite.msgServer.ApplyPendingBalance(sdk.WrapSDKContext(suite.Ctx), req)
+	_, err := suite.msgServer.ApplyPendingBalance(sdk.WrapSDKContext(suite.Ctx), req)
 	suite.Require().NoError(err, "Should not have error applying pending balance")
 
 	// Check that the account has been updated
@@ -868,7 +869,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalanceAfterWithdraw() {
 	initialState, _ := suite.SetupAccountState(testPk, DefaultTestDenom, 10, initialAvailableBalance, initialPendingBalance, 1000)
 
 	// Create an apply pending balance request
-	req, err := types.NewMsgApplyPendingBalance(
+	applyPendingBalance, _ := types.NewApplyPendingBalance(
 		*testPk,
 		testAddr.String(),
 		DefaultTestDenom,
@@ -878,10 +879,12 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalanceAfterWithdraw() {
 		initialState.PendingBalanceLo,
 		initialState.PendingBalanceHi)
 
+	req := types.NewMsgApplyPendingBalanceProto(applyPendingBalance)
+
 	// Before the pending balance is applied, a withdrawal is made, changing the available balance in the account.
 	withdrawReq, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, initialAvailableBalance/2)
 	withdrawMsg := types.NewMsgWithdrawProto(withdrawReq)
-	_, err = suite.msgServer.Withdraw(sdk.WrapSDKContext(suite.Ctx), withdrawMsg)
+	_, err := suite.msgServer.Withdraw(sdk.WrapSDKContext(suite.Ctx), withdrawMsg)
 
 	// Now execute the apply pending balance. Checks in the ApplyPendingBalance function should catch this and cause it to fail so we don't wrongly update the new decryptableAvailableBalance
 	_, err = suite.msgServer.ApplyPendingBalance(sdk.WrapSDKContext(suite.Ctx), req)
@@ -902,7 +905,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalanceAfterDeposit() {
 	initialState, _ := suite.SetupAccountState(testPk, DefaultTestDenom, 10, initialAvailableBalance, initialPendingBalance, 1000)
 
 	// Create an apply pending balance request
-	req, err := types.NewMsgApplyPendingBalance(
+	applyPendingBalance, _ := types.NewApplyPendingBalance(
 		*testPk,
 		testAddr.String(),
 		DefaultTestDenom,
@@ -912,14 +915,16 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalanceAfterDeposit() {
 		initialState.PendingBalanceLo,
 		initialState.PendingBalanceHi)
 
+	req := types.NewMsgApplyPendingBalanceProto(applyPendingBalance)
+
 	// Before the pending balance is applied, a deposit is made, changing the pending balance in the account.
 	// The same scenario happens when incoming transfers are received.
 	depositMsg := &types.MsgDeposit{
 		testAddr.String(),
 		DefaultTestDenom,
 		1000,
 	}
-	_, err = suite.msgServer.Deposit(sdk.WrapSDKContext(suite.Ctx), depositMsg)
+	_, err := suite.msgServer.Deposit(sdk.WrapSDKContext(suite.Ctx), depositMsg)
 
 	// Now execute the apply pending balance. Checks in the ApplyPendingBalance function should catch this and cause it to fail so we don't wrongly update the new decryptableAvailableBalance
 	_, err = suite.msgServer.ApplyPendingBalance(sdk.WrapSDKContext(suite.Ctx), req)
@@ -939,7 +944,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalanceNoPendingBalances
 	initialState, _ := suite.SetupAccountState(testPk, DefaultTestDenom, 0, initialAvailableBalance, uint64(0), 1000)
 
 	// Create an apply pending balance request
-	req, err := types.NewMsgApplyPendingBalance(
+	applyPendingBalance, err := types.NewApplyPendingBalance(
 		*testPk,
 		testAddr.String(),
 		DefaultTestDenom,
@@ -951,6 +956,8 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalanceNoPendingBalances
 
 	suite.Require().NoError(err, "Should not have error creating apply pending balance request")
 
+	req := types.NewMsgApplyPendingBalanceProto(applyPendingBalance)
+
 	// Execute the apply pending balance. This should fail since there are no pending balances to apply.
 	_, err = suite.msgServer.ApplyPendingBalance(sdk.WrapSDKContext(suite.Ctx), req)
 	suite.Require().Error(err, "Should have error applying pending balance on account with no pending balances")