diff --git a/qa/rpc-tests/wallet_mergetoaddress.py b/qa/rpc-tests/wallet_mergetoaddress.py index d8531319c..e5d5089a4 100755 --- a/qa/rpc-tests/wallet_mergetoaddress.py +++ b/qa/rpc-tests/wallet_mergetoaddress.py @@ -325,8 +325,10 @@ class WalletMergeToAddressTest (BitcoinTestFramework): # Verify maximum number of notes which node 0 can shield can be set by the limit parameter # Also check that we can set off a second merge before the first one is complete - result1 = self.nodes[0].z_mergetoaddress([myzaddr], myzaddr, 0, 50, 2) - result2 = self.nodes[0].z_mergetoaddress([myzaddr], myzaddr, 0, 50, 2) + + # myzaddr has 5 notes at this point + result1 = self.nodes[0].z_mergetoaddress([myzaddr], myzaddr, 0.0001, 50, 2) + result2 = self.nodes[0].z_mergetoaddress([myzaddr], myzaddr, 0.0001, 50, 2) # First merge should select from all notes assert_equal(result1["mergingUTXOs"], Decimal('0')) @@ -340,9 +342,9 @@ class WalletMergeToAddressTest (BitcoinTestFramework): assert_equal(result2["remainingUTXOs"], Decimal('0')) assert_equal(result2["mergingNotes"], Decimal('2')) assert_equal(result2["remainingNotes"], Decimal('1')) - wait_and_assert_operationid_status(self.nodes[0], result1['opid']) wait_and_assert_operationid_status(self.nodes[0], result2['opid']) + self.sync_all() self.nodes[1].generate(1) self.sync_all() diff --git a/src/wallet/asyncrpcoperation_mergetoaddress.cpp b/src/wallet/asyncrpcoperation_mergetoaddress.cpp index 30ef560df..fa823f50a 100644 --- a/src/wallet/asyncrpcoperation_mergetoaddress.cpp +++ b/src/wallet/asyncrpcoperation_mergetoaddress.cpp @@ -98,6 +98,7 @@ AsyncRPCOperation_mergetoaddress::AsyncRPCOperation_mergetoaddress( // Lock UTXOs lock_utxos(); + lock_notes(); // Enable payment disclosure if requested paymentDisclosureMode = fExperimentalMode && GetBoolArg("-paymentdisclosure", false); @@ -111,6 +112,7 @@ void AsyncRPCOperation_mergetoaddress::main() { if (isCancelled()) { unlock_utxos(); // clean up + unlock_notes(); return; } @@ -173,6 +175,7 @@ void AsyncRPCOperation_mergetoaddress::main() LogPrintf("%s", s); unlock_utxos(); // clean up + unlock_notes(); // clean up // !!! Payment disclosure START if (success && paymentDisclosureMode && paymentDisclosureData_.size() > 0) { @@ -921,3 +924,24 @@ void AsyncRPCOperation_mergetoaddress::unlock_utxos() { pwalletMain->UnlockCoin(std::get<0>(utxo)); } } + + +/** + * Lock input notes + */ + void AsyncRPCOperation_mergetoaddress::lock_notes() { + LOCK2(cs_main, pwalletMain->cs_wallet); + for (auto note : noteInputs_) { + pwalletMain->LockNote(std::get<0>(note)); + } +} + +/** + * Unlock input notes + */ +void AsyncRPCOperation_mergetoaddress::unlock_notes() { + LOCK2(cs_main, pwalletMain->cs_wallet); + for (auto note : noteInputs_) { + pwalletMain->UnlockNote(std::get<0>(note)); + } +} diff --git a/src/wallet/asyncrpcoperation_mergetoaddress.h b/src/wallet/asyncrpcoperation_mergetoaddress.h index 1619b5c97..34548a5ba 100644 --- a/src/wallet/asyncrpcoperation_mergetoaddress.h +++ b/src/wallet/asyncrpcoperation_mergetoaddress.h @@ -121,6 +121,10 @@ private: void unlock_utxos(); + void lock_notes(); + + void unlock_notes(); + // payment disclosure! std::vector paymentDisclosureData_; }; diff --git a/src/wallet/gtest/test_wallet.cpp b/src/wallet/gtest/test_wallet.cpp index 4aeaf00df..e976e4ae4 100644 --- a/src/wallet/gtest/test_wallet.cpp +++ b/src/wallet/gtest/test_wallet.cpp @@ -1046,3 +1046,36 @@ TEST(wallet_tests, MarkAffectedTransactionsDirty) { wallet.MarkAffectedTransactionsDirty(wtx2); EXPECT_FALSE(wallet.mapWallet[hash].fDebitCached); } + +TEST(wallet_tests, NoteLocking) { + TestWallet wallet; + + auto sk = libzcash::SpendingKey::random(); + wallet.AddSpendingKey(sk); + + auto wtx = GetValidReceive(sk, 10, true); + auto wtx2 = GetValidReceive(sk, 10, true); + + JSOutPoint jsoutpt {wtx.GetHash(), 0, 0}; + JSOutPoint jsoutpt2 {wtx2.GetHash(),0, 0}; + + // Test selective locking + wallet.LockNote(jsoutpt); + EXPECT_TRUE(wallet.IsLockedNote(jsoutpt.hash, jsoutpt.js, jsoutpt.n)); + EXPECT_FALSE(wallet.IsLockedNote(jsoutpt2.hash, jsoutpt2.js, jsoutpt2.n)); + + // Test selective unlocking + wallet.UnlockNote(jsoutpt); + EXPECT_FALSE(wallet.IsLockedNote(jsoutpt.hash, jsoutpt.js, jsoutpt.n)); + + // Test multiple locking + wallet.LockNote(jsoutpt); + wallet.LockNote(jsoutpt2); + EXPECT_TRUE(wallet.IsLockedNote(jsoutpt.hash, jsoutpt.js, jsoutpt.n)); + EXPECT_TRUE(wallet.IsLockedNote(jsoutpt2.hash, jsoutpt2.js, jsoutpt2.n)); + + // Test unlock all + wallet.UnlockAllNotes(); + EXPECT_FALSE(wallet.IsLockedNote(jsoutpt.hash, jsoutpt.js, jsoutpt.n)); + EXPECT_FALSE(wallet.IsLockedNote(jsoutpt2.hash, jsoutpt2.js, jsoutpt2.n)); +} diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index 841232597..60821b2db 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -3437,6 +3437,42 @@ void CWallet::ListLockedCoins(std::vector& vOutpts) } } + +// Note Locking Operations + +void CWallet::LockNote(JSOutPoint& output) +{ + AssertLockHeld(cs_wallet); // setLockedNotes + setLockedNotes.insert(output); +} + +void CWallet::UnlockNote(JSOutPoint& output) +{ + AssertLockHeld(cs_wallet); // setLockedNotes + setLockedNotes.erase(output); +} + +void CWallet::UnlockAllNotes() +{ + AssertLockHeld(cs_wallet); // setLockedNotes + setLockedNotes.clear(); +} + +bool CWallet::IsLockedNote(uint256 hash, size_t js, uint8_t n) const +{ + AssertLockHeld(cs_wallet); // setLockedNotes + JSOutPoint outpt(hash, js, n); + + return (setLockedNotes.count(outpt) > 0); +} + +std::vector CWallet::ListLockedNotes() +{ + AssertLockHeld(cs_wallet); // setLockedNotes + std::vector vOutpts(setLockedNotes.begin(), setLockedNotes.end()); + return vOutpts; +} + /** @} */ // end of Actions class CAffectedKeysVisitor : public boost::static_visitor { @@ -3725,6 +3761,11 @@ void CWallet::GetFilteredNotes( if (ignoreUnspendable && !HaveSpendingKey(pa)) { continue; } + + // skip locked notes + if (IsLockedNote(jsop.hash, jsop.js, jsop.n)) { + continue; + } int i = jsop.js; // Index into CTransaction.vjoinsplit int j = jsop.n; // Index into JSDescription.ciphertexts diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h index b43099b4b..8e3749dfe 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -883,6 +883,7 @@ public: CPubKey vchDefaultKey; std::set setLockedCoins; + std::set setLockedNotes; int64_t nTimeFirstKey; @@ -903,6 +904,14 @@ public: void UnlockAllCoins(); void ListLockedCoins(std::vector& vOutpts); + + bool IsLockedNote(uint256 hash, size_t js, uint8_t n) const; + void LockNote(JSOutPoint& output); + void UnlockNote(JSOutPoint& output); + void UnlockAllNotes(); + std::vector ListLockedNotes(); + + /** * keystore implementation * Generate a new key