wallet: Move definintely unusable TXOs to a separate container

achow101 · achow101 · commit a89fc9a90034 · 2024-09-10T16:01:24.000-04:00
Definitely unusable TXOs are those that are spent by a confirmed
transaction or were produced by a now conflicted transaction. However,
we still need them for GetDebit, so we store them in a separate
m_unusable_txos container. MarkConflicted, AbandonTransaction, and
loading (via PruneSpentTXOs) will ensure that these unusable TXOs are
properly moved.
diff --git a/src/wallet/receive.cpp b/src/wallet/receive.cpp
@@ -280,6 +280,8 @@ Balance GetBalance(const CWallet& wallet, const int min_depth, bool avoid_reuse)
 
             const bool is_trusted{CachedTxIsTrusted(wallet, txo.GetState(), outpoint.hash)};
             const int tx_depth{wallet.GetTxStateDepthInMainChain(txo.GetState())};
+            Assert(tx_depth >= 0);
+            Assert(!wallet.IsSpent(outpoint, /*min_depth=*/1));
 
             if (!wallet.IsSpent(outpoint) && (allow_used_addresses || !wallet.IsSpentKey(txo.GetTxOut().scriptPubKey))) {
                 // Get the amounts for mine and watchonly
diff --git a/src/wallet/spend.cpp b/src/wallet/spend.cpp
@@ -336,6 +336,10 @@ CoinsResult AvailableCoins(const CWallet& wallet,
         if (wallet.IsLockedCoin(outpoint) && params.skip_locked)
             continue;
 
+        int nDepth = wallet.GetTxStateDepthInMainChain(txo.GetState());
+        Assert(nDepth >= 0);
+        Assert(!wallet.IsSpent(outpoint, /*min_depth=*/1));
+
         if (wallet.IsSpent(outpoint))
             continue;
 
@@ -353,11 +357,6 @@ CoinsResult AvailableCoins(const CWallet& wallet,
 
         assert(mine != ISMINE_NO);
 
-        int nDepth = wallet.GetTxStateDepthInMainChain(txo.GetState());
-        if (nDepth < 0)
-            continue;
-
-        // Perform tx level checks if we haven't already come across outputs from this tx before.
         if (!tx_safe_cache.contains(outpoint.hash)) {
             tx_safe_cache[outpoint.hash] = {false, false};
             const CWalletTx& wtx = *wallet.GetWalletTx(outpoint.hash);
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
@@ -1171,6 +1171,20 @@ CWalletTx* CWallet::AddToWallet(CTransactionRef tx, const TxState& state, const
     // Break debit/credit balance caches:
     wtx.MarkDirty();
 
+    // Remove or add back the inputs from m_txos to match the state of this tx.
+    if (wtx.isConfirmed())
+    {
+        // When a transaction becomes confirmed, we can remove all of the txos that were spent
+        // in its inputs as they are no longer relevant.
+        for (const CTxIn& txin : wtx.tx->vin) {
+            MarkTXOUnusable(txin.prevout);
+        }
+    } else if (wtx.isInactive()) {
+        // When a transaction becomes inactive, we need to mark its inputs as usable again
+        for (const CTxIn& txin : wtx.tx->vin) {
+            MarkTXOUsable(txin.prevout);
+        }
+    }
     // Cache the outputs that belong to the wallet
     RefreshWalletTxTXOs(wtx);
 
@@ -1416,12 +1430,20 @@ void CWallet::RecursiveUpdateTxState(WalletBatch* batch, const uint256& tx_hash,
             if (batch) batch->WriteTx(wtx);
             // Iterate over all its outputs, and update those tx states as well (if applicable)
             for (unsigned int i = 0; i < wtx.tx->vout.size(); ++i) {
-                std::pair<TxSpends::const_iterator, TxSpends::const_iterator> range = mapTxSpends.equal_range(COutPoint(Txid::FromUint256(now), i));
+                COutPoint outpoint{Txid::FromUint256(now), i};
+                std::pair<TxSpends::const_iterator, TxSpends::const_iterator> range = mapTxSpends.equal_range(outpoint);
                 for (TxSpends::const_iterator iter = range.first; iter != range.second; ++iter) {
                     if (!done.count(iter->second)) {
                         todo.insert(iter->second);
                     }
                 }
+                if (wtx.state<TxStateBlockConflicted>() || wtx.state<TxStateConfirmed>()) {
+                    // If the state applied is conflicted or confirmed, the outputs are unusable
+                    MarkTXOUnusable(outpoint);
+                } else {
+                    // Otherwise make the outputs usable
+                    MarkTXOUsable(outpoint);
+                }
             }
 
             if (update_state == TxUpdate::NOTIFY_CHANGED) {
@@ -1431,6 +1453,21 @@ void CWallet::RecursiveUpdateTxState(WalletBatch* batch, const uint256& tx_hash,
             // If a transaction changes its tx state, that usually changes the balance
             // available of the outputs it spends. So force those to be recomputed
             MarkInputsDirty(wtx.tx);
+            // Make the non-conflicted inputs usable again
+            for (unsigned int i = 0; i < wtx.tx->vin.size(); ++i) {
+                const CTxIn& txin = wtx.tx->vin.at(i);
+                auto unusable_txo_it = m_unusable_txos.find(txin.prevout);
+                if (unusable_txo_it == m_unusable_txos.end()) {
+                    continue;
+                }
+
+                if (std::get_if<TxStateBlockConflicted>(&unusable_txo_it->second.GetState()) ||
+                    std::get_if<TxStateConfirmed>(&unusable_txo_it->second.GetState())) {
+                    continue;
+                }
+
+                MarkTXOUsable(txin.prevout);
+            }
         }
     }
 }
@@ -3369,6 +3406,10 @@ bool CWallet::AttachChain(const std::shared_ptr<CWallet>& walletInstance, interf
     }
     walletInstance->m_attaching_chain = false;
 
+    // Remove TXOs that have already been spent
+    // We do this here as we need to have an attached chain to figure out what has actually been spent.
+    walletInstance->PruneSpentTXOs();
+
     return true;
 }
 
@@ -4152,9 +4193,9 @@ bool CWallet::ApplyMigrationData(MigrationData& data, bilingual_str& error)
         return false;
     }
 
-    // Update m_txos to match the descriptors remaining in this wallet
+    // Clear m_txos and m_unusable_txos. These will be updated next to match the descriptors remaining in this wallet
     m_txos.clear();
-    RefreshAllTXOs();
+    m_unusable_txos.clear();
 
     // Check if the transactions in the wallet are still ours. Either they belong here, or they belong in the watchonly wallet.
     // We need to go through these in the tx insertion order so that lookups to spends works.
@@ -4180,6 +4221,9 @@ bool CWallet::ApplyMigrationData(MigrationData& data, bilingual_str& error)
         }
     }
     for (const auto& [_pos, wtx] : wtxOrdered) {
+        // First update the UTXOs
+        wtx->m_txos.clear();
+        RefreshWalletTxTXOs(*wtx);
         // Check it is the watchonly wallet's
         // solvable_wallet doesn't need to be checked because transactions for those scripts weren't being watched for
         bool is_mine = IsMine(*wtx->tx) || IsFromMe(*wtx->tx);
@@ -4193,6 +4237,7 @@ bool CWallet::ApplyMigrationData(MigrationData& data, bilingual_str& error)
                     if (!new_tx) return false;
                     ins_wtx.SetTx(to_copy_wtx.tx);
                     ins_wtx.CopyFrom(to_copy_wtx);
+                    data.watchonly_wallet->RefreshWalletTxTXOs(ins_wtx);
                     return true;
                 })) {
                     error = strprintf(_("Error: Could not add watchonly tx %s to watchonly wallet"), wtx->GetHash().GetHex());
@@ -4667,26 +4712,34 @@ std::optional<CKey> CWallet::GetKey(const CKeyID& keyid) const
     return std::nullopt;
 }
 
