Merge branch 'dev' into jl777

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));
+    }
+}

src/wallet/asyncrpcoperation_mergetoaddress.h

@@ -121,6 +121,10 @@ private:
 
     void unlock_utxos();
 
+    void lock_notes();
+
+    void unlock_notes();
+
     // payment disclosure!
     std::vector<PaymentDisclosureKeyInfo> paymentDisclosureData_;
 };

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));
+}

src/wallet/rpcdump.cpp

@@ -653,7 +653,7 @@ UniValue z_importviewingkey(const UniValue& params, bool fHelp)
     if (!EnsureWalletIsAvailable(fHelp))
         return NullUniValue;
 
-    if (fHelp || params.size() < 1 || params.size() > 2)
+    if (fHelp || params.size() < 1 || params.size() > 3)
         throw runtime_error(
             "z_importviewingkey \"vkey\" ( rescan startHeight )\n"
             "\nAdds a viewing key (as returned by z_exportviewingkey) to your wallet.\n"

src/wallet/wallet.cpp

@@ -3535,6 +3535,42 @@ void CWallet::ListLockedCoins(std::vector<COutPoint>& 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<JSOutPoint> CWallet::ListLockedNotes()
+{
+    AssertLockHeld(cs_wallet); // setLockedNotes
+    std::vector<JSOutPoint> vOutpts(setLockedNotes.begin(), setLockedNotes.end());
+    return vOutpts;
+}
+
 /** @} */ // end of Actions
 
 class CAffectedKeysVisitor : public boost::static_visitor<void> {
@@ -3825,6 +3861,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

src/wallet/wallet.h

@@ -882,6 +882,7 @@ public:
     CPubKey vchDefaultKey;
 
     std::set<COutPoint> setLockedCoins;
+    std::set<JSOutPoint> setLockedNotes;
 
     int64_t nTimeFirstKey;
 
@@ -902,6 +903,14 @@ public:
     void UnlockAllCoins();
     void ListLockedCoins(std::vector<COutPoint>& vOutpts);
 
+
+    bool IsLockedNote(uint256 hash, size_t js, uint8_t n) const;
+    void LockNote(JSOutPoint& output);
+    void UnlockNote(JSOutPoint& output);
+    void UnlockAllNotes();
+    std::vector<JSOutPoint> ListLockedNotes();
+
+
     /**
      * keystore implementation
      * Generate a new key