From 6a76c60e6c4c21c6355c5580f9f770fcc9f978ef Mon Sep 17 00:00:00 2001
From: Pieter Wuille <pieter.wuille@gmail.com>
Date: Wed, 13 Apr 2011 16:16:30 +0200
Subject: [PATCH] Automatically rescan after restoring wallet.dat

Information about the best known chain is added to wallet.dat. If this
information does not match the data in blkindex.dat, a rescan is automatically
performed, starting from the the last known block. When upgrading from a wallet
which does not have this information, no rescan is done automatically.
---
 db.h     | 12 ++++++++++++
 init.cpp | 13 ++++++++++++-
 main.cpp |  9 +++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/db.h b/db.h
index c9c40d58c..290981c06 100644
--- a/db.h
+++ b/db.h
@@ -13,6 +13,7 @@ class CAddress;
 class CWalletTx;
 class CAccount;
 class CAccountingEntry;
+class CBlockLocator;
 
 extern map<string, string> mapAddressBook;
 extern CCriticalSection cs_mapAddressBook;
@@ -405,6 +406,17 @@ public:
         return Write(make_pair(string("key"), vchPubKey), vchPrivKey, false);
     }
 
+    bool WriteBestBlock(const CBlockLocator& locator)
+    {
+        nWalletDBUpdated++;
+        return Write(string("bestblock"), locator);
+    }
+
+    bool ReadBestBlock(CBlockLocator& locator)
+    {
+        return Read(string("bestblock"), locator);
+    }
+
     bool ReadDefaultKey(vector<unsigned char>& vchPubKey)
     {
         vchPubKey.clear();
diff --git a/init.cpp b/init.cpp
index 68ed11f6b..8f7218122 100644
--- a/init.cpp
+++ b/init.cpp
@@ -372,10 +372,21 @@ bool AppInit2(int argc, char* argv[])
         strErrors += _("Error loading wallet.dat      \n");
     printf(" wallet      %15"PRI64d"ms\n", GetTimeMillis() - nStart);
 
+    CBlockIndex *pindexRescan = pindexBest;
     if (GetBoolArg("-rescan"))
+        pindexRescan = pindexGenesisBlock;
+    else
+    {
+        CWalletDB walletdb;
+        CBlockLocator locator;
+        if (walletdb.ReadBestBlock(locator))
+            pindexRescan = locator.GetBlockIndex();
+    }
+    if (pindexBest != pindexRescan)
     {
+        printf("Rescanning last %i blocks (from block %i)...\n", pindexBest->nHeight - pindexRescan->nHeight, pindexRescan->nHeight);
         nStart = GetTimeMillis();
-        ScanForWalletTransactions(pindexGenesisBlock);
+        ScanForWalletTransactions(pindexRescan);
         printf(" rescan      %15"PRI64d"ms\n", GetTimeMillis() - nStart);
     }
 
diff --git a/main.cpp b/main.cpp
index 67db0bce8..1b15d7e07 100644
--- a/main.cpp
+++ b/main.cpp
@@ -1611,6 +1611,15 @@ bool CBlock::SetBestChain(CTxDB& txdb, CBlockIndex* pindexNew)
         }
     }
 
+    // Update best block in wallet (so we can detect restored wallets)
+    if (!IsInitialBlockDownload())
+    {
+        CWalletDB walletdb;
+        const CBlockLocator locator(pindexNew);
+        if (!walletdb.WriteBestBlock(locator))
+            return error("SetBestChain() : WriteWalletBest failed");
+    }
+
     // New best block
     hashBestChain = hash;
     pindexBest = pindexNew;