From 8cf412a68bf0bf437820346bcad84875e5a5156d Mon Sep 17 00:00:00 2001 From: mj Date: Thu, 19 Dec 2024 17:32:03 +0800 Subject: [PATCH] refactor withdraw --- .../confidentialtransfers/mappings_test.go | 2 +- loadtest/main.go | 2 +- x/confidentialtransfers/client/cli/tx.go | 2 +- .../keeper/msg_server_test.go | 26 +++++++++---------- x/confidentialtransfers/types/withdraw.go | 14 +++++++++- 5 files changed, 29 insertions(+), 17 deletions(-) diff --git a/aclmapping/confidentialtransfers/mappings_test.go b/aclmapping/confidentialtransfers/mappings_test.go index a1266e06f..40559a82d 100644 --- a/aclmapping/confidentialtransfers/mappings_test.go +++ b/aclmapping/confidentialtransfers/mappings_test.go @@ -430,7 +430,7 @@ func (suite *KeeperTestSuite) TestMsgWithdrawDependencies() { initialState, _ := suite.SetupAccountState(senderPk, DefaultTestDenom, 10, 2000, 3000, 1000) withdrawAmount := new(big.Int).SetUint64(500) - withdraw, _ := types.NewWithdraw(*senderPk, + withdraw, _ := types.NewWithdrawFromPrivateKey(*senderPk, initialState.AvailableBalance, DefaultTestDenom, senderAddr.String(), diff --git a/loadtest/main.go b/loadtest/main.go index c06293dd8..a3ba99742 100644 --- a/loadtest/main.go +++ b/loadtest/main.go @@ -529,7 +529,7 @@ func (c *LoadTestClient) generateMessage(key cryptotypes.PrivKey, msgType string senderEcdsaKey, _ := crypto.HexToECDSA(senderPrivHex) address := sdk.AccAddress(key.PubKey().Address()).String() account := c.getCtAccount(address, CtDefaultDenom) - withdraw, err := cttypes.NewWithdraw( + withdraw, err := cttypes.NewWithdrawFromPrivateKey( *senderEcdsaKey, account.AvailableBalance, CtDefaultDenom, diff --git a/x/confidentialtransfers/client/cli/tx.go b/x/confidentialtransfers/client/cli/tx.go index c37a42994..ad1aff41c 100644 --- a/x/confidentialtransfers/client/cli/tx.go +++ b/x/confidentialtransfers/client/cli/tx.go @@ -321,7 +321,7 @@ func makeWithdrawCmd(cmd *cobra.Command, args []string) error { return err } - withdraw, err := types.NewWithdraw( + withdraw, err := types.NewWithdrawFromPrivateKey( *privKey, account.AvailableBalance, coin.Denom, diff --git a/x/confidentialtransfers/keeper/msg_server_test.go b/x/confidentialtransfers/keeper/msg_server_test.go index a251218f6..246e062b1 100644 --- a/x/confidentialtransfers/keeper/msg_server_test.go +++ b/x/confidentialtransfers/keeper/msg_server_test.go @@ -540,7 +540,7 @@ func (suite *KeeperTestSuite) TestMsgServer_WithdrawHappyPath() { // Create a withdraw request withdrawAmount := new(big.Int).Div(initialAvailableBalance, big.NewInt(2)) - withdrawStruct, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) + withdrawStruct, _ := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) // Execute the withdraw req := types.NewMsgWithdrawProto(withdrawStruct) @@ -598,13 +598,13 @@ func (suite *KeeperTestSuite) TestMsgServer_WithdrawSuccessive() { // Create a withdraw request with an invalid amount withdrawAmount := new(big.Int).Add(initialAvailableBalance, big.NewInt(1)) - _, err := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) + _, err := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) suite.Require().Error(err, "Cannot use client to create withdraw for more than the account balance") suite.Require().ErrorContains(err, "insufficient balance") // Create two withdraw requests for the entire balance - withdrawStruct1, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, initialAvailableBalance) - withdrawStruct2, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, initialAvailableBalance) + withdrawStruct1, _ := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, initialAvailableBalance) + withdrawStruct2, _ := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, initialAvailableBalance) // Execute the first withdraw req1 := types.NewMsgWithdrawProto(withdrawStruct1) @@ -636,7 +636,7 @@ func (suite *KeeperTestSuite) TestMsgServer_WithdrawInvalidAmount() { initialState, _ := suite.SetupAccountState(testPk, DefaultTestDenom, 50, initialAvailableBalance, big.NewInt(8000), big.NewInt(1000000000000)) // Create a withdraw request - withdrawStruct, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, initialAvailableBalance) + withdrawStruct, _ := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, initialAvailableBalance) // Manually modify the withdraw struct to have an invalid amount (since we can't do that via the client) withdrawStruct.Amount = new(big.Int).Add(initialAvailableBalance, big.NewInt(1)) @@ -681,7 +681,7 @@ func (suite *KeeperTestSuite) TestMsgServer_RepeatWithdraw() { // Create a withdraw request withdrawAmount := new(big.Int).Div(initialAvailableBalance, big.NewInt(5)) - withdrawStruct, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) + withdrawStruct, _ := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) // Execute the first withdraw req := types.NewMsgWithdrawProto(withdrawStruct) @@ -713,7 +713,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ModifiedDecryptableBalance() { // Create a withdraw request withdrawAmount := new(big.Int).Div(initialAvailableBalance, big.NewInt(5)) - withdrawStruct, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) + withdrawStruct, _ := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) // Modify the decryptable balance aesKey, _ := utils.GetAESKey(*testPk, DefaultTestDenom) @@ -728,7 +728,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ModifiedDecryptableBalance() { // At this point, the decryptable available balance is corrupted. // Any withdraw struct we create based on the decryptable balance in the account will be invalid. accountState, _ := suite.App.ConfidentialTransfersKeeper.GetAccount(suite.Ctx, testAddr.String(), DefaultTestDenom) - nextWithdrawStruct, err := types.NewWithdraw(*testPk, accountState.AvailableBalance, DefaultTestDenom, testAddr.String(), accountState.DecryptableAvailableBalance, withdrawAmount) + nextWithdrawStruct, err := types.NewWithdrawFromPrivateKey(*testPk, accountState.AvailableBalance, DefaultTestDenom, testAddr.String(), accountState.DecryptableAvailableBalance, withdrawAmount) req = types.NewMsgWithdrawProto(nextWithdrawStruct) _, err = suite.msgServer.Withdraw(sdk.WrapSDKContext(suite.Ctx), req) suite.Require().Error(err, "Should have error withdrawing since withdraw struct is invalid") @@ -751,7 +751,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ModifiedDecryptableBalance() { aesEncryptedActualBalance, _ := encryption.EncryptAESGCM(actualBalance, aesKey) // Then create the correct struct for the withdraw - correctedWithdrawStruct, err := types.NewWithdraw(*testPk, accountState.AvailableBalance, DefaultTestDenom, testAddr.String(), aesEncryptedActualBalance, withdrawAmount) + correctedWithdrawStruct, err := types.NewWithdrawFromPrivateKey(*testPk, accountState.AvailableBalance, DefaultTestDenom, testAddr.String(), aesEncryptedActualBalance, withdrawAmount) // Execute the withdraw req = types.NewMsgWithdrawProto(correctedWithdrawStruct) @@ -766,7 +766,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ModifiedDecryptableBalance() { suite.Require().Equal(decryptedAvailableBalance, newBalance, "New account value should have been updated") // Validate that I can create regular transactions with the account again - nextWithdrawStruct, err = types.NewWithdraw(*testPk, accountState.AvailableBalance, DefaultTestDenom, testAddr.String(), accountState.DecryptableAvailableBalance, big.NewInt(1)) + nextWithdrawStruct, err = types.NewWithdrawFromPrivateKey(*testPk, accountState.AvailableBalance, DefaultTestDenom, testAddr.String(), accountState.DecryptableAvailableBalance, big.NewInt(1)) req = types.NewMsgWithdrawProto(nextWithdrawStruct) _, err = suite.msgServer.Withdraw(sdk.WrapSDKContext(suite.Ctx), req) suite.Require().NoError(err, "Should not have error withdrawing since decryptable balance is no longer corrupted") @@ -791,7 +791,7 @@ func (suite *KeeperTestSuite) TestMsgServer_WithdrawFeatureDisabled() { // Create a withdraw request withdrawAmount := new(big.Int).Div(initialAvailableBalance, big.NewInt(5)) - withdrawStruct, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) + withdrawStruct, _ := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) // Disable the confidential tokens module via params suite.Ctx = suite.App.BaseApp.NewContext(false, tmproto.Header{}) @@ -893,7 +893,7 @@ func (suite *KeeperTestSuite) TestMsgServer_CloseAccountHasPendingBalance() { availableBalanceAmount, err := encryption.DecryptAESGCM(account.DecryptableAvailableBalance, aeskey) suite.Require().NoError(err, "Should be able to decrypt available balance") - withdrawStruct, _ := types.NewWithdraw(*testPk, account.AvailableBalance, DefaultTestDenom, testAddr.String(), account.DecryptableAvailableBalance, availableBalanceAmount) + withdrawStruct, _ := types.NewWithdrawFromPrivateKey(*testPk, account.AvailableBalance, DefaultTestDenom, testAddr.String(), account.DecryptableAvailableBalance, availableBalanceAmount) withdrawReq := types.NewMsgWithdrawProto(withdrawStruct) suite.msgServer.Withdraw(sdk.WrapSDKContext(suite.Ctx), withdrawReq) @@ -1067,7 +1067,7 @@ func (suite *KeeperTestSuite) TestMsgServer_ApplyPendingBalanceAfterWithdraw() { // Before the pending balance is applied, a withdrawal is made, changing the available balance in the account. withdrawAmount := new(big.Int).Div(initialAvailableBalance, big.NewInt(2)) - withdrawReq, _ := types.NewWithdraw(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) + withdrawReq, _ := types.NewWithdrawFromPrivateKey(*testPk, initialState.AvailableBalance, DefaultTestDenom, testAddr.String(), initialState.DecryptableAvailableBalance, withdrawAmount) withdrawMsg := types.NewMsgWithdrawProto(withdrawReq) _, err := suite.msgServer.Withdraw(sdk.WrapSDKContext(suite.Ctx), withdrawMsg) diff --git a/x/confidentialtransfers/types/withdraw.go b/x/confidentialtransfers/types/withdraw.go index c4ba1de5f..a450b9f71 100644 --- a/x/confidentialtransfers/types/withdraw.go +++ b/x/confidentialtransfers/types/withdraw.go @@ -33,7 +33,7 @@ type WithdrawProofs struct { RemainingBalanceEqualityProof *zkproofs.CiphertextCommitmentEqualityProof `json:"remaining_balance_equality_proof"` } -func NewWithdraw( +func NewWithdrawFromPrivateKey( privateKey ecdsa.PrivateKey, currentAvailableBalance *elgamal.Ciphertext, denom, @@ -51,6 +51,18 @@ func NewWithdraw( return &Withdraw{}, err } + return NewWithdraw(teg, keyPair, aesKey, amount, address, denom, currentAvailableBalance, currentDecryptableBalance) +} + +func NewWithdraw( + teg *elgamal.TwistedElGamal, + keyPair *elgamal.KeyPair, + aesKey []byte, + amount *big.Int, + address, + denom string, + currentAvailableBalance *elgamal.Ciphertext, + currentDecryptableBalance string) (*Withdraw, error) { currentBalance, err := encryption.DecryptAESGCM(currentDecryptableBalance, aesKey) if err != nil { return &Withdraw{}, err