wallet: collapse restored legacy hdnext

3 files changed, 93 insertions, 6 deletions

diff --git a/src/selftest.cpp b/src/selftest.cpp
index cca5335..bcbf9d3 100644
--- a/src/selftest.cpp
+++ b/src/selftest.cpp
@@ -586,6 +586,34 @@ static int RunWalletHDSelfTest()
                        "the recovery audit finds wallet.dat-only immature mining rewards") ? 0 : 1;
         nFail += Check(audit.nLegacyImmatureTx == 1,
                        "the recovery audit counts wallet.dat-only immature mining rewards") ? 0 : 1;
+
+        CKey keyLegacyHD11;
+        nSchemaForLegacyAudit = nHDKeySchema;
+        nHDKeySchema = HD_SCHEMA_LEGACY;
+        if (!DeriveHDKey(11, keyLegacyHD11, strError))
+            throw std::runtime_error("legacy compatibility index 11 derivation failed: " + strError);
+        nHDKeySchema = nSchemaForLegacyAudit;
+        if (!AddKey(keyLegacyHD11))
+            throw std::runtime_error("could not store the legacy compatibility index 11 key");
+
+        std::map<unsigned int, std::vector<unsigned char> > mapLegacyProbe;
+        mapLegacyProbe[0] = keyAuditLegacyHD.GetPubKey();
+        mapLegacyProbe[11] = keyLegacyHD11.GetPubKey();
+
+        CWalletTx wtxLegacyHD11;
+        wtxLegacyHD11.vout.push_back(CTxOut(19 * COIN, CScript() << keyLegacyHD11.GetPubKey() << OP_CHECKSIG));
+        CRITICAL_BLOCK(cs_mapWallet)
+        {
+            mapWallet.clear();
+            mapWallet[wtxLegacyHD11.GetHash()] = wtxLegacyHD11;
+        }
+        nFail += Check(WalletLastUsedPubKeyIndexNext(mapLegacyProbe) == 12,
+                       "restore collapses legacy hdnext to the last used compatibility index plus one") ? 0 : 1;
+
+        CRITICAL_BLOCK(cs_mapWallet)
+            mapWallet.clear();
+        nFail += Check(WalletLastUsedPubKeyIndexNext(mapLegacyProbe) == 0,
+                       "restore does not persist legacy hdnext at the scan depth when no compatibility key was used") ? 0 : 1;
     }
     catch (const std::exception& e)
     {
diff --git a/src/walletcmd.cpp b/src/walletcmd.cpp
index 0ea2251..e8ac4bf 100644
--- a/src/walletcmd.cpp
+++ b/src/walletcmd.cpp
@@ -54,6 +54,37 @@ bool RestoreScanReachedDepth(int nSchema,
     return nLegacyNext >= nDepth;
 }
 
+unsigned int WalletLastUsedPubKeyIndexNext(
+    const std::map<unsigned int, std::vector<unsigned char> >& mapPubKeysByIndex)
+{
+    std::map<std::vector<unsigned char>, unsigned int> mapIndexByPubKey;
+    for (std::map<unsigned int, std::vector<unsigned char> >::const_iterator it =
+             mapPubKeysByIndex.begin();
+         it != mapPubKeysByIndex.end(); ++it)
+        mapIndexByPubKey[it->second] = it->first;
+
+    unsigned int nNext = 0;
+    CRITICAL_BLOCK(cs_mapWallet)
+    {
+        for (map<uint256, CWalletTx>::iterator it = mapWallet.begin();
+             it != mapWallet.end(); ++it)
+        {
+            CWalletTx& wtx = (*it).second;
+            for (int i = 0; i < (int)wtx.vout.size(); i++)
+            {
+                vector<unsigned char> vchPubKey;
+                if (!ExtractPubKey(wtx.vout[i].scriptPubKey, false, vchPubKey))
+                    continue;
+                std::map<std::vector<unsigned char>, unsigned int>::const_iterator mi =
+                    mapIndexByPubKey.find(vchPubKey);
+                if (mi != mapIndexByPubKey.end() && mi->second + 1 > nNext)
+                    nNext = mi->second + 1;
+            }
+        }
+    }
+    return nNext;
+}
+
 int CmdNewPhrase()
 {
     AttachTerminal();
@@ -132,7 +163,10 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
     // Derive forward in batches, scanning after each, until a whole batch turns
     // up nothing. Every derived key is written to the wallet before the scan,
     // because the scan asks the wallet what belongs to it.
-    int nTotalDerived = nHDKeySchema == HD_SCHEMA_BIP44
+    const bool fRestoreBIP44 = nHDKeySchema == HD_SCHEMA_BIP44;
+    unsigned int nLegacyScanNext = nHDNext;
+    std::map<unsigned int, vector<unsigned char> > mapLegacyScanPubKeys;
+    int nTotalDerived = fRestoreBIP44
         ? (int)(nHDReceiveNext + nHDChangeNext)
         : (int)nHDNext;
     int nStopDepth = max(nMinDepth, RESTORE_MIN_SCAN);
@@ -151,7 +185,7 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
             nWalletBefore = mapWallet.size();
 
         std::string strDeriveError;
-        if (nHDKeySchema == HD_SCHEMA_BIP44)
+        if (fRestoreBIP44)
         {
             for (int i = 0; i < RESTORE_BATCH; i++)
             {
@@ -182,15 +216,15 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
             for (int i = 0; i < RESTORE_BATCH; i++)
             {
                 CKey key;
-                if (!DeriveHDKey(nHDNext, key, strDeriveError))
+                if (!DeriveHDKey(nLegacyScanNext, key, strDeriveError))
                     break;
                 if (!AddKey(key))
                     break;
-                nHDNext++;
+                mapLegacyScanPubKeys[nLegacyScanNext] = key.GetPubKey();
+                nLegacyScanNext++;
                 nTotalDerived++;
             }
             nHDKeySchema = nSavedSchema;
-            CWalletDB().WriteHDNext(nHDNext);
         }
         else
         {
@@ -209,6 +243,22 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
 
         ScanForWalletTransactions(pindexGenesisBlock);
 
+        if (fRestoreBIP44)
+        {
+            // The legacy branch is scanned only for compatibility with phrases
+            // created before BIP44. Persist the last used legacy index + 1, not
+            // the lookahead depth. Otherwise a restore with no legacy hits can
+            // leave hdnext at RESTORE_MAX and make future compatibility scans
+            // look exhausted even though no legacy key was ever used.
+            unsigned int nLegacyUsedNext =
+                WalletLastUsedPubKeyIndexNext(mapLegacyScanPubKeys);
+            if (nLegacyUsedNext != nHDNext)
+            {
+                nHDNext = nLegacyUsedNext;
+                CWalletDB().WriteHDNext(nHDNext);
+            }
+        }
+
         size_t nWalletAfter = 0;
         CRITICAL_BLOCK(cs_mapWallet)
             nWalletAfter = mapWallet.size();
@@ -239,7 +289,7 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
             RestoreScanReachedDepth(nHDKeySchema,
                                     nHDReceiveNext,
                                     nHDChangeNext,
-                                    nHDNext,
+                                    fRestoreBIP44 ? nLegacyScanNext : nHDNext,
                                     nStopDepth))
             break;
     }
diff --git a/src/walletcmd.h b/src/walletcmd.h
index 04d3fb7..007b8b5 100644
--- a/src/walletcmd.h
+++ b/src/walletcmd.h
@@ -7,7 +7,9 @@
 #ifndef BITFLASH_WALLETCMD_H
 #define BITFLASH_WALLETCMD_H
 
+#include <map>
 #include <string>
+#include <vector>
 
 // Create a phrase for a wallet that does not have one, install the seed, and
 // show the words once. Refuses if a phrase already exists. Returns 0 on success.
@@ -41,6 +43,13 @@ bool RestoreScanReachedDepth(int nSchema,
                              unsigned int nLegacyNext,
                              int nStopDepth);
 
+// Given the pubkeys derived while scanning a compatibility branch, return the
+// next index after the highest one that actually appears in wallet
+// transactions. The scan depth can be much larger than this; hdnext must record
+// use, not how far a restore looked ahead.
+unsigned int WalletLastUsedPubKeyIndexNext(
+    const std::map<unsigned int, std::vector<unsigned char> >& mapPubKeysByIndex);
+
 // Take the next address from the key pool and print it. With a recovery phrase
 // installed the address is derived, so the phrase can bring back whatever is
 // paid to it.