Keep new phrase addresses recoverable (#102)

-newphrase installed a seed but left the wallet's visible address as the pre-seed random key and the old random pool intact, so the next address handed out -- a mining payout included -- could be one the twelve words cannot reproduce, while the window said a recovery phrase existed. Installing a phrase now clears the stale pool records (private keys stay in wallet.dat), reserves a derived default receiving key, tolerates already-stored derived keys so repeated installs do not stall the pool, and -restorephrase reports what it recovered rather than the wallet's total transaction count. Self-test covers installing a phrase over a pre-existing random pool.

4 files changed, 99 insertions, 11 deletions

diff --git a/src/main.cpp b/src/main.cpp
index e5754e6..2c4e346 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -167,6 +167,20 @@ bool AddKey(const CKey& key)
     return CWalletDB().WriteKey(key.GetPubKey(), key.GetPrivKey());
 }
 
+static bool WalletAlreadyHasKey(const vector<unsigned char>& vchPubKey)
+{
+    CRITICAL_BLOCK(cs_mapKeys)
+        return mapKeys.count(vchPubKey) > 0;
+    return false;
+}
+
+static bool AddKeyIfMissing(const CKey& key)
+{
+    if (WalletAlreadyHasKey(key.GetPubKey()))
+        return true;
+    return AddKey(key);
+}
+
 vector<unsigned char> GenerateNewKey()
 {
     CKey key;
@@ -183,6 +197,23 @@ vector<unsigned char> vchHDMaster;
 vector<unsigned char> vchHDChainCode;
 unsigned int nHDNext = 0;
 
+static bool ClearKeyPoolRecords(string& strErrorRet)
+{
+    CWalletDB walletdb;
+    for (map<int64, vector<unsigned char> >::const_iterator mi = mapKeyPool.begin();
+         mi != mapKeyPool.end(); ++mi)
+    {
+        if (!walletdb.ErasePool(mi->first))
+        {
+            strErrorRet = strprintf("could not clear old key-pool entry %lld",
+                                    (long long)mi->first);
+            return false;
+        }
+    }
+    mapKeyPool.clear();
+    return true;
+}
+
 bool DeriveHDKey(unsigned int nIndex, CKey& keyRet, string& strErrorRet)
 {
     if (!HaveHDSeed())
@@ -223,9 +254,12 @@ bool SetHDSeedFromMnemonic(const string& strMnemonic, string& strErrorRet)
 
     CRITICAL_BLOCK(cs_keyPool)
     {
-        // Written before the globals change, so a failed write leaves the
-        // wallet on the seed it already had rather than on one that exists
-        // only in memory.
+        // Old private keys stay in wallet.dat, but old unused pool records must
+        // not survive the seed. Otherwise the next address can still be random
+        // while the UI says a recovery phrase exists.
+        if (!ClearKeyPoolRecords(strErrorRet))
+            return false;
+
         if (!CWalletDB().WriteHDMaster(master.privateKey, master.chainCode) ||
             !CWalletDB().WriteHDNext(0))
         {
@@ -235,6 +269,25 @@ bool SetHDSeedFromMnemonic(const string& strMnemonic, string& strErrorRet)
         vchHDMaster    = master.privateKey;
         vchHDChainCode = master.chainCode;
         nHDNext        = 0;
+
+        CKey key;
+        if (!DeriveHDKey(0, key, strErrorRet))
+            return false;
+        if (!AddKeyIfMissing(key))
+        {
+            strErrorRet = "could not write the default derived key to wallet.dat";
+            return false;
+        }
+        vector<unsigned char> vchPubKey = key.GetPubKey();
+        if (!CWalletDB().WriteDefaultKey(vchPubKey) ||
+            !SetAddressBookName(PubKeyToAddress(vchPubKey), "Your Address") ||
+            !CWalletDB().WriteHDNext(1))
+        {
+            strErrorRet = "could not set the default derived receiving key";
+            return false;
+        }
+        keyUser = key;
+        nHDNext = 1;
     }
     return true;
 }
@@ -279,7 +332,7 @@ void TopUpKeyPool()
             // AddKey writes the private key. Order matters: the key is in
             // wallet.dat before anything can point at it, so a crash between
             // the two writes costs an unused key, never a usable one.
-            if (!AddKey(key))
+            if (!AddKeyIfMissing(key))
                 return;
             vector<unsigned char> vchPubKey = key.GetPubKey();
             if (!CWalletDB().WritePool(nIndex, vchPubKey))
diff --git a/src/main.h b/src/main.h
index a04a4e3..1b0be59 100644
--- a/src/main.h
+++ b/src/main.h
@@ -159,8 +159,9 @@ extern vector<unsigned char> vchHDChainCode;  // 32 bytes (IR)
 extern unsigned int nHDNext;                  // next child index to derive
 inline bool HaveHDSeed() { return vchHDMaster.size() == 32 && vchHDChainCode.size() == 32; }
 
-// Install a seed derived from a mnemonic, replacing any existing one, and reset
-// the derivation counter. Writes to wallet.dat. Returns false and leaves the
+// Install a seed derived from a mnemonic, replacing any existing one. The
+// default receiving key is derived immediately, so the next visible address is
+// covered by the phrase. Writes to wallet.dat. Returns false and leaves the
 // wallet untouched if the phrase is not valid.
 bool SetHDSeedFromMnemonic(const string& strMnemonic, string& strErrorRet);
 
diff --git a/src/selftest.cpp b/src/selftest.cpp
index b733c60..bc9fa5a 100644
--- a/src/selftest.cpp
+++ b/src/selftest.cpp
@@ -327,10 +327,35 @@ static int RunWalletHDSelfTest()
             throw std::runtime_error("LoadWallet failed");
 
         std::string strError;
+        TopUpKeyPool();
+        int nRandomPoolSize = 0;
+        CRITICAL_BLOCK(cs_keyPool)
+            nRandomPoolSize = (int)mapKeyPool.size();
+        nFail += Check(nRandomPoolSize == KEYPOOL_SIZE,
+                       "a wallet can start with a random key pool") ? 0 : 1;
+
         nFail += Check(SetHDSeedFromMnemonic(strPhraseA, strError),
                        "a valid phrase installs a seed") ? 0 : 1;
         nFail += Check(HaveHDSeed(), "the wallet reports having a seed") ? 0 : 1;
 
+        int nPoolAfterSeed = 0;
+        CRITICAL_BLOCK(cs_keyPool)
+            nPoolAfterSeed = (int)mapKeyPool.size();
+        nFail += Check(nPoolAfterSeed == 0,
+                       "installing a phrase clears the old random key pool") ? 0 : 1;
+        nFail += Check(nHDNext == 1,
+                       "installing a phrase reserves one derived default key") ? 0 : 1;
+
+        CKey keyFirstDerived;
+        if (!DeriveHDKey(0, keyFirstDerived, strError))
+            throw std::runtime_error("default derivation failed: " + strError);
+        std::vector<unsigned char> vchDefaultKey;
+        bool fDefaultRead = CWalletDB("r").ReadDefaultKey(vchDefaultKey);
+        nFail += Check(fDefaultRead && vchDefaultKey == keyFirstDerived.GetPubKey(),
+                       "the default receiving key is derived from the phrase") ? 0 : 1;
+        nFail += Check(WalletHasPrivateKey(vchDefaultKey),
+                       "the derived default key is stored in wallet.dat") ? 0 : 1;
+
         std::vector<unsigned char> vchBefore = vchHDMaster;
         nFail += Check(!SetHDSeedFromMnemonic("not a mnemonic at all", strError),
                        "an invalid phrase is refused") ? 0 : 1;
@@ -378,8 +403,6 @@ static int RunWalletHDSelfTest()
         // With a seed installed the pool must be derived from it, and the
         // counter must move exactly once per key.
         SetHDSeedFromMnemonic(strPhraseA, strError);
-        CRITICAL_BLOCK(cs_keyPool)
-            mapKeyPool.clear();
         unsigned int nNextBefore = nHDNext;
         TopUpKeyPool();
 
diff --git a/src/walletcmd.cpp b/src/walletcmd.cpp
index 94e9d72..8eb3592 100644
--- a/src/walletcmd.cpp
+++ b/src/walletcmd.cpp
@@ -94,13 +94,17 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
     if (!CanScanWalletTransactions(strErrorRet))
         return false;
 
+    size_t nWalletStart = 0;
+    CRITICAL_BLOCK(cs_mapWallet)
+        nWalletStart = mapWallet.size();
+
     if (!SetHDSeedFromMnemonic(strMnemonic, strErrorRet))
         return false;
 
     // 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 = 0;
+    int nTotalDerived = (int)nHDNext;
     while (nTotalDerived < RESTORE_MAX)
     {
         // Count wallet transactions, not scan hits.
@@ -134,7 +138,11 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
             nWalletAfter = mapWallet.size();
 
         if (fnProgress)
-            fnProgress(pArg, nTotalDerived, (int)nWalletAfter);
+        {
+            size_t nRecoveredNow = nWalletAfter > nWalletStart ?
+                                   nWalletAfter - nWalletStart : 0;
+            fnProgress(pArg, nTotalDerived, (int)nRecoveredNow);
+        }
 
         // A whole batch with nothing in it means far enough -- unless the
         // caller asked to look deeper anyway.
@@ -148,8 +156,11 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
         mapKeyPool.clear();
     TopUpKeyPool();
 
+    size_t nWalletEnd = 0;
     CRITICAL_BLOCK(cs_mapWallet)
-        nRecoveredRet = (int)mapWallet.size();
+        nWalletEnd = mapWallet.size();
+    nRecoveredRet = (int)(nWalletEnd > nWalletStart ?
+                          nWalletEnd - nWalletStart : 0);
     nDerivedRet = nTotalDerived;
     return true;
 }