From c978c6d39cdeb78fc4720767b943d03d6a9a36d8 Mon Sep 17 00:00:00 2001
From: Andrew Chow <achow101-github@achow101.com>
Date: Mon, 18 Apr 2022 14:26:26 -0400
Subject: [PATCH] walletdb: refactor active spkm loading

Instead of loading active spkm records as we come across them when
iterating the database, load them explicitly.

Due to exception handling changes, deserialization errors are now
treated as critical.
---
 src/wallet/walletdb.cpp | 54 ++++++++++++++++++++++++-----------------
 1 file changed, 32 insertions(+), 22 deletions(-)

diff --git a/src/wallet/walletdb.cpp b/src/wallet/walletdb.cpp
index f3a82734c1..20df4f589f 100644
--- a/src/wallet/walletdb.cpp
+++ b/src/wallet/walletdb.cpp
@@ -301,8 +301,6 @@ bool WalletBatch::EraseLockedUTXO(const COutPoint& output)
 class CWalletScanState {
 public:
     unsigned int m_unknown_records{0};
-    std::map<OutputType, uint256> m_active_external_spks;
-    std::map<OutputType, uint256> m_active_internal_spks;
 
     CWalletScanState() = default;
 };
@@ -495,18 +493,6 @@ ReadKeyValue(CWallet* pwallet, DataStream& ssKey, CDataStream& ssValue,
             strErr = "Found unsupported 'wkey' record, try loading with version 0.18";
             return false;
         } else if (strType == DBKeys::ACTIVEEXTERNALSPK || strType == DBKeys::ACTIVEINTERNALSPK) {
-            uint8_t type;
-            ssKey >> type;
-            uint256 id;
-            ssValue >> id;
-
-            bool internal = strType == DBKeys::ACTIVEINTERNALSPK;
-            auto& spk_mans = internal ? wss.m_active_internal_spks : wss.m_active_external_spks;
-            if (spk_mans.count(static_cast<OutputType>(type)) > 0) {
-                strErr = "Multiple ScriptPubKeyMans specified for a single type";
-                return false;
-            }
-            spk_mans[static_cast<OutputType>(type)] = id;
         } else if (strType == DBKeys::WALLETDESCRIPTOR) {
         } else if (strType == DBKeys::WALLETDESCRIPTORCACHE) {
         } else if (strType == DBKeys::WALLETDESCRIPTORLHCACHE) {
@@ -1146,6 +1132,35 @@ static DBErrors LoadTxRecords(CWallet* pwallet, DatabaseBatch& batch, std::vecto
     return result;
 }
 
+static DBErrors LoadActiveSPKMs(CWallet* pwallet, DatabaseBatch& batch) EXCLUSIVE_LOCKS_REQUIRED(pwallet->cs_wallet)
+{
+    AssertLockHeld(pwallet->cs_wallet);
+    DBErrors result = DBErrors::LOAD_OK;
+
+    // Load spk records
+    std::set<std::pair<OutputType, bool>> seen_spks;
+    for (const auto& spk_key : {DBKeys::ACTIVEEXTERNALSPK, DBKeys::ACTIVEINTERNALSPK}) {
+        LoadResult spkm_res = LoadRecords(pwallet, batch, spk_key,
+            [&seen_spks, &spk_key] (CWallet* pwallet, DataStream& key, CDataStream& value, std::string& strErr) {
+            uint8_t output_type;
+            key >> output_type;
+            uint256 id;
+            value >> id;
+
+            bool internal = spk_key == DBKeys::ACTIVEINTERNALSPK;
+            auto [it, insert] = seen_spks.emplace(static_cast<OutputType>(output_type), internal);
+            if (!insert) {
+                strErr = "Multiple ScriptpubKeyMans specified for a single type";
+                return DBErrors::CORRUPT;
+            }
+            pwallet->LoadActiveScriptPubKeyMan(id, static_cast<OutputType>(output_type), /*internal=*/internal);
+            return DBErrors::LOAD_OK;
+        });
+        result = std::max(result, spkm_res.m_result);
+    }
+    return result;
+}
+
 DBErrors WalletBatch::LoadWallet(CWallet* pwallet)
 {
     CWalletScanState wss;
@@ -1191,6 +1206,9 @@ DBErrors WalletBatch::LoadWallet(CWallet* pwallet)
         // Load tx records
         result = std::max(LoadTxRecords(pwallet, *m_batch, upgraded_txs, any_unordered), result);
 
+        // Load SPKMs
+        result = std::max(LoadActiveSPKMs(pwallet, *m_batch), result);
+
         // Get cursor
         std::unique_ptr<DatabaseCursor> cursor = m_batch->GetNewCursor();
         if (!cursor)
@@ -1236,14 +1254,6 @@ DBErrors WalletBatch::LoadWallet(CWallet* pwallet)
         result = DBErrors::CORRUPT;
     }
 
-    // Set the active ScriptPubKeyMans
-    for (auto spk_man_pair : wss.m_active_external_spks) {
-        pwallet->LoadActiveScriptPubKeyMan(spk_man_pair.second, spk_man_pair.first, /*internal=*/false);
-    }
-    for (auto spk_man_pair : wss.m_active_internal_spks) {
-        pwallet->LoadActiveScriptPubKeyMan(spk_man_pair.second, spk_man_pair.first, /*internal=*/true);
-    }
-
     if (fNoncriticalErrors && result == DBErrors::LOAD_OK) {
         result = DBErrors::NONCRITICAL_ERROR;
     }