Merge #17373: wallet: Various fixes and cleanup to keypool handling in LegacyScriptPubKeyMan and CWallet

886f1731bec4393dd342403ac34069a3a4f95eea Key pool: Fix omitted pre-split count in GetKeyPoolSize (Andrew Chow)
386a994b853bc5b3a2ed0d812673465b8ffa4849 Key pool: Change ReturnDestination interface to take address instead of key (Andrew Chow)
ba41aa4969169cd73d6b4f57444ed7d8d875de10 Key pool: Move LearnRelated and GetDestination calls (Andrew Chow)
65833a74076cddf986037c6eb3b29a9b9dbe31c5 Add OutputType and CPubKey parameters to KeepDestination (Andrew Chow)
9fcf8ce7ae02bf170b9bf0c2887fd709d752cbf7 Rename Keep/ReturnKey to Keep/ReturnDestination and remove the wrapper (Andrew Chow)
596f6460f9fd8273665c8754ccd673d93a4f25f0 Key pool: Move CanGetAddresses call (Andrew Chow)

Pull request description:

  * The `pwallet->CanGetAddresses()` call in `ReserveDestination::GetReservedDestination` to `LegacyScriptPubKeyMan::GetReservedDestination` so that the sanity check results in a failure when a `ScriptPubKeyMan` individually cannot get a destination, not when any of the `ScriptPubKeyMan`s can't.
  * `ScriptPubKeyMan::GetReservedDestination` is changed to return the destination so that future `ScriptPubKeyMan`s can return destinations constructed in other ways. This is implemented for `LegacyScriptPubKeyMan` by moving key-to-destination code from `CWallet` to `LegacyScriptPubKeyMan`
  * In order for `ScriptPubKeyMan` to be generic and work with future `ScriptPubKeyMan`s, `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.
  * A bug is fixed in how the total keypool size is calculated as it was omitting `set_pre_split_keypool` which is a bug.

  Split from #17261

ACKs for top commit:
  ryanofsky:
    Code review ACK 886f1731bec4393dd342403ac34069a3a4f95eea. Only change is moving earlier fix to a better commit (same end result).
  promag:
    Code review ACK 886f1731bec4393dd342403ac34069a3a4f95eea.
  instagibbs:
    code review re-ACK https://github.com/bitcoin/bitcoin/pull/17373/commits/886f1731bec4393dd342403ac34069a3a4f95eea
  Sjors:
    Code review re-ACK 886f1731bec4393dd342403ac34069a3a4f95eea

Tree-SHA512: f4be290759f63fdc920d5c02bd0d09acc4b06a5f053787d4afcd3c921b2e35d2bd97617fadae015da853dc189f559fb8d2c6e58d53e4cabfac9af151cd97ad19

4 files changed, 36 insertions, 43 deletions

diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index 3eaaf3786c..2b9e8fbf2a 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -18,7 +18,7 @@ bool LegacyScriptPubKeyMan::GetNewDestination(const OutputType type, CTxDestinat
 
     // 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;
     }
@@ -262,24 +262,19 @@ bool LegacyScriptPubKeyMan::EncryptKeys(CKeyingMaterial& vMasterKeyIn)
     return true;
 }
 
-bool LegacyScriptPubKeyMan::GetReservedDestination(const OutputType type, bool internal, int64_t& index, CKeyPool& keypool)
+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);
@@ -460,7 +455,7 @@ size_t LegacyScriptPubKeyMan::KeypoolCountExternalKeys()
 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
@@ -1092,15 +1087,20 @@ void LegacyScriptPubKeyMan::AddKeypoolPubkeyWithDB(const CPubKey& pubkey, const
     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
     {
@@ -1112,13 +1112,15 @@ void LegacyScriptPubKeyMan::ReturnKey(int64_t nIndex, bool fInternal, const CPub
         } 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;
@@ -1134,7 +1136,7 @@ bool LegacyScriptPubKeyMan::GetKeyFromPool(CPubKey& result, bool internal)
             result = GenerateNewKey(batch, internal);
             return true;
         }
-        KeepKey(nIndex);
+        KeepDestination(nIndex, type);
         result = keypool.vchPubKey;
     }
     return true;
@@ -1179,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 4f17156792..6ed9a4787a 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -150,9 +150,9 @@ public:
     virtual bool GetNewDestination(const OutputType type, CTxDestination& dest, std::string& error) { return false; }
     virtual isminetype IsMine(const CScript& script) const { return ISMINE_NO; }
 
-    virtual bool GetReservedDestination(const OutputType type, bool internal, int64_t& index, CKeyPool& keypool) { return false; }
-    virtual void KeepDestination(int64_t index) {}
-    virtual void ReturnDestination(int64_t index, bool internal, const CPubKey& pubkey) {}
+    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) {}
+    virtual void ReturnDestination(int64_t index, bool internal, const CTxDestination& addr) {}
 
     virtual bool TopUp(unsigned int size = 0) { return false; }
 
@@ -246,9 +246,11 @@ private:
     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, 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
@@ -266,9 +268,6 @@ private:
      */
     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;
     isminetype IsMine(const CScript& script) const override;
@@ -276,9 +275,9 @@ public:
     //! will encrypt previously unencrypted keys
     bool EncryptKeys(CKeyingMaterial& vMasterKeyIn);
 
-    bool GetReservedDestination(const OutputType type, bool internal, int64_t& index, CKeyPool& keypool) override;
-    void KeepDestination(int64_t index) override;
-    void ReturnDestination(int64_t index, bool internal, const CPubKey& pubkey) override;
+    bool GetReservedDestination(const OutputType type, bool internal, 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 CTxDestination&) override;
 
     bool TopUp(unsigned int size = 0) override;
 
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 0a8ce7caf1..abee497c1d 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -3295,21 +3295,15 @@ bool ReserveDestination::GetReservedDestination(CTxDestination& dest, bool inter
         return false;
     }
 
-    if (!pwallet->CanGetAddresses(internal)) {
-        return false;
-    }
 
     if (nIndex == -1)
     {
         CKeyPool keypool;
-        if (!m_spk_man->GetReservedDestination(type, internal, nIndex, keypool)) {
+        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;
 }
@@ -3317,21 +3311,18 @@ bool ReserveDestination::GetReservedDestination(CTxDestination& dest, bool inter
 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();
 }
 
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 51df4a9a8b..b02e092f0a 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -139,12 +139,11 @@ class ReserveDestination
 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