add ct to load test framework (#1959)

* add ct to load test framework

* comment out receiver account creation for now

* - add ct account population script
- add ct query client
- add 4 types of ct messages for load test

* - script updates

* - add error handling

* script updates

* formatting

* disable state writing for transfers

* add metrics

* Revert "add metrics"

This reverts commit a7814b1.

* self review fixes

loadtest/config.json

@@ -61,7 +61,11 @@
       "wasm_occ_iterator_range",
       "wasm_occ_parallel_write",
       "evm",
-      "univ2"
+      "univ2",
+      "confidential_transfers_transfer",
+      "confidential_transfers_deposit",
+      "confidential_transfers_apply_pending_balance",
+      "confidential_transfers_withdraw"
     ],
     "evm_use_eip1559_txs": true,
     "run_oracle": false,

loadtest/loadtest_client.go

@@ -10,6 +10,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	cttypes "github.com/sei-protocol/sei-chain/x/confidentialtransfers/types"
+
 	"golang.org/x/exp/slices"
 	"golang.org/x/time/rate"
 
@@ -35,6 +37,7 @@ type LoadTestClient struct {
 	SignerClient       *SignerClient
 	ChainID            string
 	GrpcConns          []*grpc.ClientConn
+	CtQueryClient      cttypes.QueryClient
 	StakingQueryClient stakingtypes.QueryClient
 	// Staking specific variables
 	Validators []stakingtypes.Validator
@@ -76,6 +79,7 @@ func NewLoadTestClient(config Config) *LoadTestClient {
 		SignerClient:                  signerClient,
 		ChainID:                       config.ChainID,
 		GrpcConns:                     grpcConns,
+		CtQueryClient:                 cttypes.NewQueryClient(grpcConns[0]),
 		StakingQueryClient:            stakingtypes.NewQueryClient(grpcConns[0]),
 		DelegationMap:                 map[string]map[string]int{},
 		TokenFactoryDenomOwner:        map[string]string{},

loadtest/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"bytes"
 	"context"
+	"encoding/hex"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -21,6 +22,9 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/sei-protocol/sei-cryptography/pkg/encryption/elgamal"
+
 	wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/codec/types"
@@ -39,6 +43,7 @@ import (
 
 	"github.com/sei-protocol/sei-chain/app"
 	"github.com/sei-protocol/sei-chain/utils/metrics"
+	cttypes "github.com/sei-protocol/sei-chain/x/confidentialtransfers/types"
 	tokenfactorytypes "github.com/sei-protocol/sei-chain/x/tokenfactory/types"
 )
 
@@ -512,6 +517,94 @@ func (c *LoadTestClient) generateMessage(key cryptotypes.PrivKey, msgType string
 			Contract: contract,
 			Msg:      wasmtypes.RawContractMessage([]byte(fmt.Sprintf("{\"test_occ_parallelism\":{\"value\": %d}}", value))),
 		}}
+	case ConfidentialTransfersDeposit:
+		depositMsg := &cttypes.MsgDeposit{
+			FromAddress: sdk.AccAddress(key.PubKey().Address()).String(),
+			Denom:       CtDefaultDenom,
+			Amount:      1,
+		}
+		msgs = append(msgs, depositMsg)
+	case ConfidentialTransfersWithdraw:
+		senderPrivHex := hex.EncodeToString(key.Bytes())
+		senderEcdsaKey, _ := crypto.HexToECDSA(senderPrivHex)
+		address := sdk.AccAddress(key.PubKey().Address()).String()
+		account := c.getCtAccount(address, CtDefaultDenom)
+		withdraw, err := cttypes.NewWithdraw(
+			*senderEcdsaKey,
+			account.AvailableBalance,
+			CtDefaultDenom,
+			sdk.AccAddress(key.PubKey().Address()).String(),
+			account.DecryptableAvailableBalance,
+			1,
+		)
+		if err != nil {
+			panic(fmt.Sprintf("error for senderAddress %s: %s\n", address, err.Error()))
+		}
+		withdrawMsg := cttypes.NewMsgWithdrawProto(withdraw)
+		msgs = append(msgs, withdrawMsg)
+	case ConfidentialTransfersApplyPendingBalance:
+		senderPrivHex := hex.EncodeToString(key.Bytes())
+		senderEcdsaKey, _ := crypto.HexToECDSA(senderPrivHex)
+		account := c.getCtAccount(sdk.AccAddress(key.PubKey().Address()).String(), CtDefaultDenom)
+
+		address := sdk.AccAddress(key.PubKey().Address()).String()
+		applyPendingBalance, err := cttypes.NewApplyPendingBalance(
+			*senderEcdsaKey,
+			address,
+			CtDefaultDenom,
+			account.DecryptableAvailableBalance,
+			account.PendingBalanceCreditCounter,
+			account.AvailableBalance,
+			account.PendingBalanceLo,
+			account.PendingBalanceHi)
+		if err != nil {
+			panic(fmt.Sprintf("error for senderAddress %s: %s\n", address, err.Error()))
+		}
+
+		applyPendingBalanceMsg := cttypes.NewMsgApplyPendingBalanceProto(applyPendingBalance)
+		msgs = append(msgs, applyPendingBalanceMsg)
+	case ConfidentialTransfersTransfer:
+		accountKeys := c.AccountKeys
+		// get a random key for receiver and if it's same as the current key, get another one
+		if len(accountKeys) < 2 {
+			panic("Need at least 2 accounts to transfer")
+		}
+		receiverKey := accountKeys[rand.Intn(len(accountKeys))]
+		for receiverKey.PubKey().Equals(key.PubKey()) {
+			receiverKey = accountKeys[rand.Intn(len(accountKeys))]
+		}
+
+		senderPrivHex := hex.EncodeToString(key.Bytes())
+		senderEcdsaKey, _ := crypto.HexToECDSA(senderPrivHex)
+		receiverPrivHex := hex.EncodeToString(receiverKey.Bytes())
+		receiverEcdsaKey, _ := crypto.HexToECDSA(receiverPrivHex)
+
+		teg := elgamal.NewTwistedElgamal()
+		receiverKeyPair, _ := teg.KeyGen(*receiverEcdsaKey, CtDefaultDenom)
+		senderAddress := sdk.AccAddress(key.PubKey().Address()).String()
+		receiverAddress := sdk.AccAddress(receiverKey.PubKey().Address()).String()
+		senderAccount := c.getCtAccount(senderAddress, CtDefaultDenom)
+		if senderAccount == nil {
+			panic(fmt.Sprintf("Sender account not found for address %s\n", senderAddress))
+		}
+
+		transfer, err := cttypes.NewTransfer(
+			senderEcdsaKey,
+			senderAddress,
+			receiverAddress,
+			CtDefaultDenom,
+			senderAccount.DecryptableAvailableBalance,
+			senderAccount.AvailableBalance,
+			1,
+			&receiverKeyPair.PublicKey,
+			[]cttypes.AuditorInput{},
+		)
+		if err != nil {
+			panic(fmt.Sprintf("error for address %s: %s\n", senderAddress, err.Error()))
+		}
+
+		transferMsg := cttypes.NewMsgTransferProto(transfer)
+		msgs = append(msgs, transferMsg)
 	default:
 		fmt.Printf("Unrecognized message type %s", msgType)
 	}
