Never remove database files on shutdown, it caused unreadable wallets on some testers...

[novacoin.git] / src / wallet.cpp

diff --git a/src/wallet.cpp b/src/wallet.cpp
index 1daec98..43fb6b6 100644 (file)
--- a/src/wallet.cpp
+++ b/src/wallet.cpp
@@ -40,6 +40,7 @@ bool CWallet::AddCryptedKey(const vector<unsigned char> &vchPubKey, const vector
         else
             return CWalletDB(strWalletFile).WriteCryptedKey(vchPubKey, vchCryptedSecret);
     }
+    return false;
 }
 
 bool CWallet::Unlock(const string& strWalletPassphrase)
@@ -187,6 +188,13 @@ bool CWallet::EncryptWallet(const string& strWalletPassphrase)
         }
 
         Lock();
+        Unlock(strWalletPassphrase);
+        NewKeyPool();
+        Lock();
+
+        // Need to completely rewrite the wallet file; if we don't, bdb might keep
+        // bits of the unencrypted private key in slack space in the database file.
+        CDB::Rewrite(strWalletFile);
     }
 
     return true;
@@ -268,8 +276,12 @@ bool CWallet::AddToWallet(const CWalletTx& wtxIn)
         {
             if (txout.scriptPubKey == scriptDefaultKey)
             {
-                SetDefaultKey(GetOrReuseKeyFromPool());
-                SetAddressBookName(CBitcoinAddress(vchDefaultKey), "");
+                std::vector<unsigned char> newDefaultKey;
+                if (GetKeyFromPool(newDefaultKey, false))
+                {
+                    SetDefaultKey(newDefaultKey);
+                    SetAddressBookName(CBitcoinAddress(vchDefaultKey), "");
+                }
             }
         }
 
@@ -1117,6 +1129,18 @@ int CWallet::LoadWallet(bool& fFirstRunRet)
         return false;
     fFirstRunRet = false;
     int nLoadWalletRet = CWalletDB(strWalletFile,"cr+").LoadWallet(this);
+    if (nLoadWalletRet == DB_NEED_REWRITE)
+    {
+        if (CDB::Rewrite(strWalletFile, "\x04pool"))
+        {
+            setKeyPool.clear();
+            // Note: can't top-up keypool here, because wallet is locked.
+            // User will be prompted to unlock wallet the next operation
+            // the requires a new key.
+        }
+        nLoadWalletRet = DB_NEED_REWRITE;
+    }
+
     if (nLoadWalletRet != DB_LOAD_OK)
         return nLoadWalletRet;
     fFirstRunRet = vchDefaultKey.empty();
@@ -1126,7 +1150,10 @@ int CWallet::LoadWallet(bool& fFirstRunRet)
         // Create new keyUser and set as default key
         RandAddSeedPerfmon();
 
-        SetDefaultKey(GetOrReuseKeyFromPool());
+        std::vector<unsigned char> newDefaultKey;
+        if (!GetKeyFromPool(newDefaultKey, false))
+            return DB_LOAD_FAIL;
+        SetDefaultKey(newDefaultKey);
         if (!SetAddressBookName(CBitcoinAddress(vchDefaultKey), ""))
             return DB_LOAD_FAIL;
     }
@@ -1199,6 +1226,34 @@ bool GetWalletFile(CWallet* pwallet, string &strWalletFileOut)
     return true;
 }
 
+//
+// Mark old keypool keys as used,
+// and generate all new keys
+//
+bool CWallet::NewKeyPool()
+{
+    CRITICAL_BLOCK(cs_wallet)
+    {
+        CWalletDB walletdb(strWalletFile);
+        BOOST_FOREACH(int64 nIndex, setKeyPool)
+            walletdb.ErasePool(nIndex);
+        setKeyPool.clear();
+
+        if (IsLocked())
+            return false;
+
+        int64 nKeys = max(GetArg("-keypool", 100), (int64)0);
+        for (int i = 0; i < nKeys; i++)
+        {
+            int64 nIndex = i+1;
+            walletdb.WritePool(nIndex, CKeyPool(GenerateNewKey()));
+            setKeyPool.insert(nIndex);
+        }
+        printf("CWallet::NewKeyPool wrote %"PRI64d" new keys\n", nKeys);
+    }
+    return true;
+}
+
 bool CWallet::TopUpKeyPool()
 {
     CRITICAL_BLOCK(cs_wallet)
@@ -1269,15 +1324,28 @@ void CWallet::ReturnKey(int64 nIndex)
     printf("keypool return %"PRI64d"\n", nIndex);
 }
 
-vector<unsigned char> CWallet::GetOrReuseKeyFromPool()
+bool CWallet::GetKeyFromPool(vector<unsigned char>& result, bool fAllowReuse)
 {
     int64 nIndex = 0;
     CKeyPool keypool;
-    ReserveKeyFromKeyPool(nIndex, keypool);
-    if(nIndex == -1)
-        return vchDefaultKey;
-    KeepKey(nIndex);
-    return keypool.vchPubKey;
+    CRITICAL_BLOCK(cs_wallet)
+    {
+        ReserveKeyFromKeyPool(nIndex, keypool);
+        if (nIndex == -1)
+        {
+            if (fAllowReuse && !vchDefaultKey.empty())
+            {
+                result = vchDefaultKey;
+                return true;
+            }
+            if (IsLocked()) return false;
+            result = GenerateNewKey();
+            return true;
+        }
+        KeepKey(nIndex);
+        result = keypool.vchPubKey;
+    }
+    return true;
 }
 
 int64 CWallet::GetOldestKeyPoolTime()