test: add case for full askfor set.

furszy · furszy · commit 96c64ec84675 · 2021-10-29T11:29:34.000-03:00
diff --git a/test/functional/p2p_invalid_messages.py b/test/functional/p2p_invalid_messages.py
@@ -12,10 +12,13 @@
     P2PDataStore,
     P2PInterface,
 )
+from test_framework.messages import CTxIn, COutPoint, msg_mnping
 from test_framework.test_framework import PivxTestFramework
 from test_framework.util import (
+    assert_equal,
     hex_str_to_bytes,
 )
+from random import getrandbits
 
 MSG_LIMIT = 2 * 1024 * 1024  # 2MB, per MAX_PROTOCOL_MESSAGE_LENGTH
 VALID_DATA_LIMIT = MSG_LIMIT - 5  # Account for the 5-byte length prefix
@@ -39,6 +42,19 @@ class SenderOfAddrV2(P2PInterface):
     def wait_for_sendaddrv2(self):
         self.wait_until(lambda: 'sendaddrv2' in self.last_message)
 
+class InvReceiver(P2PInterface):
+
+    def __init__(self):
+        super().__init__()
+        self.vec_mnp = {}
+        self.getdata_count = 0
+
+    def on_getdata(self, message):
+        for inv in message.inv:
+            if inv.type == 15: # MNPING
+                self.send_message(self.vec_mnp[inv.hash])
+                self.getdata_count+=1
+
 class InvalidMessagesTest(PivxTestFramework):
     def set_test_params(self):
         self.num_nodes = 1
@@ -55,6 +71,7 @@ def run_test(self):
         self.test_addrv2_unrecognized_network()
         self.test_large_inv()
         self.test_resource_exhaustion()
+        self.test_fill_askfor()
 
     def test_magic_bytes(self):
         conn = self.nodes[0].add_p2p_connection(P2PDataStore())
@@ -187,6 +204,33 @@ def test_large_inv(self):
             conn.send_and_ping(msg)
         self.nodes[0].disconnect_p2ps()
 
+    def test_fill_askfor(self):
+        self.nodes[0].generate(1) # IBD
+        conn = self.nodes[0].add_p2p_connection(InvReceiver())
+        invs = []
+        blockhash = int(self.nodes[0].getbestblockhash(), 16)
+        for _ in range(50000):
+            mnp = msg_mnping(CTxIn(COutPoint(getrandbits(256))), blockhash, int(time.time()))
+            conn.vec_mnp[mnp.get_hash()] = mnp
+            invs.append(messages.CInv(15, mnp.get_hash()))
+        assert_equal(len(conn.vec_mnp), 50000)
+        assert_equal(len(invs), 50000)
+        msg = messages.msg_inv(invs)
+        conn.send_message(msg)
+
+        time.sleep(20) # wait a bit
+        assert_equal(conn.getdata_count, 50000)
+
+        # Prior #2611 the node was blocking any follow-up request.
+        mnp = msg_mnping(CTxIn(COutPoint(getrandbits(256))), getrandbits(256), int(time.time()))
+        conn.vec_mnp[mnp.get_hash()] = mnp
+        msg = messages.msg_inv([messages.CInv(15, mnp.get_hash())])
+        conn.send_and_ping(msg)
+        time.sleep(3)
+
+        assert_equal(conn.getdata_count, 50001)
+        self.nodes[0].disconnect_p2ps()
+
     def test_resource_exhaustion(self):
         conn = self.nodes[0].add_p2p_connection(P2PDataStore())
         conn2 = self.nodes[0].add_p2p_connection(P2PDataStore())
diff --git a/test/functional/test_framework/messages.py b/test/functional/test_framework/messages.py
@@ -1149,6 +1149,41 @@ def __repr__(self):
         return "msg_getblocks(locator=%s hashstop=%064x)" \
             % (repr(self.locator), self.hashstop)
 
+class msg_mnping:
+    command = b"mnp"
+
+    def __init__(self, _vin, _blockhash, _sigtime):
+        self.vin = _vin
+        self.blockhash = _blockhash
+        self.sigtime = _sigtime
+        self.vch_sig = b""
+        self.mess_version = 1
+        self.cached_hash = b""
+
+    def deserialize(self, f):
+        self.vin.deserialize(f)
+        self.blockhash = deser_uint256(f)
+        self.sigtime = struct.unpack("<q", f.read(8))[0]
+        self.vch_sig = deser_string(f)
+        self.mess_version = struct.unpack("<i", f.read(4))[0]
+
+    def serialize(self):
+        r = b""
+        r += self.vin.serialize()
+        r += ser_uint256(self.blockhash)
+        r += struct.pack("<q", self.sigtime)
+        r += ser_string(self.vch_sig)
+        r += struct.pack("<I", self.mess_version)
+        return r
+
+    def get_hash(self):
+        if self.cached_hash == b"":
+            self.cached_hash = uint256_from_str(hash256(self.serialize()))
+        return self.cached_hash
+
+    def __repr__(self):
+        return "msg_mnping(vin=%s blockhash=%064x)" \
+               % (repr(self.vin), self.blockhash)
 
 class msg_tx:
     command = b"tx"
