From c89fd12d48c5fa77e7f607ec9857544ffc640405 Mon Sep 17 00:00:00 2001
From: Konstantin Akimov <knstqq@gmail.com>
Date: Thu, 16 Feb 2023 15:56:10 +0700
Subject: [PATCH] fix: follow-up bitcoin#17260 for dash specific code

Changes in this commit are required as a preparation to bitcoin#17261
Method GenerateNewHDChainEncrypted moved back from LegacyScriptManager to CWallet
This methods should not be moved before in #17260.

Also added 2 new methods in interface WalletStorage: NewKeyPoolCallback and KeepDestinationCallback
---
 src/wallet/rpcwallet.cpp       |  2 +-
 src/wallet/scriptpubkeyman.cpp | 95 ++++------------------------------
 src/wallet/scriptpubkeyman.h   |  9 ++--
 src/wallet/wallet.cpp          | 93 ++++++++++++++++++++++++++++++++-
 src/wallet/wallet.h            |  9 ++++
 5 files changed, 117 insertions(+), 91 deletions(-)

diff --git a/src/wallet/rpcwallet.cpp b/src/wallet/rpcwallet.cpp
index e11cb1bc41..ca8b5e1161 100644
--- a/src/wallet/rpcwallet.cpp
+++ b/src/wallet/rpcwallet.cpp
@@ -2719,7 +2719,7 @@ static UniValue upgradetohd(const JSONRPCRequest& request)
     pwallet->SetMinVersion(FEATURE_HD);
 
     if (prev_encrypted) {
-        if (!spk_man->GenerateNewHDChainEncrypted(secureMnemonic, secureMnemonicPassphrase, secureWalletPassphrase)) {
+        if (!pwallet->GenerateNewHDChainEncrypted(secureMnemonic, secureMnemonicPassphrase, secureWalletPassphrase)) {
             throw JSONRPCError(RPC_WALLET_ERROR, "Failed to generate encrypted HD wallet");
         }
     } else {
diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
index 430c21fa76..bae908c3ac 100644
--- a/src/wallet/scriptpubkeyman.cpp
+++ b/src/wallet/scriptpubkeyman.cpp
@@ -4,9 +4,9 @@
 
 #include <key_io.h>
 #include <chainparams.h>
-#include <coinjoin/client.h>
 #include <logging.h>
 #include <script/descriptor.h>
+#include <ui_interface.h>
 #include <util/bip32.h>
 #include <util/strencodings.h>
 #include <util/system.h>
@@ -265,7 +265,7 @@ bool LegacyScriptPubKeyMan::GetReservedDestination(bool internal, CTxDestination
     return true;
 }
 
-void LegacyScriptPubKeyMan::MarkUnusedAddresses(WalletBatch &batch, const CScript& script, const uint256& hashBlock)
+void LegacyScriptPubKeyMan::MarkUnusedAddresses(WalletBatch &batch, const CScript& script, const std::optional<int64_t>& block_time)
 {
     AssertLockHeld(cs_wallet);
     // extract addresses and check if they match with an unused keypool key
@@ -279,18 +279,15 @@ void LegacyScriptPubKeyMan::MarkUnusedAddresses(WalletBatch &batch, const CScrip
                 WalletLogPrintf("%s: Topping up keypool failed (locked wallet)\n", __func__);
             }
         }
-        if (!hashBlock.IsNull()) {
-            int64_t block_time;
-            bool found_block = m_wallet.chain().findBlock(hashBlock, nullptr /* block */, &block_time);
-            assert(found_block);
-            if (mapKeyMetadata[keyid].nCreateTime > block_time) {
+        if (block_time) {
+            if (mapKeyMetadata[keyid].nCreateTime > *block_time) {
                 WalletLogPrintf("%s: Found a key which appears to be used earlier than we expected, updating metadata\n", __func__);
                 CPubKey vchPubKey;
                 bool res = GetPubKey(keyid, vchPubKey);
                 assert(res); // this should never fail
-                mapKeyMetadata[keyid].nCreateTime = block_time;
+                mapKeyMetadata[keyid].nCreateTime = *block_time;
                 batch.WriteKeyMetadata(mapKeyMetadata[keyid], vchPubKey, true);
-                UpdateTimeFirstKey(block_time);
+                UpdateTimeFirstKey(*block_time);
             }
         }
     }
@@ -370,73 +367,6 @@ void LegacyScriptPubKeyMan::GenerateNewHDChain(const SecureString& secureMnemoni
     }
 }
 
-bool LegacyScriptPubKeyMan::GenerateNewHDChainEncrypted(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, const SecureString& secureWalletPassphrase)
-{
-    assert(!m_storage.IsWalletFlagSet(WALLET_FLAG_DISABLE_PRIVATE_KEYS));
-    LOCK(cs_wallet);
-
-    if (!m_storage.HasEncryptionKeys()) {
-        return false;
-    }
-
-    CCrypter crypter;
-    CKeyingMaterial vMasterKey;
-    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();
-    hdChainTmp.Debug(__func__);
-
-    for (const CWallet::MasterKeyMap::value_type& pMasterKey : m_wallet.mapMasterKeys) {
-        if (!crypter.SetKeyFromPassphrase(secureWalletPassphrase, pMasterKey.second.vchSalt, pMasterKey.second.nDeriveIterations, pMasterKey.second.nDerivationMethod)) {
-            return false;
-        }
-        // get vMasterKey to encrypt new hdChain
-        if (!crypter.Decrypt(pMasterKey.second.vchCryptedKey, vMasterKey)) {
-            continue; // try another master key
-        }
-
-        bool res = EncryptHDChain(vMasterKey, hdChainTmp);
-        assert(res);
-
-        CHDChain hdChainCrypted;
-        res = GetHDChain(hdChainCrypted);
-        assert(res);
-
-        DBG(
-            tfm::format(std::cout, "GenerateNewHDChainEncrypted -- current seed: '%s'\n", HexStr(hdChainTmp.GetSeed()));
-            tfm::format(std::cout, "GenerateNewHDChainEncrypted -- crypted seed: '%s'\n", HexStr(hdChainCrypted.GetSeed()));
-        );
-
-        // ids should match, seed hashes should not
-        assert(hdChainTmp.GetID() == hdChainCrypted.GetID());
-        assert(hdChainTmp.GetSeedHash() != hdChainCrypted.GetSeedHash());
-
-        hdChainCrypted.Debug(__func__);
-
-        if (SetCryptedHDChainSingle(hdChainCrypted, false)) {
-            m_wallet.Lock();
-            if (!m_wallet.Unlock(secureWalletPassphrase)) {
-                // this should never happen
-                throw std::runtime_error(std::string(__func__) + ": Unlock failed");
-            }
-            if (!NewKeyPool()) {
-                throw std::runtime_error(std::string(__func__) + ": NewKeyPool failed");
-            }
-            m_wallet.Lock();
-            return true;
-        }
-    }
-
-    return false;
-}
-
 bool LegacyScriptPubKeyMan::SetHDChain(WalletBatch &batch, const CHDChain& chain, bool memonly)
 {
     LOCK(cs_wallet);
@@ -1337,12 +1267,8 @@ bool LegacyScriptPubKeyMan::NewKeyPool()
             batch.ErasePool(nIndex);
         }
         setExternalKeyPool.clear();
-        auto it = coinJoinClientManagers.find(m_wallet.GetName());
-        if (it != coinJoinClientManagers.end()) {
-            it->second->StopMixing();
-        }
-        m_wallet.nKeysLeftSinceAutoBackup = 0;
 
+        m_storage.NewKeyPoolCallback();
         m_pool_key_to_index.clear();
 
         if (!TopUp()) {
@@ -1443,11 +1369,8 @@ void LegacyScriptPubKeyMan::KeepDestination(int64_t nIndex)
     {
         LOCK(cs_wallet);
         WalletBatch batch(m_storage.GetDatabase());
-        if (batch.ErasePool(nIndex))
-            --m_wallet.nKeysLeftSinceAutoBackup;
-        if (!nWalletBackups)
-            m_wallet.nKeysLeftSinceAutoBackup = 0;
-
+        bool erased = batch.ErasePool(nIndex);
+        m_storage.KeepDestinationCallback(erased);
         CPubKey pubkey;
         bool have_pk = GetPubKey(m_index_to_reserved_key.at(nIndex), pubkey);
         assert(have_pk);
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
index 4b6789aab1..30a95c4607 100644
--- a/src/wallet/scriptpubkeyman.h
+++ b/src/wallet/scriptpubkeyman.h
@@ -33,6 +33,10 @@ public:
     virtual const CKeyingMaterial& GetEncryptionKey() const = 0;
     virtual bool HasEncryptionKeys() const = 0;
     virtual bool IsLocked(bool fForMixing = false) const = 0;
+
+    // methods below are unique from Dash due to different implementation of HD
+    virtual void NewKeyPoolCallback() = 0;
+    virtual void KeepDestinationCallback(bool erased) = 0;
 };
 
 //! Default for -keypool
@@ -163,7 +167,7 @@ public:
     virtual bool TopUp(unsigned int size = 0) { return false; }
 
     //! Mark unused addresses as being used
-    virtual void MarkUnusedAddresses(WalletBatch &batch, const CScript& script, const uint256& hashBlock) {}
+    virtual void MarkUnusedAddresses(WalletBatch &batch, const CScript& script, const std::optional<int64_t>& block_time) {}
 
     /* Returns true if HD is enabled */
     virtual bool IsHDEnabled() const { return false; }
@@ -282,7 +286,7 @@ public:
 
     bool TopUp(unsigned int size = 0) override;
 
-    void MarkUnusedAddresses(WalletBatch &batch, const CScript& script, const uint256& hashBlock) override;
+    void MarkUnusedAddresses(WalletBatch &batch, const CScript& script, const std::optional<int64_t>& block_time) override;
 
     //! Upgrade stored CKeyMetadata objects to store key origin info as KeyOriginInfo
     void UpgradeKeyMetadata() EXCLUSIVE_LOCKS_REQUIRED(cs_wallet);
@@ -410,7 +414,6 @@ public:
 
     /* Generates a new HD chain */
     void GenerateNewHDChain(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase);
-    bool GenerateNewHDChainEncrypted(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, const SecureString& secureWalletPassphrase);
 
     /**
      * 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 5bb81e470e..b8115b6708 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -1010,7 +1010,14 @@ bool CWallet::AddToWalletIfInvolvingMe(const CTransactionRef& ptx, CWalletTx::Co
             // loop though all outputs
             for (const CTxOut& txout: tx.vout) {
                 if (auto spk_man = m_spk_man.get()) {
-                    spk_man->MarkUnusedAddresses(batch, txout.scriptPubKey, confirm.hashBlock);
+                    std::optional<int64_t> block_time;
+                    if (!confirm.hashBlock.IsNull()) {
+                        int64_t block_time_tmp;
+                        bool found_block = chain().findBlock(confirm.hashBlock, nullptr /* block */, &block_time_tmp);
+                        assert(found_block);
+                        block_time = block_time_tmp;
+                    }
+                    spk_man->MarkUnusedAddresses(batch, txout.scriptPubKey, block_time);
                 }
             }
 
