From 96f34cd5c4d76459917b29b15aa9f4b7e2a6cec1 Mon Sep 17 00:00:00 2001
From: Matt Corallo <matt@bluematt.me>
Date: Fri, 8 Jul 2011 15:08:27 +0200
Subject: [PATCH] Use DB Transactions when encrypting wallet. This speeds up
 the encryption process significantly.

---
 src/wallet.cpp | 31 +++++++++++++++++++++++++------
 src/wallet.h   |  4 ++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/wallet.cpp b/src/wallet.cpp
index 216ddc40..9f3701a8 100644
--- a/src/wallet.cpp
+++ b/src/wallet.cpp
@@ -31,7 +31,13 @@ bool CWallet::AddCryptedKey(const vector<unsigned char> &vchPubKey, const vector
         return false;
     if (!fFileBacked)
         return true;
-    return CWalletDB(strWalletFile).WriteCryptedKey(vchPubKey, vchCryptedSecret);
+    CRITICAL_BLOCK(cs_pwalletdbEncryption)
+    {
+        if (pwalletdbEncryption)
+            return pwalletdbEncryption->WriteCryptedKey(vchPubKey, vchCryptedSecret);
+        else
+            return CWalletDB(strWalletFile).WriteCryptedKey(vchPubKey, vchCryptedSecret);
+    }
 }
 
 bool CWallet::Unlock(const string& strWalletPassphrase)
@@ -104,10 +110,10 @@ bool CWallet::ChangeWalletPassphrase(const string& strOldWalletPassphrase, const
 
 bool CWallet::EncryptWallet(const string& strWalletPassphrase)
 {
-    //TODO: use db commits
     CRITICAL_BLOCK(cs_mapPubKeys)
     CRITICAL_BLOCK(cs_KeyStore)
     CRITICAL_BLOCK(cs_vMasterKey)
+    CRITICAL_BLOCK(cs_pwalletdbEncryption)
     {
         if (IsCrypted())
             return false;
@@ -146,13 +152,26 @@ bool CWallet::EncryptWallet(const string& strWalletPassphrase)
         mapMasterKeys[++nMasterKeyMaxID] = kMasterKey;
         if (fFileBacked)
         {
-            DBFlush(false);
-            CWalletDB(strWalletFile).WriteMasterKey(nMasterKeyMaxID, kMasterKey);
-            DBFlush(false);
+            pwalletdbEncryption = new CWalletDB(strWalletFile);
+            pwalletdbEncryption->TxnBegin();
+            pwalletdbEncryption->WriteMasterKey(nMasterKeyMaxID, kMasterKey);
         }
 
         if (!EncryptKeys(vMasterKey))
-            exit(1); //We now probably have half of our keys encrypted, and half not...die and let the user ask someone with experience to recover their wallet.
+        {
+            if (fFileBacked)
+                pwalletdbEncryption->TxnAbort();
+            exit(1); //We now probably have half of our keys encrypted in memory, and half not...die and let the user reload their unencrypted wallet.
+        }
+
+        if (fFileBacked)
+        {
+            if (!pwalletdbEncryption->TxnCommit())
+                exit(1); //We now have keys encrypted in memory, but no on disk...die to avoid confusion and let the user reload their unencrypted wallet.
+
+            pwalletdbEncryption->Close();
+            pwalletdbEncryption = NULL;
+        }
 
         Lock();
     }
diff --git a/src/wallet.h b/src/wallet.h
index c0ee24f0..d336d387 100644
--- a/src/wallet.h
+++ b/src/wallet.h
@@ -18,6 +18,8 @@ private:
     bool SelectCoinsMinConf(int64 nTargetValue, int nConfMine, int nConfTheirs, std::set<std::pair<const CWalletTx*,unsigned int> >& setCoinsRet, int64& nValueRet) const;
     bool SelectCoins(int64 nTargetValue, std::set<std::pair<const CWalletTx*,unsigned int> >& setCoinsRet, int64& nValueRet) const;
 
+    CWalletDB *pwalletdbEncryption;
+    CCriticalSection cs_pwalletdbEncryption;
 
 public:
     bool fFileBacked;
@@ -34,12 +36,14 @@ public:
     {
         fFileBacked = false;
         nMasterKeyMaxID = 0;
+        pwalletdbEncryption = NULL;
     }
     CWallet(std::string strWalletFileIn)
     {
         strWalletFile = strWalletFileIn;
         fFileBacked = true;
         nMasterKeyMaxID = 0;
+        pwalletdbEncryption = NULL;
     }
 
     mutable CCriticalSection cs_mapWallet;