Add query decrypted tx method to confidential transfers client (#1950)

* queries work

* comments

* update doc

* use decryptor instead of from flag

* validate decryptor

* lint

* merge

aclmapping/confidentialtransfers/mappings_test.go

@@ -3,6 +3,7 @@ package confidentialtransfers_test
 import (
 	"crypto/ecdsa"
 	"fmt"
+
 	"github.com/cosmos/cosmos-sdk/simapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkacltypes "github.com/cosmos/cosmos-sdk/types/accesscontrol"
@@ -21,10 +22,11 @@ import (
 	"github.com/stretchr/testify/require"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 
+	"testing"
+
 	"github.com/sei-protocol/sei-chain/x/confidentialtransfers/keeper"
 	"github.com/sei-protocol/sei-chain/x/confidentialtransfers/types"
 	"github.com/stretchr/testify/suite"
-	"testing"
 )
 
 const (
@@ -272,7 +274,7 @@ func (suite *KeeperTestSuite) TestMsgApplyPendingBalanceDependencies() {
 	// Initialize an account
 	initialState, _ := suite.SetupAccountState(senderPk, DefaultTestDenom, 10, 2000, 3000, 1000)
 
-	applyPendingBalanceMsg, _ := types.NewMsgApplyPendingBalance(
+	applyPendingBalance, _ := types.NewApplyPendingBalance(
 		*senderPk,
 		senderAddr.String(),
 		DefaultTestDenom,
@@ -282,6 +284,8 @@ func (suite *KeeperTestSuite) TestMsgApplyPendingBalanceDependencies() {
 		initialState.PendingBalanceLo,
 		initialState.PendingBalanceHi)
 
+	applyPendingBalanceMsg := types.NewMsgApplyPendingBalanceProto(applyPendingBalance)
+
 	tests := []struct {
 		name          string
 		expectedError error

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,23 +1,31 @@
 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"
 	"github.com/spf13/cobra"
 )
 
 const decryptAvailableBalanceFlag = "decrypt-available-balance"
+const decryptorFlag = "decryptor"
 
 // 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 +35,7 @@ func GetQueryCmd() *cobra.Command {
 	confidentialTransfersQueryCmd.AddCommand(
 		GetCmdQueryAccount(),
 		GetCmdQueryAllAccount(),
+		GetCmdQueryTx(),
 	)
 
 	return confidentialTransfersQueryCmd
@@ -38,14 +47,14 @@ func GetCmdQueryAccount() *cobra.Command {
 		Use:   "account [denom] [address] [flags]",
 		Short: "Query the account state",
 		Long: "Queries the account state associated with the address and denom." +
-			"Pass the --from flag to decrypt the account" +
+			"Pass the --decryptor flag to decrypt the account" +
 			"Pass the --decrypt-available-balance flag to attempt to decrypt the available balance.",
 		Args: cobra.ExactArgs(2),
 		RunE: queryAccount,
 	}
 
 	flags.AddQueryFlagsToCmd(cmd)
-	cmd.Flags().String(flags.FlagFrom, "", "Name or address of private key to decrypt the account")
+	cmd.Flags().String(decryptorFlag, "", "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
 }
@@ -72,11 +81,11 @@ func queryAccount(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("invalid address: %v", err)
 	}
 
-	from, err := cmd.Flags().GetString(flags.FlagFrom)
+	decryptorAccount, err := cmd.Flags().GetString(decryptorFlag)
 	if err != nil {
 		return err
 	}
-	fromAddr, _, _, err := client.GetFromFields(clientCtx, clientCtx.Keyring, from)
+	decryptorAddr, name, _, err := client.GetFromFields(clientCtx, clientCtx.Keyring, decryptorAccount)
 	if err != nil {
 		return err
 	}
@@ -89,13 +98,13 @@ func queryAccount(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	// If the --from flag passed matches the queried address, attempt to decrypt the contents
-	if fromAddr.String() == address {
+	// If the decryptor flag passed matches the queried address, attempt to decrypt the contents
+	if decryptorAddr.String() == address {
 		account, err := res.Account.FromProto()
 		if err != nil {
 			return err
 		}
-		privateKey, err := getPrivateKey(cmd)
+		privateKey, err := getPrivateKey(cmd, name)
 		if err != nil {
 			return err
 		}
@@ -178,3 +187,145 @@ 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's decrypted state by decrypting with the senders private key.
+func GetCmdQueryTx() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "tx [hash] --decryptor [decryptor] [flags]",
+		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 --decryptor. For decryption to work, the decryptor should be of a sender, receiver or auditor." +
+			"Pass the --decrypt-available-balance flag to attempt to decrypt the available balance. (This is an expensive operation and may not succeed even if the private key provided is valid)",
+		Args: cobra.ExactArgs(1),
+		RunE: queryDecryptedTx,
+	}
+
+	flags.AddQueryFlagsToCmd(cmd)
+	cmd.Flags().String(decryptorFlag, "", "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
+	}
+
+	// Get the transaction hash
+	txHashHex := args[0]
+
+	decryptorAccount, err := cmd.Flags().GetString(decryptorFlag)
+	if err != nil {
+		return err
+	}
+
+	if decryptorAccount == "" {
+		return fmt.Errorf("--decryptor flag must be set since we need the private key to decrypt the transaction")
+	}
+
+	fromAddr, name, _, err := client.GetFromFields(clientCtx, clientCtx.Keyring, decryptorAccount)
+	if err != nil {
+		return err
+	}
+	clientCtx = clientCtx.WithFrom(decryptorAccount).WithFromAddress(fromAddr).WithFromName(name)
+
+	decryptAvailableBalance, err := cmd.Flags().GetBool(decryptAvailableBalanceFlag)
+	if err != nil {
+		return err
+	}
+	if decryptAvailableBalance {
+		err = clientCtx.PrintString("--decrypt-available-balance set to true. This operation could take a long time and may not succeed even if the private key provided is valid\n")
+		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, name)
+
+	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
+	}
+
+	var result proto.Message
+	switch message := sdkmsg.(type) {
+	case *types.MsgInitializeAccount:
+		result, err = message.Decrypt(decryptor, *privKey, decryptAvailableBalance)
+	case *types.MsgWithdraw:
+		result, err = message.Decrypt(decryptor, *privKey, decryptAvailableBalance)
+	case *types.MsgApplyPendingBalance:
+		result, err = message.Decrypt(decryptor, *privKey, decryptAvailableBalance)
+	case *types.MsgTransfer:
+		result, err = message.Decrypt(decryptor, *privKey, decryptAvailableBalance, address)
+	case *types.MsgDeposit:
+		result = message
+	case *types.MsgCloseAccount:
+		result = message
+	default:
+		return nil, false, nil
+	}
+
+	return result, true, err
+}