@@ -1603,6 +1610,21 @@ void CWallet::UnsetBlankWalletFlag(WalletBatch& batch)
     UnsetWalletFlagWithDB(batch, WALLET_FLAG_BLANK_WALLET);
 }
 
+void CWallet::NewKeyPoolCallback()
+{
+    auto it = coinJoinClientManagers.find(GetName());
+    if (it != coinJoinClientManagers.end()) {
+        it->second->StopMixing();
+    }
+    nKeysLeftSinceAutoBackup = 0;
+}
+
+void CWallet::KeepDestinationCallback(bool erased)
+{
+    if (erased) --nKeysLeftSinceAutoBackup;
+    if (!nWalletBackups) nKeysLeftSinceAutoBackup = 0;
+}
+
 bool CWallet::IsWalletFlagSet(uint64_t flag) const
 {
     return (m_wallet_flags & flag);
@@ -5036,3 +5058,72 @@ bool CWallet::HasEncryptionKeys() const
 {
     return !mapMasterKeys.empty();
 }
+
+bool CWallet::GenerateNewHDChainEncrypted(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, const SecureString& secureWalletPassphrase)
+{
+    assert(m_spk_man);
+
+    assert(!IsWalletFlagSet(WALLET_FLAG_DISABLE_PRIVATE_KEYS));
+    LOCK(cs_wallet);
+
+    if (!HasEncryptionKeys()) {
+        return false;
+    }
+
+    CCrypter crypter;
+    CKeyingMaterial vMasterKey;
+    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();
+    hdChainTmp.Debug(__func__);
+
+    for (const CWallet::MasterKeyMap::value_type& pMasterKey : mapMasterKeys) {
+        if (!crypter.SetKeyFromPassphrase(secureWalletPassphrase, pMasterKey.second.vchSalt, pMasterKey.second.nDeriveIterations, pMasterKey.second.nDerivationMethod)) {
+            return false;
+        }
+        // get vMasterKey to encrypt new hdChain
+        if (!crypter.Decrypt(pMasterKey.second.vchCryptedKey, vMasterKey)) {
+            continue; // try another master key
+        }
+
+        bool res = m_spk_man->EncryptHDChain(vMasterKey, hdChainTmp);
+        assert(res);
+
+        CHDChain hdChainCrypted;
+        res = m_spk_man->GetHDChain(hdChainCrypted);
+        assert(res);
+
+        DBG(
+            tfm::format(std::cout, "GenerateNewHDChainEncrypted -- current seed: '%s'\n", HexStr(hdChainTmp.GetSeed()));
+            tfm::format(std::cout, "GenerateNewHDChainEncrypted -- crypted seed: '%s'\n", HexStr(hdChainCrypted.GetSeed()));
+        );
+
+        // ids should match, seed hashes should not
+        assert(hdChainTmp.GetID() == hdChainCrypted.GetID());
+        assert(hdChainTmp.GetSeedHash() != hdChainCrypted.GetSeedHash());
+
+        hdChainCrypted.Debug(__func__);
+
+        if (m_spk_man->SetCryptedHDChainSingle(hdChainCrypted, false)) {
+            Lock();
+            if (!Unlock(secureWalletPassphrase)) {
+                // this should never happen
+                throw std::runtime_error(std::string(__func__) + ": Unlock failed");
+            }
+            if (!m_spk_man->NewKeyPool()) {
+                throw std::runtime_error(std::string(__func__) + ": NewKeyPool failed");
+            }
+            Lock();
+            return true;
+        }
+    }
+
+    return false;
+}
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 43e70399e1..ad5902313b 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -726,6 +726,12 @@ private:
     //! Unset the blank wallet flag and saves it to disk
     void UnsetBlankWalletFlag(WalletBatch& batch) override;
 
+    // Reset coinjoin and reset key counter
+    void NewKeyPoolCallback() override;
+
+    // Decreases amount of nKeysLeftSinceAutoBackup after KeepDestination
+    void KeepDestinationCallback(bool erased) override;
+
     /** Interface for accessing chain state. */
     interfaces::Chain* m_chain;
 
@@ -1194,6 +1200,9 @@ public:
     /* Returns true if HD is enabled */
     bool IsHDEnabled() const;
 
+    /* Generates a new HD chain */
+    bool GenerateNewHDChainEncrypted(const SecureString& secureMnemonic, const SecureString& secureMnemonicPassphrase, const SecureString& secureWalletPassphrase);
+
     /* Returns true if the wallet can give out new addresses. This means it has keys in the keypool or can generate new keys */
     bool CanGetAddresses(bool internal = false);