From 58e4a76e60fdfe0920b94e08638fad8a15f24791 Mon Sep 17 00:00:00 2001
From: Larry Ruane <larry@z.cash>
Date: Mon, 4 Oct 2021 23:27:10 -0600
Subject: [PATCH] better wallet network info error handling

---
 src/wallet/gtest/test_wallet.cpp | 10 ++++++++--
 src/wallet/wallet.cpp            | 14 ++++++--------
 src/wallet/wallet.h              |  2 +-
 src/wallet/walletdb.cpp          | 19 ++++++++++++-------
 src/wallet/walletdb.h            |  3 ++-
 5 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/src/wallet/gtest/test_wallet.cpp b/src/wallet/gtest/test_wallet.cpp
index 4a47804b3..edc060585 100644
--- a/src/wallet/gtest/test_wallet.cpp
+++ b/src/wallet/gtest/test_wallet.cpp
@@ -138,9 +138,15 @@ TEST(WalletTests, WalletNetworkSerialization) {
     // create a new testnet wallet and generate a seed
     CWallet wallet(Params(), "wallet.dat");
     wallet.InitLoadWallet(Params(), true);
-    wallet.GenerateNewSeed();
     wallet.Flush();
 
+    // Stay on TESTNET, make sure wallet can be successfully loaded.
+    {
+        CWallet restored(Params(), "wallet.dat");
+        bool fFirstRunRet;
+        EXPECT_EQ(restored.LoadWallet(fFirstRunRet), DB_LOAD_OK);
+    }
+
     // now, switch to mainnet and attempt to restore the wallet
     // using the same wallet.dat
     SelectParams(CBaseChainParams::MAIN);
@@ -148,7 +154,7 @@ TEST(WalletTests, WalletNetworkSerialization) {
 
     // load should fail due to being associated with the wrong network
     bool fFirstRunRet;
-    EXPECT_NE(restored.LoadWallet(fFirstRunRet), DB_LOAD_OK);
+    EXPECT_EQ(restored.LoadWallet(fFirstRunRet), DB_WRONG_NETWORK);
 }
 
 TEST(WalletTests, SproutNoteDataSerialisation) {
diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 7d26262c4..f8fc1d30f 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -2295,17 +2295,11 @@ void CWallet::SetHDChain(const CHDChain& chain, bool memonly)
     hdChain = chain;
 }
 
-void CWallet::CheckNetworkInfo(std::pair<std::string, std::string> readNetworkInfo)
+bool CWallet::CheckNetworkInfo(std::pair<std::string, std::string> readNetworkInfo)
 {
     LOCK(cs_wallet);
     std::pair<string, string> networkInfo(PACKAGE_NAME, networkIdString);
-    if (readNetworkInfo != networkInfo)
-        throw std::runtime_error(
-                strprintf("%s: this wallet is for a different network (%s, %s) than the node is configured for (%s, %s)",
-                          std::string(__func__),
-                          readNetworkInfo.first, readNetworkInfo.second,
-                          networkInfo.first, networkInfo.second)
-        );
+    return readNetworkInfo == networkInfo;
 }
 
 uint32_t CWallet::BIP44CoinType() {
@@ -4773,6 +4767,10 @@ bool CWallet::InitLoadWallet(const CChainParams& params, bool clearWitnessCaches
         {
             return UIError(strprintf(_("Wallet needed to be rewritten: restart %s to complete"), _(PACKAGE_NAME)));
         }
+        else if (nLoadWalletRet == DB_WRONG_NETWORK)
+        {
+            return UIError(strprintf(_("Wallet %s is not for %s %s network"), walletFile, _(PACKAGE_NAME), params.NetworkIDString()));
+        }
         else
             return UIError(strprintf(_("Error loading %s"), walletFile));
     }
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 1bd9f1682..b0d5c2df7 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -1298,7 +1298,7 @@ public:
     void SetHDChain(const CHDChain& chain, bool memonly);
     const CHDChain& GetHDChain() const { return hdChain; }
 
-    void CheckNetworkInfo(std::pair<std::string, std::string> networkInfo);
+    bool CheckNetworkInfo(std::pair<std::string, std::string> networkInfo);
     uint32_t BIP44CoinType();
 
     /* Set the current HD seed, without saving it to disk (used by LoadWallet) */
diff --git a/src/wallet/walletdb.cpp b/src/wallet/walletdb.cpp
index 7ac13f32a..2793a9cf5 100644
--- a/src/wallet/walletdb.cpp
+++ b/src/wallet/walletdb.cpp
@@ -728,7 +728,10 @@ ReadKeyValue(CWallet* pwallet, CDataStream& ssKey, CDataStream& ssValue,
         {
             std::pair<std::string, std::string> networkInfo;
             ssValue >> networkInfo;
-            pwallet->CheckNetworkInfo(networkInfo);
+            if (!pwallet->CheckNetworkInfo(networkInfo)) {
+                strErr = "Error in wallet database: unexpected network";
+                return false;
+            }
         }
     } catch (...)
     {
@@ -790,17 +793,19 @@ DBErrors CWalletDB::LoadWallet(CWallet* pwallet)
             string strType, strErr;
             if (!ReadKeyValue(pwallet, ssKey, ssValue, wss, strType, strErr))
             {
-                // losing keys is considered a catastrophic error, anything else
-                // we assume the user can live with:
-                if (IsKeyType(strType))
+                if (strType == "networkinfo") {
+                    // example: running mainnet, but this wallet.dat is from testnet
+                    result = DB_WRONG_NETWORK;
+                } else if (result != DB_WRONG_NETWORK && IsKeyType(strType)) {
+                    // losing keys is considered a catastrophic error
                     result = DB_CORRUPT;
-                else
-                {
+                } else {
                     // Leave other errors alone, if we try to fix them we might make things worse.
                     fNoncriticalErrors = true; // ... but do warn the user there is something wrong.
-                    if (strType == "tx")
+                    if (strType == "tx") {
                         // Rescan if there is a bad transaction record:
                         SoftSetBoolArg("-rescan", true);
+                    }
                 }
             }
             if (!strErr.empty())
diff --git a/src/wallet/walletdb.h b/src/wallet/walletdb.h
index 5d0a2d0de..5518e38d9 100644
--- a/src/wallet/walletdb.h
+++ b/src/wallet/walletdb.h
@@ -37,7 +37,8 @@ enum DBErrors
     DB_NONCRITICAL_ERROR,
     DB_TOO_NEW,
     DB_LOAD_FAIL,
-    DB_NEED_REWRITE
+    DB_NEED_REWRITE,
+    DB_WRONG_NETWORK,
 };
 
 /* simple hd chain data model */