Function to rollback, eventually disconnect a block by recreating the…

… witness cache

src/rpc/blockchain.cpp

@@ -8,8 +8,9 @@
 #include "budget/budgetmanager.h"
 #include "checkpoints.h"
 #include "clientversion.h"
-#include "core_io.h"
 #include "consensus/upgrades.h"
+#include "core_io.h"
+#include "hash.h"
 #include "kernel.h"
 #include "key_io.h"
 #include "masternodeman.h"
@@ -22,7 +23,7 @@
 #include "util/system.h"
 #include "utilmoneystr.h"
 #include "utilstrencodings.h"
-#include "hash.h"
+#include "validation.h"
 #include "validationinterface.h"
 #include "wallet/wallet.h"
 #include "warnings.h"
@@ -1228,6 +1229,15 @@ UniValue invalidateblock(const JSONRPCRequest& request)
             throw JSONRPCError(RPC_INVALID_ADDRESS_OR_KEY, "Block not found");
 
         CBlockIndex* pblockindex = mapBlockIndex[hash];
+        //For each walllet in your wallet list
+        for (auto pwallet : vpwallets) {
+            //Do we need to recreate the witnesscache or is the current one enough?
+            if (pwallet->GetSaplingScriptPubKeyMan()->nWitnessCacheSize <= (chainActive.Height() - pblockindex->nHeight + 1)) {
+                if (!pwallet->GetSaplingScriptPubKeyMan()->BuildWitnessChain(pblockindex)) {
+                    throw JSONRPCError(RPC_DATABASE_ERROR, "Cannot load default note values");
+                }
+            }
+        }
         InvalidateBlock(state, Params(), pblockindex);
     }
 

src/sapling/saplingscriptpubkeyman.cpp

@@ -6,7 +6,14 @@
 #include "sapling/saplingscriptpubkeyman.h"
 
 #include "chain.h" // for CBlockIndex
+#include "primitives/block.h"
+#include "sapling/incrementalmerkletree.h"
+#include "uint256.h"
 #include "validation.h" // for ReadBlockFromDisk()
+#include "wallet/wallet.h"
+#include <algorithm>
+#include <map>
+#include <vector>
 
 void SaplingScriptPubKeyMan::AddToSaplingSpends(const uint256& nullifier, const uint256& wtxid)
 {
@@ -198,6 +205,93 @@ void UpdateWitnessHeights(NoteDataMap& noteDataMap, int indexHeight, int64_t nWi
     }
 }
 
+bool SaplingScriptPubKeyMan::BuildWitnessChain(const CBlockIndex* pTargetBlock)
+{
+    LOCK2(cs_main, wallet->cs_wallet);
+    // Target is the last block we want to invalidate
+    rollbackTargetHeight = pTargetBlock->nHeight;
+    cachedWitnessMap.clear();
+
+    // Find the oldest sapling note
+    int minHeight = INT_MAX;
+    for (auto& it : wallet->mapWallet) {
+        CWalletTx& wtx = it.second;
+        if (wtx.mapSaplingNoteData.empty()) continue;
+        // Skip abandoned and conflicted txs for which the block_height is not defined (more precisely it it set to 0 by default)
+        if (wtx.m_confirm.status != CWalletTx::CONFIRMED) continue;
+        minHeight = std::min(wtx.m_confirm.block_height, minHeight);
+    }
+
+    // For the moment as a maximum rollback span we use 1 month or 43200 blocks
+    if ((chainActive.Height() - minHeight) > 43200) {
+        cachedWitnessMap.clear();
+        rollbackTargetHeight = -1;
+        return false;
+    }
+
+    // Read blocks from the disk from chaintip to the minimum found height
+    std::vector<CBlock> cblocks;
+    const CBlockIndex* pIndex = GetChainTip();
+    int currentHeight = GetChainTip()->nHeight;
+    while (currentHeight >= minHeight) {
+        CBlock cblock;
+        ReadBlockFromDisk(cblock, pIndex);
+        cblocks.insert(cblocks.begin(), cblock);
+        pIndex = pIndex->pprev;
+        currentHeight = pIndex->nHeight;
+    }
+
+    // Load the SaplingMerkleTree for the block before the oldest note
+    SaplingMerkleTree initialSaplingTree;
+    if (!pcoinsTip->GetSaplingAnchorAt(pIndex->hashFinalSaplingRoot, initialSaplingTree)) {
+        return false;
+    }
+    // Finally build the witness cache for each sapling note of your wallet
+    for (CBlock block : cblocks) {
+        // Finally build the witness cache for each sapling note
+        std::vector<uint256> noteCommitments;
+        std::vector<SaplingNoteData*> inBlockArrivingNotes;
+        for (const auto& tx : block.vtx) {
+            const auto& hash = tx->GetHash();
+            auto it = wallet->mapWallet.find(hash);
+            bool txIsOurs = it != wallet->mapWallet.end();
+
+            if (!tx->IsShieldedTx()) continue;
+            for (uint32_t i = 0; i < tx->sapData->vShieldedOutput.size(); i++) {
+                const auto& cmu = tx->sapData->vShieldedOutput[i].cmu;
+                noteCommitments.emplace_back(cmu);
+                for (auto& item : inBlockArrivingNotes) {
+                    item->witnesses.front().append(cmu);
+                }
+                initialSaplingTree.append(cmu);
+                if (txIsOurs) {
+                    CWalletTx* wtx = &it->second;
+                    auto ndIt = wtx->mapSaplingNoteData.find({hash, i});
+                    if (ndIt != wtx->mapSaplingNoteData.end()) {
+                        SaplingNoteData* nd = &ndIt->second;
+                        nd->witnesses.push_front(initialSaplingTree.witness());
+                        inBlockArrivingNotes.emplace_back(nd);
+                    }
+                }
+            }
+        }
+        for (auto& it2 : cachedWitnessMap) {
+            it2.second.emplace_front(it2.second.front());
+            for (auto& noteComm : noteCommitments) {
+                it2.second.front().append(noteComm);
+            }
+        }
+        for (auto nd : inBlockArrivingNotes) {
+            if (nd->nullifier) {
+                std::list<SaplingWitness> witnesses;
+                witnesses.push_front(nd->witnesses.front());
+                cachedWitnessMap.emplace(*(nd->nullifier), witnesses);
+            }
+        }
+    }
+    return true;
+}
+
 void SaplingScriptPubKeyMan::IncrementNoteWitnesses(const CBlockIndex* pindex,
                                                     const CBlock* pblock,
                                                     SaplingMerkleTree& saplingTreeRes)
