From 31404e0332e97f0cb2536ea224e622a7ef5ef23c Mon Sep 17 00:00:00 2001
From: David Burkett <davidburkett38@gmail.com>
Date: Sat, 29 Jan 2022 19:54:50 -0500
Subject: [PATCH] MWEB: Block undo changes

---
 src/undo.h         | 30 ++++++++++++++++++++++++++++--
 src/validation.cpp | 18 +++++++++++++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/src/undo.h b/src/undo.h
index a98f046735..0085adacf9 100644
--- a/src/undo.h
+++ b/src/undo.h
@@ -9,6 +9,7 @@
 #include <coins.h>
 #include <compressor.h>
 #include <consensus/consensus.h>
+#include <mw/models/block/BlockUndo.h>
 #include <primitives/transaction.h>
 #include <serialize.h>
 #include <version.h>
@@ -42,7 +43,7 @@ struct TxInUndoFormatter
             // Old versions stored the version number for the last spend of
             // a transaction's outputs. Non-final spends were indicated with
             // height = 0.
-            unsigned int nVersionDummy;
+            unsigned int nVersionDummy = 0;
             ::Unserialize(s, VARINT(nVersionDummy));
         }
         ::Unserialize(s, Using<TxOutCompression>(txout.out));
@@ -64,8 +65,33 @@ class CBlockUndo
 {
 public:
     std::vector<CTxUndo> vtxundo; // for all but the coinbase
+    mw::BlockUndo::CPtr mwundo;
 
-    SERIALIZE_METHODS(CBlockUndo, obj) { READWRITE(obj.vtxundo); }
+    SERIALIZE_METHODS(CBlockUndo, obj)
+	{
+		READWRITE(obj.vtxundo);
+		
+		if (obj.mwundo !=  nullptr) {
+			READWRITE(obj.mwundo);
+		}
+	}
 };
 
+template <typename Stream>
+inline void UnserializeBlockUndo(CBlockUndo& blockundo, Stream& s, const unsigned int num_bytes)
+{
+    const uint64_t num_txs = ::ReadCompactSize(s);
+    blockundo.vtxundo.reserve(num_txs);
+    for (uint64_t i = 0; i < num_txs; i++) {
+        CTxUndo txundo;
+        s >> txundo;
+        blockundo.vtxundo.emplace_back(std::move(txundo));
+    }
+
+    if (::GetSerializeSize(blockundo) < num_bytes) {
+        // There's more data to read. MWEB rewind data must be available.
+        s >> blockundo.mwundo;
+    }
+}
+
 #endif // BITCOIN_UNDO_H
diff --git a/src/validation.cpp b/src/validation.cpp
index ffd134b55f..057957930a 100644
--- a/src/validation.cpp
+++ b/src/validation.cpp
@@ -1656,17 +1656,24 @@ bool UndoReadFromDisk(CBlockUndo& blockundo, const CBlockIndex* pindex)
         return error("%s: no undo data available", __func__);
     }
 
+    // Rewind 4 bytes in order to read the size
+    pos.nPos -= 4;
+
     // Open history file to read
     CAutoFile filein(OpenUndoFile(pos, true), SER_DISK, CLIENT_VERSION);
     if (filein.IsNull())
         return error("%s: OpenUndoFile failed", __func__);
 
+    // Read undo size
+    unsigned int undo_size = 0;
+    filein >> undo_size;
+
     // Read block
     uint256 hashChecksum;
     CHashVerifier<CAutoFile> verifier(&filein); // We need a CHashVerifier as reserializing may lose data
     try {
         verifier << pindex->pprev->GetBlockHash();
-        verifier >> blockundo;
+        UnserializeBlockUndo(blockundo, verifier, undo_size);
         filein >> hashChecksum;
     }
     catch (const std::exception& e) {
@@ -1787,6 +1794,15 @@ DisconnectResult CChainState::DisconnectBlock(const CBlock& block, const CBlockI
         }
     }
 
+    if (blockUndo.mwundo != nullptr) {
+        try {
+            view.GetMWEBCacheView()->UndoBlock(blockUndo.mwundo);
+        } catch (const std::exception& e) {
+            error("DisconnectBlock(): Failed to disconnect MWEB block");
+            return DISCONNECT_FAILED;
+        }
+    }
+
     // move best block pointer to prevout block
     view.SetBestBlock(pindex->pprev->GetBlockHash());