Add RAII wallet rescan reserver

jonasschnelli · random-zebra · commit 4d056725a271 · 2021-04-11T11:33:36.000+02:00
diff --git a/src/wallet/rpcdump.cpp b/src/wallet/rpcdump.cpp
@@ -221,6 +221,12 @@ UniValue importaddress(const JSONRPCRequest& request)
     const std::string strLabel = (request.params.size() > 1 ? request.params[1].get_str() : "");
     // Whether to perform rescan after import
     const bool fRescan = (request.params.size() > 2 ? request.params[2].get_bool() : true);
+
+    WalletRescanReserver reserver(pwalletMain);
+    if (fRescan && !reserver.reserve()) {
+        throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
+    }
+
     // Whether to import a p2sh version, too
     const bool fP2SH = (request.params.size() > 3 ? request.params[3].get_bool() : false);
 
@@ -276,6 +282,11 @@ UniValue importpubkey(const JSONRPCRequest& request)
     // Whether to perform rescan after import
     const bool fRescan = (request.params.size() > 2 ? request.params[2].get_bool() : true);
 
+    WalletRescanReserver reserver(pwalletMain);
+    if (fRescan && !reserver.reserve()) {
+        throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
+    }
+
     if (!IsHex(request.params[0].get_str()))
         throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Pubkey must be a hex string");
     std::vector<unsigned char> data(ParseHex(request.params[0].get_str()));
@@ -317,19 +328,23 @@ UniValue importwallet(const JSONRPCRequest& request)
             "\nImport using the json rpc call\n" +
             HelpExampleRpc("importwallet", "\"test\""));
 
-
     std::ifstream file;
     file.open(request.params[0].get_str().c_str(), std::ios::in | std::ios::ate);
     if (!file.is_open())
         throw JSONRPCError(RPC_INVALID_PARAMETER, "Cannot open wallet dump file");
 
-    int64_t nTimeBegin = chainActive.Tip()->GetBlockTime();
+    WalletRescanReserver reserver(pwalletMain);
+    if (!reserver.reserve()) {
+        throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
+    }
 
+    int64_t nTimeBegin = 0;
     bool fGood = true;
     {
         LOCK2(cs_main, pwalletMain->cs_wallet);
         EnsureWalletIsUnlocked();
 
+        nTimeBegin = chainActive.Tip()->GetBlockTime();
         int64_t nFilesize = std::max((int64_t)1, (int64_t)file.tellg());
         file.seekg(0, file.beg);
 
@@ -973,6 +988,11 @@ UniValue importmulti(const JSONRPCRequest& mainRequest)
         }
     }
 
+    WalletRescanReserver reserver(pwalletMain);
+    if (fRescan && !reserver.reserve()) {
+        throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
+    }
+
     int64_t now = 0;
     bool fRunScan = false;
     int64_t nLowestTimestamp = 0;
diff --git a/src/wallet/rpcwallet.cpp b/src/wallet/rpcwallet.cpp
@@ -4207,7 +4207,8 @@ UniValue rescanblockchain(const JSONRPCRequest& request)
     }
 
     EnsureWallet();
-    if (pwalletMain->IsScanning()) {
+    WalletRescanReserver reserver(pwalletMain);
+    if (!reserver.reserve()) {
         throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");
     }
 
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
@@ -44,7 +44,8 @@
 #include <utility>
 #include <vector>
 
-extern CWallet* pwalletMain;
+typedef CWallet* CWalletRef;
+extern CWalletRef pwalletMain;
 
 /**
  * Settings
@@ -564,6 +565,9 @@ class CWallet : public CCryptoKeyStore, public CValidationInterface
     static std::atomic<bool> fFlushScheduled;
     std::atomic<bool> fAbortRescan;
     std::atomic<bool> fScanningWallet;
+    std::mutex mutexScanning;
+    friend class WalletRescanReserver;
+
 
     //! keeps track of whether Unlock has run a thorough check before
     bool fDecryptionThoroughlyChecked{false};
@@ -1229,4 +1233,34 @@ class CStakeableOutput : public COutput
 
 };
 
+/** RAII object to check and reserve a wallet rescan */
+class WalletRescanReserver
+{
+private:
+    CWalletRef m_wallet;
+    bool m_could_reserve;
+public:
+    explicit WalletRescanReserver(CWalletRef w) : m_wallet(w), m_could_reserve(false) {}
+
+    bool reserve()
+    {
+        assert(!m_could_reserve);
+        std::lock_guard<std::mutex> lock(m_wallet->mutexScanning);
+        if (m_wallet->fScanningWallet) {
+            return false;
+        }
+        m_wallet->fScanningWallet = true;
+        m_could_reserve = true;
+        return true;
+    }
+
+    ~WalletRescanReserver()
+    {
+        std::lock_guard<std::mutex> lock(m_wallet->mutexScanning);
+        if (m_could_reserve) {
+            m_wallet->fScanningWallet = false;
+        }
+    }
+};
+
 #endif // BITCOIN_WALLET_H

Original file line number	Diff line number	Diff line change
`@@ -4207,7 +4207,8 @@ UniValue rescanblockchain(const JSONRPCRequest& request)`
`4207`	`4207`	`}`
`4208`	`4208`
`4209`	`4209`	`EnsureWallet();`
`4210`		`- if (pwalletMain->IsScanning()) {`
	`4210`	`+ WalletRescanReserver reserver(pwalletMain);`
	`4211`	`+ if (!reserver.reserve()) {`
`4211`	`4212`	`throw JSONRPCError(RPC_WALLET_ERROR, "Wallet is currently rescanning. Abort existing rescan or wait.");`
`4212`	`4213`	`}`
`4213`	`4214`