@@ -281,6 +375,35 @@ void SaplingScriptPubKeyMan::IncrementNoteWitnesses(const CBlockIndex* pindex,
     // CWallet::SetBestChain() (which also ensures that overall consistency
     // of the wallet.dat is maintained).
 }
+/*
+ *
+ * Clear and eventrually reset each witness of noteDataMap with the corresponding front-value of cachedWitnessMap, indexHeight is the blockHeight being invalidated
+ */
+void ResetNoteWitnesses(std::map<SaplingOutPoint, SaplingNoteData>& noteDataMap, std::map<uint256, std::list<SaplingWitness>>& cachedWitnessMap, int indexHeight)
+{
+    // For each note that you own:
+    for (auto& item : noteDataMap) {
+        auto* nd = &(item.second);
+        // skip externally sent notes
+        if (!nd->IsMyNote()) continue;
+        // Clear the cache
+        nd->witnesses.clear();
+        // The withnessHeight must be EITHER -1 or equal to the block indexHeight
+        // The case in which indexHeight > witnessHeight is due to conflicted notes, which are irrelevant
+        // TODO: Allow invalidating blocks only if there are not conflicted txs?
+        if (nd->witnessHeight <= indexHeight) {
+            assert((nd->witnessHeight == -1) || (nd->witnessHeight == indexHeight));
+        }
+        // Decrease the witnessHeight
+        nd->witnessHeight = indexHeight - 1;
+        if (nd->nullifier && cachedWitnessMap.at(*nd->nullifier).size() > 0) {
+            // Update the witness value with the cached one
+            nd->witnesses.push_front(cachedWitnessMap.at(*nd->nullifier).front());
+            cachedWitnessMap.at(*nd->nullifier).pop_front();
+        }
+    }
+}
+
 
 template<typename NoteDataMap>
 void DecrementNoteWitnesses(NoteDataMap& noteDataMap, int indexHeight, int64_t nWitnessCacheSize)
@@ -324,9 +447,30 @@ void DecrementNoteWitnesses(NoteDataMap& noteDataMap, int indexHeight, int64_t n
     }
 }
 
