From 386a994b853bc5b3a2ed0d812673465b8ffa4849 Mon Sep 17 00:00:00 2001 From: Andrew Chow Date: Mon, 7 Oct 2019 14:11:34 -0400 Subject: [PATCH] Key pool: Change ReturnDestination interface to take address instead of key In order for ScriptPubKeyMan to be generic and work with future ScriptPubKeyMans, ScriptPubKeyMan::ReturnDestination is changed to take a CTxDestination instead of a CPubKey. Since LegacyScriptPubKeyMan still deals with keys internally, a new map m_reserved_key_to_index is added in order to track the keypool indexes that have been reserved. The CPubKey argument of KeepDestination is also removed so that it is more generic. Instead of taking a CPubKey or a CTxDestination, we just use the nIndex given to find the pubkey. --- src/wallet/scriptpubkeyman.cpp | 16 ++++++++++++---- src/wallet/scriptpubkeyman.h | 10 ++++++---- src/wallet/wallet.cpp | 8 ++------ src/wallet/wallet.h | 2 -- 4 files changed, 20 insertions(+), 16 deletions(-) diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp index 4986871fb..e5b45a81a 100644 --- a/src/wallet/scriptpubkeyman.cpp +++ b/src/wallet/scriptpubkeyman.cpp @@ -1087,16 +1087,20 @@ void LegacyScriptPubKeyMan::AddKeypoolPubkeyWithDB(const CPubKey& pubkey, const m_pool_key_to_index[pubkey.GetID()] = index; } -void LegacyScriptPubKeyMan::KeepDestination(int64_t nIndex, const OutputType& type, const CPubKey& pubkey) +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::ReturnDestination(int64_t nIndex, bool fInternal, const CPubKey& pubkey) +void LegacyScriptPubKeyMan::ReturnDestination(int64_t nIndex, bool fInternal, const CTxDestination&) { // Return to key pool { @@ -1108,7 +1112,9 @@ void LegacyScriptPubKeyMan::ReturnDestination(int64_t nIndex, bool fInternal, co } 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); @@ -1130,7 +1136,7 @@ bool LegacyScriptPubKeyMan::GetKeyFromPool(CPubKey& result, const OutputType typ result = GenerateNewKey(batch, internal); return true; } - KeepDestination(nIndex, type, keypool.vchPubKey); + KeepDestination(nIndex, type); result = keypool.vchPubKey; } return true; @@ -1175,6 +1181,8 @@ bool LegacyScriptPubKeyMan::ReserveKeyFromKeyPool(int64_t& nIndex, CKeyPool& key throw std::runtime_error(std::string(__func__) + ": 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/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h index bcc62ece1..6ed9a4787 100644 --- a/src/wallet/scriptpubkeyman.h +++ b/src/wallet/scriptpubkeyman.h @@ -151,8 +151,8 @@ public: virtual isminetype IsMine(const CScript& script) const { return ISMINE_NO; } virtual bool GetReservedDestination(const OutputType type, bool internal, CTxDestination& address, int64_t& index, CKeyPool& keypool) { return false; } - virtual void KeepDestination(int64_t index, const OutputType& type, const CPubKey& pubkey) {} - virtual void ReturnDestination(int64_t index, bool internal, const CPubKey& pubkey) {} + virtual void KeepDestination(int64_t index, const OutputType& type) {} + virtual void ReturnDestination(int64_t index, bool internal, const CTxDestination& addr) {} virtual bool TopUp(unsigned int size = 0) { return false; } @@ -246,6 +246,8 @@ private: 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, const OutputType type, bool internal = false); @@ -274,8 +276,8 @@ public: bool EncryptKeys(CKeyingMaterial& vMasterKeyIn); bool GetReservedDestination(const OutputType type, bool internal, CTxDestination& address, int64_t& index, CKeyPool& keypool) override; - void KeepDestination(int64_t index, const OutputType& type, const CPubKey& pubkey) override; - void ReturnDestination(int64_t index, bool internal, const CPubKey& pubkey) override; + void KeepDestination(int64_t index, const OutputType& type) override; + void ReturnDestination(int64_t index, bool internal, const CTxDestination&) override; bool TopUp(unsigned int size = 0) override; diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp index 2bdec16f2..e50dd1c12 100644 --- a/src/wallet/wallet.cpp +++ b/src/wallet/wallet.cpp @@ -3305,10 +3305,8 @@ bool ReserveDestination::GetReservedDestination(CTxDestination& dest, bool inter if (!m_spk_man->GetReservedDestination(type, internal, address, nIndex, keypool)) { return false; } - vchPubKey = keypool.vchPubKey; fInternal = keypool.fInternal; } - assert(vchPubKey.IsValid()); dest = address; return true; } @@ -3316,20 +3314,18 @@ bool ReserveDestination::GetReservedDestination(CTxDestination& dest, bool inter void ReserveDestination::KeepDestination() { if (nIndex != -1) { - m_spk_man->KeepDestination(nIndex, type, vchPubKey); + 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(); } diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h index f6cbe6f6d..db5788ebb 100644 --- a/src/wallet/wallet.h +++ b/src/wallet/wallet.h @@ -146,8 +146,6 @@ protected: 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