Auto-rescan wallets on reorg if note witnesses get out of sync

src/main.cpp

@@ -4048,7 +4048,10 @@ static bool ActivateBestChainStep(CValidationState &state, CBlockIndex *pindexMo
         CheckForkWarningConditionsOnNewFork(vpindexToConnect.back());
     else
         CheckForkWarningConditions();
-    
+
+    // if we need to, rescan wallets
+    RescanWallets();
+
     return true;
 }
 

src/validationinterface.cpp

@@ -17,6 +17,7 @@ void RegisterValidationInterface(CValidationInterface* pwalletIn) {
     g_signals.SyncTransaction.connect(boost::bind(&CValidationInterface::SyncTransaction, pwalletIn, _1, _2));
     g_signals.EraseTransaction.connect(boost::bind(&CValidationInterface::EraseFromWallet, pwalletIn, _1));
     g_signals.UpdatedTransaction.connect(boost::bind(&CValidationInterface::UpdatedTransaction, pwalletIn, _1));
+    g_signals.RescanWallet.connect(boost::bind(&CValidationInterface::RescanWallet, pwalletIn));
     g_signals.ChainTip.connect(boost::bind(&CValidationInterface::ChainTip, pwalletIn, _1, _2, _3, _4, _5));
     g_signals.SetBestChain.connect(boost::bind(&CValidationInterface::SetBestChain, pwalletIn, _1));
     g_signals.Inventory.connect(boost::bind(&CValidationInterface::Inventory, pwalletIn, _1));
@@ -33,6 +34,7 @@ void UnregisterValidationInterface(CValidationInterface* pwalletIn) {
     g_signals.UpdatedTransaction.disconnect(boost::bind(&CValidationInterface::UpdatedTransaction, pwalletIn, _1));
     g_signals.EraseTransaction.disconnect(boost::bind(&CValidationInterface::EraseFromWallet, pwalletIn, _1));
     g_signals.SyncTransaction.disconnect(boost::bind(&CValidationInterface::SyncTransaction, pwalletIn, _1, _2));
+    g_signals.RescanWallet.disconnect(boost::bind(&CValidationInterface::RescanWallet, pwalletIn));
     g_signals.UpdatedBlockTip.disconnect(boost::bind(&CValidationInterface::UpdatedBlockTip, pwalletIn, _1));
 }
 
@@ -45,6 +47,7 @@ void UnregisterAllValidationInterfaces() {
     g_signals.UpdatedTransaction.disconnect_all_slots();
     g_signals.EraseTransaction.disconnect_all_slots();
     g_signals.SyncTransaction.disconnect_all_slots();
+    g_signals.RescanWallet.disconnect_all_slots();
     g_signals.UpdatedBlockTip.disconnect_all_slots();
 }
 
@@ -55,3 +58,7 @@ void SyncWithWallets(const CTransaction &tx, const CBlock *pblock) {
 void EraseFromWallets(const uint256 &hash) {
     g_signals.EraseTransaction(hash);
 }
+
+void RescanWallets() {
+    g_signals.RescanWallet();
+}

src/validationinterface.h

@@ -30,12 +30,15 @@ void UnregisterAllValidationInterfaces();
 void SyncWithWallets(const CTransaction& tx, const CBlock* pblock = NULL);
 /** Erase a transaction from all registered wallets */
 void EraseFromWallets(const uint256 &hash);
+/** Rescan all registered wallets */
+void RescanWallets();
 
 class CValidationInterface {
 protected:
     virtual void UpdatedBlockTip(const CBlockIndex *pindex) {}
     virtual void SyncTransaction(const CTransaction &tx, const CBlock *pblock) {}
     virtual void EraseFromWallet(const uint256 &hash) {}
+    virtual void RescanWallet() {}
     virtual void ChainTip(const CBlockIndex *pindex, const CBlock *pblock, SproutMerkleTree sproutTree, SaplingMerkleTree saplingTree, bool added) {}
     virtual void SetBestChain(const CBlockLocator &locator) {}
     virtual void UpdatedTransaction(const uint256 &hash) {}
@@ -52,8 +55,10 @@ struct CMainSignals {
     boost::signals2::signal<void (const CBlockIndex *)> UpdatedBlockTip;
     /** Notifies listeners of updated transaction data (transaction, and optionally the block it is found in. */
     boost::signals2::signal<void (const CTransaction &, const CBlock *)> SyncTransaction;
-    /** Notifies listeners of an erased transaction (currently disabled, requires transaction replacement). */
+    /** Notifies listeners of an erased transaction. */
     boost::signals2::signal<void (const uint256 &)> EraseTransaction;
+    /** Notifies listeners of the need to rescan the wallet. */
+    boost::signals2::signal<void ()> RescanWallet;
     /** Notifies listeners of an updated transaction without new data (for now: a coinbase potentially becoming visible). */
     boost::signals2::signal<void (const uint256 &)> UpdatedTransaction;
     /** Notifies listeners of a change to the tip of the active block chain. */

src/wallet/wallet.cpp

@@ -1123,7 +1123,7 @@ void CWallet::IncrementNoteWitnesses(const CBlockIndex* pindex,
 }
 
 template<typename NoteDataMap>
-void DecrementNoteWitnesses(NoteDataMap& noteDataMap, int indexHeight, int64_t nWitnessCacheSize)
+bool DecrementNoteWitnesses(NoteDataMap& noteDataMap, int indexHeight, int64_t nWitnessCacheSize)
 {
     extern int32_t KOMODO_REWIND;
 
@@ -1138,7 +1138,11 @@ void DecrementNoteWitnesses(NoteDataMap& noteDataMap, int indexHeight, int64_t n
             // Witnesses being decremented should always be either -1
             // (never incremented or decremented) or equal to the height
             // of the block being removed (indexHeight)
-            assert((nd->witnessHeight == -1) || (nd->witnessHeight == indexHeight));
+            if (!((nd->witnessHeight == -1) || (nd->witnessHeight == indexHeight)))
+            {
+                printf("at height %d\n", indexHeight);
+                return false;
+            }
             if (nd->witnesses.size() > 0) {
                 nd->witnesses.pop_front();
             }
@@ -1163,14 +1167,17 @@ void DecrementNoteWitnesses(NoteDataMap& noteDataMap, int indexHeight, int64_t n
         }
     }
     assert(KOMODO_REWIND != 0 || nWitnessCacheSize > 0);
+    return true;
 }
 
 void CWallet::DecrementNoteWitnesses(const CBlockIndex* pindex)
 {
     LOCK(cs_wallet);
     for (std::pair<const uint256, CWalletTx>& wtxItem : mapWallet) {
-        ::DecrementNoteWitnesses(wtxItem.second.mapSproutNoteData, pindex->GetHeight(), nWitnessCacheSize);
-        ::DecrementNoteWitnesses(wtxItem.second.mapSaplingNoteData, pindex->GetHeight(), nWitnessCacheSize);
+        if (!::DecrementNoteWitnesses(wtxItem.second.mapSproutNoteData, pindex->GetHeight(), nWitnessCacheSize))
+            needsRescan = true;
+        if (!::DecrementNoteWitnesses(wtxItem.second.mapSaplingNoteData, pindex->GetHeight(), nWitnessCacheSize))
+            needsRescan = true;
     }
     nWitnessCacheSize -= 1;
     // TODO: If nWitnessCache is zero, we need to regenerate the caches (#1302)
@@ -1838,6 +1845,17 @@ void CWallet::EraseFromWallet(const uint256 &hash)
     return;
 }
 
+void CWallet::RescanWallet()
+{
+    if (needsRescan)
+    {
+        CBlockIndex *start = chainActive.Height() > 0 ? chainActive[1] : NULL;
+        if (start)
+            ScanForWalletTransactions(start, true);
+        needsRescan = false;
+    }
+}
+
 
 /**
  * Returns a nullifier if the SpendingKey is available
@@ -3001,6 +3019,7 @@ std::vector<uint256> CWallet::ResendWalletTransactionsBefore(int64_t nTime)
     // Sort them in chronological order
     multimap<unsigned int, CWalletTx*> mapSorted;
     uint32_t now = (uint32_t)time(NULL);
+    std::vector<uint256> vwtxh;
     BOOST_FOREACH(PAIRTYPE(const uint256, CWalletTx)& item, mapWallet)
     {
         CWalletTx& wtx = item.second;
@@ -3012,8 +3031,7 @@ std::vector<uint256> CWallet::ResendWalletTransactionsBefore(int64_t nTime)
             if ( wtx.nLockTime >= LOCKTIME_THRESHOLD && wtx.nLockTime < now-KOMODO_MAXMEMPOOLTIME )
             {
                 LogPrintf("skip Relaying wtx %s nLockTime %u vs now.%u\n", wtx.GetHash().ToString(),(uint32_t)wtx.nLockTime,now);
-                //TODO: EraseFromWallet(wtx.GetHash()); //should be erased, but this creates issues, likely better to create
-                // vector and do it outside of this loop, but for later
+                vwtxh.push_back(wtx.GetHash());
                 continue;
             }
         }
@@ -3028,6 +3046,10 @@ std::vector<uint256> CWallet::ResendWalletTransactionsBefore(int64_t nTime)
                 result.push_back(wtx.GetHash());
         }
     }
+    for (auto hash : vwtxh)
+    {
+        EraseFromWallet(hash);
+    }
     return result;
 }
 

src/wallet/wallet.h

@@ -781,6 +781,7 @@ public:
      * incremental witness cache in any transaction in mapWallet.
      */
     int64_t nWitnessCacheSize;
+    bool needsRescan = false;
 
     void ClearNoteWitnessCache();
 
@@ -1120,6 +1121,7 @@ public:
     bool AddToWallet(const CWalletTx& wtxIn, bool fFromLoadWallet, CWalletDB* pwalletdb);
     void EraseFromWallet(const uint256 &hash);
     void SyncTransaction(const CTransaction& tx, const CBlock* pblock);
+    void RescanWallet();
     bool AddToWalletIfInvolvingMe(const CTransaction& tx, const CBlock* pblock, bool fUpdate);
     void WitnessNoteCommitment(
          std::vector<uint256> commitments,