-void SaplingScriptPubKeyMan::DecrementNoteWitnesses(int nChainHeight)
+void SaplingScriptPubKeyMan::DecrementNoteWitnesses(const CBlockIndex* pindex)
 {
+    assert(pindex);
     LOCK(wallet->cs_wallet);
+    int nChainHeight = pindex->nHeight;
+    // if the targetHeight is different from -1 we have a cache to use
+    if (rollbackTargetHeight != -1) {
+        for (std::pair<const uint256, CWalletTx>& wtxItem : wallet->mapWallet) {
+            if (!wtxItem.second.mapSaplingNoteData.empty()) {
+                // For each sapling note that you own reset the current witness with the cached one
+                ResetNoteWitnesses(wtxItem.second.mapSaplingNoteData, cachedWitnessMap, nChainHeight);
+            }
+        }
+        nWitnessCacheSize = 1;
+        nWitnessCacheNeedsUpdate = true;
+        // If we reached the target height empty the cache and reset the target height to -1
+        // Remember that the targetHeight is indeed the last block we want to invalidate
+        if (rollbackTargetHeight == pindex->nHeight) {
+            cachedWitnessMap.clear();
+            rollbackTargetHeight = -1;
+        }
+        return;
+    }
+
     for (std::pair<const uint256, CWalletTx>& wtxItem : wallet->mapWallet) {
         ::DecrementNoteWitnesses(wtxItem.second.mapSaplingNoteData, nChainHeight, nWitnessCacheSize);
     }
@@ -504,7 +648,6 @@ void SaplingScriptPubKeyMan::GetFilteredNotes(
         for (const auto& it : wtx.mapSaplingNoteData) {
             const SaplingOutPoint& op = it.first;
             const SaplingNoteData& nd = it.second;
-
             // skip sent notes
             if (!nd.IsMyNote()) continue;
 

src/sapling/saplingscriptpubkeyman.h

@@ -6,11 +6,14 @@
 #define PIVX_SAPLINGSCRIPTPUBKEYMAN_H
 
 #include "consensus/consensus.h"
+#include "sapling/incrementalmerkletree.h"
 #include "sapling/note.h"
+#include "uint256.h"
 #include "wallet/hdchain.h"
+#include "wallet/scriptpubkeyman.h"
 #include "wallet/wallet.h"
 #include "wallet/walletdb.h"
-#include "sapling/incrementalmerkletree.h"
+#include <map>
 
 //! Size of witness cache
 //  Should be large enough that we can expect not to reorg beyond our cache
@@ -47,6 +50,7 @@ class SaplingNoteData
 
     /* witnesses/ivk: only for own (received) outputs */
     std::list<SaplingWitness> witnesses;
+
     Optional<libzcash::SaplingIncomingViewingKey> ivk {nullopt};
     inline bool IsMyNote() const { return ivk != nullopt; }
 
@@ -153,16 +157,21 @@ class SaplingScriptPubKeyMan {
     void AddToSaplingSpends(const uint256& nullifier, const uint256& wtxid);
     bool IsSaplingSpent(const uint256& nullifier) const;
 
+    /**
+     * Build the old witness chain.
+     */
+    bool BuildWitnessChain(const CBlockIndex* pTargetBlock);
+
     /**
      * pindex is the new tip being connected.
      */
     void IncrementNoteWitnesses(const CBlockIndex* pindex,
                                 const CBlock* pblock,
                                 SaplingMerkleTree& saplingTree);
     /**
-     * nChainHeight is the old tip height being disconnected.
+     * pindex is the old tip being disconnected.
      */
-    void DecrementNoteWitnesses(int nChainHeight);
+    void DecrementNoteWitnesses(const CBlockIndex* pindex);
 
     /**
      * Update mapSaplingNullifiersToNotes
@@ -405,6 +414,9 @@ class SaplingScriptPubKeyMan {
     std::map<uint256, SaplingOutPoint> mapSaplingNullifiersToNotes;
 
 private:
+    /* Map hash nullifiers, list Sapling Witness*/
+    std::map<uint256, std::list<SaplingWitness>> cachedWitnessMap;
+    int rollbackTargetHeight = -1;
     /* Parent wallet */
     CWallet* wallet{nullptr};
     /* the HD chain data model (external/internal chain counters) */

src/wallet/wallet.cpp

@@ -5,6 +5,7 @@
 // Distributed under the MIT software license, see the accompanying
 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
+#include "validation.h"
 #if defined(HAVE_CONFIG_H)
 #include "config/pivx-config.h"
 #endif
@@ -1363,7 +1364,7 @@ void CWallet::BlockDisconnected(const std::shared_ptr<const CBlock>& pblock, con
 
     if (Params().GetConsensus().NetworkUpgradeActive(nBlockHeight, Consensus::UPGRADE_V5_0)) {
         // Update Sapling cached incremental witnesses
-        m_sspk_man->DecrementNoteWitnesses(nBlockHeight);
+        m_sspk_man->DecrementNoteWitnesses(mapBlockIndex[blockHash]);
         m_sspk_man->UpdateSaplingNullifierNoteMapForBlock(pblock.get());
     }
 }
@@ -4630,7 +4631,7 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
                             const CBlock* pblock,
                             SaplingMerkleTree& saplingTree) { m_sspk_man->IncrementNoteWitnesses(pindex, pblock, saplingTree); }
 
-void CWallet::DecrementNoteWitnesses(const CBlockIndex* pindex) { m_sspk_man->DecrementNoteWitnesses(pindex->nHeight); }
+void CWallet::DecrementNoteWitnesses(const CBlockIndex* pindex) { m_sspk_man->DecrementNoteWitnesses(pindex); }
 
 void CWallet::ClearNoteWitnessCache() { m_sspk_man->ClearNoteWitnessCache(); }