@@ -522,6 +615,17 @@ func (c *LoadTestClient) generateMessage(key cryptotypes.PrivKey, msgType string
 	return msgs, false, signer, gas, int64(fee)
 }
 
+func (c *LoadTestClient) getCtAccount(address string, denom string) *cttypes.Account {
+	getCtAccountReq := &cttypes.GetCtAccountRequest{
+		Address: address,
+		Denom:   denom,
+	}
+	getCtAccountRes, _ := c.CtQueryClient.GetCtAccount(context.Background(), getCtAccountReq)
+	ctAccount := getCtAccountRes.GetAccount()
+	account, _ := ctAccount.FromProto()
+	return account
+}
+
 func (c *LoadTestClient) generateStakingMsg(delegatorAddr string, chosenValidator string, srcAddr string) sdk.Msg {
 	c.mtx.Lock()
 	defer c.mtx.Unlock()

loadtest/scripts/populate_ct_accounts.py

@@ -0,0 +1,51 @@
+import os
+import json
+import subprocess
+import logging
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+def read_files_and_run_command(directory, max_threads=10):
+    with ThreadPoolExecutor(max_workers=max_threads) as executor:
+        futures = []
+        for filename in os.listdir(directory):
+            file_path = os.path.join(directory, filename)
+            if os.path.isfile(file_path) and filename.endswith('.json'):
+                with open(file_path, 'r') as file:
+                    data = json.load(file)
+                    mnemonic = data.get('mnemonic')
+                    if mnemonic:
+                        file_prefix = filename.split('.')[0]
+                        logging.info(f"Processing file: {filename}")
+                        futures.append(executor.submit(run_commands, mnemonic, file_prefix))
+
+        for future in as_completed(futures):
+            try:
+                future.result()
+            except Exception as e:
+                logging.error(f"Error occurred: {e}")
+
+def run_commands(mnemonic, file_prefix):
+    commands = [
+        f"printf '{mnemonic}\n12345678\n' | ~/go/bin/seid keys add {file_prefix} --recover",
+        f"printf '12345678\n' | ~/go/bin/seid tx ct init-account usei --from {file_prefix} --fees 20000usei -y -b block",
+        f"printf '12345678\n' | ~/go/bin/seid tx ct deposit usei 1000000000 --from {file_prefix} --fees 20000usei -y -b block",
+        f"printf '12345678\n' | ~/go/bin/seid tx ct apply-pending-balance usei --from {file_prefix} --fees 20000usei -y",
+        f"printf '12345678\n' | ~/go/bin/seid keys delete {file_prefix} -y"
+    ]
+    for cmd in commands:
+        logging.info(f"Running command: {cmd}")
+        process = subprocess.Popen(cmd, stdin=subprocess.PIPE, text=True, shell=True, start_new_session=True)
+        stdout, stderr = process.communicate()
+        if process.returncode == 0:
+            logging.info(f"Command succeeded: {cmd}")
+        else:
+            logging.error(f"Command failed: {cmd}\nError: {stderr}")
+
+if __name__ == "__main__":
+    directory = '/root/test_accounts'
+    logging.info(f"Starting to process directory: {directory}")
+    read_files_and_run_command(directory, max_threads=50)
+    logging.info("Finished processing directory")

loadtest/types.go

@@ -16,28 +16,33 @@ import (
 )
 
 const (
-	Bank                 string = "bank"
-	EVM                  string = "evm"
-	ERC20                string = "erc20"
-	ERC721               string = "erc721"
-	CollectRewards       string = "collect_rewards"
-	DistributeRewards    string = "distribute_rewards"
-	FailureBankMalformed string = "failure_bank_malformed"
-	FailureBankInvalid   string = "failure_bank_invalid"
-	FailureDexMalformed  string = "failure_dex_malformed"
-	FailureDexInvalid    string = "failure_dex_invalid"
-	Dex                  string = "dex"
-	Staking              string = "staking"
-	Tokenfactory         string = "tokenfactory"
-	Limit                string = "limit"
-	Market               string = "market"
-	WasmMintNft          string = "wasm_mint_nft"
-	UNIV2                string = "univ2"
-	Vortex               string = "vortex"
-	WasmInstantiate      string = "wasm_instantiate"
-	WasmOccIteratorWrite string = "wasm_occ_iterator_write"
-	WasmOccIteratorRange string = "wasm_occ_iterator_range"
-	WasmOccParallelWrite string = "wasm_occ_parallel_write"
+	Bank                                     string = "bank"
+	EVM                                      string = "evm"
+	ERC20                                    string = "erc20"
+	ERC721                                   string = "erc721"
+	CollectRewards                           string = "collect_rewards"
+	DistributeRewards                        string = "distribute_rewards"
+	FailureBankMalformed                     string = "failure_bank_malformed"
+	FailureBankInvalid                       string = "failure_bank_invalid"
+	FailureDexMalformed                      string = "failure_dex_malformed"
+	FailureDexInvalid                        string = "failure_dex_invalid"
+	Dex                                      string = "dex"
+	Staking                                  string = "staking"
+	Tokenfactory                             string = "tokenfactory"
+	Limit                                    string = "limit"
+	Market                                   string = "market"
+	WasmMintNft                              string = "wasm_mint_nft"
+	UNIV2                                    string = "univ2"
+	Vortex                                   string = "vortex"
+	WasmInstantiate                          string = "wasm_instantiate"
+	WasmOccIteratorWrite                     string = "wasm_occ_iterator_write"
+	WasmOccIteratorRange                     string = "wasm_occ_iterator_range"
+	WasmOccParallelWrite                     string = "wasm_occ_parallel_write"
+	ConfidentialTransfersDeposit             string = "confidential_transfers_deposit"
+	ConfidentialTransfersTransfer            string = "confidential_transfers_transfer"
+	ConfidentialTransfersApplyPendingBalance string = "confidential_transfers_apply_pending_balance"
+	ConfidentialTransfersWithdraw            string = "confidential_transfers_withdraw"
+	CtDefaultDenom                           string = "usei"
 )
 
 type WasmIteratorWriteMsg struct {

x/confidentialtransfers/client/cli/tx.go

@@ -467,10 +467,6 @@ func makeApplyPendingBalanceCmd(cmd *cobra.Command, args []string) error {
 
 	msg := types.NewMsgApplyPendingBalanceProto(applyPendingBalance)
 
-	if err != nil {
-		return err
-	}
-
 	if err = msg.ValidateBasic(); err != nil {
 		return err
 	}

x/confidentialtransfers/keeper/keeper_test.go

@@ -87,11 +87,7 @@ func (suite *KeeperTestSuite) SetupAccountState(privateKey *ecdsa.PrivateKey, de
 	// Extract the next 32 bits (from bit 16 to bit 47)
 	pendingBalanceHi := uint32((pendingBalance >> 16) & 0xFFFFFFFF)
 
-	if err != nil {
-		return types.Account{}, err
-	}
-
-	availableBalanceCipherText, _, err := teg.Encrypt(keypair.PublicKey, uint64(availableBalance))
+	availableBalanceCipherText, _, err := teg.Encrypt(keypair.PublicKey, availableBalance)
 	if err != nil {
 		return types.Account{}, err
 	}
@@ -106,7 +102,7 @@ func (suite *KeeperTestSuite) SetupAccountState(privateKey *ecdsa.PrivateKey, de
 		return types.Account{}, err
 	}
 
-	decryptableAvailableBalance, err := encryption.EncryptAESGCM(uint64(availableBalance), aesKey)
+	decryptableAvailableBalance, err := encryption.EncryptAESGCM(availableBalance, aesKey)
 	if err != nil {
 		return types.Account{}, err
 	}
@@ -121,7 +117,10 @@ func (suite *KeeperTestSuite) SetupAccountState(privateKey *ecdsa.PrivateKey, de
 	}
 
 	addr := privkeyToAddress(privateKey)
-	suite.App.ConfidentialTransfersKeeper.SetAccount(suite.Ctx, addr.String(), denom, initialAccountState)
+	err = suite.App.ConfidentialTransfersKeeper.SetAccount(suite.Ctx, addr.String(), denom, initialAccountState)
+	if err != nil {
+		return types.Account{}, err
+	}
 
 	bankModuleTokens := sdk.NewCoins(sdk.Coin{Amount: sdk.NewInt(int64(bankAmount)), Denom: denom})