From 163d31861c14d37c0a407d70d10f847050b67e37 Mon Sep 17 00:00:00 2001
From: Kittywhiskers Van Gogh <63189531+kwvg@users.noreply.github.com>
Date: Wed, 17 Jul 2024 16:31:30 +0000
Subject: [PATCH] wallet: unify HD chain generation in LegacyScriptPubKeyMan

most of the steps are overlapping except for the encryption sections,
which we can make contingent on supplying of the keying material. `res`
isn't used anywhere since we plan on hard crashing if they fail anyway,
so if we got that far, we've already succeeded.

we will validate vMasterKey's size before starting the encryption
routine to prevent a value outside of our control from hard-crashing
the client.
---
 src/wallet/scriptpubkeyman.cpp | 64 +++++++++++++---------------------
 src/wallet/scriptpubkeyman.h   |  3 +-
 src/wallet/wallet.cpp          |  5 +--
 3 files changed, 27 insertions(+), 45 deletions(-)

diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index 678022e1f2..631fcac566 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -377,56 +377,42 @@ void LegacyScriptPubKeyMan::UpgradeKeyMetadata()
     }
 }
 
-void LegacyScriptPubKeyMan::GenerateNewCryptedHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, CKeyingMaterial vMasterKey)
-{
-    assert(!m_storage.IsWalletFlagSet(WALLET_FLAG_DISABLE_PRIVATE_KEYS));
-
-
-    CHDChain hdChainTmp;
-
-    // NOTE: an empty mnemonic means "generate a new one for me"
-    // NOTE: default mnemonic passphrase is an empty string
-    if (!hdChainTmp.SetMnemonic(secureMnemonic, secureMnemonicPassphrase, true)) {
-        throw std::runtime_error(std::string(__func__) + ": SetMnemonic failed");
-    }
-
-    // add default account
-    hdChainTmp.AddAccount();
-
-    // We need to safe chain for validation further
-    CHDChain hdChainPrev = hdChainTmp;
-    bool res = EncryptHDChain(vMasterKey, hdChainTmp);
-    assert(res);
-    res = LoadHDChain(hdChainTmp);
-    assert(res);
-
-    CHDChain hdChainCrypted;
-    res = GetHDChain(hdChainCrypted);
-    assert(res);
-
-    // ids should match, seed hashes should not
-    assert(hdChainPrev.GetID() == hdChainCrypted.GetID());
-    assert(hdChainPrev.GetSeedHash() != hdChainCrypted.GetSeedHash());
-
-    if (!AddHDChainSingle(hdChainCrypted)) {
-        throw std::runtime_error(std::string(__func__) + ": AddHDChainSingle failed");
-    }
-}
-
-void LegacyScriptPubKeyMan::GenerateNewHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase)
+void LegacyScriptPubKeyMan::GenerateNewHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, std::optional<CKeyingMaterial> vMasterKeyOpt)
 {
     assert(!m_storage.IsWalletFlagSet(WALLET_FLAG_DISABLE_PRIVATE_KEYS));
     CHDChain newHdChain;
 
     // NOTE: an empty mnemonic means "generate a new one for me"
     // NOTE: default mnemonic passphrase is an empty string
-    if (!newHdChain.SetMnemonic(secureMnemonic, secureMnemonicPassphrase, true)) {
+    if (!newHdChain.SetMnemonic(secureMnemonic, secureMnemonicPassphrase, /* fUpdateID = */ true)) {
         throw std::runtime_error(std::string(__func__) + ": SetMnemonic failed");
     }
 
-    // add default account
+    // Add default account
     newHdChain.AddAccount();
 
+    // Encryption routine if vMasterKey has been supplied
+    if (vMasterKeyOpt.has_value()) {
+        auto vMasterKey = vMasterKeyOpt.value();
+        if (vMasterKey.size() != WALLET_CRYPTO_KEY_SIZE) {
+            throw std::runtime_error(strprintf("%s : invalid vMasterKey size, got %zd (expected %lld)", __func__, vMasterKey.size(), WALLET_CRYPTO_KEY_SIZE));
+        }
+
+        // Maintain an unencrypted copy of the chain for sanity checking
+        CHDChain prevHdChain{newHdChain};
+
+        bool res = EncryptHDChain(vMasterKey, newHdChain);
+        assert(res);
+        res = LoadHDChain(newHdChain);
+        assert(res);
+        res = GetHDChain(newHdChain);
+        assert(res);
+
+        // IDs should match, seed hashes should not
+        assert(prevHdChain.GetID() == newHdChain.GetID());
+        assert(prevHdChain.GetSeedHash() != newHdChain.GetSeedHash());
+    }
+
     if (!AddHDChainSingle(newHdChain)) {
         throw std::runtime_error(std::string(__func__) + ": AddHDChainSingle failed");
     }
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index 9c639b73eb..5ebb757628 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -463,8 +463,7 @@ public:
     bool GetDecryptedHDChain(CHDChain& hdChainRet);
 
     /* Generates a new HD chain */
-    void GenerateNewCryptedHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, CKeyingMaterial vMasterKey);
-    void GenerateNewHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase);
+    void GenerateNewHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, std::optional<CKeyingMaterial> vMasterKey = std::nullopt);
 
     /**
      * Explicitly make the wallet learn the related scripts for outputs to the
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 29d6104270..68e1919955 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -5677,16 +5677,13 @@ bool CWallet::GenerateNewHDChainEncrypted(const SecureString& secureMnemonic, co
 
     }
 
-    spk_man->GenerateNewCryptedHDChain(secureMnemonic, secureMnemonicPassphrase, vMasterKey);
+    spk_man->GenerateNewHDChain(secureMnemonic, secureMnemonicPassphrase, vMasterKey);
 
     Lock();
     if (!Unlock(secureWalletPassphrase)) {
         // this should never happen
         throw std::runtime_error(std::string(__func__) + ": Unlock failed");
     }
-    if (!spk_man->NewKeyPool()) {
-        throw std::runtime_error(std::string(__func__) + ": NewKeyPool failed");
-    }
     return true;
 }