From cd10562c73a773e7926118d677e7fafb34ff1b87 Mon Sep 17 00:00:00 2001
From: Simon <simon@bitcartel.com>
Date: Thu, 11 Oct 2018 21:25:53 -0700
Subject: [PATCH 1/2] Resolves Sapling nullifier persistence issue when
 importing a key.

During a rescan, a CWalletTx was persisted to disk before it had its
note data set.  This meant that upon restart, the CWalletTx would
potentially be missing its nullifiers causing the wallet's balance
to include notes which had already been spent.

The resolution is to ensure that after a rescan, a CWalletTx is
persisted after it has had its nullifiers set correctly.

Co-authored-by: Eirik Ogilvie-Wigley <eirik@z.cash>
---
 src/wallet/wallet.cpp | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 3690ed166..b9da65454 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -2398,6 +2398,9 @@ int CWallet::ScanForWalletTransactions(CBlockIndex* pindexStart, bool fUpdate)
     const CChainParams& chainParams = Params();
 
     CBlockIndex* pindex = pindexStart;
+
+    std::vector<uint256> myTxHashes;
+
     {
         LOCK2(cs_main, cs_wallet);
 
@@ -2418,8 +2421,10 @@ int CWallet::ScanForWalletTransactions(CBlockIndex* pindexStart, bool fUpdate)
             ReadBlockFromDisk(block, pindex);
             BOOST_FOREACH(CTransaction& tx, block.vtx)
             {
-                if (AddToWalletIfInvolvingMe(tx, &block, fUpdate))
+                if (AddToWalletIfInvolvingMe(tx, &block, fUpdate)) {
+                    myTxHashes.push_back(tx.GetHash());
                     ret++;
+                }
             }
 
             SproutMerkleTree sproutTree;
@@ -2441,6 +2446,19 @@ int CWallet::ScanForWalletTransactions(CBlockIndex* pindexStart, bool fUpdate)
                 LogPrintf("Still rescanning. At block %d. Progress=%f\n", pindex->nHeight, Checkpoints::GuessVerificationProgress(chainParams.Checkpoints(), pindex));
             }
         }
+
+        // After rescanning, persist Sapling note data that might have changed, e.g. nullifiers.
+        // Do not flush the wallet here for performance reasons.
+        CWalletDB walletdb(strWalletFile, "r+", false);
+        for (auto hash : myTxHashes) {
+            CWalletTx wtx = mapWallet[hash];
+            if (!wtx.mapSaplingNoteData.empty()) {
+                if (!wtx.WriteToDisk(&walletdb)) {
+                    LogPrintf("Rescanning... WriteToDisk failed to update Sapling note data for: %s\n", hash.ToString());
+                }
+            }
+        }
+
         ShowProgress(_("Rescanning..."), 100); // hide progress dialog in GUI
     }
     return ret;

From 01918e8742028846259f0b1389b17179a8df92f9 Mon Sep 17 00:00:00 2001
From: Simon <simon@bitcartel.com>
Date: Thu, 11 Oct 2018 21:52:56 -0700
Subject: [PATCH 2/2] Update test to verify Sapling nullifiers and witnesses
 persist correctly.

---
 qa/rpc-tests/wallet_persistence.py | 36 +++++++++++++++++++++++++++---
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/qa/rpc-tests/wallet_persistence.py b/qa/rpc-tests/wallet_persistence.py
index 612226275..f6c227d2a 100755
--- a/qa/rpc-tests/wallet_persistence.py
+++ b/qa/rpc-tests/wallet_persistence.py
@@ -16,15 +16,16 @@ class WalletPersistenceTest (BitcoinTestFramework):
 
     def setup_chain(self):
         print("Initializing test directory " + self.options.tmpdir)
-        initialize_chain_clean(self.options.tmpdir, 2)
+        initialize_chain_clean(self.options.tmpdir, 3)
 
     def setup_network(self, split=False):
-        self.nodes = start_nodes(2, self.options.tmpdir,
+        self.nodes = start_nodes(3, self.options.tmpdir,
             extra_args=[[
                 '-nuparams=5ba81b19:100', # Overwinter
                 '-nuparams=76b809bb:201', # Sapling
-            ]] * 2)
+            ]] * 3)
         connect_nodes_bi(self.nodes,0,1)
+        connect_nodes_bi(self.nodes,1,2)
         self.is_network_split=False
         self.sync_all()
 
@@ -97,5 +98,34 @@ class WalletPersistenceTest (BitcoinTestFramework):
         assert_equal(self.nodes[0].z_getbalance(sapling_addr), Decimal('5'))
         assert_equal(self.nodes[1].z_getbalance(dest_addr), Decimal('15'))
 
+        # Verify importing a spending key will update and persist the nullifiers and witnesses correctly
+        sk0 = self.nodes[0].z_exportkey(sapling_addr)
+        self.nodes[2].z_importkey(sk0, "yes")
+        assert_equal(self.nodes[2].z_getbalance(sapling_addr), Decimal('5'))
+
+        # Restart the nodes
+        stop_nodes(self.nodes)
+        wait_bitcoinds()
+        self.setup_network()
+
+        # Verify nullifiers persisted correctly by checking balance
+        # Prior to PR #3590, there will be an error as spent notes are considered unspent:
+        #    Assertion failed: expected: <25.00000000> but was: <5>
+        assert_equal(self.nodes[2].z_getbalance(sapling_addr), Decimal('5'))
+
+        # Verity witnesses persisted correctly by sending shielded funds
+        recipients = []
+        recipients.append({"address": dest_addr, "amount": Decimal('1')})
+        myopid = self.nodes[2].z_sendmany(sapling_addr, recipients, 1, 0)
+        wait_and_assert_operationid_status(self.nodes[2], myopid)
+
+        self.sync_all()
+        self.nodes[0].generate(1)
+        self.sync_all()
+
+        # Verify balances
+        assert_equal(self.nodes[2].z_getbalance(sapling_addr), Decimal('4'))
+        assert_equal(self.nodes[1].z_getbalance(dest_addr), Decimal('16'))
+
 if __name__ == '__main__':
     WalletPersistenceTest().main()
\ No newline at end of file