wallet: restore BDB wallets from SQLite exports

4 files changed, 176 insertions, 0 deletions

diff --git a/src/main_gui.cpp b/src/main_gui.cpp
index caacf13..cf96822 100644
--- a/src/main_gui.cpp
+++ b/src/main_gui.cpp
@@ -185,6 +185,8 @@ static void PrintUsage()
     printf("                              experimental SQLite store, then exit)\n");
     printf("  /walletsqliteverify=FILE   (compare wallet.dat with a SQLite export,\n");
     printf("                              printing counts only, then exit)\n");
+    printf("  /walletsqliterestore=FILE  (rebuild wallet.dat from a SQLite export in\n");
+    printf("                              an empty data directory, then exit)\n");
     printf("  /rescan                    (walk the chain for coins this wallet owns\n");
     printf("                              but never recorded, then exit)\n");
     printf("\n");
@@ -551,6 +553,15 @@ int main(int argc, char* argv[])
         return nRet;
     }
 
+    string strWalletSQLiteRestore =
+        argval2(argc, argv, "/walletsqliterestore", "-walletsqliterestore");
+    if (!strWalletSQLiteRestore.empty())
+    {
+        int nRet = CmdWalletSQLiteRestore(strWalletSQLiteRestore);
+        DBFlush(true);
+        return nRet;
+    }
+
     // Berkeley DB reports failure by throwing, and CDB's constructor lets it
     // through. Nothing on this path caught anything, so an unreadable
     // blkindex.dat or wallet.dat -- one written by another platform's Berkeley
diff --git a/src/selftest.cpp b/src/selftest.cpp
index 4ba90c3..7d3a96e 100644
--- a/src/selftest.cpp
+++ b/src/selftest.cpp
@@ -1230,6 +1230,41 @@ static int RunWalletSQLiteMigrationSelfTest()
                        FileContainsText(strVerifyOut, "verification:              ok"),
                        "SQLite wallet export verifies against wallet.dat") ? 0 : 1;
 
+        string strRestoredDir = tmp + "/restored-wallet";
+        nFail += Check(MakeDirLocal(strRestoredDir),
+                       "restored wallet directory can be created") ? 0 : 1;
+        string strRestoreOut = tmp + "/sqlite-restore.txt";
+        vector<string> vRestoreArgs;
+        vRestoreArgs.push_back("-datadir=" + strRestoredDir);
+        vRestoreArgs.push_back("-nomanagedtor");
+        vRestoreArgs.push_back("-nogui");
+        vRestoreArgs.push_back("-walletsqliterestore=" + strSQLite);
+        int nRestoreRet = RunBitflashChild(strExe, vRestoreArgs, NULL, &strRestoreOut);
+        nFail += Check(nRestoreRet == 0 &&
+                       FileContainsText(strRestoreOut, "records restored:"),
+                       "SQLite wallet export restores into a fresh wallet.dat") ? 0 : 1;
+
+        string strVerifyRestoredOut = tmp + "/sqlite-verify-restored.txt";
+        vector<string> vVerifyRestoredArgs;
+        vVerifyRestoredArgs.push_back("-datadir=" + strRestoredDir);
+        vVerifyRestoredArgs.push_back("-nomanagedtor");
+        vVerifyRestoredArgs.push_back("-nogui");
+        vVerifyRestoredArgs.push_back("-walletsqliteverify=" + strSQLite);
+        int nVerifyRestoredRet =
+            RunBitflashChild(strExe, vVerifyRestoredArgs, NULL, &strVerifyRestoredOut);
+        nFail += Check(nVerifyRestoredRet == 0 &&
+                       FileContainsText(strVerifyRestoredOut,
+                                        "verification:              ok"),
+                       "restored wallet.dat verifies against the SQLite export") ? 0 : 1;
+
+        string strRestoreOverwriteOut = tmp + "/sqlite-restore-overwrite.txt";
+        int nRestoreOverwriteRet =
+            RunBitflashChild(strExe, vRestoreArgs, NULL, &strRestoreOverwriteOut);
+        nFail += Check(nRestoreOverwriteRet == 1 &&
+                       FileContainsText(strRestoreOverwriteOut,
+                                        "Refusing to overwrite existing wallet.dat"),
+                       "SQLite wallet restore refuses to overwrite wallet.dat") ? 0 : 1;
+
         string strOverwriteOut = tmp + "/sqlite-export-overwrite.txt";
         int nOverwriteRet =
             RunBitflashChild(strExe, vExportArgs, NULL, &strOverwriteOut);
diff --git a/src/walletcmd.cpp b/src/walletcmd.cpp
index 36f6783..d1bad98 100644
--- a/src/walletcmd.cpp
+++ b/src/walletcmd.cpp
@@ -844,6 +844,52 @@ public:
     }
 };
 
