Added "mwebleafset" getdata type and CMWEBLeafsetMsg response

src/net_processing.cpp

@@ -102,6 +102,8 @@ static const unsigned int MAX_HEADERS_RESULTS = 2000;
 static const int MAX_CMPCTBLOCK_DEPTH = 5;
 /** Maximum depth of blocks we're willing to respond to GETBLOCKTXN requests for. */
 static const int MAX_BLOCKTXN_DEPTH = 10;
+/** Maximum depth of blocks we're willing to serve MWEB leafsets for. */
+static const int MAX_MWEB_LEAFSET_DEPTH = 10;
 /** Size of the "block download window": how far ahead of our current height do we fetch?
  *  Larger windows tolerate larger download speed differences between peer, but increase the potential
  *  degree of disordering of blocks on disk (which make reindexing and pruning harder). We'll probably
@@ -1536,22 +1538,8 @@ static void RelayAddress(const CAddress& addr, bool fReachable, const CConnman&
     connman.ForEachNodeThen(std::move(sortfunc), std::move(pushfunc));
 }
 
-void static ProcessGetBlockData(CNode& pfrom, const CChainParams& chainparams, const CInv& inv, CConnman& connman)
+static void ActivateBestChainIfNeeded(const CChainParams& chainparams, const CInv& inv)
 {
-    bool send = false;
-    std::shared_ptr<const CBlock> a_recent_block;
-    std::shared_ptr<const CBlockHeaderAndShortTxIDs> a_recent_compact_block;
-    bool fWitnessesPresentInARecentCompactBlock;
-    bool fMWEBPresentInARecentCompactBlock;
-    const Consensus::Params& consensusParams = chainparams.GetConsensus();
-    {
-        LOCK(cs_most_recent_block);
-        a_recent_block = most_recent_block;
-        a_recent_compact_block = most_recent_compact_block;
-        fWitnessesPresentInARecentCompactBlock = fWitnessesPresentInMostRecentCompactBlock;
-        fMWEBPresentInARecentCompactBlock = fMWEBPresentInMostRecentCompactBlock;
-    }
-
     bool need_activate_chain = false;
     {
         LOCK(cs_main);
@@ -1568,12 +1556,40 @@ void static ProcessGetBlockData(CNode& pfrom, const CChainParams& chainparams, c
             }
         }
     } // release cs_main before calling ActivateBestChain
+
     if (need_activate_chain) {
+        // Grab the current most_recent_block and pass it to ActivateBestChain
+        // which hopefully will prevent needing to load blocks from disk.
+        std::shared_ptr<const CBlock> a_recent_block;
+        {
+            LOCK(cs_most_recent_block);
+            a_recent_block = most_recent_block;
+        }
+
         BlockValidationState state;
         if (!ActivateBestChain(state, chainparams, a_recent_block)) {
             LogPrint(BCLog::NET, "failed to activate chain (%s)\n", state.ToString());
         }
     }
+}
+
+void static ProcessGetBlockData(CNode& pfrom, const CChainParams& chainparams, const CInv& inv, CConnman& connman)
+{
+    bool send = false;
+    std::shared_ptr<const CBlock> a_recent_block;
+    std::shared_ptr<const CBlockHeaderAndShortTxIDs> a_recent_compact_block;
+    bool fWitnessesPresentInARecentCompactBlock;
+    bool fMWEBPresentInARecentCompactBlock;
+    const Consensus::Params& consensusParams = chainparams.GetConsensus();
+    {
+        LOCK(cs_most_recent_block);
+        a_recent_block = most_recent_block;
+        a_recent_compact_block = most_recent_compact_block;
+        fWitnessesPresentInARecentCompactBlock = fWitnessesPresentInMostRecentCompactBlock;
+        fMWEBPresentInARecentCompactBlock = fMWEBPresentInMostRecentCompactBlock;
+    }
+
+    ActivateBestChainIfNeeded(chainparams, inv);
 
     LOCK(cs_main);
     const CBlockIndex* pindex = LookupBlockIndex(inv.hash);
@@ -1702,6 +1718,63 @@ void static ProcessGetBlockData(CNode& pfrom, const CChainParams& chainparams, c
     }
 }
 
+class CMWEBLeafsetMsg
+{
+public:
+    CMWEBLeafsetMsg() = default;
+    CMWEBLeafsetMsg(uint256 block_hash_in, BitSet leafset_in)
+        : block_hash(std::move(block_hash_in)), leafset(std::move(leafset_in)) { }
+
+    SERIALIZE_METHODS(CMWEBLeafsetMsg, obj) { READWRITE(obj.block_hash, obj.leafset); }
+
+    uint256 block_hash;
+    BitSet leafset;
+};
+
+static void ProcessGetMWEBLeafset(CNode& pfrom, const CChainParams& chainparams, const CInv& inv, CConnman& connman)
+{
+    ActivateBestChainIfNeeded(chainparams, inv);
+
+    LOCK(cs_main);
+    CBlockIndex* pindex = LookupBlockIndex(inv.hash);
+    if (!pindex) {
+        return;
+    }
+
+    // TODO: Add an outbound limit
+
+    // Avoid leaking prune-height by never sending blocks below the NODE_NETWORK_LIMITED threshold
+    if (::ChainActive().Tip()->nHeight - pindex->nHeight > MAX_MWEB_LEAFSET_DEPTH) {
+        LogPrint(BCLog::NET, "Ignore block request below MAX_MWEB_LEAFSET_DEPTH threshold from peer=%d\n", pfrom.GetId());
+
+        // disconnect node and prevent it from stalling (would otherwise wait for the MWEB leafset)
+        if (!pfrom.HasPermission(PF_NOBAN)) {
+            pfrom.fDisconnect = true;
+        }
+
+        return;
+    }
+
+    // Pruned nodes may have deleted the block, so check whether
+    // it's available before trying to send.
+    if (pindex->nStatus & BLOCK_HAVE_DATA && pindex->nStatus & BLOCK_HAVE_MWEB) {
+        // Rewind leafset to block height
+        BlockValidationState state;
+        CCoinsViewCache temp_view(&::ChainstateActive().CoinsTip());
+        if (!ActivateArbitraryChain(state, temp_view, chainparams, pindex)) {
+            pfrom.fDisconnect = true;
+            return;
+        }
+
+        // Serve leafset to peer
+        const CNetMsgMaker msgMaker(pfrom.GetCommonVersion());
+        auto pLeafset = temp_view.GetMWEBCacheView()->GetLeafSet();
+
+        CMWEBLeafsetMsg leafset_msg(pindex->GetBlockHash(), pLeafset->ToBitSet());
+        connman.PushMessage(&pfrom, msgMaker.Make(NetMsgType::MWEBLEAFSET, leafset_msg));
+    }
+}
+
 //! Determine whether or not a peer can request a transaction, and return it (or nullptr if not found or not allowed).
 static CTransactionRef FindTxForGetData(const CTxMemPool& mempool, const CNode& peer, const GenTxid& gtxid, const std::chrono::seconds mempool_req, const std::chrono::seconds now) LOCKS_EXCLUDED(cs_main)
 {
@@ -1798,6 +1871,8 @@ void static ProcessGetData(CNode& pfrom, Peer& peer, const CChainParams& chainpa
         const CInv &inv = *it++;
         if (inv.IsGenBlkMsg()) {
             ProcessGetBlockData(pfrom, chainparams, inv, connman);
+        } else if (inv.IsMsgMWEBLeafset()) {
+            ProcessGetMWEBLeafset(pfrom, chainparams, inv, connman);
         }
         // else: If the first item on the queue is an unknown type, we erase it
         // and continue processing the queue on the next call.

src/protocol.cpp

@@ -47,6 +47,7 @@ const char *GETCFCHECKPT="getcfcheckpt";
 const char *CFCHECKPT="cfcheckpt";
 const char *WTXIDRELAY="wtxidrelay";
 const char *MWEBHEADER="mwebheader";
+const char *MWEBLEAFSET="mwebleafset";
 } // namespace NetMsgType
 
 /** All known message types. Keep this in the same order as the list of
@@ -88,6 +89,7 @@ const static std::string allNetMessageTypes[] = {
     NetMsgType::CFCHECKPT,
     NetMsgType::WTXIDRELAY,
     NetMsgType::MWEBHEADER,
+    NetMsgType::MWEBLEAFSET,
 };
 const static std::vector<std::string> allNetMessageTypesVec(allNetMessageTypes, allNetMessageTypes+ARRAYLEN(allNetMessageTypes));
 
@@ -179,6 +181,7 @@ std::string CInv::GetCommand() const
     case MSG_FILTERED_BLOCK: return cmd.append(NetMsgType::MERKLEBLOCK);
     case MSG_CMPCT_BLOCK:    return cmd.append(NetMsgType::CMPCTBLOCK);
     case MSG_MWEB_HEADER:    return cmd.append(NetMsgType::MWEBHEADER);
+    case MSG_MWEB_LEAFSET:    return cmd.append(NetMsgType::MWEBLEAFSET);
     default:
         throw std::out_of_range(strprintf("CInv::GetCommand(): type=%d unknown type", type));
     }

src/protocol.h

@@ -267,6 +267,13 @@ extern const char* WTXIDRELAY;
  * @since protocol version 70017 as described by LIP-0007
  */
 extern const char* MWEBHEADER;
