From f48649a61867e0fb523fa408b30f5eb39cb9913a Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Tue, 20 Feb 2024 11:48:33 -0500 Subject: [PATCH 1/7] wallet: Consolidate generation setup callers into one function --- src/wallet/wallet.cpp | 94 +++++++++++++++++-------------------------- src/wallet/wallet.h | 3 ++ 2 files changed, 41 insertions(+), 56 deletions(-) diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index 96c43975049..c5032ed1221 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -423,6 +423,11 @@ std::shared_ptr CreateWallet(WalletContext& context, const std::string& return nullptr; } + // Unset the blank flag if not specified by the user + if (!create_blank) { + wallet->UnsetWalletFlag(WALLET_FLAG_BLANK_WALLET); + } + // Encrypt the wallet if (!passphrase.empty() && !(wallet_creation_flags & WALLET_FLAG_DISABLE_PRIVATE_KEYS)) { if (!wallet->EncryptWallet(passphrase)) { @@ -430,33 +435,6 @@ std::shared_ptr CreateWallet(WalletContext& context, const std::string& status = DatabaseStatus::FAILED_ENCRYPT; return nullptr; } - if (!create_blank) { - // Unlock the wallet - if (!wallet->Unlock(passphrase)) { - error = Untranslated("Error: Wallet was encrypted but could not be unlocked"); - status = DatabaseStatus::FAILED_ENCRYPT; - return nullptr; - } - - // Set a seed for the wallet - { - LOCK(wallet->cs_wallet); - if (wallet->IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS)) { - wallet->SetupDescriptorScriptPubKeyMans(); - } else { - for (auto spk_man : wallet->GetActiveScriptPubKeyMans()) { - if (!spk_man->SetupGeneration()) { - error = Untranslated("Unable to generate initial keys"); - status = DatabaseStatus::FAILED_CREATE; - return nullptr; - } - } - } - } - - // Relock the wallet - wallet->Lock(); - } } NotifyWalletLoaded(context, wallet); @@ -861,16 +839,9 @@ bool CWallet::EncryptWallet(const SecureString& strWalletPassphrase) Lock(); Unlock(strWalletPassphrase); - // If we are using descriptors, make new descriptors with a new seed - if (IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS) && !IsWalletFlagSet(WALLET_FLAG_BLANK_WALLET)) { - SetupDescriptorScriptPubKeyMans(); - } else if (auto spk_man = GetLegacyScriptPubKeyMan()) { - // if we are using HD, replace the HD seed with a new one - if (spk_man->IsHDEnabled()) { - if (!spk_man->SetupGeneration(true)) { - return false; - } - } + // Generate new descriptors or seed if not blank or disable private keys + if (!SetupWalletGeneration()) { + return false; } Lock(); @@ -3013,25 +2984,9 @@ std::shared_ptr CWallet::Create(WalletContext& context, const std::stri walletInstance->InitWalletFlags(wallet_creation_flags); - // Only create LegacyScriptPubKeyMan when not descriptor wallet - if (!walletInstance->IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS)) { - walletInstance->SetupLegacyScriptPubKeyMan(); - } - - if ((wallet_creation_flags & WALLET_FLAG_EXTERNAL_SIGNER) || !(wallet_creation_flags & (WALLET_FLAG_DISABLE_PRIVATE_KEYS | WALLET_FLAG_BLANK_WALLET))) { - LOCK(walletInstance->cs_wallet); - if (walletInstance->IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS)) { - walletInstance->SetupDescriptorScriptPubKeyMans(); - // SetupDescriptorScriptPubKeyMans already calls SetupGeneration for us so we don't need to call SetupGeneration separately - } else { - // Legacy wallets need SetupGeneration here. - for (auto spk_man : walletInstance->GetActiveScriptPubKeyMans()) { - if (!spk_man->SetupGeneration()) { - error = _("Unable to generate initial keys"); - return nullptr; - } - } - } + if (!walletInstance->SetupWalletGeneration()) { + error = _("Unable to generate initial keys"); + return nullptr; } if (chain) { @@ -3755,6 +3710,33 @@ void CWallet::SetupDescriptorScriptPubKeyMans() } } +bool CWallet::SetupWalletGeneration() +{ + LOCK(cs_wallet); + // Skip if blank or no privkeys + if (!IsWalletFlagSet(WALLET_FLAG_EXTERNAL_SIGNER) && + (IsWalletFlagSet(WALLET_FLAG_BLANK_WALLET) || IsWalletFlagSet(WALLET_FLAG_DISABLE_PRIVATE_KEYS))) { + // Just make sure that there's a legacy spkm if this wallet is legacy + if (!IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS)) { + SetupLegacyScriptPubKeyMan(); + } + return true; + } + // If we are using descriptors, make new descriptors with a new seed + if (IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS)) { + SetupDescriptorScriptPubKeyMans(); + } else if (auto spk_man = GetOrCreateLegacyScriptPubKeyMan()) { + // if we are using HD, set or replace the HD seed with a new one + // non-HD always has key generation enabled by default + if (CanSupportFeature(FEATURE_HD)) { + if (!spk_man->SetupGeneration(true)) { + return false; + } + } + } + return true; +} + void CWallet::AddActiveScriptPubKeyMan(uint256 id, OutputType type, bool internal) { WalletBatch batch(GetDatabase()); diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h index b49b5a7d0d8..683867a937d 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -1020,6 +1020,9 @@ public: void SetupDescriptorScriptPubKeyMans(const CExtKey& master_key) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet); void SetupDescriptorScriptPubKeyMans() EXCLUSIVE_LOCKS_REQUIRED(cs_wallet); + //! Setup new descriptors or seed for new address generation + bool SetupWalletGeneration(); + //! Return the DescriptorScriptPubKeyMan for a WalletDescriptor if it is already in the wallet DescriptorScriptPubKeyMan* GetDescriptorScriptPubKeyMan(const WalletDescriptor& desc) const; From a1ec92ca70bd25850ed2ae95b4406ff23bf2e157 Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Tue, 20 Feb 2024 11:48:36 -0500 Subject: [PATCH 2/7] wallet: Remove SetupGeneration from ScriptPubKeyMan interface --- src/wallet/scriptpubkeyman.h | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h index 4575881d96a..4fa7c246716 100644 --- a/src/wallet/scriptpubkeyman.h +++ b/src/wallet/scriptpubkeyman.h @@ -204,12 +204,6 @@ public: */ virtual std::vector MarkUnusedAddresses(const CScript& script) { return {}; } - /** Sets up the key generation stuff, i.e. generates new HD seeds and sets them as active. - * Returns false if already setup or setup fails, true if setup is successful - * Set force=true to make it re-setup if already setup, used for upgrades - */ - virtual bool SetupGeneration(bool force = false) { return false; } - /* Returns true if HD is enabled */ virtual bool IsHDEnabled() const { return false; } @@ -398,7 +392,11 @@ public: bool IsHDEnabled() const override; - bool SetupGeneration(bool force = false) override; + /** Sets up the key generation stuff, i.e. generates new HD seeds and sets them as active. + * Returns false if already setup or setup fails, true if setup is successful + * Set force=true to make it re-setup if already setup, used for upgrades + */ + bool SetupGeneration(bool force = false); bool Upgrade(int prev_version, int new_version, bilingual_str& error) override; From d81da95b9499886cc4a28d8bd145eaddeba5488b Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Tue, 20 Feb 2024 11:48:38 -0500 Subject: [PATCH 3/7] 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 constructor when loading. --- src/wallet/external_signer_scriptpubkeyman.h | 4 +-- src/wallet/scriptpubkeyman.cpp | 34 +++++++++----------- src/wallet/scriptpubkeyman.h | 25 +++++++------- src/wallet/wallet.cpp | 7 ++-- src/wallet/wallet.h | 2 +- src/wallet/walletdb.cpp | 30 ++++++++--------- 6 files changed, 49 insertions(+), 53 deletions(-) diff --git a/src/wallet/external_signer_scriptpubkeyman.h b/src/wallet/external_signer_scriptpubkeyman.h index c052ce61298..c9e6f39c221 100644 --- a/src/wallet/external_signer_scriptpubkeyman.h +++ b/src/wallet/external_signer_scriptpubkeyman.h @@ -13,8 +13,8 @@ namespace wallet { class ExternalSignerScriptPubKeyMan : public DescriptorScriptPubKeyMan { public: - ExternalSignerScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size) - : DescriptorScriptPubKeyMan(storage, descriptor, keypool_size) + ExternalSignerScriptPubKeyMan(WalletStorage& storage, const uint256& id, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys) + : DescriptorScriptPubKeyMan(storage, id, descriptor, keypool_size, keys, ckeys) {} ExternalSignerScriptPubKeyMan(WalletStorage& storage, int64_t keypool_size) : DescriptorScriptPubKeyMan(storage, keypool_size) diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp index 59171f6db76..fe669d54823 100644 --- a/src/wallet/scriptpubkeyman.cpp +++ b/src/wallet/scriptpubkeyman.cpp @@ -2003,6 +2003,22 @@ bool LegacyScriptPubKeyMan::DeleteRecords() return batch.EraseRecords(DBKeys::LEGACY_TYPES); } +DescriptorScriptPubKeyMan::DescriptorScriptPubKeyMan(WalletStorage& storage, const uint256& id, 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 (!m_map_keys.empty() && !m_map_crypted_keys.empty()) { + throw std::runtime_error("Error: Wallet contains both unencrypted and encrypted keys"); + } + if (id != GetID()) { + throw std::runtime_error("The descriptor ID calculated by the wallet differs from the one in DB"); + } + SetCache(m_wallet_descriptor.cache); +} + 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 @@ -2640,24 +2656,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 4fa7c246716..0e1e9986c8f 100644 --- a/src/wallet/scriptpubkeyman.h +++ b/src/wallet/scriptpubkeyman.h @@ -271,18 +271,22 @@ static const std::unordered_set LEGACY_OUTPUT_TYPES { class DescriptorScriptPubKeyMan; +using WatchOnlySet = std::set; +using WatchKeyMap = std::map; +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 + class LegacyScriptPubKeyMan : public ScriptPubKeyMan, public FillableSigningProvider { private: //! keeps track of whether Unlock has run a thorough check before bool fDecryptionThoroughlyChecked = true; - using WatchOnlySet = std::set; - using WatchKeyMap = std::map; WalletBatch *encrypted_batch GUARDED_BY(cs_KeyStore) = nullptr; - using CryptedKeyMap = std::map>>; CryptedKeyMap mapCryptedKeys GUARDED_BY(cs_KeyStore); WatchOnlySet setWatchOnly GUARDED_BY(cs_KeyStore); @@ -558,11 +562,6 @@ public: class DescriptorScriptPubKeyMan : public ScriptPubKeyMan { 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; @@ -589,6 +588,8 @@ 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 SetCache(const DescriptorCache& cache); + protected: WalletDescriptor m_wallet_descriptor GUARDED_BY(cs_desc_man); @@ -596,11 +597,14 @@ protected: 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), m_wallet_descriptor(descriptor) {} + //! Create a DescriptorScriptPubKeyMan from existing data (i.e. during loading) + DescriptorScriptPubKeyMan(WalletStorage& storage, const uint256& id, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys); DescriptorScriptPubKeyMan(WalletStorage& storage, int64_t keypool_size) : ScriptPubKeyMan(storage), m_keypool_size(keypool_size) @@ -654,11 +658,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; void UpdateWalletDescriptor(WalletDescriptor& descriptor); bool CanUpdateToWalletDescriptor(const WalletDescriptor& descriptor, std::string& error); diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index c5032ed1221..8527b7d0250 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -3605,16 +3605,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; if (IsWalletFlagSet(WALLET_FLAG_EXTERNAL_SIGNER)) { - spk_manager = new ExternalSignerScriptPubKeyMan(*this, desc, m_keypool_size); + spk_manager = new ExternalSignerScriptPubKeyMan(*this, id, desc, m_keypool_size, keys, ckeys); } else { - spk_manager = new DescriptorScriptPubKeyMan(*this, desc, m_keypool_size); + spk_manager = new DescriptorScriptPubKeyMan(*this, id, desc, m_keypool_size, keys, ckeys); } AddScriptPubKeyMan(id, std::unique_ptr(spk_manager)); - return *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 683867a937d..c89990e3861 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -994,7 +994,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 b1ce7ee4e70..f23c9ceefee 100644 --- a/src/wallet/walletdb.cpp +++ b/src/wallet/walletdb.cpp @@ -808,13 +808,6 @@ 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()) { - strErr = "The descriptor ID calculated by the wallet differs from the one in DB"; - return DBErrors::CORRUPT; - } DescriptorCache cache; @@ -870,15 +863,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 + std::map 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; @@ -913,16 +905,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 + std::map>> 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; @@ -936,12 +929,19 @@ static DBErrors LoadDescriptorWalletRecords(CWallet* pwallet, DatabaseBatch& bat std::vector privkey; value >> privkey; - spk_man->AddCryptedKey(pubkey.GetID(), pubkey, privkey); + ckeys[pubkey.GetID()] = 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; }); From dae9767bf63bd425ff666703bacf682a5ab68758 Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Wed, 3 Apr 2024 18:38:51 -0400 Subject: [PATCH 4/7] fuzz: Skip adding descriptor to wallet if it cannot be expanded --- src/wallet/test/fuzz/scriptpubkeyman.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/wallet/test/fuzz/scriptpubkeyman.cpp b/src/wallet/test/fuzz/scriptpubkeyman.cpp index 228e9629edd..63d09d66cae 100644 --- a/src/wallet/test/fuzz/scriptpubkeyman.cpp +++ b/src/wallet/test/fuzz/scriptpubkeyman.cpp @@ -72,6 +72,11 @@ static std::optional> CreateWal std::unique_ptr parsed_desc{Parse(desc_str.value(), keys, error, false)}; if (!parsed_desc) return std::nullopt; + FlatSigningProvider out_keys; + std::vector scripts_temp; + DescriptorCache temp_cache; + if (!parsed_desc->Expand(0, keys, scripts_temp, out_keys, &temp_cache)) return std::nullopt; + WalletDescriptor w_desc{std::move(parsed_desc), /*creation_time=*/0, /*range_start=*/0, /*range_end=*/1, /*next_index=*/1}; return std::make_pair(w_desc, keys); } From 4f066b63870f50183c79831a43193a64eb6f3338 Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Tue, 20 Feb 2024 11:48:42 -0500 Subject: [PATCH 5/7] wallet: include keys when constructing DescriptorSPKM during import When importing a descriptor, all of the descriptor data should be provided at the same time in the constructor. --- src/wallet/scriptpubkeyman.cpp | 42 +++++++++++++++++------- src/wallet/scriptpubkeyman.h | 12 +++---- src/wallet/test/fuzz/scriptpubkeyman.cpp | 9 ----- src/wallet/wallet.cpp | 16 ++------- 4 files changed, 38 insertions(+), 41 deletions(-) diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp index fe669d54823..37159227400 100644 --- a/src/wallet/scriptpubkeyman.cpp +++ b/src/wallet/scriptpubkeyman.cpp @@ -1811,9 +1811,8 @@ std::optional LegacyScriptPubKeyMan::MigrateToDescriptor() WalletDescriptor w_desc(std::move(desc), creation_time, 0, 0, 0); // Make the DescriptorScriptPubKeyMan and get the scriptPubKeys - auto desc_spk_man = std::unique_ptr(new DescriptorScriptPubKeyMan(m_storage, w_desc, m_keypool_size)); - desc_spk_man->AddDescriptorKey(key, key.GetPubKey()); - desc_spk_man->TopUp(); + keys.keys.emplace(key.GetPubKey().GetID(), key); + auto desc_spk_man = std::unique_ptr(new DescriptorScriptPubKeyMan(m_storage, w_desc, m_keypool_size, keys)); auto desc_spks = desc_spk_man->GetScriptPubKeys(); // Remove the scriptPubKeys from our current set @@ -1856,9 +1855,8 @@ std::optional LegacyScriptPubKeyMan::MigrateToDescriptor() WalletDescriptor w_desc(std::move(desc), 0, 0, chain_counter, 0); // Make the DescriptorScriptPubKeyMan and get the scriptPubKeys - auto desc_spk_man = std::unique_ptr(new DescriptorScriptPubKeyMan(m_storage, w_desc, m_keypool_size)); - desc_spk_man->AddDescriptorKey(master_key.key, master_key.key.GetPubKey()); - desc_spk_man->TopUp(); + keys.keys.emplace(master_key.key.GetPubKey().GetID(), master_key.key); + auto desc_spk_man = std::unique_ptr(new DescriptorScriptPubKeyMan(m_storage, w_desc, m_keypool_size, keys)); auto desc_spks = desc_spk_man->GetScriptPubKeys(); // Remove the scriptPubKeys from our current set @@ -1917,16 +1915,15 @@ std::optional LegacyScriptPubKeyMan::MigrateToDescriptor() desc->Expand(0, provider, desc_spks, provider); } else { // Make the DescriptorScriptPubKeyMan and get the scriptPubKeys - WalletDescriptor w_desc(std::move(desc), creation_time, 0, 0, 0); - auto desc_spk_man = std::unique_ptr(new DescriptorScriptPubKeyMan(m_storage, w_desc, m_keypool_size)); for (const auto& keyid : privkeyids) { CKey key; if (!GetKey(keyid, key)) { continue; } - desc_spk_man->AddDescriptorKey(key, key.GetPubKey()); + keys.keys.emplace(key.GetPubKey().GetID(), key); } - desc_spk_man->TopUp(); + WalletDescriptor w_desc(std::move(desc), creation_time, 0, 0, 0); + auto desc_spk_man = std::unique_ptr(new DescriptorScriptPubKeyMan(m_storage, w_desc, m_keypool_size, keys)); auto desc_spks_set = desc_spk_man->GetScriptPubKeys(); desc_spks.insert(desc_spks.end(), desc_spks_set.begin(), desc_spks_set.end()); @@ -2003,6 +2000,14 @@ bool LegacyScriptPubKeyMan::DeleteRecords() return batch.EraseRecords(DBKeys::LEGACY_TYPES); } +DescriptorScriptPubKeyMan::DescriptorScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const FlatSigningProvider& provider) + : ScriptPubKeyMan(storage), + m_keypool_size(keypool_size), + m_wallet_descriptor(descriptor) +{ + UpdateWithSigningProvider(provider); +} + DescriptorScriptPubKeyMan::DescriptorScriptPubKeyMan(WalletStorage& storage, const uint256& id, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys) : ScriptPubKeyMan(storage), m_map_keys(keys), @@ -2744,7 +2749,7 @@ void DescriptorScriptPubKeyMan::UpgradeDescriptorCache() } } -void DescriptorScriptPubKeyMan::UpdateWalletDescriptor(WalletDescriptor& descriptor) +void DescriptorScriptPubKeyMan::UpdateWalletDescriptor(WalletDescriptor& descriptor, const FlatSigningProvider& provider) { LOCK(cs_desc_man); std::string error; @@ -2757,9 +2762,24 @@ void DescriptorScriptPubKeyMan::UpdateWalletDescriptor(WalletDescriptor& descrip m_max_cached_index = -1; m_wallet_descriptor = descriptor; + UpdateWithSigningProvider(provider); NotifyFirstKeyTimeChanged(this, m_wallet_descriptor.creation_time); } +void DescriptorScriptPubKeyMan::UpdateWithSigningProvider(const FlatSigningProvider& signing_provider) +{ + // Add the private keys to the descriptor + for (const auto& entry : signing_provider.keys) { + const CKey& key = entry.second; + AddDescriptorKey(key, key.GetPubKey()); + } + + // Top up key pool, to generate scriptPubKeys + if (!TopUp()) { + throw std::runtime_error("Could not top up scriptPubKeys"); + } +} + bool DescriptorScriptPubKeyMan::CanUpdateToWalletDescriptor(const WalletDescriptor& descriptor, std::string& error) { LOCK(cs_desc_man); diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h index 0e1e9986c8f..c3aa0779289 100644 --- a/src/wallet/scriptpubkeyman.h +++ b/src/wallet/scriptpubkeyman.h @@ -590,6 +590,9 @@ private: void SetCache(const DescriptorCache& cache); + void AddDescriptorKey(const CKey& key, const CPubKey &pubkey); + void UpdateWithSigningProvider(const FlatSigningProvider& signing_provider); + protected: WalletDescriptor m_wallet_descriptor GUARDED_BY(cs_desc_man); @@ -598,11 +601,7 @@ protected: 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), - m_wallet_descriptor(descriptor) - {} + DescriptorScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const FlatSigningProvider& provider); //! Create a DescriptorScriptPubKeyMan from existing data (i.e. during loading) DescriptorScriptPubKeyMan(WalletStorage& storage, const uint256& id, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys); DescriptorScriptPubKeyMan(WalletStorage& storage, int64_t keypool_size) @@ -659,9 +658,8 @@ public: uint256 GetID() const override; bool HasWalletDescriptor(const WalletDescriptor& desc) const; - void UpdateWalletDescriptor(WalletDescriptor& descriptor); + void UpdateWalletDescriptor(WalletDescriptor& descriptor, const FlatSigningProvider& provider); bool CanUpdateToWalletDescriptor(const WalletDescriptor& descriptor, std::string& error); - void AddDescriptorKey(const CKey& key, const CPubKey &pubkey); void WriteDescriptor(); WalletDescriptor GetWalletDescriptor() const EXCLUSIVE_LOCKS_REQUIRED(cs_desc_man); diff --git a/src/wallet/test/fuzz/scriptpubkeyman.cpp b/src/wallet/test/fuzz/scriptpubkeyman.cpp index 63d09d66cae..c082289c9d9 100644 --- a/src/wallet/test/fuzz/scriptpubkeyman.cpp +++ b/src/wallet/test/fuzz/scriptpubkeyman.cpp @@ -145,15 +145,6 @@ FUZZ_TARGET(scriptpubkeyman, .init = initialize_spkm) } } }, - [&] { - CKey key{ConsumePrivateKey(fuzzed_data_provider, /*compressed=*/fuzzed_data_provider.ConsumeBool())}; - if (!key.IsValid()) { - good_data = false; - return; - } - spk_manager->AddDescriptorKey(key, key.GetPubKey()); - spk_manager->TopUp(); - }, [&] { std::string descriptor; (void)spk_manager->GetDescriptorString(descriptor, /*priv=*/fuzzed_data_provider.ConsumeBool()); diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index 8527b7d0250..ec1472de369 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -3845,9 +3845,9 @@ ScriptPubKeyMan* CWallet::AddWalletDescriptor(WalletDescriptor& desc, const Flat auto spk_man = GetDescriptorScriptPubKeyMan(desc); if (spk_man) { WalletLogPrintf("Update existing descriptor: %s\n", desc.descriptor->ToString()); - spk_man->UpdateWalletDescriptor(desc); + spk_man->UpdateWalletDescriptor(desc, signing_provider); } else { - auto new_spk_man = std::unique_ptr(new DescriptorScriptPubKeyMan(*this, desc, m_keypool_size)); + auto new_spk_man = std::unique_ptr(new DescriptorScriptPubKeyMan(*this, desc, m_keypool_size, signing_provider)); spk_man = new_spk_man.get(); // Save the descriptor to memory @@ -3855,18 +3855,6 @@ ScriptPubKeyMan* CWallet::AddWalletDescriptor(WalletDescriptor& desc, const Flat AddScriptPubKeyMan(id, std::move(new_spk_man)); } - // Add the private keys to the descriptor - for (const auto& entry : signing_provider.keys) { - const CKey& key = entry.second; - spk_man->AddDescriptorKey(key, key.GetPubKey()); - } - - // Top up key pool, the manager will generate new scriptPubKeys internally - if (!spk_man->TopUp()) { - WalletLogPrintf("Could not top up scriptPubKeys\n"); - return nullptr; - } - // Apply the label if necessary // Note: we disable labels for ranged descriptors if (!desc.descriptor->IsRange()) { From 7b817279e3c385dbdf65364f3fe0e9b49e9fbaa0 Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Tue, 20 Feb 2024 11:48:44 -0500 Subject: [PATCH 6/7] wallet: Construct ExternalSignerSPKM with the new descriptor Instead of constructing then setting the descriptor with SetupDescriptor, just pass in that descriptor to the constructor. --- src/wallet/external_signer_scriptpubkeyman.cpp | 4 ++-- src/wallet/external_signer_scriptpubkeyman.h | 15 +++++---------- src/wallet/wallet.cpp | 3 +-- 3 files changed, 8 insertions(+), 14 deletions(-) diff --git a/src/wallet/external_signer_scriptpubkeyman.cpp b/src/wallet/external_signer_scriptpubkeyman.cpp index a71f8f9fbc0..dcbdd32a122 100644 --- a/src/wallet/external_signer_scriptpubkeyman.cpp +++ b/src/wallet/external_signer_scriptpubkeyman.cpp @@ -16,7 +16,8 @@ #include namespace wallet { -bool ExternalSignerScriptPubKeyMan::SetupDescriptor(WalletBatch& batch, std::unique_ptr desc) +ExternalSignerScriptPubKeyMan::ExternalSignerScriptPubKeyMan(WalletStorage& storage, WalletBatch& batch, int64_t keypool_size, std::unique_ptr desc) + : DescriptorScriptPubKeyMan(storage, keypool_size) { LOCK(cs_desc_man); assert(m_storage.IsWalletFlagSet(WALLET_FLAG_DESCRIPTORS)); @@ -37,7 +38,6 @@ bool ExternalSignerScriptPubKeyMan::SetupDescriptor(WalletBatch& batch, std::uni TopUpWithDB(batch); m_storage.UnsetBlankWalletFlag(batch); - return true; } ExternalSigner ExternalSignerScriptPubKeyMan::GetExternalSigner() { diff --git a/src/wallet/external_signer_scriptpubkeyman.h b/src/wallet/external_signer_scriptpubkeyman.h index c9e6f39c221..cfdc4715e49 100644 --- a/src/wallet/external_signer_scriptpubkeyman.h +++ b/src/wallet/external_signer_scriptpubkeyman.h @@ -12,18 +12,13 @@ namespace wallet { class ExternalSignerScriptPubKeyMan : public DescriptorScriptPubKeyMan { - public: - ExternalSignerScriptPubKeyMan(WalletStorage& storage, const uint256& id, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys) +public: + //! Create an ExternalSPKM from existing wallet data + ExternalSignerScriptPubKeyMan(WalletStorage& storage, const uint256& id, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys) : DescriptorScriptPubKeyMan(storage, id, descriptor, keypool_size, keys, ckeys) {} - ExternalSignerScriptPubKeyMan(WalletStorage& storage, int64_t keypool_size) - : DescriptorScriptPubKeyMan(storage, keypool_size) - {} - - /** Provide a descriptor at setup time - * Returns false if already setup or setup fails, true if setup is successful - */ - bool SetupDescriptor(WalletBatch& batch, std::unique_ptrdesc); + //! Create a new ExternalSPKM from just a descriptor + ExternalSignerScriptPubKeyMan(WalletStorage& storage, WalletBatch& batch, int64_t keypool_size, std::unique_ptr desc); static ExternalSigner GetExternalSigner(); diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index ec1472de369..383f1ffcd22 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -3696,8 +3696,7 @@ void CWallet::SetupDescriptorScriptPubKeyMans() continue; } OutputType t = *desc->GetOutputType(); - auto spk_manager = std::unique_ptr(new ExternalSignerScriptPubKeyMan(*this, m_keypool_size)); - spk_manager->SetupDescriptor(batch, std::move(desc)); + auto spk_manager = std::unique_ptr(new ExternalSignerScriptPubKeyMan(*this, batch, m_keypool_size, std::move(desc))); uint256 id = spk_manager->GetID(); AddScriptPubKeyMan(id, std::move(spk_manager)); AddActiveScriptPubKeyManWithDb(batch, id, t, internal); From e432a938e81089e3966f128671cb74326830a713 Mon Sep 17 00:00:00 2001 From: Ava Chow Date: Tue, 20 Feb 2024 11:48:46 -0500 Subject: [PATCH 7/7] wallet: Setup new autogenerated descriptors on construction Instead of having a caller use SetupDescriptorGeneration, just have a constructor that takes those arguments and sets up the descriptor with the autogenerated key. --- src/wallet/scriptpubkeyman.cpp | 7 +++++++ src/wallet/scriptpubkeyman.h | 17 ++++++++++------- src/wallet/wallet.cpp | 3 +-- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp index 37159227400..d2e79d07073 100644 --- a/src/wallet/scriptpubkeyman.cpp +++ b/src/wallet/scriptpubkeyman.cpp @@ -2024,6 +2024,13 @@ DescriptorScriptPubKeyMan::DescriptorScriptPubKeyMan(WalletStorage& storage, con SetCache(m_wallet_descriptor.cache); } +DescriptorScriptPubKeyMan::DescriptorScriptPubKeyMan(WalletStorage& storage, WalletBatch& batch, int64_t keypool_size, const CExtKey& master_key, OutputType addr_type, bool internal) + : ScriptPubKeyMan(storage), + m_keypool_size(keypool_size) +{ + SetupDescriptorGeneration(batch, master_key, addr_type, internal); +} + 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 diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h index c3aa0779289..4619871b086 100644 --- a/src/wallet/scriptpubkeyman.h +++ b/src/wallet/scriptpubkeyman.h @@ -593,7 +593,15 @@ private: void AddDescriptorKey(const CKey& key, const CPubKey &pubkey); void UpdateWithSigningProvider(const FlatSigningProvider& signing_provider); + //! Setup descriptors based on the given CExtkey + bool SetupDescriptorGeneration(WalletBatch& batch, const CExtKey& master_key, OutputType addr_type, bool internal); + protected: + DescriptorScriptPubKeyMan(WalletStorage& storage, int64_t keypool_size) + : ScriptPubKeyMan(storage), + m_keypool_size(keypool_size) + {} + WalletDescriptor m_wallet_descriptor GUARDED_BY(cs_desc_man); //! Same as 'TopUp' but designed for use within a batch transaction context @@ -604,10 +612,8 @@ public: DescriptorScriptPubKeyMan(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const FlatSigningProvider& provider); //! Create a DescriptorScriptPubKeyMan from existing data (i.e. during loading) DescriptorScriptPubKeyMan(WalletStorage& storage, const uint256& id, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys); - DescriptorScriptPubKeyMan(WalletStorage& storage, int64_t keypool_size) - : ScriptPubKeyMan(storage), - m_keypool_size(keypool_size) - {} + //! Create an automatically generated DescriptorScriptPubKeyMan + DescriptorScriptPubKeyMan(WalletStorage& storage, WalletBatch& batch, int64_t keypool_size, const CExtKey& master_key, OutputType addr_type, bool internal); mutable RecursiveMutex cs_desc_man; @@ -630,9 +636,6 @@ public: bool IsHDEnabled() const override; - //! Setup descriptors based on the given CExtkey - bool SetupDescriptorGeneration(WalletBatch& batch, const CExtKey& master_key, OutputType addr_type, bool internal); - bool HavePrivateKeys() const override; bool HasPrivKey(const CKeyID& keyid) const EXCLUSIVE_LOCKS_REQUIRED(cs_desc_man); //! Retrieve the particular key if it is available. Returns nullopt if the key is not in the wallet, or if the wallet is locked. diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index 383f1ffcd22..65662a25ac9 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -3619,7 +3619,7 @@ void CWallet::LoadDescriptorScriptPubKeyMan(uint256 id, WalletDescriptor& desc, DescriptorScriptPubKeyMan& CWallet::SetupDescriptorScriptPubKeyMan(WalletBatch& batch, const CExtKey& master_key, const OutputType& output_type, bool internal) { AssertLockHeld(cs_wallet); - auto spk_manager = std::unique_ptr(new DescriptorScriptPubKeyMan(*this, m_keypool_size)); + auto spk_manager = std::unique_ptr(new DescriptorScriptPubKeyMan(*this, batch, m_keypool_size, master_key, output_type, internal)); if (IsCrypted()) { if (IsLocked()) { throw std::runtime_error(std::string(__func__) + ": Wallet is locked, cannot setup new descriptors"); @@ -3628,7 +3628,6 @@ DescriptorScriptPubKeyMan& CWallet::SetupDescriptorScriptPubKeyMan(WalletBatch& throw std::runtime_error(std::string(__func__) + ": Could not encrypt new descriptors"); } } - spk_manager->SetupDescriptorGeneration(batch, master_key, output_type, internal); DescriptorScriptPubKeyMan* out = spk_manager.get(); uint256 id = spk_manager->GetID(); AddScriptPubKeyMan(id, std::move(spk_manager));