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.
This commit is contained in:
Ava Chow
2024-02-20 11:48:38 -05:00
parent f713fd1725
commit 80b0c25992
7 changed files with 68 additions and 54 deletions

View File

@@ -21,6 +21,11 @@
using common::PSBTError;
namespace wallet {
std::unique_ptr<ExternalSignerScriptPubKeyMan> ExternalSignerScriptPubKeyMan::LoadFromStorage(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys)
{
return std::unique_ptr<ExternalSignerScriptPubKeyMan>(new ExternalSignerScriptPubKeyMan(storage, descriptor, keypool_size, keys, ckeys));
}
bool ExternalSignerScriptPubKeyMan::SetupDescriptor(WalletBatch& batch, std::unique_ptr<Descriptor> desc)
{
LOCK(cs_desc_man);

View File

@@ -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<ExternalSignerScriptPubKeyMan> 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
*/

View File

@@ -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> DescriptorScriptPubKeyMan::LoadFromStorage(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys)
{
return std::unique_ptr<DescriptorScriptPubKeyMan>(new DescriptorScriptPubKeyMan(storage, descriptor, keypool_size, keys, ckeys));
}
util::Result<CTxDestination> 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<CScript> 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<CScript> 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<unsigned char>& 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);

View File

@@ -166,14 +166,18 @@ static const std::unordered_set<OutputType> LEGACY_OUTPUT_TYPES {
OutputType::BECH32,
};
using KeyMap = std::map<CKeyID, CKey>;
using CryptedKeyMap = std::map<CKeyID, std::pair<CPubKey, std::vector<unsigned char>>>;
using ScriptPubKeyMap = std::map<CScript, int32_t>; // Map of scripts to descriptor range index
using PubKeyMap = std::map<CPubKey, int32_t>; // 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<CScript>;
using WatchKeyMap = std::map<CKeyID, CPubKey>;
using CryptedKeyMap = std::map<CKeyID, std::pair<CPubKey, std::vector<unsigned char>>>;
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<CScript, int32_t>; // Map of scripts to descriptor range index
using PubKeyMap = std::map<CPubKey, int32_t>; // Map of pubkeys involved in scripts to descriptor range index
using CryptedKeyMap = std::map<CKeyID, std::pair<CPubKey, std::vector<unsigned char>>>;
using KeyMap = std::map<CKeyID, CKey>;
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<FlatSigningProvider> 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<DescriptorScriptPubKeyMan> LoadFromStorage(WalletStorage& storage, WalletDescriptor& descriptor, int64_t keypool_size, const KeyMap& keys, const CryptedKeyMap& ckeys);
mutable RecursiveMutex cs_desc_man;
util::Result<CTxDestination> 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<unsigned char>& crypted_key);
bool HasWalletDescriptor(const WalletDescriptor& desc) const;
util::Result<void> UpdateWalletDescriptor(WalletDescriptor& descriptor);
bool CanUpdateToWalletDescriptor(const WalletDescriptor& descriptor, std::string& error);

View File

@@ -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<DescriptorScriptPubKeyMan> 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<ScriptPubKeyMan>(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)

View File

@@ -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

View File

@@ -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<unsigned char> 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;
});