wblocktxn receiving test, mempool insertion

instagibbs · Mar 8, 2024 · ecdaaf8 · ecdaaf8
1 parent 107d97e
commit ecdaaf8
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 4 deletions.
diff --git a/src/net_processing.cpp b/src/net_processing.cpp
@@ -4523,9 +4523,50 @@ void PeerManagerImpl::ProcessMessage(CNode& pfrom, const std::string& msg_type,
 
     if (msg_type == NetMsgType::WBLOCKTXN)
     {
-        // Reconstruct block, cache it
+        LogPrint(BCLog::NET, "Received wblocktxn from peer %d\n", pfrom.GetId());
 
-        // Send off weak compact blocks over hb links
+        BlockTransactions resp;
+        vRecv >> resp;
+
+        if (pfrom.GetId() != m_weak_partial_block_source) {
+            LogPrint(BCLog::NET, "Didn't ask for wblocktxn from peer %d; ignoring\n", pfrom.GetId());
+            return;
+        }
+
+        // Check that it matches what we most recently asked for
+        if (!m_weak_partial_block || m_weak_partial_block->header.GetHash() != resp.blockhash) {
+            LogPrint(BCLog::NET, "Wrong wblocktxn header %s from peer %d; ignoring\n", resp.blockhash.ToString(), pfrom.GetId());
+            return;
+        }
+
+        std::shared_ptr<CBlock> pblock = std::make_shared<CBlock>();
+        ReadStatus status = m_weak_partial_block->FillBlock(*pblock, resp.txn, /*check_pow=*/false);
+        if (status == READ_STATUS_OK) {
+            LogPrint(BCLog::NET, "Reconstructed weak compact block %s from peer %d\n", pblock->GetHash().ToString(), pfrom.GetId());
+            for (const auto& tx : pblock->vtx) {
+                AddToWeakBlockTransactions(tx);
+
+                // We try to insert everything we hear about into mempool
+                {
+                    LOCK(cs_main);
+                    const MempoolAcceptResult result = m_chainman.ProcessTransaction(tx);
+                    if (result.m_result_type == MempoolAcceptResult::ResultType::VALID) {
+                        LogPrint(BCLog::MEMPOOL, "AcceptToMemoryPool: peer=%d: accepted %s (wtxid=%s) (poolsz %u txn, %u kB)\n",
+                            pfrom.GetId(),
+                            tx->GetHash().ToString(),
+                            tx->GetWitnessHash().ToString(),
+                            m_mempool.size(), m_mempool.DynamicMemoryUsage() / 1000);
+                    }
+                }
+            }
+            // FIXME send to all peers who requested hb connections
+        } else {
+            LogPrint(BCLog::NET, "Unable to reconstruct weak compact block %s from peer %d\n", resp.blockhash.ToString(), pfrom.GetId());
+        }
+        m_weak_partial_block.reset();
+        m_weak_partial_block_source = -1;
+
+        return;
     }
 
     if (msg_type == NetMsgType::CMPCTBLOCK)

diff --git a/test/functional/p2p_weakblocks.py b/test/functional/p2p_weakblocks.py
@@ -31,7 +31,7 @@
     PrefilledTransaction,
     calculate_shortid,
     msg_block,
-    msg_blocktxn,
+    msg_wblocktxn,
     msg_cmpctblock,
     msg_getblocktxn,
     msg_getdata,
@@ -313,7 +313,18 @@ def test_getwcmpctblock_message(self):
         with p2p_lock:
             assert "getwblocktxn" in self.segwit_node.last_message
 
-        # TODO respond to it
+        # Respond to it wtih WBLOCKTXN, look for reconstruction
+        msg = msg_wblocktxn()
+        msg.weak_block_transactions.blockhash = block.sha256
+        msg.weak_block_transactions.transactions = block.vtx[1:]
+
+        assert_equal(self.nodes[0].getrawmempool(), [])
+
+        with self.nodes[0].assert_debug_log(expected_msgs=["Reconstructed weak compact block"]):
+            self.segwit_node.send_and_ping(msg)
+
+        assert_equal(self.nodes[0].getrawmempool(), [block.vtx[-1].rehash()])
+
 
     def test_basic_weakcmpctblock_message(self):
         self.generate(self.nodes[0], COINBASE_MATURITY + 1)

diff --git a/test/functional/test_framework/messages.py b/test/functional/test_framework/messages.py
@@ -1728,6 +1728,25 @@ def __repr__(self):
         return "msg_getwblocktxn(weak_block_txn_request=%s)" % (repr(self.weak_block_txn_request))
 
 
+class msg_wblocktxn:
+    __slots__ = ("weak_block_transactions",)
+    msgtype = b"wblocktxn"
+
+    def __init__(self):
+        self.weak_block_transactions = BlockTransactions()
+
+    def deserialize(self, f):
+        self.weak_block_transactions.deserialize(f)
+
+    def serialize(self):
+        r = b""
+        r += self.weak_block_transactions.serialize()
+        return r
+
+    def __repr__(self):
+        return "msg_wblocktxn(weak_block_transactions=%s)" % (repr(self.weak_block_transactions))
+
+
 class msg_getblocktxn:
     __slots__ = ("block_txn_request",)
     msgtype = b"getblocktxn"