Add msg_server methods (#1923)

* new types

* add deposit to msgs

* add import issues

* revert evm/params to old state

* withdraw with tests

* small changes

* refactor and add tests

* comment

* refactor initialize type

* types

* non module account types

* msg server methods with module accounts

* tests scaffold

* add msg server tests

* codeql

* test issues

* test improvements

* comments

* todo

* app

app/app.go

@@ -228,7 +228,9 @@ var (
 		wasm.ModuleName:                {authtypes.Burner},
 		evmtypes.ModuleName:            {authtypes.Minter, authtypes.Burner},
 		tokenfactorytypes.ModuleName:   {authtypes.Minter, authtypes.Burner},
-		cttypes.ModuleName:             {},
+		// Confidential Transfers module is not live yet, but we add the permissions for testing
+		cttypes.ModuleName: nil,
+
 		// this line is used by starport scaffolding # stargate/app/maccPerms
 	}
 
@@ -562,7 +564,8 @@ func New(
 		appCodec,
 		app.keys[(cttypes.StoreKey)],
 		app.GetSubspace(cttypes.ModuleName),
-		app.AccountKeeper)
+		app.AccountKeeper,
+		app.BankKeeper)
 
 	// The last arguments can contain custom message handlers, and custom query handlers,
 	// if we want to allow any custom callbacks

app/apptesting/test_suite.go

