From 20e6710fc67df89ea24b2b13b0f75bbf7716aa26 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Wed, 20 Jul 2022 19:32:23 +0000
Subject: [PATCH] wallet: Enforce an assumption about how wallet data evolves

We never delete IVKs from the wallet, so the amount of data that can be
decrypted should strictly increase.
---
 src/wallet/wallet.cpp | 30 ++++++++++++++++++------------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 0ba5df8dd..852e47b73 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -3375,13 +3375,16 @@ bool CWallet::UpdatedNoteData(const CWalletTx& wtxIn, CWalletTx& wtx)
         auto tmp = wtxIn.mapSproutNoteData;
         // Ensure we keep any cached witnesses we may already have
         for (const std::pair <JSOutPoint, SproutNoteData> nd : wtx.mapSproutNoteData) {
-            if (tmp.count(nd.first)) {
-                if (nd.second.witnesses.size() > 0) {
-                    tmp.at(nd.first).witnesses.assign(
-                            nd.second.witnesses.cbegin(), nd.second.witnesses.cend());
-                }
-                tmp.at(nd.first).witnessHeight = nd.second.witnessHeight;
+            // Require that wtxIn's data is a superset of wtx's data. This holds
+            // because viewing keys are _never_ deleted from the wallet, so the
+            // number of detected notes can only increase.
+            assert(tmp.count(nd.first) == 1);
+
+            if (nd.second.witnesses.size() > 0) {
+                tmp.at(nd.first).witnesses.assign(
+                        nd.second.witnesses.cbegin(), nd.second.witnesses.cend());
             }
+            tmp.at(nd.first).witnessHeight = nd.second.witnessHeight;
         }
         // Now copy over the updated note data
         wtx.mapSproutNoteData = tmp;
@@ -3393,13 +3396,16 @@ bool CWallet::UpdatedNoteData(const CWalletTx& wtxIn, CWalletTx& wtx)
         // Ensure we keep any cached witnesses we may already have
 
         for (const std::pair <SaplingOutPoint, SaplingNoteData> nd : wtx.mapSaplingNoteData) {
-            if (tmp.count(nd.first)) {
-                if (nd.second.witnesses.size() > 0) {
-                    tmp.at(nd.first).witnesses.assign(
-                            nd.second.witnesses.cbegin(), nd.second.witnesses.cend());
-                }
-                tmp.at(nd.first).witnessHeight = nd.second.witnessHeight;
+            // Require that wtxIn's data is a superset of wtx's data. This holds
+            // because viewing keys are _never_ deleted from the wallet, so the
+            // number of detected notes can only increase.
+            assert(tmp.count(nd.first) == 1);
+
+            if (nd.second.witnesses.size() > 0) {
+                tmp.at(nd.first).witnesses.assign(
+                        nd.second.witnesses.cbegin(), nd.second.witnesses.cend());
             }
+            tmp.at(nd.first).witnessHeight = nd.second.witnessHeight;
         }
 
         // Now copy over the updated note data