From 80b0c259921d03754320754f14fec7a21bce3126 Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Tue, 20 Feb 2024 11:48:38 -0500 Subject: [PATCH] wallet: Load everything into DescSPKM on construction Instead of creating a DescSPKM that is then progressively loaded, we should instead create it all at once in a factory function when loading. --- .../external_signer_scriptpubkeyman.cpp | 5 +++ src/wallet/external_signer_scriptpubkeyman.h | 12 ++++-- src/wallet/scriptpubkeyman.cpp | 39 +++++++++---------- src/wallet/scriptpubkeyman.h | 26 +++++++------ src/wallet/wallet.cpp | 11 +++--- src/wallet/wallet.h | 2 +- src/wallet/walletdb.cpp | 27 +++++++------ 7 files changed, 68 insertions(+), 54 deletions(-) diff --git a/src/wallet/external_signer_scriptpubkeyman.cpp b/src/wallet/external_signer_scriptpubkeyman.cpp index f66686254a4..189f7698b50 100644 --- a/src/wallet/external_signer_scriptpubkeyman.cpp +++ b/src/wallet/external_signer_scriptpubkeyman.cpp @@ -21,6 +21,11 @@ using common::PSBTError; namespace wallet { +std::unique_ptr ExternalSignerScriptPubKeyMan::LoadFromStorage(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys) +{ + return std::unique_ptr(new ExternalSignerScriptPubKeyMan(storage, descriptor, keypool_size, keys, ckeys)); +} + bool ExternalSignerScriptPubKeyMan::SetupDescriptor(WalletBatch& batch, std::unique_ptr desc) { LOCK(cs_desc_man); diff --git a/src/wallet/external_signer_scriptpubkeyman.h b/src/wallet/external_signer_scriptpubkeyman.h index 0ccc243c256..1c3324f1915 100644 --- a/src/wallet/external_signer_scriptpubkeyman.h +++ b/src/wallet/external_signer_scriptpubkeyman.h @@ -15,14 +15,18 @@ struct bilingual_str; namespace wallet { class ExternalSignerScriptPubKeyMan : public DescriptorScriptPubKeyMan { - public: - ExternalSignerScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size) - : DescriptorScriptPubKeyMan(storage, descriptor, keypool_size) - {} +private: + ExternalSignerScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys) + : DescriptorScriptPubKeyMan(storage, descriptor, keypool_size, keys, ckeys) + {} + +public: ExternalSignerScriptPubKeyMan(WalletStorage& storage, int64_t keypool_size) : DescriptorScriptPubKeyMan(storage, keypool_size) {} + static std::unique_ptr LoadFromStorage(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys); + /** Provide a descriptor at setup time * Returns false if already setup or setup fails, true if setup is successful */ diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp index a65dee60ac0..5ec51a737cd 100644 --- a/src/wallet/scriptpubkeyman.cpp +++ b/src/wallet/scriptpubkeyman.cpp @@ -821,6 +821,24 @@ bool LegacyDataSPKM::DeleteRecordsWithDB(WalletBatch& batch) return batch.EraseRecords(DBKeys::LEGACY_TYPES); } +DescriptorScriptPubKeyMan::DescriptorScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys) + : ScriptPubKeyMan(storage), + m_map_keys(keys), + m_map_crypted_keys(ckeys), + m_keypool_size(keypool_size), + m_wallet_descriptor(descriptor) +{ + if (!keys.empty() && !ckeys.empty()) { + throw std::runtime_error("Wallet contains both unencrypted and encrypted keys"); + } + Load(); +} + +std::unique_ptr DescriptorScriptPubKeyMan::LoadFromStorage(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys) +{ + return std::unique_ptr(new DescriptorScriptPubKeyMan(storage, descriptor, keypool_size, keys, ckeys)); +} + util::Result DescriptorScriptPubKeyMan::GetNewDestination(const OutputType type) { // Returns true if this descriptor supports getting new addresses. Conditions where we may be unable to fetch them (e.g. locked) are caught later @@ -1426,11 +1444,10 @@ uint256 DescriptorScriptPubKeyMan::GetID() const return m_wallet_descriptor.id; } -void DescriptorScriptPubKeyMan::SetCache(const DescriptorCache& cache) +void DescriptorScriptPubKeyMan::Load() { LOCK(cs_desc_man); std::set new_spks; - m_wallet_descriptor.cache = cache; for (int32_t i = m_wallet_descriptor.range_start; i < m_wallet_descriptor.range_end; ++i) { FlatSigningProvider out_keys; std::vector scripts_temp; @@ -1460,24 +1477,6 @@ void DescriptorScriptPubKeyMan::SetCache(const DescriptorCache& cache) m_storage.TopUpCallback(new_spks, this); } -bool DescriptorScriptPubKeyMan::AddKey(const CKeyID& key_id, const CKey& key) -{ - LOCK(cs_desc_man); - m_map_keys[key_id] = key; - return true; -} - -bool DescriptorScriptPubKeyMan::AddCryptedKey(const CKeyID& key_id, const CPubKey& pubkey, const std::vector& crypted_key) -{ - LOCK(cs_desc_man); - if (!m_map_keys.empty()) { - return false; - } - - m_map_crypted_keys[key_id] = make_pair(pubkey, crypted_key); - return true; -} - bool DescriptorScriptPubKeyMan::HasWalletDescriptor(const WalletDescriptor& desc) const { LOCK(cs_desc_man); diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h index cf024294bf3..2762ebc2e61 100644 --- a/src/wallet/scriptpubkeyman.h +++ b/src/wallet/scriptpubkeyman.h @@ -166,14 +166,18 @@ static const std::unordered_set LEGACY_OUTPUT_TYPES { OutputType::BECH32, }; +using KeyMap = std::map; +using CryptedKeyMap = std::map>>; +using ScriptPubKeyMap = std::map; // Map of scripts to descriptor range index +using PubKeyMap = std::map; // Map of pubkeys involved in scripts to descriptor range index + // Manages the data for a LegacyScriptPubKeyMan. // This is the minimum necessary to load a legacy wallet so that it can be migrated. class LegacyDataSPKM : public ScriptPubKeyMan, public FillableSigningProvider { -private: +protected: using WatchOnlySet = std::set; using WatchKeyMap = std::map; - using CryptedKeyMap = std::map>>; CryptedKeyMap mapCryptedKeys GUARDED_BY(cs_KeyStore); WatchOnlySet setWatchOnly GUARDED_BY(cs_KeyStore); @@ -273,11 +277,6 @@ class DescriptorScriptPubKeyMan : public ScriptPubKeyMan { friend class LegacyDataSPKM; private: - using ScriptPubKeyMap = std::map; // Map of scripts to descriptor range index - using PubKeyMap = std::map; // Map of pubkeys involved in scripts to descriptor range index - using CryptedKeyMap = std::map>>; - using KeyMap = std::map; - ScriptPubKeyMap m_map_script_pub_keys GUARDED_BY(cs_desc_man); PubKeyMap m_map_pubkeys GUARDED_BY(cs_desc_man); int32_t m_max_cached_index = -1; @@ -315,13 +314,19 @@ private: // Fetch the SigningProvider for a given index and optionally include private keys. Called by the above functions. std::unique_ptr GetSigningProvider(int32_t index, bool include_private = false) const EXCLUSIVE_LOCKS_REQUIRED(cs_desc_man); + void Load(); + protected: + //! Create a DescriptorScriptPubKeyMan from existing data (i.e. during loading) + DescriptorScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys); + WalletDescriptor m_wallet_descriptor GUARDED_BY(cs_desc_man); //! Same as 'TopUp' but designed for use within a batch transaction context bool TopUpWithDB(WalletBatch& batch, unsigned int size = 0); public: + //! Create a new DescriptorScriptPubKeyMan from an existing descriptor (i.e. from an import) DescriptorScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size) : ScriptPubKeyMan(storage), m_keypool_size(keypool_size), @@ -332,6 +337,8 @@ public: m_keypool_size(keypool_size) {} + static std::unique_ptr LoadFromStorage(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys); + mutable RecursiveMutex cs_desc_man; util::Result GetNewDestination(OutputType type) override; @@ -383,11 +390,6 @@ public: uint256 GetID() const override; - void SetCache(const DescriptorCache& cache); - - bool AddKey(const CKeyID& key_id, const CKey& key); - bool AddCryptedKey(const CKeyID& key_id, const CPubKey& pubkey, const std::vector& crypted_key); - bool HasWalletDescriptor(const WalletDescriptor& desc) const; util::Result UpdateWalletDescriptor(WalletDescriptor& descriptor); bool CanUpdateToWalletDescriptor(const WalletDescriptor& descriptor, std::string& error); diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index c337c45d63a..7b94b659700 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -3539,16 +3539,15 @@ void CWallet::ConnectScriptPubKeyManNotifiers() } } -DescriptorScriptPubKeyMan& CWallet::LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc) +void CWallet::LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc, const KeyMap& keys, const CryptedKeyMap& ckeys) { - DescriptorScriptPubKeyMan* spk_manager; + std::unique_ptr spk_manager; if (IsWalletFlagSet(WALLET_FLAG_EXTERNAL_SIGNER)) { - spk_manager = new ExternalSignerScriptPubKeyMan(*this, desc, m_keypool_size); + spk_manager = ExternalSignerScriptPubKeyMan::LoadFromStorage(*this, desc, m_keypool_size, keys, ckeys); } else { - spk_manager = new DescriptorScriptPubKeyMan(*this, desc, m_keypool_size); + spk_manager = DescriptorScriptPubKeyMan::LoadFromStorage(*this, desc, m_keypool_size, keys, ckeys); } - AddScriptPubKeyMan(id, std::unique_ptr(spk_manager)); - return *spk_manager; + AddScriptPubKeyMan(id, std::move(spk_manager)); } DescriptorScriptPubKeyMan& CWallet::SetupDescriptorScriptPubKeyMan(WalletBatch& batch, const CExtKey& master_key, const OutputType& output_type, bool internal) diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h index 360b7315b47..4b5df5c5006 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -1004,7 +1004,7 @@ public: void ConnectScriptPubKeyManNotifiers(); //! Instantiate a descriptor ScriptPubKeyMan from the WalletDescriptor and load it - DescriptorScriptPubKeyMan& LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc); + void LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc, const KeyMap& keys, const CryptedKeyMap& ckeys); //! Adds the active ScriptPubKeyMan for the specified type and internal. Writes it to the wallet file //! @param[in] id The unique id for the ScriptPubKeyMan diff --git a/src/wallet/walletdb.cpp b/src/wallet/walletdb.cpp index c8b7ae53cd5..5ad29970d20 100644 --- a/src/wallet/walletdb.cpp +++ b/src/wallet/walletdb.cpp @@ -771,10 +771,8 @@ static DBErrors LoadDescriptorWalletRecords(CWallet* pwallet, DatabaseBatch& bat strErr = strprintf("%s\nDetails: %s", strErr, e.what()); return DBErrors::UNKNOWN_DESCRIPTOR; } - DescriptorScriptPubKeyMan& spkm = pwallet->LoadDescriptorScriptPubKeyMan(id, desc); - // Prior to doing anything with this spkm, verify ID compatibility - if (id != spkm.GetID()) { + if (id != desc.id) { strErr = "The descriptor ID calculated by the wallet differs from the one in DB"; return DBErrors::CORRUPT; } @@ -833,15 +831,14 @@ static DBErrors LoadDescriptorWalletRecords(CWallet* pwallet, DatabaseBatch& bat }); result = std::max(result, lh_cache_res.m_result); - // Set the cache for this descriptor - auto spk_man = (DescriptorScriptPubKeyMan*)pwallet->GetScriptPubKeyMan(id); - assert(spk_man); - spk_man->SetCache(cache); + // Set the cache to the WalletDescriptor + desc.cache = cache; // Get unencrypted keys + KeyMap keys; prefix = PrefixStream(DBKeys::WALLETDESCRIPTORKEY, id); LoadResult key_res = LoadRecords(pwallet, batch, DBKeys::WALLETDESCRIPTORKEY, prefix, - [&id, &spk_man] (CWallet* pwallet, DataStream& key, DataStream& value, std::string& strErr) { + [&id, &keys] (CWallet* pwallet, DataStream& key, DataStream& value, std::string& strErr) { uint256 desc_id; CPubKey pubkey; key >> desc_id; @@ -873,16 +870,17 @@ static DBErrors LoadDescriptorWalletRecords(CWallet* pwallet, DatabaseBatch& bat strErr = "Error reading wallet database: descriptor unencrypted key CPrivKey corrupt"; return DBErrors::CORRUPT; } - spk_man->AddKey(pubkey.GetID(), privkey); + keys[pubkey.GetID()] = privkey; return DBErrors::LOAD_OK; }); result = std::max(result, key_res.m_result); num_keys = key_res.m_records; // Get encrypted keys + CryptedKeyMap ckeys; prefix = PrefixStream(DBKeys::WALLETDESCRIPTORCKEY, id); LoadResult ckey_res = LoadRecords(pwallet, batch, DBKeys::WALLETDESCRIPTORCKEY, prefix, - [&id, &spk_man] (CWallet* pwallet, DataStream& key, DataStream& value, std::string& err) { + [&id, &ckeys] (CWallet* pwallet, DataStream& key, DataStream& value, std::string& err) { uint256 desc_id; CPubKey pubkey; key >> desc_id; @@ -896,12 +894,19 @@ static DBErrors LoadDescriptorWalletRecords(CWallet* pwallet, DatabaseBatch& bat std::vector privkey; value >> privkey; - spk_man->AddCryptedKey(pubkey.GetID(), pubkey, privkey); + ckeys[pubkey.GetID()] = std::make_pair(pubkey, privkey); return DBErrors::LOAD_OK; }); result = std::max(result, ckey_res.m_result); num_ckeys = ckey_res.m_records; + try { + pwallet->LoadDescriptorScriptPubKeyMan(id, desc, keys, ckeys); + } catch (std::runtime_error& e) { + strErr = e.what(); + return DBErrors::CORRUPT; + } + return result; });