+/**
+ * Contains a block hash and its serialized leafset.
+ * Sent in response to a getdata message which requested
+ * data using the inventory type MSG_MWEB_LEAFSET.
+ * @since protocol version 70017 as described by LIP-0007
+ */
+extern const char* MWEBLEAFSET;
 }; // namespace NetMsgType
 
 /* Get a vector of all valid message types (see above) */
@@ -435,6 +442,7 @@ enum GetDataMsg : uint32_t {
     MSG_MWEB_BLOCK = MSG_WITNESS_BLOCK | MSG_MWEB_FLAG,
     MSG_MWEB_TX = MSG_WITNESS_TX | MSG_MWEB_FLAG,
     MSG_MWEB_HEADER = 8 | MSG_MWEB_FLAG,            //!< Defined in LIP-0007
+    MSG_MWEB_LEAFSET = 9 | MSG_MWEB_FLAG,           //!< Defined in LIP-0007
 };
 
 /** inv message data */
@@ -460,6 +468,7 @@ public:
     bool IsMsgWitnessBlk() const { return type == MSG_WITNESS_BLOCK; }
     bool IsMsgMWEBBlk() const { return type == MSG_MWEB_BLOCK; }
     bool IsMsgMWEBHeader() const { return type == MSG_MWEB_HEADER; }
