wallet: route rewrite copies through record scan

2 files changed, 54 insertions, 55 deletions

diff --git a/docs/storage.md b/docs/storage.md
index e533589..e4a66d2 100644
--- a/docs/storage.md
+++ b/docs/storage.md
@@ -43,7 +43,9 @@ level schema migration is intentionally introduced.
 `wallet.dat` is accessed through `CWalletDB`.
 Read-only tooling should go through `ScanWalletRecords()` instead of opening a
 Berkeley DB cursor directly. That keeps migration checks, audits, and future
-backends behind one record-stream boundary.
+backends behind one record-stream boundary. Wallet rewrite code should also use
+that boundary when it needs to copy existing records into a newly written
+wallet file.
 
 | Key shape | Value | Secret? | Meaning |
 |---|---|---:|---|
diff --git a/src/db.cpp b/src/db.cpp
index 4fa72af..b90219e 100644
--- a/src/db.cpp
+++ b/src/db.cpp
@@ -1354,70 +1354,68 @@ public:
     {
         return Write(make_pair(string("pool"), nIndex), vchPubKey);
     }
+};
+
+static bool IsWalletEncryptionPrivateRecord(const string& strType)
+{
+    return strType == "key" ||
+           strType == "pool" ||
+           strType == "hdmaster" ||
+           strType == "hdchaincode" ||
+           strType == "mkey" ||
+           strType == "ckey" ||
+           strType == "cryptedhdmaster" ||
+           strType == "cryptedhdchaincode" ||
+           strType == "walletminversion";
+}
 
-    bool CopyPublicRecordsTo(CWalletRewriteDB& dbTo, int& nCopiedRet, string& strErrorRet)
+class CWalletPublicRecordCopyVisitor : public CWalletRecordVisitor
+{
+public:
+    CWalletRewriteDB& dbTo;
+    int& nCopied;
+
+    CWalletPublicRecordCopyVisitor(CWalletRewriteDB& dbToIn, int& nCopiedIn)
+        : dbTo(dbToIn), nCopied(nCopiedIn) { }
+
+    bool VisitWalletRecord(const CDataStream& ssKey,
+                           const CDataStream& ssValue,
+                           string& strErrorRet)
     {
-        nCopiedRet = 0;
-        Dbc* pcursor = GetCursor();
-        if (!pcursor)
+        string strType;
+        try
         {
-            strErrorRet = "could not open a wallet cursor";
-            return false;
+            CDataStream ssType = ssKey;
+            ssType >> strType;
         }
-
-        unsigned int fFlags = DB_NEXT;
-        loop
+        catch (...)
         {
-            CDataStream ssKey;
-            CDataStream ssValue;
-            int ret = ReadAtCursor(pcursor, ssKey, ssValue, fFlags);
-            if (ret == DB_NOTFOUND)
-                break;
-            if (ret != 0)
-            {
-                pcursor->close();
-                strErrorRet = strprintf("could not read wallet record: Berkeley DB error %d", ret);
-                return false;
-            }
-
-            string strType;
-            try
-            {
-                CDataStream ssType = ssKey;
-                ssType >> strType;
-            }
-            catch (...)
-            {
-                pcursor->close();
-                strErrorRet = "wallet.dat has a record whose key type cannot be read";
-                return false;
-            }
+            strErrorRet = "wallet.dat has a record whose key type cannot be read";
+            return false;
+        }
 
-            if (strType == "key" ||
-                strType == "pool" ||
-                strType == "hdmaster" ||
-                strType == "hdchaincode" ||
-                strType == "mkey" ||
-                strType == "ckey" ||
-                strType == "cryptedhdmaster" ||
-                strType == "cryptedhdchaincode" ||
-                strType == "walletminversion")
-                continue;
+        if (IsWalletEncryptionPrivateRecord(strType))
+            return true;
 
-            if (!dbTo.WriteRaw(ssKey, ssValue))
-            {
-                pcursor->close();
-                strErrorRet = strprintf("could not copy wallet record '%s'", strType.c_str());
-                return false;
-            }
-            nCopiedRet++;
+        if (!dbTo.WriteRaw(ssKey, ssValue))
+        {
+            strErrorRet = strprintf("could not copy wallet record '%s'", strType.c_str());
+            return false;
         }
-
-        pcursor->close();
+        nCopied++;
         return true;
     }
 };
 
+static bool CopyPublicWalletRecordsTo(CWalletRewriteDB& dbTo,
+                                      int& nCopiedRet,
+                                      string& strErrorRet)
+{
+    nCopiedRet = 0;
+    CWalletPublicRecordCopyVisitor visitor(dbTo, nCopiedRet);
+    return ScanWalletRecords(visitor, strErrorRet);
+}
+
 static string WalletEncryptBackupPath()
 {
     string strBase = GetAppDir() + "/wallet.dat.before-encrypt.";
@@ -1594,9 +1592,8 @@ bool EncryptWallet(const string& strPassphrase, string& strBackupRet, string& st
 
     int nCopied = 0;
     {
-        CWalletRewriteDB dbSource("wallet.dat", "r");
         CWalletRewriteDB dbTarget(strTempFile.c_str(), "c+");
-        if (!dbSource.CopyPublicRecordsTo(dbTarget, nCopied, strErrorRet))
+        if (!CopyPublicWalletRecordsTo(dbTarget, nCopied, strErrorRet))
         {
             remove(strTempPath.c_str());
             return false;