wallet: clean up failed SQLite migrations

2 files changed, 85 insertions, 0 deletions

diff --git a/src/selftest.cpp b/src/selftest.cpp
index 7d3a96e..343d9c5 100644
--- a/src/selftest.cpp
+++ b/src/selftest.cpp
@@ -1162,6 +1162,39 @@ public:
 
 } // namespace
 
+static bool CreateSQLiteWalletWithEmptyKey(const string& strPath)
+{
+    sqlite3* pdb = NULL;
+    int ret = sqlite3_open_v2(strPath.c_str(), &pdb,
+                              SQLITE_OPEN_READWRITE |
+                              SQLITE_OPEN_CREATE |
+                              SQLITE_OPEN_FULLMUTEX,
+                              NULL);
+    if (ret != SQLITE_OK || !pdb)
+    {
+        if (pdb)
+            sqlite3_close(pdb);
+        return false;
+    }
+
+    char* pszErr = NULL;
+    bool fOk =
+        sqlite3_exec(pdb,
+                     "CREATE TABLE wallet_records ("
+                     "key BLOB PRIMARY KEY NOT NULL,"
+                     "value BLOB NOT NULL"
+                     ");",
+                     NULL, NULL, &pszErr) == SQLITE_OK &&
+        sqlite3_exec(pdb,
+                     "INSERT INTO wallet_records(key, value) "
+                     "VALUES(x'', x'01');",
+                     NULL, NULL, &pszErr) == SQLITE_OK;
+    if (pszErr)
+        sqlite3_free(pszErr);
+    sqlite3_close(pdb);
+    return fOk;
+}
+
 static int RunWalletSQLiteMigrationSelfTest()
 {
     fflush(stdout);
@@ -1265,6 +1298,26 @@ static int RunWalletSQLiteMigrationSelfTest()
                                         "Refusing to overwrite existing wallet.dat"),
                        "SQLite wallet restore refuses to overwrite wallet.dat") ? 0 : 1;
 
+        string strBadSQLite = tmp + "/bad-empty-key.sqlite";
+        nFail += Check(CreateSQLiteWalletWithEmptyKey(strBadSQLite),
+                       "bad SQLite export with an empty key can be created") ? 0 : 1;
+        string strBadRestoreDir = tmp + "/bad-restore-wallet";
+        nFail += Check(MakeDirLocal(strBadRestoreDir),
+                       "bad restore wallet directory can be created") ? 0 : 1;
+        string strBadRestoreOut = tmp + "/sqlite-restore-bad.txt";
+        vector<string> vBadRestoreArgs;
+        vBadRestoreArgs.push_back("-datadir=" + strBadRestoreDir);
+        vBadRestoreArgs.push_back("-nomanagedtor");
+        vBadRestoreArgs.push_back("-nogui");
+        vBadRestoreArgs.push_back("-walletsqliterestore=" + strBadSQLite);
+        int nBadRestoreRet =
+            RunBitflashChild(strExe, vBadRestoreArgs, NULL, &strBadRestoreOut);
+        nFail += Check(nBadRestoreRet == 1 &&
+                       FileContainsText(strBadRestoreOut,
+                                        "Cannot restore SQLite wallet export") &&
+                       !FileExists((strBadRestoreDir + "/wallet.dat").c_str()),
+                       "failed SQLite wallet restore removes incomplete 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 d1bad98..9e4a981 100644
--- a/src/walletcmd.cpp
+++ b/src/walletcmd.cpp
@@ -798,6 +798,19 @@ static vector<unsigned char> DataStreamBytesLocal(const CDataStream& ss)
 
 namespace {
 
+static void RemoveSQLiteExportFiles(const string& strPath)
+{
+    remove(strPath.c_str());
+    remove((strPath + "-wal").c_str());
+    remove((strPath + "-shm").c_str());
+}
+
+static void RemoveFailedRestoreWallet(const string& strWalletPath)
+{
+    DBFlush(true);
+    remove(strWalletPath.c_str());
+}
+
 class CWalletSQLiteExportVisitor : public CWalletRecordVisitor
 {
 public:
@@ -912,11 +925,14 @@ int CmdWalletSQLiteExport(const std::string& strDest)
     CWalletDBSQLite db;
     if (!db.Open(strDest, strError))
     {
+        RemoveSQLiteExportFiles(strDest);
         fprintf(stderr, "Cannot open SQLite wallet export: %s\n", strError.c_str());
         return 1;
     }
     if (!db.BeginTransaction(strError))
     {
+        db.Close();
+        RemoveSQLiteExportFiles(strDest);
         fprintf(stderr, "Cannot begin SQLite wallet export: %s\n", strError.c_str());
         return 1;
     }
@@ -927,16 +943,22 @@ int CmdWalletSQLiteExport(const std::string& strDest)
     {
         string strRollbackError;
         db.RollbackTransaction(strRollbackError);
+        db.Close();
+        RemoveSQLiteExportFiles(strDest);
         fprintf(stderr, "Cannot export wallet.dat to SQLite: %s\n", strError.c_str());
         return 1;
     }
     if (!db.CommitTransaction(strError))
     {
+        db.Close();
+        RemoveSQLiteExportFiles(strDest);
         fprintf(stderr, "Cannot commit SQLite wallet export: %s\n", strError.c_str());
         return 1;
     }
     if (!db.Checkpoint(strError))
     {
+        db.Close();
+        RemoveSQLiteExportFiles(strDest);
         fprintf(stderr, "Cannot checkpoint SQLite wallet export: %s\n", strError.c_str());
         return 1;
     }
@@ -944,12 +966,16 @@ int CmdWalletSQLiteExport(const std::string& strDest)
     int nCount = 0;
     if (!db.CountRecords(nCount, strError))
     {
+        db.Close();
+        RemoveSQLiteExportFiles(strDest);
         fprintf(stderr, "Cannot verify SQLite wallet export count: %s\n",
                 strError.c_str());
         return 1;
     }
     if (nCount != nCopied)
     {
+        db.Close();
+        RemoveSQLiteExportFiles(strDest);
         fprintf(stderr, "SQLite wallet export count mismatch: copied %d, stored %d\n",
                 nCopied, nCount);
         return 1;
@@ -1097,6 +1123,8 @@ int CmdWalletSQLiteRestore(const std::string& strPath)
     if (!sqlite.ScanRecords(visitor, strError))
     {
         bdb.TxnAbort();
+        bdb.Close();
+        RemoveFailedRestoreWallet(strWalletPath);
         fprintf(stderr, "Cannot restore SQLite wallet export: %s\n",
                 strError.c_str());
         return 1;
@@ -1104,12 +1132,16 @@ int CmdWalletSQLiteRestore(const std::string& strPath)
     if (nCopied != nSQLiteRecords)
     {
         bdb.TxnAbort();
+        bdb.Close();
+        RemoveFailedRestoreWallet(strWalletPath);
         fprintf(stderr, "SQLite wallet restore copied %d records, expected %d.\n",
                 nCopied, nSQLiteRecords);
         return 2;
     }
     if (!bdb.TxnCommit())
     {
+        bdb.Close();
+        RemoveFailedRestoreWallet(strWalletPath);
         fprintf(stderr, "Cannot commit Berkeley DB restore transaction.\n");
         return 1;
     }