Refactor: Move SetupGeneration code out of CWallet

This commit does not change behavior.

3 files changed, 39 insertions, 16 deletions

diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index 9d90edeea1..14e4163be6 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -339,6 +339,19 @@ void LegacyScriptPubKeyMan::UpgradeKeyMetadata()
     }
 }
 
+bool LegacyScriptPubKeyMan::SetupGeneration(bool force)
+{
+    if ((CanGenerateKeys() && !force) || m_storage.IsLocked()) {
+        return false;
+    }
+
+    SetHDSeed(GenerateNewSeed());
+    if (!NewKeyPool()) {
+        return false;
+    }
+    return true;
+}
+
 bool LegacyScriptPubKeyMan::IsHDEnabled() const
 {
     return !hdChain.seed_id.IsNull();
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index 927401be20..0e38cbf657 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -159,6 +159,12 @@ public:
     //! Mark unused addresses as being used
     virtual void MarkUnusedAddresses(const CScript& script) {}
 
+    /** Sets up the key generation stuff, i.e. generates new HD seeds and sets them as active.
+      * Returns false if already setup or setup fails, true if setup is successful
+      * Set force=true to make it re-setup if already setup, used for upgrades
+      */
+    virtual bool SetupGeneration(bool force = false) { return false; }
+
     /* Returns true if HD is enabled */
     virtual bool IsHDEnabled() const { return false; }
 
@@ -276,6 +282,8 @@ public:
 
     bool IsHDEnabled() const override;
 
+    bool SetupGeneration(bool force = false) override;
+
     bool Upgrade(int prev_version, std::string& error) override;
 
     bool HavePrivateKeys() const override;
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index e6ef932ac2..1b27077d67 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -210,9 +210,14 @@ WalletCreationStatus CreateWallet(interfaces::Chain& chain, const SecureString&
             }
 
             // Set a seed for the wallet
-            CPubKey master_pub_key = wallet->m_spk_man->GenerateNewSeed();
-            wallet->m_spk_man->SetHDSeed(master_pub_key);
-            wallet->m_spk_man->NewKeyPool();
+            {
+                if (auto spk_man = wallet->m_spk_man.get()) {
+                    if (!spk_man->SetupGeneration()) {
+                        error = "Unable to generate initial keys";
+                        return WalletCreationStatus::CREATION_FAILED;
+                    }
+                }
+            }
 
             // Relock the wallet
             wallet->Lock();
@@ -565,11 +570,11 @@ bool CWallet::EncryptWallet(const SecureString& strWalletPassphrase)
         Unlock(strWalletPassphrase);
 
         // if we are using HD, replace the HD seed with a new one
-        if (m_spk_man->IsHDEnabled()) {
-            m_spk_man->SetHDSeed(m_spk_man->GenerateNewSeed());
+        if (auto spk_man = m_spk_man.get()) {
+            if (spk_man->IsHDEnabled()) {
+                spk_man->SetupGeneration(true);
+            }
         }
-
-        m_spk_man->NewKeyPool();
         Lock();
 
         // Need to completely rewrite the wallet file; if we don't, bdb might keep
@@ -3630,15 +3635,12 @@ std::shared_ptr<CWallet> CWallet::CreateWalletFromFile(interfaces::Chain& chain,
 
         walletInstance->SetWalletFlags(wallet_creation_flags, false);
         if (!(wallet_creation_flags & (WALLET_FLAG_DISABLE_PRIVATE_KEYS | WALLET_FLAG_BLANK_WALLET))) {
-            // generate a new seed
-            CPubKey seed = walletInstance->m_spk_man->GenerateNewSeed();
-            walletInstance->m_spk_man->SetHDSeed(seed);
-        }
-
-        // Top up the keypool
-        if (walletInstance->m_spk_man->CanGenerateKeys() && !walletInstance->m_spk_man->TopUp()) {
-            error = _("Unable to generate initial keys").translated;
-            return nullptr;
+            if (auto spk_man = walletInstance->m_spk_man.get()) {
+                if (!spk_man->SetupGeneration()) {
+                    error = _("Unable to generate initial keys").translated;
+                    return nullptr;
+                }
+            }
         }
 
         auto locked_chain = chain.lock();