+class CWalletRawRestoreDB : public CWalletDB
+{
+public:
+    CWalletRawRestoreDB(const char* pszMode="cr+", bool fTxn=true)
+        : CWalletDB(pszMode, fTxn) { }
+
+    bool WriteRaw(CDataStream ssKey, CDataStream ssValue)
+    {
+        if (!pdb || ssKey.empty())
+            return false;
+
+        Dbt datKey((void*)&ssKey[0], ssKey.size());
+        Dbt datValue(ssValue.empty() ? NULL : (void*)&ssValue[0],
+                     ssValue.size());
+        int ret = pdb->put(GetTxn(), &datKey, &datValue, 0);
+
+        memset(datKey.get_data(), 0, datKey.get_size());
+        if (datValue.get_data())
+            memset(datValue.get_data(), 0, datValue.get_size());
+        return ret == 0;
+    }
+};
+
+class CWalletSQLiteRestoreVisitor : public CWalletRecordVisitor
+{
+public:
+    CWalletRawRestoreDB& db;
+    int& nCopied;
+
+    CWalletSQLiteRestoreVisitor(CWalletRawRestoreDB& dbIn, int& nCopiedIn)
+        : db(dbIn), nCopied(nCopiedIn) { }
+
+    bool VisitWalletRecord(const CDataStream& ssKey,
+                           const CDataStream& ssValue,
+                           string& strErrorRet)
+    {
+        if (!db.WriteRaw(ssKey, ssValue))
+        {
+            strErrorRet = "could not write raw wallet record";
+            return false;
+        }
+        nCopied++;
+        return true;
+    }
+};
+
 } // namespace
 
 int CmdWalletSQLiteExport(const std::string& strDest)
@@ -994,6 +1040,86 @@ int CmdWalletSQLiteVerify(const std::string& strPath)
     return 2;
 }
 
+int CmdWalletSQLiteRestore(const std::string& strPath)
+{
+    AttachTerminal();
+
+    if (strPath.empty())
+    {
+        fprintf(stderr, "Missing SQLite wallet export path.\n");
+        return 1;
+    }
+    if (!FileExists(strPath.c_str()))
+    {
+        fprintf(stderr, "SQLite wallet export does not exist: %s\n",
+                strPath.c_str());
+        return 1;
+    }
+
+    string strWalletPath = GetAppDir() + "/wallet.dat";
+    if (FileExists(strWalletPath.c_str()))
+    {
+        fprintf(stderr, "Refusing to overwrite existing wallet.dat: %s\n",
+                strWalletPath.c_str());
+        return 1;
+    }
+
+    string strError;
+    CWalletDBSQLite sqlite;
+    if (!sqlite.Open(strPath, strError))
+    {
+        fprintf(stderr, "Cannot open SQLite wallet export: %s\n", strError.c_str());
+        return 1;
+    }
+
+    int nSQLiteRecords = 0;
+    if (!sqlite.CountRecords(nSQLiteRecords, strError))
+    {
+        fprintf(stderr, "Cannot count SQLite wallet export records: %s\n",
+                strError.c_str());
+        return 1;
+    }
+    if (nSQLiteRecords <= 0)
+    {
+        fprintf(stderr, "SQLite wallet export contains no records.\n");
+        return 2;
+    }
+
+    CWalletRawRestoreDB bdb("cr+", true);
+    if (!bdb.TxnBegin())
+    {
+        fprintf(stderr, "Cannot begin Berkeley DB restore transaction.\n");
+        return 1;
+    }
+
+    int nCopied = 0;
+    CWalletSQLiteRestoreVisitor visitor(bdb, nCopied);
+    if (!sqlite.ScanRecords(visitor, strError))
+    {
+        bdb.TxnAbort();
+        fprintf(stderr, "Cannot restore SQLite wallet export: %s\n",
+                strError.c_str());
+        return 1;
+    }
+    if (nCopied != nSQLiteRecords)
+    {
+        bdb.TxnAbort();
+        fprintf(stderr, "SQLite wallet restore copied %d records, expected %d.\n",
+                nCopied, nSQLiteRecords);
+        return 2;
+    }
+    if (!bdb.TxnCommit())
+    {
+        fprintf(stderr, "Cannot commit Berkeley DB restore transaction.\n");
+        return 1;
+    }
+
+    printf("SQLite wallet export restored to wallet.dat\n");
+    printf("  records restored:          %d\n", nCopied);
+    fflush(stdout);
+    return 0;
+}
+
 int CmdEncryptWallet(const std::string& strPassphrase)
 {
     AttachTerminal();
diff --git a/src/walletcmd.h b/src/walletcmd.h
index 784c74c..50dc69e 100644
--- a/src/walletcmd.h
+++ b/src/walletcmd.h
@@ -79,6 +79,10 @@ int CmdWalletSQLiteExport(const std::string& strDest);
 // wallet.dat raw records. Prints counts only, never record keys or values.
 int CmdWalletSQLiteVerify(const std::string& strPath);
 
+// Diagnostic/migration staging: rebuild wallet.dat from a SQLite export in an
+// empty data directory. Refuses to overwrite an existing wallet.dat.
+int CmdWalletSQLiteRestore(const std::string& strPath);
+
 // Rewrite wallet.dat so private keys and the HD seed are encrypted with the
 // given passphrase. Returns 0 on success and exits without starting the node.
 int CmdEncryptWallet(const std::string& strPassphrase);