From 96e7a89c5e0bf5e1a4497d8296e8841edd7ebbf1 Mon Sep 17 00:00:00 2001
From: Ava Chow <github@achow101.com>
Date: Tue, 20 Feb 2024 11:50:00 -0500
Subject: [PATCH] wallet: Recalculate the wallet's txos after any imports

---
 src/wallet/rpc/addresses.cpp | 1 +
 src/wallet/rpc/backup.cpp    | 1 +
 src/wallet/wallet.cpp        | 8 ++++++++
 src/wallet/wallet.h          | 2 ++
 4 files changed, 12 insertions(+)

diff --git a/src/wallet/rpc/addresses.cpp b/src/wallet/rpc/addresses.cpp
index 39df2166bef..eece0815111 100644
--- a/src/wallet/rpc/addresses.cpp
+++ b/src/wallet/rpc/addresses.cpp
@@ -250,6 +250,7 @@ RPCHelpMan keypoolrefill()
     if (pwallet->GetKeyPoolSize() < kpSize) {
         throw JSONRPCError(RPC_WALLET_ERROR, "Error refreshing keypool.");
     }
+    pwallet->RefreshAllTXOs();
 
     return UniValue::VNULL;
 },
diff --git a/src/wallet/rpc/backup.cpp b/src/wallet/rpc/backup.cpp
index f7d37239c6c..edd1fb86c99 100644
--- a/src/wallet/rpc/backup.cpp
+++ b/src/wallet/rpc/backup.cpp
@@ -394,6 +394,7 @@ RPCHelpMan importdescriptors()
             }
         }
         pwallet->ConnectScriptPubKeyManNotifiers();
+        pwallet->RefreshAllTXOs();
     }
 
     // Rescan the blockchain using the lowest timestamp
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 49f1cf22bc5..ecbddf341f4 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -4451,4 +4451,12 @@ void CWallet::RefreshTXOsFromTx(const CWalletTx& wtx)
         }
     }
 }
+
+void CWallet::RefreshAllTXOs()
+{
+    AssertLockHeld(cs_wallet);
+    for (const auto& [_, wtx] : mapWallet) {
+        RefreshTXOsFromTx(wtx);
+    }
+}
 } // namespace wallet
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 4d7a369baf2..2b52842236e 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -525,6 +525,8 @@ public:
 
     /** Cache outputs that belong to the wallet from a single transaction */
     void RefreshTXOsFromTx(const CWalletTx& wtx) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
+    /** Cache outputs that belong to the wallet for all transactions in the wallet */
+    void RefreshAllTXOs() EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
 
     /**
      * Return depth of transaction in blockchain: