diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h --- a/src/wallet/scriptpubkeyman.h +++ b/src/wallet/scriptpubkeyman.h @@ -165,12 +165,13 @@ virtual isminetype IsMine(const CScript &script) const { return ISMINE_NO; } virtual bool GetReservedDestination(const OutputType type, bool internal, - int64_t &index, CKeyPool &keypool) { + CTxDestination &address, int64_t &index, + CKeyPool &keypool) { return false; } - virtual void KeepDestination(int64_t index) {} + virtual void KeepDestination(int64_t index, const OutputType &type) {} virtual void ReturnDestination(int64_t index, bool internal, - const CPubKey &pubkey) {} + const CTxDestination &addr) {} virtual bool TopUp(unsigned int size = 0) { return false; } @@ -282,9 +283,13 @@ std::set set_pre_split_keypool GUARDED_BY(cs_wallet); int64_t m_max_keypool_index GUARDED_BY(cs_wallet) = 0; std::map m_pool_key_to_index; + // Tracks keypool indexes to CKeyIDs of keys that have been taken out of the + // keypool but may be returned to it + std::map m_index_to_reserved_key; //! Fetches a key from the keypool - bool GetKeyFromPool(CPubKey &key, bool internal = false); + bool GetKeyFromPool(CPubKey &key, const OutputType type, + bool internal = false); /** * Reserves a key from the keypool and sets nIndex to its index @@ -303,9 +308,6 @@ bool ReserveKeyFromKeyPool(int64_t &nIndex, CKeyPool &keypool, bool fRequestedInternal); - void KeepKey(int64_t nIndex); - void ReturnKey(int64_t nIndex, bool fInternal, const CPubKey &pubkey); - public: bool GetNewDestination(const OutputType type, CTxDestination &dest, std::string &error) override; @@ -315,10 +317,11 @@ bool EncryptKeys(CKeyingMaterial &vMasterKeyIn); bool GetReservedDestination(const OutputType type, bool internal, - int64_t &index, CKeyPool &keypool) override; - void KeepDestination(int64_t index) override; + CTxDestination &address, int64_t &index, + CKeyPool &keypool) override; + void KeepDestination(int64_t index, const OutputType &type) override; void ReturnDestination(int64_t index, bool internal, - const CPubKey &pubkey) override; + const CTxDestination &) override; bool TopUp(unsigned int size = 0) override; diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp --- a/src/wallet/scriptpubkeyman.cpp +++ b/src/wallet/scriptpubkeyman.cpp @@ -20,7 +20,7 @@ // Generate a new key that is added to wallet CPubKey new_key; - if (!GetKeyFromPool(new_key)) { + if (!GetKeyFromPool(new_key, type)) { error = "Error: Keypool ran out, please call keypoolrefill first"; return false; } @@ -215,23 +215,20 @@ bool LegacyScriptPubKeyMan::GetReservedDestination(const OutputType type, bool internal, + CTxDestination &address, int64_t &index, CKeyPool &keypool) { + if (!CanGetAddresses(internal)) { + return false; + } + if (!ReserveKeyFromKeyPool(index, keypool, internal)) { return false; } + address = GetDestinationForKey(keypool.vchPubKey, type); return true; } -void LegacyScriptPubKeyMan::KeepDestination(int64_t index) { - KeepKey(index); -} - -void LegacyScriptPubKeyMan::ReturnDestination(int64_t index, bool internal, - const CPubKey &pubkey) { - ReturnKey(index, internal, pubkey); -} - void LegacyScriptPubKeyMan::MarkUnusedAddresses(const CScript &script) { AssertLockHeld(cs_wallet); // extract addresses and check if they match with an unused keypool key @@ -416,7 +413,8 @@ unsigned int LegacyScriptPubKeyMan::GetKeyPoolSize() const { AssertLockHeld(cs_wallet); - return setInternalKeyPool.size() + setExternalKeyPool.size(); + return setInternalKeyPool.size() + setExternalKeyPool.size() + + set_pre_split_keypool.size(); } int64_t LegacyScriptPubKeyMan::GetTimeFirstKey() const { @@ -1081,15 +1079,21 @@ m_pool_key_to_index[pubkey.GetID()] = index; } -void LegacyScriptPubKeyMan::KeepKey(int64_t nIndex) { +void LegacyScriptPubKeyMan::KeepDestination(int64_t nIndex, + const OutputType &type) { // Remove from key pool. WalletBatch batch(m_storage.GetDatabase()); batch.ErasePool(nIndex); + CPubKey pubkey; + bool have_pk = GetPubKey(m_index_to_reserved_key.at(nIndex), pubkey); + assert(have_pk); + LearnRelatedScripts(pubkey, type); + m_index_to_reserved_key.erase(nIndex); WalletLogPrintf("keypool keep %d\n", nIndex); } -void LegacyScriptPubKeyMan::ReturnKey(int64_t nIndex, bool fInternal, - const CPubKey &pubkey) { +void LegacyScriptPubKeyMan::ReturnDestination(int64_t nIndex, bool fInternal, + const CTxDestination &) { // Return to key pool { LOCK(cs_wallet); @@ -1100,14 +1104,18 @@ } else { setExternalKeyPool.insert(nIndex); } - m_pool_key_to_index[pubkey.GetID()] = nIndex; + CKeyID &pubkey_id = m_index_to_reserved_key.at(nIndex); + m_pool_key_to_index[pubkey_id] = nIndex; + m_index_to_reserved_key.erase(nIndex); NotifyCanGetAddressesChanged(); } WalletLogPrintf("keypool return %d\n", nIndex); } -bool LegacyScriptPubKeyMan::GetKeyFromPool(CPubKey &result, bool internal) { +bool LegacyScriptPubKeyMan::GetKeyFromPool(CPubKey &result, + const OutputType type, + bool internal) { if (!CanGetAddresses(internal)) { return false; } @@ -1125,7 +1133,7 @@ return true; } - KeepKey(nIndex); + KeepDestination(nIndex, type); result = keypool.vchPubKey; return true; @@ -1179,6 +1187,8 @@ ": keypool entry invalid"); } + assert(m_index_to_reserved_key.count(nIndex) == 0); + m_index_to_reserved_key[nIndex] = keypool.vchPubKey.GetID(); m_pool_key_to_index.erase(keypool.vchPubKey.GetID()); WalletLogPrintf("keypool reserve %d\n", nIndex); } diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -141,12 +141,12 @@ protected: //! The wallet to reserve from CWallet *const pwallet; - LegacyScriptPubKeyMan *m_spk_man{nullptr}; + //! The ScriptPubKeyMan to reserve from. Based on type when + //! GetReservedDestination is called + ScriptPubKeyMan *m_spk_man{nullptr}; OutputType const type; //! The index of the address's key in the keypool int64_t nIndex{-1}; - //! The public key for the address - CPubKey vchPubKey; //! The destination CTxDestination address; //! Whether this is from the internal (change output) keypool diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -3695,43 +3695,32 @@ return false; } - if (!pwallet->CanGetAddresses(internal)) { - return false; - } - if (nIndex == -1) { CKeyPool keypool; - if (!m_spk_man->GetReservedDestination(type, internal, nIndex, + if (!m_spk_man->GetReservedDestination(type, internal, address, nIndex, keypool)) { return false; } - - vchPubKey = keypool.vchPubKey; fInternal = keypool.fInternal; } - assert(vchPubKey.IsValid()); - address = GetDestinationForKey(vchPubKey, type); dest = address; return true; } void ReserveDestination::KeepDestination() { if (nIndex != -1) { - m_spk_man->KeepDestination(nIndex); - m_spk_man->LearnRelatedScripts(vchPubKey, type); + m_spk_man->KeepDestination(nIndex, type); } nIndex = -1; - vchPubKey = CPubKey(); address = CNoDestination(); } void ReserveDestination::ReturnDestination() { if (nIndex != -1) { - m_spk_man->ReturnDestination(nIndex, fInternal, vchPubKey); + m_spk_man->ReturnDestination(nIndex, fInternal, address); } nIndex = -1; - vchPubKey = CPubKey(); address = CNoDestination(); }