@@ -2,6 +2,8 @@ package apptesting
 
 import (
 	"context"
+	"crypto/ecdsa"
+	"github.com/ethereum/go-ethereum/crypto"
 	"time"
 
 	"github.com/cosmos/cosmos-sdk/baseapp"
@@ -176,6 +178,20 @@ func CreateRandomAccounts(numAccts int) []sdk.AccAddress {
 	return testAddrs
 }
 
+// CreateRandomAccountKeys is a function return a list of randomly generated private keys
+func CreateRandomAccountKeys(numAccts int) []*ecdsa.PrivateKey {
+	testAccts := make([]*ecdsa.PrivateKey, numAccts)
+	for i := 0; i < numAccts; i++ {
+		pk, err := crypto.GenerateKey()
+		if err != nil {
+			panic(err)
+		}
+		testAccts[i] = pk
+	}
+
+	return testAccts
+}
+
 func GenerateTestAddrs() (string, string) {
 	pk1 := ed25519.GenPrivKey().PubKey()
 	validAddr := sdk.AccAddress(pk1.Address()).String()

proto/confidentialtransfers/tx.proto

@@ -27,7 +27,6 @@ service Msg {
 
   // CloseAccount defines a method for closing an account.
   rpc CloseAccount(MsgCloseAccount) returns (MsgCloseAccountResponse);
-
 }
 
 // MsgTransfer represents a message to send coins from one account to another.
@@ -70,7 +69,10 @@ message MsgInitializeAccount {
   string   denom = 2 [(gogoproto.moretags) = "yaml:\"denom\""];
   bytes public_key = 3 [(gogoproto.moretags) = "yaml:\"public_key\""];
   string decryptable_balance = 4 [(gogoproto.moretags) = "yaml:\"decryptable_balance\""];
-  InitializeAccountMsgProofs proofs = 5 [(gogoproto.moretags) = "yaml:\"proofs\""];
+  Ciphertext pending_balance_lo = 5 [(gogoproto.moretags) = "yaml:\"pending_balance_lo\""];
+  Ciphertext pending_balance_hi = 6 [(gogoproto.moretags) = "yaml:\"pending_balance_hi\""];
+  Ciphertext available_balance = 7 [(gogoproto.moretags) = "yaml:\"available_balance\""];
+  InitializeAccountMsgProofs proofs = 8 [(gogoproto.moretags) = "yaml:\"proofs\""];
 }
 
 // MsgInitializeAccountResponse defines the Msg/Send response type.
@@ -110,9 +112,9 @@ message MsgApplyPendingBalance {
   option (gogoproto.equal) = false;
   option (gogoproto.goproto_getters) = false;
 
-  string address = 1;
-  string denom = 2;
-  string new_decryptable_available_balance = 3;
+  string address = 1 [(gogoproto.moretags) = "yaml:\"address\""];
+  string denom = 2 [(gogoproto.moretags) = "yaml:\"denom\""];
+  string new_decryptable_available_balance = 3 [(gogoproto.moretags) = "yaml:\"new_decryptable_available_balance\""];
 }
 
 message MsgApplyPendingBalanceResponse {}

proto/confidentialtransfers/zk.proto

@@ -18,6 +18,9 @@ message TransferMsgProofs {
 
 message InitializeAccountMsgProofs {
   PubkeyValidityProof pubkey_validity_proof = 1;
+  ZeroBalanceProof zero_pending_balance_lo_proof = 2;
+  ZeroBalanceProof zero_pending_balance_hi_proof = 3;
+  ZeroBalanceProof zero_available_balance_proof = 4;
 }
 
 message WithdrawMsgProofs {

x/confidentialtransfers/keeper/genesis.go

@@ -2,13 +2,15 @@ package keeper
 
 import (
 	"fmt"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/query"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	"github.com/sei-protocol/sei-chain/x/confidentialtransfers/types"
 )
 
 func (k BaseKeeper) InitGenesis(ctx sdk.Context, gs *types.GenesisState) {
+	moduleAcc := authtypes.NewEmptyModuleAccount(types.ModuleName)
+	k.accountKeeper.SetModuleAccount(ctx, moduleAcc)
 	k.SetParams(ctx, gs.Params)
 	store := k.getAccountStore(ctx)
 	for i := range gs.Accounts {

x/confidentialtransfers/keeper/keeper.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"fmt"
+	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
 
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
@@ -11,8 +12,6 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-
-	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
 )
 
 var _ Keeper = (*BaseKeeper)(nil)
@@ -24,9 +23,14 @@ type Keeper interface {
 	GetAccount(ctx sdk.Context, addrString string, denom string) (types.Account, bool)
 	SetAccount(ctx sdk.Context, addrString string, denom string, account types.Account) error
 
+	DeleteAccount(ctx sdk.Context, addrString string, denom string) error
 	GetParams(ctx sdk.Context) types.Params
 	SetParams(ctx sdk.Context, params types.Params)
 
+	// TODO: See if there's a way to put this somewhere else
+	SendTokens(ctx sdk.Context, to sdk.AccAddress, amount sdk.Coins) error
+	ReceiveTokens(ctx sdk.Context, from sdk.AccAddress, amount sdk.Coins) error
+
 	CreateModuleAccount(ctx sdk.Context)
 
 	types.QueryServer
@@ -39,6 +43,7 @@ type BaseKeeper struct {
 
 	paramSpace    paramtypes.Subspace
 	accountKeeper types.AccountKeeper
+	bankKeeper    types.BankKeeper
 }
 
 // NewKeeper returns a new instance of the x/confidentialtransfers keeper
@@ -47,6 +52,7 @@ func NewKeeper(
 	storeKey sdk.StoreKey,
 	paramSpace paramtypes.Subspace,
 	accountKeeper types.AccountKeeper,
+	bankKeeper types.BankKeeper,
 ) Keeper {
 
 	if !paramSpace.HasKeyTable() {
@@ -57,6 +63,7 @@ func NewKeeper(
 		cdc:           codec,
 		storeKey:      storeKey,
 		accountKeeper: accountKeeper,
+		bankKeeper:    bankKeeper,
 		paramSpace:    paramSpace,
 	}
 }
@@ -77,6 +84,17 @@ func (k BaseKeeper) GetAccount(ctx sdk.Context, address string, denom string) (t
 	return *account, true
 }
 
+func (k BaseKeeper) DeleteAccount(ctx sdk.Context, addrString, denom string) error {
+	address, err := sdk.AccAddressFromBech32(addrString)
+	if err != nil {
+		return err
+	}
+
+	store := k.getAccountStoreForAddress(ctx, address)
+	store.Delete([]byte(denom)) // Store the serialized account under the key
+	return nil
+}
+
 func (k BaseKeeper) getCtAccount(ctx sdk.Context, address sdk.AccAddress, denom string) (types.CtAccount, bool) {
 	store := k.getAccountStoreForAddress(ctx, address)
 	key := []byte(denom)
@@ -139,6 +157,15 @@ func (k BaseKeeper) GetAccountsForAddress(ctx sdk.Context, address sdk.AccAddres
 	return accounts, nil
 }
 
+// CreateModuleAccount creates the module account for confidentialtransfers
+func (k BaseKeeper) CreateModuleAccount(ctx sdk.Context) {
+	account := k.accountKeeper.GetModuleAccount(ctx, types.ModuleName)
+	if account == nil {
+		moduleAcc := authtypes.NewEmptyModuleAccount(types.ModuleName)
+		k.accountKeeper.SetModuleAccount(ctx, moduleAcc)
+	}
+}
+
 func (k BaseKeeper) GetParams(ctx sdk.Context) (params types.Params) {
 	k.paramSpace.GetParamSet(ctx, &params)
 	return params
@@ -149,9 +176,12 @@ func (k BaseKeeper) SetParams(ctx sdk.Context, params types.Params) {
 	k.paramSpace.SetParamSet(ctx, &params)
 }
 
-func (k BaseKeeper) CreateModuleAccount(ctx sdk.Context) {
-	moduleAcc := authtypes.NewEmptyModuleAccount(types.ModuleName, authtypes.Minter, authtypes.Burner)
-	k.accountKeeper.SetModuleAccount(ctx, moduleAcc)
+func (k BaseKeeper) SendTokens(ctx sdk.Context, to sdk.AccAddress, amount sdk.Coins) error {
+	return k.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, to, amount)
+}
+
+func (k BaseKeeper) ReceiveTokens(ctx sdk.Context, from sdk.AccAddress, amount sdk.Coins) error {
+	return k.bankKeeper.SendCoinsFromAccountToModule(ctx, from, types.ModuleName, amount)
 }
 
 func (k BaseKeeper) getAccountStore(ctx sdk.Context) prefix.Store {

x/confidentialtransfers/keeper/keeper_test.go

@@ -1,20 +1,29 @@
 package keeper_test
 
 import (
+	"crypto/ecdsa"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/sei-protocol/sei-chain/app/apptesting"
 	"github.com/sei-protocol/sei-chain/x/confidentialtransfers/keeper"
 	"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/stretchr/testify/suite"
 	"testing"
 )
 
+const DefaultTestDenom = "test"
+
 type KeeperTestSuite struct {
 	apptesting.KeeperTestHelper
 
 	queryClient types.QueryClient
 	msgServer   types.MsgServer
 	// defaultDenom is on the suite, as it depends on the creator test address.
 	defaultDenom string
+
+	PrivKeys []*ecdsa.PrivateKey
 }
 
 func TestKeeperTestSuite(t *testing.T) {
@@ -27,4 +36,98 @@ func (suite *KeeperTestSuite) SetupTest() {
 	suite.queryClient = types.NewQueryClient(suite.QueryHelper)
 	suite.msgServer = keeper.NewMsgServerImpl(suite.App.ConfidentialTransfersKeeper)
 
+	// TODO: remove this once the app initializes confidentialtransfers keeper
+	suite.App.ConfidentialTransfersKeeper = keeper.NewKeeper(
+		suite.App.AppCodec(),
+		suite.App.GetKey(types.StoreKey),
+		suite.App.GetSubspace(types.ModuleName),
+		suite.App.AccountKeeper,
+		suite.App.BankKeeper,
+	)
+	suite.msgServer = keeper.NewMsgServerImpl(suite.App.ConfidentialTransfersKeeper)
+	suite.PrivKeys = apptesting.CreateRandomAccountKeys(3)
+}
+
+func (suite *KeeperTestSuite) SetupAccount() {
+	suite.queryClient = types.NewQueryClient(suite.QueryHelper)
+	// TODO: remove this once the app initializes confidentialtransfers keeper
+	suite.App.ConfidentialTransfersKeeper = keeper.NewKeeper(
+		suite.App.AppCodec(),
+		suite.App.GetKey(types.StoreKey),
+		suite.App.GetSubspace(types.ModuleName),
+		suite.App.AccountKeeper,
+		suite.App.BankKeeper,
+	)
+	suite.msgServer = keeper.NewMsgServerImpl(suite.App.ConfidentialTransfersKeeper)
+	suite.PrivKeys = apptesting.CreateRandomAccountKeys(4)
+}
+
+func (suite *KeeperTestSuite) SetupAccountState(privateKey *ecdsa.PrivateKey, denom string, pendingBalanceCreditCounter uint16, initialAvailableBalance, initialPendingBalance, bankAmount uint64) (types.Account, error) {
+	aesKey, err := encryption.GetAESKey(*privateKey, denom)
+	if err != nil {
+		return types.Account{}, err
+	}
+
+	teg := elgamal.NewTwistedElgamal()
+	keypair, err := teg.KeyGen(*privateKey, denom)
+	if err != nil {
+		return types.Account{}, err
+	}
+
+	availableBalance := initialAvailableBalance
+	pendingBalance := initialPendingBalance
+
+	// Extract the bottom 16 bits (rightmost 16 bits)
+	pendingBalanceLo := uint16(pendingBalance & 0xFFFF)
+
+	// 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))
+	if err != nil {
+		return types.Account{}, err
+	}
+
+	pendingBalanceLoCipherText, _, err := teg.Encrypt(keypair.PublicKey, uint64(pendingBalanceLo))
+	if err != nil {
+		return types.Account{}, err
+	}
+
+	pendingBalanceHiCipherText, _, err := teg.Encrypt(keypair.PublicKey, uint64(pendingBalanceHi))
+	if err != nil {
+		return types.Account{}, err
+	}
+
+	decryptableAvailableBalance, err := encryption.EncryptAESGCM(uint64(availableBalance), aesKey)
+	if err != nil {
+		return types.Account{}, err
+	}
+
+	initialAccountState := types.Account{
+		PublicKey:                   keypair.PublicKey,
+		PendingBalanceLo:            pendingBalanceLoCipherText,
+		PendingBalanceHi:            pendingBalanceHiCipherText,
+		PendingBalanceCreditCounter: pendingBalanceCreditCounter,
+		AvailableBalance:            availableBalanceCipherText,
+		DecryptableAvailableBalance: decryptableAvailableBalance,
+	}
+
+	addr := privkeyToAddress(privateKey)
+	suite.App.ConfidentialTransfersKeeper.SetAccount(suite.Ctx, addr.String(), denom, initialAccountState)
+
+	bankModuleTokens := sdk.NewCoins(sdk.Coin{Amount: sdk.NewInt(int64(bankAmount)), Denom: denom})
+
+	suite.FundAcc(addr, bankModuleTokens)
+
+	return initialAccountState, nil
+}
+
+func privkeyToAddress(privateKey *ecdsa.PrivateKey) sdk.AccAddress {
+	publicKeyBytes := crypto.FromECDSAPub(&privateKey.PublicKey)
+	testAddr := sdk.AccAddress(crypto.Keccak256(publicKeyBytes[1:])[12:])
+	return testAddr
 }