Ct queries (#1927)

* add genesis init/export tests

* refactor store

* confidential transfers queries

* working draft test

* more tests

* refactor tests

* all accounts query

* formatting

* add pagination to request/response

* update implementation to use paginated response instead

* formatting

* remove redundant param

* all accounts with denoms

* clean up commented code

* return GetAccount back to keeper

* formatting

app/app.go

@@ -116,6 +116,7 @@ import (
 	"github.com/sei-protocol/sei-chain/utils"
 	"github.com/sei-protocol/sei-chain/utils/metrics"
 	"github.com/sei-protocol/sei-chain/wasmbinding"
+	ctmodule "github.com/sei-protocol/sei-chain/x/confidentialtransfers"
 	ctkeeper "github.com/sei-protocol/sei-chain/x/confidentialtransfers/keeper"
 	cttypes "github.com/sei-protocol/sei-chain/x/confidentialtransfers/types"
 	epochmodule "github.com/sei-protocol/sei-chain/x/epoch"
@@ -209,6 +210,7 @@ var (
 		wasm.AppModuleBasic{},
 		epochmodule.AppModuleBasic{},
 		tokenfactorymodule.AppModuleBasic{},
+		ctmodule.AppModuleBasic{},
 		// this line is used by starport scaffolding # stargate/app/moduleBasic
 	)
 
@@ -226,6 +228,7 @@ var (
 		wasm.ModuleName:                {authtypes.Burner},
 		evmtypes.ModuleName:            {authtypes.Minter, authtypes.Burner},
 		tokenfactorytypes.ModuleName:   {authtypes.Minter, authtypes.Burner},
+		cttypes.ModuleName:             {},
 		// this line is used by starport scaffolding # stargate/app/maccPerms
 	}
 
@@ -549,6 +552,12 @@ func New(
 		app.DistrKeeper,
 	)
 
+	app.ConfidentialTransfersKeeper = ctkeeper.NewKeeper(
+		appCodec,
+		app.keys[(cttypes.StoreKey)],
+		app.GetSubspace(cttypes.ModuleName),
+		app.AccountKeeper)
+
 	// The last arguments can contain custom message handlers, and custom query handlers,
 	// if we want to allow any custom callbacks
 	supportedFeatures := "iterator,staking,stargate,sei"
@@ -752,6 +761,7 @@ func New(
 		epochModule,
 		tokenfactorymodule.NewAppModule(app.TokenFactoryKeeper, app.AccountKeeper, app.BankKeeper),
 		authzmodule.NewAppModule(appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
+		ctmodule.NewAppModule(app.ConfidentialTransfersKeeper),
 		// this line is used by starport scaffolding # stargate/app/appModule
 	)
 
@@ -783,6 +793,7 @@ func New(
 		evmtypes.ModuleName,
 		wasm.ModuleName,
 		tokenfactorytypes.ModuleName,
+		cttypes.ModuleName,
 		acltypes.ModuleName,
 	)
 
@@ -814,6 +825,7 @@ func New(
 		evmtypes.ModuleName,
 		wasm.ModuleName,
 		tokenfactorytypes.ModuleName,
+		cttypes.ModuleName,
 		acltypes.ModuleName,
 	)
 
@@ -843,6 +855,7 @@ func New(
 		feegrant.ModuleName,
 		oracletypes.ModuleName,
 		tokenfactorytypes.ModuleName,
+		cttypes.ModuleName,
 		epochmoduletypes.ModuleName,
 		wasm.ModuleName,
 		evmtypes.ModuleName,
@@ -874,6 +887,7 @@ func New(
 		transferModule,
 		epochModule,
 		tokenfactorymodule.NewAppModule(app.TokenFactoryKeeper, app.AccountKeeper, app.BankKeeper),
+		ctmodule.NewAppModule(app.ConfidentialTransfersKeeper),
 		// this line is used by starport scaffolding # stargate/app/appModule
 	)
 	app.sm.RegisterStoreDecoders()

proto/confidentialtransfers/confidential.proto

@@ -15,4 +15,9 @@ message CtAccount {
   uint32   pending_balance_credit_counter = 4;
   Ciphertext   available_balance = 5; // elgamal encoded balance
   string   decryptable_available_balance = 6; // aes encoded balance
+}
+
+message CtAccountWithDenom {
+  string denom = 1;
+  CtAccount account = 2 [(gogoproto.nullable) = false];
 }

proto/confidentialtransfers/query.proto

@@ -1,28 +1,50 @@
 syntax = "proto3";
 package seiprotocol.seichain.confidentialtransfers;
 
+import "cosmos/base/query/v1beta1/pagination.proto";
 import "gogoproto/gogo.proto";
 import "google/api/annotations.proto";
+import "confidentialtransfers/confidential.proto";
 
 option go_package = "github.com/sei-protocol/sei-chain/x/confidentialtransfers/types";
 
-// TODO: Define any query messages here
 // Query defines the gRPC querier service.
 service Query {
-  // TODO: This is a mock method. Remove when we add real queries
-  rpc TestQuery (TestQueryRequest) returns (TestQueryResponse) {
+  rpc GetCtAccount (GetCtAccountRequest) returns (GetCtAccountResponse) {
     option (google.api.http) = {
-      get: "/seichain/confidentialtransfers/test_query"
+      get: "/seichain/confidentialtransfers/account/{address}/{denom}"
     };
   }
+
+  rpc GetAllCtAccounts (GetAllCtAccountsRequest) returns (GetAllCtAccountsResponse) {
+    option (google.api.http) = {
+      get: "/seichain/confidentialtransfers/account/{address}"
+    };
+  }
+}
+
+message GetCtAccountRequest {
+  string address = 1;
+  string denom = 2;
 }
 
-// TODO: This is a mock method. Remove when we add real queries
-message TestQueryRequest {
-  string test_field = 1;
+message GetCtAccountResponse {
+  CtAccount account = 1;
 }
 
-// TODO: This is a mock method. Remove when we add real queries
-message TestQueryResponse {
-  string test_field = 1;
+message GetAllCtAccountsRequest {
+  option (gogoproto.equal)           = false;
+  option (gogoproto.goproto_getters) = false;
+
+  string address = 1;
+
+  // pagination defines an optional pagination for the request.
+  cosmos.base.query.v1beta1.PageRequest pagination = 2;
+}
+
+message GetAllCtAccountsResponse {
+  repeated CtAccountWithDenom accounts = 1 [(gogoproto.nullable) = false];
+
+  // pagination defines the pagination in the response.
+  cosmos.base.query.v1beta1.PageResponse pagination = 2;
 }

x/confidentialtransfers/keeper/genesis.go

@@ -3,17 +3,17 @@ package keeper
 import (
 	"fmt"
 
-	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/query"
 	"github.com/sei-protocol/sei-chain/x/confidentialtransfers/types"
 )
 
 func (k BaseKeeper) InitGenesis(ctx sdk.Context, gs *types.GenesisState) {
 	k.SetParams(ctx, gs.Params)
+	store := k.getAccountStore(ctx)
 	for i := range gs.Accounts {
 		genesisCtAccount := gs.Accounts[i]
-		store := ctx.KVStore(k.storeKey)
+
 		bz := k.cdc.MustMarshal(&genesisCtAccount.Account) // Marshal the Account object into bytes
 		store.Set(genesisCtAccount.Key, bz)
 	}
@@ -31,11 +31,10 @@ func (k BaseKeeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
 }
 
 func (k BaseKeeper) GetPaginatedAccounts(ctx sdk.Context, pagination *query.PageRequest) ([]types.GenesisCtAccount, *query.PageResponse, error) {
-	store := ctx.KVStore(k.storeKey)
-	supplyStore := prefix.NewStore(store, types.AccountsKey)
+	store := k.getAccountStore(ctx)
 
 	genesisAccounts := make([]types.GenesisCtAccount, 0)
-	pageRes, err := query.Paginate(supplyStore, pagination, func(key, value []byte) error {
+	pageRes, err := query.Paginate(store, pagination, func(key, value []byte) error {
 		var ctAccount types.CtAccount
 		err := ctAccount.Unmarshal(value)
 		if err != nil {

x/confidentialtransfers/keeper/genesis_test.go

@@ -1,7 +1,12 @@
 package keeper_test
 
 import (
+	"crypto/ecdsa"
+	"fmt"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"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"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 )
 
@@ -16,3 +21,63 @@ func (suite *KeeperTestSuite) TestDefaultGenesisState() {
 	suite.Require().NotNil(exportedGenesis)
 	suite.Require().Equal(genesisState, exportedGenesis)
 }
+
+func (suite *KeeperTestSuite) TestGenesisExportImportState() {
+	pk1, _ := encryption.GenerateKey()
+	pk2, _ := encryption.GenerateKey()
+	addr1 := sdk.AccAddress("addr1")
+	addr2 := sdk.AccAddress("addr2")
+	testDenom1 := fmt.Sprintf("factory/%s/TEST1", addr1.String())
+	testDenom2 := fmt.Sprintf("factory/%s/TEST2", addr2.String())
+
+	ctAcc1 := generateCtAccount(pk1, testDenom1, 1000)
+	ctAcc2 := generateCtAccount(pk2, testDenom2, 2000)
+
+	accounts := []types.GenesisCtAccount{
+		{
+			Key:     []byte(addr1.String() + testDenom1),
+			Account: ctAcc1,
+		},
+		{
+			Key:     []byte(addr2.String() + testDenom2),
+			Account: ctAcc2,
+		},
+	}
+	genesisState := types.NewGenesisState(types.DefaultParams(), accounts)
+	app := suite.App
+	suite.Ctx = app.BaseApp.NewContext(false, tmproto.Header{})
+
+	suite.App.ConfidentialTransfersKeeper.InitGenesis(suite.Ctx, genesisState)
+
+	exportedGenesis := suite.App.ConfidentialTransfersKeeper.ExportGenesis(suite.Ctx)
+	suite.Require().NotNil(exportedGenesis)
+	suite.Require().Equal(genesisState, exportedGenesis)
+}
+
+func generateCtAccount(pk *ecdsa.PrivateKey, testDenom string, balance uint64) types.CtAccount {
+	eg := elgamal.NewTwistedElgamal()
+	keyPair, _ := eg.KeyGen(*pk, testDenom)
+
+	aesPK1, _ := encryption.GetAESKey(*pk, testDenom)
+
+	amountLo := uint64(100)
+	amountHi := uint64(0)
+
+	decryptableBalance, _ := encryption.EncryptAESGCM(balance, aesPK1)
+
+	ciphertextLo, _, _ := eg.Encrypt(keyPair.PublicKey, amountLo)
+	ciphertextHi, _, _ := eg.Encrypt(keyPair.PublicKey, amountHi)
+
+	zeroCiphertextAvailable, _, _ := eg.Encrypt(keyPair.PublicKey, amountLo)
+
+	account := &types.Account{
+		PublicKey:                   keyPair.PublicKey,
+		PendingBalanceLo:            ciphertextLo,
+		PendingBalanceHi:            ciphertextHi,
+		PendingBalanceCreditCounter: 1,
+		AvailableBalance:            zeroCiphertextAvailable,
+		DecryptableAvailableBalance: decryptableBalance,
+	}
+
+	return *types.NewCtAccount(account)
+}

x/confidentialtransfers/keeper/grpc_query.go

@@ -0,0 +1,80 @@
+package keeper
+
+import (
+	"context"
+
+	"github.com/cosmos/cosmos-sdk/types/query"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/sei-protocol/sei-chain/x/confidentialtransfers/types"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+var _ types.QueryServer = BaseKeeper{}
+
+func (k BaseKeeper) GetCtAccount(ctx context.Context, req *types.GetCtAccountRequest) (*types.GetCtAccountResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "empty request")
+	}
+
+	if req.Address == "" {
+		return nil, status.Error(codes.InvalidArgument, "address cannot be empty")
+	}
+
+	address, err := sdk.AccAddressFromBech32(req.Address)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid address: %s", err.Error())
+	}
+
+	if req.Denom == "" {
+		return nil, status.Error(codes.InvalidArgument, "invalid denom")
+	}
+
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+
+	ctAccount, found := k.getCtAccount(sdkCtx, address, req.Denom)
+	if !found {
+		return nil, status.Errorf(codes.NotFound, "account not found for account %s and denom %s",
+			req.Address, req.Denom)
+	}
+
+	return &types.GetCtAccountResponse{Account: &ctAccount}, nil
+}
+
+func (k BaseKeeper) GetAllCtAccounts(ctx context.Context, req *types.GetAllCtAccountsRequest) (*types.GetAllCtAccountsResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "empty request")
+	}
+
+	if req.Address == "" {
+		return nil, status.Error(codes.InvalidArgument, "address cannot be empty")
+	}
+
+	address, err := sdk.AccAddressFromBech32(req.Address)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid address: %s", err.Error())
+	}
+
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+
+	store := k.getAccountStoreForAddress(sdkCtx, address)
+	accounts := make([]types.CtAccountWithDenom, 0)
+	pageRes, err := query.Paginate(store, req.Pagination, func(denom, value []byte) error {
+
+		var ctAccount types.CtAccount
+		err = k.cdc.Unmarshal(value, &ctAccount)
+		if err != nil {
+			return err
+		}
+		accounts = append(accounts, types.CtAccountWithDenom{Denom: string(denom), Account: ctAccount})
+		return nil
+	})
+
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "paginate: %v", err)
+	}
+
+	return &types.GetAllCtAccountsResponse{Accounts: accounts, Pagination: pageRes}, nil
+
+}