+using TXOMap = std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher>;
 void CWallet::RefreshWalletTxTXOs(const CWalletTx& wtx)
 {
     AssertLockHeld(cs_wallet);
     for (uint32_t i = 0; i < wtx.tx->vout.size(); ++i) {
         const CTxOut& txout = wtx.tx->vout.at(i);
         COutPoint outpoint(wtx.GetHash(), i);
 
-        auto it = m_txos.find(outpoint);
-
         isminetype ismine = IsMine(txout);
         if (ismine == ISMINE_NO) {
             continue;
         }
 
-        if (it != m_txos.end()) {
+        auto it = wtx.m_txos.find(i);
+        if (it != wtx.m_txos.end()) {
             it->second.SetIsMine(ismine);
             it->second.SetState(wtx.GetState());
         } else {
-            auto [txo_it, _] = m_txos.emplace(outpoint, WalletTXO{txout, ismine, wtx.GetState(), wtx.IsCoinBase(), wtx.m_from_me, wtx.GetTxTime()});
-            wtx.m_txos.emplace(i, txo_it->second);
+            TXOMap::iterator txo_it;
+            bool txos_inserted = false;
+            if (m_last_block_processed_height >= 0 && IsSpent(outpoint, /*min_depth=*/1)) {
+                std::tie(txo_it, txos_inserted) = m_unusable_txos.emplace(outpoint, WalletTXO{txout, ismine, wtx.GetState(), wtx.IsCoinBase(), wtx.m_from_me, wtx.GetTxTime()});
+                assert(txos_inserted);
+            } else {
+                std::tie(txo_it, txos_inserted) = m_txos.emplace(outpoint, WalletTXO{txout, ismine, wtx.GetState(), wtx.IsCoinBase(), wtx.m_from_me, wtx.GetTxTime()});
+            }
+            auto [_, wtx_inserted] = wtx.m_txos.emplace(i, txo_it->second);
+            assert(wtx_inserted);
         }
     }
 }
@@ -4703,9 +4756,58 @@ std::optional<WalletTXO> CWallet::GetTXO(const COutPoint& outpoint) const
 {
     AssertLockHeld(cs_wallet);
     const auto& it = m_txos.find(outpoint);
-    if (it == m_txos.end()) {
-        return std::nullopt;
+    if (it != m_txos.end()) {
+        return it->second;
     }
-    return it->second;
+    const auto& u_it = m_unusable_txos.find(outpoint);
+    if (u_it != m_unusable_txos.end()) {
+        return u_it->second;
+    }
+    return std::nullopt;
+}
+
+void CWallet::PruneSpentTXOs()
+{
+    AssertLockHeld(cs_wallet);
+    auto it = m_txos.begin();
+    while (it != m_txos.end()) {
+        if (std::get_if<TxStateBlockConflicted>(&it->second.GetState()) || IsSpent(it->first, /*min_depth=*/1)) {
+            it = MarkTXOUnusable(it->first).first;
+        } else {
+            it++;
+        }
+    }
+}
+
+std::pair<TXOMap::iterator, TXOMap::iterator> CWallet::MarkTXOUnusable(const COutPoint& outpoint)
+{
+    AssertLockHeld(cs_wallet);
+    auto txos_it = m_txos.find(outpoint);
+    auto unusable_txos_it = m_unusable_txos.end();
+    if (txos_it != m_txos.end()) {
+        auto next_txo_it = std::next(txos_it);
+        auto nh = m_txos.extract(txos_it);
+        txos_it = next_txo_it;
+        auto [position, inserted, _] = m_unusable_txos.insert(std::move(nh));
+        unusable_txos_it = position;
+        assert(inserted);
+    }
+    return {txos_it, unusable_txos_it};
+}
+
+std::pair<TXOMap::iterator, TXOMap::iterator> CWallet::MarkTXOUsable(const COutPoint& outpoint)
+{
+    AssertLockHeld(cs_wallet);
+    auto txos_it = m_txos.end();
+    auto unusable_txos_it = m_unusable_txos.find(outpoint);
+    if (unusable_txos_it != m_unusable_txos.end()) {
+        auto next_unusable_txo_it = std::next(unusable_txos_it);
+        auto nh = m_unusable_txos.extract(unusable_txos_it);
+        unusable_txos_it = next_unusable_txo_it;
+        auto [position, inserted, _] = m_txos.insert(std::move(nh));
+        assert(inserted);
+        txos_it = position;
+    }
+    return {unusable_txos_it, txos_it};
 }
 } // namespace wallet
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
@@ -426,7 +426,11 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati
     std::unordered_map<CScript, std::vector<ScriptPubKeyMan*>, SaltedSipHasher> m_cached_spks;
 
     //! Set of both spent and unspent transaction outputs owned by this wallet
-    std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher> m_txos GUARDED_BY(cs_wallet);
+    using TXOMap = std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher>;
+    TXOMap m_txos GUARDED_BY(cs_wallet);
+    //! Set of transaction outputs that are definitely no longer usuable
+    //! These outputs may already be spent in a confirmed tx, or are the outputs of a conflicted tx
+    TXOMap m_unusable_txos GUARDED_BY(cs_wallet);
 
     /**
      * Catch wallet up to current chain, scanning new blocks, updating the best
@@ -507,11 +511,14 @@ class CWallet final : public WalletStorage, public interfaces::Chain::Notificati
 
     std::set<uint256> GetTxConflicts(const CWalletTx& wtx) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
-    const std::unordered_map<COutPoint, WalletTXO, SaltedOutpointHasher>& GetTXOs() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet) { AssertLockHeld(cs_wallet); return m_txos; };
+    const TXOMap& GetTXOs() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet) { AssertLockHeld(cs_wallet); return m_txos; };
     std::optional<WalletTXO> GetTXO(const COutPoint& outpoint) const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
     void RefreshWalletTxTXOs(const CWalletTx& wtx) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
     void RefreshAllTXOs() EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    void PruneSpentTXOs() EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    std::pair<TXOMap::iterator, TXOMap::iterator> MarkTXOUnusable(const COutPoint& outpoint) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    std::pair<TXOMap::iterator, TXOMap::iterator> MarkTXOUsable(const COutPoint& outpoint) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
     /**
      * Return depth of transaction in blockchain:
diff --git a/test/functional/wallet_balance.py b/test/functional/wallet_balance.py
@@ -13,6 +13,7 @@
     assert_equal,
     assert_is_hash_string,
     assert_raises_rpc_error,
+    find_vout_for_address,
 )
 from test_framework.wallet_util import get_generate_key
 
@@ -398,5 +399,35 @@ def test_balances(*, fee_node_1=0):
         balances = wallet.getbalances()
         assert_equal(balances["mine"]["trusted"], amount * 2)
 
+        wallet.unloadwallet()
+
+        self.log.info("Test that the balance is unchanged by an import that makes an already spent output in an existing tx \"mine\"")
+        self.nodes[0].createwallet("importalreadyspent")
+        wallet = self.nodes[0].get_wallet_rpc("importalreadyspent")
+
+        import_change_key = get_generate_key()
+        addr1 = wallet.getnewaddress()
+        addr2 = wallet.getnewaddress()
+
+        default.importprivkey(privkey=import_change_key.privkey, rescan=False)
+
+        res = default.send(outputs=[{addr1: amount}], options={"change_address": import_change_key.p2wpkh_addr})
+        inputs = [{"txid":res["txid"], "vout": find_vout_for_address(default, res["txid"], import_change_key.p2wpkh_addr)}]
+        default.send(outputs=[{addr2: amount}], options={"inputs": inputs, "add_inputs": True})
+
+        # Mock the time forward by another day so that "now" will exclude the block we just mined
+        self.nodes[0].setmocktime(int(time.time()) + 86400 * 2)
+        # Mine 11 blocks to move the MTP past the block we just mined
+        self.generate(self.nodes[0], 11, sync_fun=self.no_op)
+
+        balances = wallet.getbalances()
+        assert_equal(balances["mine"]["trusted"], amount * 2)
+
+        # Don't rescan to make sure that the import updates the wallet txos
+        # The balance should not change because the output for this key is already spent
+        wallet.importprivkey(privkey=import_change_key.privkey, rescan=False)
+        balances = wallet.getbalances()
+        assert_equal(balances["mine"]["trusted"], amount * 2)
+
 if __name__ == '__main__':
     WalletTest(__file__).main()
