[Wallet] Upgrade to HD flow with -upgradewallet flag.

furszy · furszy · commit cc1a60b5d9f6 · 2020-03-04T11:22:52.000-03:00
* Essentially marks non HD wallet keys as pre_split keys, stop creating pre_split keys and creates a new HD seed and HD keypool to be used from there on.
diff --git a/src/init.cpp b/src/init.cpp
@@ -1671,6 +1671,7 @@ bool AppInit2()
                 strErrors << _("Error loading wallet.dat") << "\n";
         }
 
+        int prev_version = pwalletMain->GetVersion();
         if (GetBoolArg("-upgradewallet", fFirstRun)) {
             int nMaxVersion = GetArg("-upgradewallet", 0);
             if (nMaxVersion == 0) // the -upgradewallet without argument case
@@ -1685,6 +1686,12 @@ bool AppInit2()
             pwalletMain->SetMaxVersion(nMaxVersion);
         }
 
+        // Upgrade to HD if explicit upgrade was requested.
+        std::string upgradeError;
+        if (!pwalletMain->Upgrade(upgradeError, prev_version)) {
+            strErrors << upgradeError << "\n";
+        }
+
         if (fFirstRun) {
             // Create new HD Wallet
             LogPrintf("Creating HD Wallet\n");
diff --git a/src/wallet/scriptpubkeyman.cpp b/src/wallet/scriptpubkeyman.cpp
@@ -20,6 +20,37 @@ bool ScriptPubKeyMan::SetupGeneration(bool force)
     return true;
 }
 
+bool ScriptPubKeyMan::Upgrade(const int& prev_version, std::string& error)
+{
+    LOCK(wallet->cs_KeyStore);
+    error = "";
+    bool hd_upgrade = false;
+    bool split_upgrade = false;
+    if (wallet->CanSupportFeature(FEATURE_HD_SPLIT) && !IsHDEnabled()) {
+        LogPrintf("Upgrading wallet to use HD chain split\n");
+        wallet->SetMinVersion(FEATURE_PRE_SPLIT_KEYPOOL);
+
+        // generate a new master key
+        CPubKey masterPubKey = GenerateNewSeed();
+        SetHDSeed(masterPubKey);
+        hd_upgrade = true;
+        split_upgrade = FEATURE_HD_SPLIT > prev_version;
+    }
+
+    // Mark all keys currently in the keypool as pre-split
+    if (split_upgrade) {
+        MarkPreSplitKeys();
+    }
+    // Regenerate the keypool if upgraded to HD
+    if (hd_upgrade) {
+        if (!TopUp()) {
+            error = _("Unable to generate keys");
+            return false;
+        }
+    }
+    return true;
+}
+
 bool ScriptPubKeyMan::CanGenerateKeys()
 {
     // A wallet can generate keys if it has an HD seed (IsHDEnabled) or it is a non-HD wallet (pre FEATURE_HD)
@@ -231,6 +262,24 @@ void ScriptPubKeyMan::MarkUnusedAddresses(const CScript& script)
     }
 }
 
+void ScriptPubKeyMan::MarkPreSplitKeys()
+{
+    CWalletDB batch(wallet->strWalletFile);
+    for (auto it = setExternalKeyPool.begin(); it != setExternalKeyPool.end();) {
+        int64_t index = *it;
+        CKeyPool keypool;
+        if (!batch.ReadPool(index, keypool)) {
+            throw std::runtime_error(std::string(__func__) + ": read keypool entry failed");
+        }
+        keypool.m_pre_split = true;
+        if (!batch.WritePool(index, keypool)) {
+            throw std::runtime_error(std::string(__func__) + ": writing modified keypool entry failed");
+        }
+        set_pre_split_keypool.insert(index);
+        it = setExternalKeyPool.erase(it);
+    }
+}
+
 /**
  * Mark old keypool keys as used,
  * and generate all new keys
diff --git a/src/wallet/scriptpubkeyman.h b/src/wallet/scriptpubkeyman.h
@@ -37,6 +37,9 @@ class ScriptPubKeyMan {
       */
     bool SetupGeneration(bool force = false);
 
+    /** Upgrades the wallet to the specified version */
+    bool Upgrade(const int& prev_version, std::string& error);
+
     /* Returns true if the wallet can generate new keys */
     bool CanGenerateKeys();
 
@@ -71,6 +74,8 @@ class ScriptPubKeyMan {
     void LoadKeyPool(int64_t nIndex, const CKeyPool &keypool);
     //! Key pool
     bool NewKeyPool();
+    //! Update pre HD keys in db with the pre-split flag enabled.
+    void MarkPreSplitKeys();
 
     /** Fills internal address pool. Use within ScriptPubKeyMan implementations should be used sparingly and only
       * when something from the address pool is removed, excluding GetNewDestination and GetReservedDestination.
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
@@ -1542,6 +1542,24 @@ bool CWalletTx::WriteToDisk(CWalletDB *pwalletdb)
     return CWalletDB(pwallet->strWalletFile).WriteTx(GetHash(), *this);
 }
 
+bool CWallet::Upgrade(std::string& error, const int& prevVersion)
+{
+    // Upgrade to HD if explicit upgrade
+    if (GetBoolArg("-upgradewallet", false)) {
+        LOCK(cs_wallet);
+
+        // Do not upgrade versions to any version between HD_SPLIT and FEATURE_PRE_SPLIT_KEYPOOL unless already supporting HD_SPLIT
+        int max_version = GetVersion();
+        if (!CanSupportFeature(FEATURE_HD_SPLIT) && max_version >= FEATURE_HD_SPLIT && max_version < FEATURE_PRE_SPLIT_KEYPOOL) {
+            error = _("Cannot upgrade a non HD split wallet without upgrading to support pre split keypool. Please use -upgradewallet=169900 or -upgradewallet with no version specified.");
+            return false;
+        }
+
+        return m_spk_man->Upgrade(prevVersion, error);
+    }
+    return true;
+}
+
 /**
  * Scan the block chain (starting in pindexStart) for transactions
  * from or to us. If fUpdate is true, found transactions that already
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
@@ -86,7 +86,9 @@ enum WalletFeature {
 
     FEATURE_HD_SPLIT = 139900, // Wallet with HD chain split (change outputs will use m/0'/1'/k)
 
-    FEATURE_LATEST = FEATURE_HD_SPLIT
+    FEATURE_PRE_SPLIT_KEYPOOL = 169900, // Upgraded to HD SPLIT and can have a pre-split keypool
+
+    FEATURE_LATEST = FEATURE_PRE_SPLIT_KEYPOOL
 };
 
 enum AvailableCoinsType {
@@ -408,6 +410,12 @@ class CWallet : public CCryptoKeyStore, public CValidationInterface
     void SyncTransaction(const CTransaction& tx, const CBlock* pblock);
     bool AddToWalletIfInvolvingMe(const CTransaction& tx, const CBlock* pblock, bool fUpdate);
     void EraseFromWallet(const uint256& hash);
+
+    /**
+     * Upgrade wallet to HD if needed. Does nothing if not.
+     */
+    bool Upgrade(std::string& error, const int& prevVersion);
+
     int ScanForWalletTransactions(CBlockIndex* pindexStart, bool fUpdate = false, bool fromStartup = false);
     void ReacceptWalletTransactions(bool fFirstLoad = false);
     void ResendWalletTransactions();
