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 | 58 +++++++++++++---------------------
 src/wallet/scriptpubkeyman.h   |  3 +-
 src/wallet/wallet.cpp          |  5 +--
 3 files changed, 24 insertions(+), 42 deletions(-)

diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index 678022e1f22e7..631fcac566cbe 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -377,55 +377,41 @@ void LegacyScriptPubKeyMan::UpgradeKeyMetadata()
     }
 }
 
-void LegacyScriptPubKeyMan::GenerateNewCryptedHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, CKeyingMaterial vMasterKey)
+void LegacyScriptPubKeyMan::GenerateNewHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, std::optional<CKeyingMaterial> vMasterKeyOpt)
 {
     assert(!m_storage.IsWalletFlagSet(WALLET_FLAG_DISABLE_PRIVATE_KEYS));
-
-
-    CHDChain hdChainTmp;
+    CHDChain newHdChain;
 
     // 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)) {
+    if (!newHdChain.SetMnemonic(secureMnemonic, secureMnemonicPassphrase, /* fUpdateID = */ true)) {
         throw std::runtime_error(std::string(__func__) + ": SetMnemonic failed");
     }
 
-    // add default account
-    hdChainTmp.AddAccount();
+    // Add default account
+    newHdChain.AddAccount();
 
-    // We need to safe chain for validation further
-    CHDChain hdChainPrev = hdChainTmp;
-    bool res = EncryptHDChain(vMasterKey, hdChainTmp);
-    assert(res);
-    res = LoadHDChain(hdChainTmp);
-    assert(res);
+    // 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));
+        }
 
-    CHDChain hdChainCrypted;
-    res = GetHDChain(hdChainCrypted);
-    assert(res);
+        // Maintain an unencrypted copy of the chain for sanity checking
+        CHDChain prevHdChain{newHdChain};
 
-    // ids should match, seed hashes should not
-    assert(hdChainPrev.GetID() == hdChainCrypted.GetID());
-    assert(hdChainPrev.GetSeedHash() != hdChainCrypted.GetSeedHash());
+        bool res = EncryptHDChain(vMasterKey, newHdChain);
+        assert(res);
+        res = LoadHDChain(newHdChain);
+        assert(res);
+        res = GetHDChain(newHdChain);
+        assert(res);
 
-    if (!AddHDChainSingle(hdChainCrypted)) {
-        throw std::runtime_error(std::string(__func__) + ": AddHDChainSingle failed");
+        // IDs should match, seed hashes should not
+        assert(prevHdChain.GetID() == newHdChain.GetID());
+        assert(prevHdChain.GetSeedHash() != newHdChain.GetSeedHash());
     }
-}
-
-void LegacyScriptPubKeyMan::GenerateNewHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase)
-{
-    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)) {
-        throw std::runtime_error(std::string(__func__) + ": SetMnemonic failed");
-    }
-
-    // add default account
-    newHdChain.AddAccount();
 
     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 9c639b73ebfb7..5ebb757628c89 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -463,8 +463,7 @@ class LegacyScriptPubKeyMan : public ScriptPubKeyMan, public FillableSigningProv
     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 29d6104270a67..68e1919955b1e 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;
 }