[Sapling] Try to recover corruption of notes cache during send operation

src/sapling/sapling_operation.cpp

@@ -339,9 +339,45 @@ OperationResult SaplingOperation::loadUtxos(TxValues& txValues, const std::vecto
     return OperationResult(true);
 }
 
+/*
+ * Check that the witness and nullifier of a sapling note (about to be spent) have been
+ * correctly cached. If the witness is missing, return an error. If the nullifier is missing,
+ * recover it from the note (now that we have the spending key).
+ */
+enum CacheCheckResult {OK, SPENT, INVALID};
+static CacheCheckResult CheckCachedNote(const SaplingNoteEntry& t, const libzcash::SaplingExpandedSpendingKey& expsk)
+{
+    auto sspkm = pwalletMain->GetSaplingScriptPubKeyMan();
+    CWalletTx& prevTx = pwalletMain->mapWallet.at(t.op.hash);
+    SaplingNoteData& nd = prevTx.mapSaplingNoteData.at(t.op);
+    if (nd.witnesses.empty()) {
+        return CacheCheckResult::INVALID;
+    }
+    if (nd.nullifier == nullopt) {
+        const std::string& noteStr = t.op.ToString();
+        LogPrintf("WARNING: nullifier not cached for note %s. Updating...\n", noteStr);
+        // get the nullifier from the note and update the cache
+        const auto& witness = nd.witnesses.front();
+        const Optional<uint256> nf = t.note.nullifier(expsk.full_viewing_key(), witness.position());
+        // check that it's valid
+        if (nf == nullopt) {
+            LogPrintf("ERROR: Unable to recover nullifier for note %s.\n", noteStr);
+            return CacheCheckResult::INVALID;
+        }
+        WITH_LOCK(pwalletMain->cs_wallet, sspkm->UpdateSaplingNullifierNoteMap(nd, t.op, nf));
+        // re-check the spent status
+        if (sspkm->IsSaplingSpent(*(nd.nullifier))) {
+            LogPrintf("Removed note %s as it appears to be already spent.\n", noteStr);
+            return CacheCheckResult::SPENT;
+        }
+    }
+    return CacheCheckResult::OK;
+}
+
 OperationResult SaplingOperation::loadUnspentNotes(TxValues& txValues, uint256& ovk)
 {
     shieldedInputs.clear();
+    auto sspkm = pwalletMain->GetSaplingScriptPubKeyMan();
     // if we already have selected the notes, let's directly set them.
     bool hasCoinControl = coinControl && coinControl->HasSelected();
     if (hasCoinControl) {
@@ -355,18 +391,18 @@ OperationResult SaplingOperation::loadUnspentNotes(TxValues& txValues, uint256&
             vSaplingOutpoints.emplace_back(outpoint.outPoint.hash, outpoint.outPoint.n);
         }
 
-        pwalletMain->GetSaplingScriptPubKeyMan()->GetNotes(vSaplingOutpoints, shieldedInputs);
+        sspkm->GetNotes(vSaplingOutpoints, shieldedInputs);
 
         if (shieldedInputs.empty()) {
             return errorOut("Insufficient funds, no available notes to spend");
         }
     } else {
         // If we don't have coinControl then let's find the notes
-        pwalletMain->GetSaplingScriptPubKeyMan()->GetFilteredNotes(shieldedInputs, fromAddress.fromSapAddr, mindepth);
+        sspkm->GetFilteredNotes(shieldedInputs, fromAddress.fromSapAddr, mindepth);
         if (shieldedInputs.empty()) {
             // Just to notify the user properly, check if the wallet has notes with less than the min depth
             std::vector<SaplingNoteEntry> _shieldedInputs;
-            pwalletMain->GetSaplingScriptPubKeyMan()->GetFilteredNotes(_shieldedInputs, fromAddress.fromSapAddr, 0);
+            sspkm->GetFilteredNotes(_shieldedInputs, fromAddress.fromSapAddr, 0);
             return errorOut(_shieldedInputs.empty() ?
                     "Insufficient funds, no available notes to spend" :
                     "Insufficient funds, shielded PIV need at least 5 confirmations");
@@ -392,14 +428,26 @@ OperationResult SaplingOperation::loadUnspentNotes(TxValues& txValues, uint256&
         uint256 ovkIn;
         auto resLoadKeys = loadKeysFromShieldedFrom(t.address, expsk, ovkIn);
         if (!resLoadKeys) return resLoadKeys;
-        spendingKeys.emplace_back(expsk);
+
+        // If the noteData is not properly cached, for whatever reason,
+        // try to update it here, now that we have the spending key.
+        CacheCheckResult res = CheckCachedNote(t, expsk);
+        if (res == CacheCheckResult::INVALID) {
+            // This should never happen. User would be forced to zap.
+            LogPrintf("ERROR: Witness/Nullifier invalid for note %s. Restart with --zapwallettxes\n", t.op.ToString());
+            return errorOut("Note cache corrupt. Try \"Recover transactions\" (Settings-->Debug-->\"wallet repair\")");
+        } else if (res == CacheCheckResult::SPENT) {
+            // note was already spent, don't include it in the inputs
+            continue;
+        }
 
         // Return ovk to be used in the outputs
         if (ovk.IsNull()) {
             ovk = ovkIn;
         }
 
         // Load data
+        spendingKeys.emplace_back(expsk);
         ops.emplace_back(t.op);
         notes.emplace_back(t.note);
         txValues.shieldedInTotal += t.note.value();

src/sapling/saplingscriptpubkeyman.cpp

@@ -43,22 +43,30 @@ void SaplingScriptPubKeyMan::UpdateSaplingNullifierNoteMapForBlock(const CBlock
     }
 }
 
+// Updates noteData and mapSaplingNullifiersToNotes directly
+void SaplingScriptPubKeyMan::UpdateSaplingNullifierNoteMap(SaplingNoteData& nd, const SaplingOutPoint& op, const Optional<uint256>& nullifier)
+{
+    AssertLockHeld(wallet->cs_wallet);
+    nd.nullifier = nullifier;
+    if (nullifier) mapSaplingNullifiersToNotes[*nullifier] = op;
+}
+
 /**
  * Update mapSaplingNullifiersToNotes, computing the nullifier from a cached witness if necessary.
  */
 void SaplingScriptPubKeyMan::UpdateSaplingNullifierNoteMapWithTx(CWalletTx& wtx) {
     LOCK(wallet->cs_wallet);
 
     for (mapSaplingNoteData_t::value_type &item : wtx.mapSaplingNoteData) {
-        SaplingOutPoint op = item.first;
-        SaplingNoteData nd = item.second;
+        const SaplingOutPoint& op = item.first;
+        SaplingNoteData& nd = item.second;
 
         if (nd.witnesses.empty() || !nd.IsMyNote()) {
             // If there are no witnesses, erase the nullifier and associated mapping.
-            if (item.second.nullifier) {
+            if (nd.nullifier) {
                 mapSaplingNullifiersToNotes.erase(item.second.nullifier.get());
             }
-            item.second.nullifier = boost::none;
+            nd.nullifier = boost::none;
         } else {
             const libzcash::SaplingIncomingViewingKey& ivk = *(nd.ivk);
             uint64_t position = nd.witnesses.front().position();
@@ -79,9 +87,7 @@ void SaplingScriptPubKeyMan::UpdateSaplingNullifierNoteMapWithTx(CWalletTx& wtx)
                 // This should not happen.  If it does, maybe the position has been corrupted or miscalculated?
                 assert(false);
             }
-            uint256 nullifier = optNullifier.get();
-            mapSaplingNullifiersToNotes[nullifier] = op;
-            item.second.nullifier = nullifier;
+            UpdateSaplingNullifierNoteMap(nd, op, optNullifier.get());
         }
     }
 }

src/sapling/saplingscriptpubkeyman.h

@@ -173,6 +173,12 @@ class SaplingScriptPubKeyMan {
      */
     void UpdateNullifierNoteMapWithTx(const CWalletTx& wtx);
 
+    /**
+     *  Update mapSaplingNullifiersToNotes, and NoteData of a specific outpoint,
+     *  directly with nullifier provided by the caller.
+     */
+    void UpdateSaplingNullifierNoteMap(SaplingNoteData& nd, const SaplingOutPoint& op, const Optional<uint256>& nullifier);
+
     /**
      *  Update mapSaplingNullifiersToNotes, computing the nullifier
      *  from a cached witness if necessary.

src/test/librust/sapling_rpc_wallet_tests.cpp

@@ -27,9 +27,6 @@
 
 extern UniValue CallRPC(std::string args); // Implemented in rpc_tests.cpp
 
-// Remember: this method will be moved to an utility file in the short future. For now, it's in sapling_keystore_tests.cpp
-extern libzcash::SaplingExtendedSpendingKey GetTestMasterSaplingSpendingKey();
-
 namespace {
 
     /** Set the working directory for the duration of the scope. */