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.
dashpay · Jul 17, 2024 · 163d318 · 163d318
1 parent 1394c41
commit 163d318
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 42 deletions.
diff --git 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
@@ -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
@@ -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;
 }