+    bool IsMsgMWEBLeafset() const { return type == MSG_MWEB_LEAFSET; }
 
     // Combined-message helper methods
     bool IsGenTxMsg() const

src/validation.cpp

@@ -3047,6 +3047,52 @@ bool ActivateBestChain(BlockValidationState &state, const CChainParams& chainpar
     return ::ChainstateActive().ActivateBestChain(state, chainparams, std::move(pblock));
 }
 
+bool ActivateArbitraryChain(BlockValidationState& state, CCoinsViewCache& view, const CChainParams& chainparams, CBlockIndex* pindex)
+{
+    AssertLockHeld(cs_main);
+
+    const CBlockIndex* pindexFork = ::ChainstateActive().m_chain.FindFork(pindex);
+    CBlockIndex* pindexTip = ::ChainstateActive().m_chain.Tip();
+
+    // Disconnect blocks from view until we reach the fork block
+    while (pindexTip->GetBlockHash() != pindexFork->GetBlockHash()) {
+        CBlock block;
+        if (!ReadBlockFromDisk(block, pindexTip, chainparams.GetConsensus())) {
+            return error("ActivateArbitraryChain(): Failed to read block %s", pindexTip->GetBlockHash().ToString());
+        }
+
+        if (::ChainstateActive().DisconnectBlock(block, pindexTip, view) != DISCONNECT_OK) {
+            return error("ActivateArbitraryChain(): DisconnectBlock %s failed", pindexTip->GetBlockHash().ToString());
+        }
+
+        pindexTip = pindexTip->pprev;
+    }
+
+    // Build list of new blocks to connect.
+    std::vector<CBlockIndex*> vpindexToConnect;
+    vpindexToConnect.reserve(pindex->nHeight - pindexTip->nHeight);
+
+    CBlockIndex* pindexIter = pindex;
+    while (pindexIter && pindexIter->nHeight != pindexTip->nHeight) {
+        vpindexToConnect.push_back(pindexIter);
+        pindexIter = pindexIter->pprev;
+    }
+
+    // Connect the new blocks
+    for (CBlockIndex* pindexConnect : reverse_iterate(vpindexToConnect)) {
+        CBlock block;
+        if (!ReadBlockFromDisk(block, pindexConnect, chainparams.GetConsensus())) {
+            return error("ActivateArbitraryChain(): Failed to read block %s", pindexConnect->GetBlockHash().ToString());
+        }
+
+        if (!::ChainstateActive().ConnectBlock(block, state, pindexConnect, view, chainparams, true)) {
+            return error("ActivateArbitraryChain(): ConnectBlock %s failed", pindexConnect->GetBlockHash().ToString());
+        }
+    }
+
+    return true;
+}
+
 bool CChainState::PreciousBlock(BlockValidationState& state, const CChainParams& params, CBlockIndex *pindex)
 {
     {

src/validation.h

@@ -182,6 +182,15 @@ CTransactionRef GetTransaction(const CBlockIndex* const block_index, const CTxMe
  * validationinterface callback.
  */
 bool ActivateBestChain(BlockValidationState& state, const CChainParams& chainparams, std::shared_ptr<const CBlock> pblock = std::shared_ptr<const CBlock>());
+
+/**
+ * Make the provided index the tip of the chain, regardless of the amount of work.
+ * 
+ * Unlike ActivateBestChain, this only updates the provided coins view, not the active chain state.
+ * No calls to any validationinterface callbacks will be made.
+ */
+bool ActivateArbitraryChain(BlockValidationState& state, CCoinsViewCache& view, const CChainParams& chainparams, CBlockIndex* pindex) EXCLUSIVE_LOCKS_REQUIRED(cs_main);
+
 CAmount GetBlockSubsidy(int nHeight, const Consensus::Params& consensusParams);
 
 /** Guess verification progress (as a fraction between 0.0=genesis and 1.0=current tip). */

test/functional/mweb_p2p.py

@@ -7,6 +7,9 @@ Test LIP-0007
 
 1. Test getdata 'mwebheader' *before* MWEB activation
 2. Test getdata 'mwebheader' *after* MWEB activation
+3. Test getdata 'mwebleafset' *before* MWEB activation
+4. Test getdata 'mwebleafset' *after* MWEB activation
+   - Request from earlier block (not tip) to make sure rewind works
 """
 
 from test_framework.messages import (
@@ -16,6 +19,7 @@ from test_framework.messages import (
     hash256,
     msg_getdata,
     MSG_MWEB_HEADER,
+    MSG_MWEB_LEAFSET,
 )
 from test_framework.p2p import P2PInterface, p2p_lock
 from test_framework.script import MAX_SCRIPT_ELEMENT_SIZE
@@ -29,6 +33,7 @@ class MockLightClient(P2PInterface):
         super().__init__()
         self.merkle_blocks_with_mweb = {}
         self.block_headers = {}
+        self.leafsets = {}
 
     def request_mweb_header(self, block_hash):
         want = msg_getdata([CInv(MSG_MWEB_HEADER, int(block_hash, 16))])
@@ -36,6 +41,13 @@ class MockLightClient(P2PInterface):
 
     def on_mwebheader(self, message):
         self.merkle_blocks_with_mweb[message.header_hash()] = message.merkleblockwithmweb
+        
+    def request_mweb_leafset(self, block_hash):
+        want = msg_getdata([CInv(MSG_MWEB_LEAFSET, int(block_hash, 16))])
+        self.send_message(want)
+
+    def on_mwebleafset(self, message):
+        self.leafsets[message.block_hash] = message.leafset
 
     def on_block(self, message):
         message.block.calc_sha256()
@@ -48,33 +60,7 @@ class MWEBP2PTest(BitcoinTestFramework):
         self.extra_args = [['-whitelist=noban@127.0.0.1']]  # immediate tx relay
         self.num_nodes = 1
 
-    def run_test(self):
-        node = self.nodes[0]
-        light_client = node.add_p2p_connection(MockLightClient())
-
-        self.log.info("Generate pre-MWEB blocks")
-        pre_mweb_block_hash = node.generate(FIRST_MWEB_HEIGHT - 1)[-1]
-
-        self.log.info("Request mweb_header for pre-MWEB block '{}'".format(pre_mweb_block_hash))
-        light_client.request_mweb_header(pre_mweb_block_hash)
-        
-        self.log.info("Activate MWEB")
-        node.sendtoaddress(node.getnewaddress(address_type='mweb'), 1)
-        post_mweb_block_hash = node.generate(1)[0]
-
-        self.log.info("Request mweb_header for block '{}'".format(post_mweb_block_hash))
-        light_client.request_mweb_header(post_mweb_block_hash)
-
-        self.log.info("Waiting for mweb_header")
-        light_client.wait_for_mwebheader(post_mweb_block_hash, 5)
-        light_client.wait_for_block(int(post_mweb_block_hash, 16), 5)
-
-        self.log.info("Assert results")
-
-        # Before MWEB activation, no merkle block should be returned
-        assert pre_mweb_block_hash not in light_client.merkle_blocks_with_mweb
-
-        # After MWEB activation, the requested merkle block should be returned
+    def assert_mweb_header(self, node, light_client, post_mweb_block_hash):
         assert post_mweb_block_hash in light_client.merkle_blocks_with_mweb
         merkle_block_with_mweb = light_client.merkle_blocks_with_mweb[post_mweb_block_hash]
 
@@ -104,5 +90,55 @@ class MWEBP2PTest(BitcoinTestFramework):
         merkle_root_bytes = hash256(left_hash.serialize() + right_branch_bytes)
         assert_equal(Hash.from_byte_arr(merkle_root_bytes), Hash(block_header.hashMerkleRoot))
 
+    def run_test(self):
+        node = self.nodes[0]
+        light_client = node.add_p2p_connection(MockLightClient())
+
+        self.log.info("Generate pre-MWEB blocks")
+        pre_mweb_block_hash = node.generate(FIRST_MWEB_HEIGHT - 1)[-1]
+
+        self.log.info("Request 'mwebheader' and 'mwebleafset' for pre-MWEB block '{}'".format(pre_mweb_block_hash))
+        light_client.request_mweb_header(pre_mweb_block_hash)
+        light_client.request_mweb_leafset(pre_mweb_block_hash)
+        
+        self.log.info("Activate MWEB")
+        node.sendtoaddress(node.getnewaddress(address_type='mweb'), 1)
+        post_mweb_block_hash = node.generate(1)[0]
+
+        self.log.info("Waiting for mweb_header")
+        light_client.wait_for_block(int(post_mweb_block_hash, 16), 5)
+        
+        self.log.info("Pegin some additional coins")
+        node.sendtoaddress(node.getnewaddress(address_type='mweb'), 10)
+        post_mweb_block_hash2 = node.generate(1)[0]
+        light_client.wait_for_block(int(post_mweb_block_hash2, 16), 5)
+        
+        self.log.info("Request 'mwebheader' and 'mwebleafset' for block '{}'".format(post_mweb_block_hash))
+        light_client.request_mweb_header(post_mweb_block_hash)
+        light_client.request_mweb_leafset(post_mweb_block_hash)
+
+        self.log.info("Waiting for 'mwebheader' and 'mwebleafset'")
+        light_client.wait_for_mwebheader(post_mweb_block_hash, 5)
+        light_client.wait_for_mwebleafset(post_mweb_block_hash, 5)
+
+        self.log.info("Assert results")
+
+        # Before MWEB activation, no merkle block should be returned
+        assert pre_mweb_block_hash not in light_client.merkle_blocks_with_mweb
+
+        # After MWEB activation, the requested merkle block should be returned
+        self.assert_mweb_header(node, light_client, post_mweb_block_hash)
+
+        # Before MWEB activation, no leafset should be returned
+        assert pre_mweb_block_hash not in light_client.leafsets
+
+        # After MWEB activation, the leafset should be returned
+        # Only 2 outputs should be in the UTXO set (the pegin and its change)
+        # That's '11' and then padded to the right with 0's, then serialized in big endian.
+        # So we expect the serialized leafset to be 0b11000000 or 0xc0
+        assert Hash.from_hex(post_mweb_block_hash) in light_client.leafsets
+        leafset = light_client.leafsets[Hash.from_hex(post_mweb_block_hash)]
+        assert_equal([0xc0], leafset)
+
 if __name__ == '__main__':
     MWEBP2PTest().main()

test/functional/test_framework/messages.py

@@ -72,6 +72,7 @@ MSG_WITNESS_TX = MSG_TX | MSG_WITNESS_FLAG
 MSG_MWEB_BLOCK = MSG_BLOCK | MSG_WITNESS_FLAG | MSG_MWEB_FLAG
 MSG_MWEB_TX = MSG_WITNESS_TX | MSG_MWEB_FLAG
 MSG_MWEB_HEADER = 8 | MSG_MWEB_FLAG
+MSG_MWEB_LEAFSET = 9 | MSG_MWEB_FLAG
 
 FILTER_TYPE_BASIC = 0
 
@@ -354,7 +355,8 @@ class CInv:
         MSG_FILTERED_BLOCK: "filtered Block",
         MSG_CMPCT_BLOCK: "CompactBlock",
         MSG_WTX: "WTX",
-        MSG_MWEB_HEADER: "MWEB Header"
+        MSG_MWEB_HEADER: "MWEB Header",
+        MSG_MWEB_LEAFSET: "MWEB Leafset"
     }
 
     def __init__(self, t=0, h=0):
@@ -2045,7 +2047,7 @@ class MWEBInput:
         self.hash = None
 
     def deserialize(self, f):
-        self.features = f.read(1)
+        self.features = struct.unpack("B", f.read(1))[0]
         self.output_id = Hash.deserialize(f)
         self.commitment = deser_pubkey(f)
         self.output_pubkey = deser_pubkey(f)
@@ -2060,14 +2062,14 @@ class MWEBInput:
     def serialize(self):
         r = b""
         r += struct.pack("B", self.features)
-        r += ser_uint256(self.output_id)
+        r += self.output_id.serialize()
         r += ser_pubkey(self.commitment)
         r += ser_pubkey(self.output_pubkey)
         if self.features & 1:
             r += ser_pubkey(self.input_pubkey)
         if self.features & 2:
-            r += ser_compact_size(self.extradata)
-            r += ser_fixed_bytes(self.extradata, len(self.extradata))#extradata
+            r += ser_compact_size(len(self.extradata))
+            r += ser_fixed_bytes(self.extradata, len(self.extradata))
         r += ser_signature(self.signature)
         return r
     
@@ -2404,4 +2406,26 @@ class msg_mwebheader:
         return self.merkleblockwithmweb.merkle.header.hash
 
     def __repr__(self):
-        return "msg_mwebheader(merkleblockwithmweb=%s)" % (repr(self.merkleblockwithmweb))
+        return "msg_mwebheader(merkleblockwithmweb=%s)" % (repr(self.merkleblockwithmweb))
+
+class msg_mwebleafset:
+    __slots__ = ("block_hash", "leafset")
+    msgtype = b"mwebleafset"
+
+    def __init__(self, block_hash=None, leafset=None):
+        self.block_hash = block_hash
+        self.leafset = leafset
+
+    def deserialize(self, f):
+        self.block_hash = Hash.deserialize(f)
+        self.leafset = deser_fixed_bytes(f, deser_compact_size(f))
+
+    def serialize(self):
+        r = b""
+        r += ser_compact_size(len(self.leafset))
+        r += ser_fixed_bytes(self.leafset, len(self.leafset))
+        return r
+
+    def __repr__(self):
+        leafset_hex = ser_fixed_bytes(self.leafset, len(self.leafset)).hex() #encode(self.leafset, 'hex_codec').decode('ascii')
+        return "msg_mwebleafset(block_hash=%s, leafset=%s%s)" % (repr(self.block_hash), repr(leafset_hex)[:50], "..." if len(leafset_hex) > 50 else "")

test/functional/test_framework/p2p.py

@@ -30,6 +30,7 @@ import threading
 
 from test_framework.messages import (
     CBlockHeader,
+    Hash,
     MAX_HEADERS_RESULTS,
     MIN_VERSION_SUPPORTED,
     msg_addr,
@@ -55,6 +56,7 @@ from test_framework.messages import (
     msg_mempool,
     msg_merkleblock,
     msg_mwebheader,
+    msg_mwebleafset,
     msg_notfound,
     msg_ping,
     msg_pong,
@@ -104,6 +106,7 @@ MESSAGEMAP = {
     b"mempool": msg_mempool,
     b"merkleblock": msg_merkleblock,
     b"mwebheader": msg_mwebheader,
+    b"mwebleafset": msg_mwebleafset,
     b"notfound": msg_notfound,
     b"ping": msg_ping,
     b"pong": msg_pong,
@@ -490,6 +493,18 @@ class P2PInterface(P2PConnection):
 
         self.wait_until(test_function, timeout=timeout)
 
+    def wait_for_mwebleafset(self, blockhash, timeout=60):
+        """Waits for an mwebleafset message
+        
+        The hash of the block header must match the provided blockhash"""
+        def test_function():
+            last_mwebleafset = self.last_message.get('mwebleafset')
+            if not last_mwebleafset:
+                return False
+            return last_mwebleafset.block_hash == Hash(int(blockhash, 16))
+
+        self.wait_until(test_function, timeout=timeout)
+
     def wait_for_getdata(self, hash_list, timeout=60):
         """Waits for a getdata message.