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.

Author: achow101

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);
     }

src/wallet/scriptpubkeyman.h

@@ -151,8 +151,8 @@ class ScriptPubKeyMan
     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 @@ class LegacyScriptPubKeyMan : public ScriptPubKeyMan, public FillableSigningProv
     std::set<int64_t> set_pre_split_keypool GUARDED_BY(cs_wallet);
     int64_t m_max_keypool_index GUARDED_BY(cs_wallet) = 0;
     std::map<CKeyID, int64_t> 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<int64_t, CKeyID> 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 @@ class LegacyScriptPubKeyMan : public ScriptPubKeyMan, public FillableSigningProv
     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;
 

src/wallet/wallet.cpp

@@ -3305,31 +3305,27 @@ 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;
 }
 
 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();
 }
 

src/wallet/wallet.h

@@ -146,8 +146,6 @@ class ReserveDestination
     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