test: add coverage for regular peers disconnecting after receiving an…

… authenticated connection.

test/functional/p2p_quorum_connect.py

@@ -6,17 +6,24 @@
 
 import time
 
+from random import getrandbits
 from test_framework.test_framework import PivxTestFramework
+from test_framework.mininode import P2PInterface
+from test_framework.messages import msg_version
 from test_framework.util import (
     assert_equal,
     assert_true,
     bytes_to_hex_str,
-    disconnect_nodes,
+    connect_nodes_clique,
     hash256,
     hex_str_to_bytes,
     wait_until,
 )
 
+class TestP2PConn(P2PInterface):
+    def on_version(self, message):
+        pass
+
 class DMNConnectionTest(PivxTestFramework):
 
     def set_test_params(self):
@@ -47,7 +54,7 @@ def check_mn_enabled_count(self, enabled, total):
     def wait_until_mnsync_completed(self):
         SYNC_FINISHED = [999] * self.num_nodes
         synced = [-1] * self.num_nodes
-        timeout = time.time() + 160
+        timeout = time.time() + 240
         while synced != SYNC_FINISHED and time.time() < timeout:
             synced = [node.mnsync("status")["RequestedMasternodeAssets"]
                       for node in self.nodes]
@@ -57,6 +64,12 @@ def wait_until_mnsync_completed(self):
             raise AssertionError("Unable to complete mnsync: %s" % str(synced))
 
     def setup_phase(self):
+        # Disconnect every node and generate connections to all of them manually.
+        for node in self.nodes:
+            self.disconnect_peers(node)
+        connect_nodes_clique(self.nodes)
+        self.check_peers_count(14)
+
         # Mine 110 blocks
         self.log.info("Mining...")
         self.miner.generate(110)
@@ -83,11 +96,16 @@ def setup_phase(self):
         self.log.info("All masternodes ready")
 
     def disconnect_peers(self, node):
-        for i in range(0, 7):
-            disconnect_nodes(node, i)
+        node.setnetworkactive(False)
+        time.sleep(3)
         assert_equal(len(node.getpeerinfo()), 0)
+        node.setnetworkactive(True)
 
-    def check_peer_info(self, peer_info, mn, is_iqr_conn, inbound = False):
+    def check_peers_count(self, expected_count):
+        for i in range(0, 7):
+            assert_equal(len(self.nodes[i].getpeerinfo()), expected_count)
+
+    def check_peer_info(self, peer_info, mn, is_iqr_conn, inbound=False):
         assert_equal(peer_info["masternode"], True)
         assert_equal(peer_info["verif_mn_proreg_tx_hash"], mn.proTx)
         assert_equal(peer_info["verif_mn_operator_pubkey_hash"], bytes_to_hex_str(hash256(hex_str_to_bytes(mn.operator_pk))))
@@ -96,7 +114,7 @@ def check_peer_info(self, peer_info, mn, is_iqr_conn, inbound = False):
         if not inbound:
             assert_equal(peer_info["addr"], mn.ipport)
 
-    def check_peers_info(self, peers_info, quorum_members, is_iqr_conn, inbound = False):
+    def check_peers_info(self, peers_info, quorum_members, is_iqr_conn, inbound=False):
         for quorum_node in quorum_members:
             found = False
             for peer in peers_info:
@@ -180,7 +198,7 @@ def run_test(self):
         ###########################################
         # 4) Now test the connections probe process
         ###########################################
-        self.log.info("3) Testing MN probe connection process..")
+        self.log.info("4) Testing MN probe connection process..")
         # Take mn6, disconnect all the nodes and try to probe connection to one of them
         mn6_node = self.nodes[mn6.idx]
         self.disconnect_peers(mn6_node)
@@ -190,6 +208,22 @@ def run_test(self):
             time.sleep(10) # wait a bit until the connection gets established
         self.log.info("Completed MN connection probe!")
 
+        ###############################################################################
+        # 5) Now test regular node disconnecting after receiving an auth DMN connection
+        ###############################################################################
+        self.log.info("5) Testing regular node disconnection after receiving an auth DMN connection..")
+        self.disconnect_peers(self.miner)
+        no_version_node = self.miner.add_p2p_connection(TestP2PConn(), send_version=False, wait_for_verack=False)
+        assert_equal(len(self.miner.getpeerinfo()), 1)
+        # send the version as it would be a MN
+        mn_challenge = getrandbits(256)
+        with self.miner.assert_debug_log(["but we're not a masternode, disconnecting"]):
+            no_version_node.send_message(msg_version(mn_challenge))
+            time.sleep(2)
+        # as the miner is not a DMN, the miner should had dropped the connection.
+        assert_equal(len(self.miner.getpeerinfo()), 0)
+        self.log.info("Regular node disconnected auth connection successfully")
+
 
 if __name__ == '__main__':
     DMNConnectionTest().main()

test/functional/test_framework/messages.py

@@ -959,7 +959,7 @@ def __repr__(self):
 class msg_version:
     command = b"version"
 
-    def __init__(self):
+    def __init__(self, mn_auth_challenge=0):
         self.nVersion = MY_VERSION
         self.nServices = NODE_NETWORK
         self.nTime = int(time.time())
@@ -969,6 +969,7 @@ def __init__(self):
         self.strSubVer = MY_SUBVERSION
         self.nStartingHeight = -1
         self.nRelay = MY_RELAY
+        self.mn_auth_challenge = mn_auth_challenge
 
     def deserialize(self, f):
         self.nVersion = struct.unpack("<i", f.read(4))[0]
@@ -1004,6 +1005,13 @@ def deserialize(self, f):
         else:
             self.nRelay = 0
 
+        if self.nVersion >= 70925:
+            try:
+                self.mn_auth_challenge = deser_uint256(f)
+            except:
+                self.mn_auth_challenge = 0
+
+
     def serialize(self):
         r = b""
         r += struct.pack("<i", self.nVersion)
@@ -1015,6 +1023,8 @@ def serialize(self):
         r += ser_string(self.strSubVer.encode('utf-8'))
         r += struct.pack("<i", self.nStartingHeight)
         r += struct.pack("<b", self.nRelay)
+        if self.mn_auth_challenge != 0:
+            r += ser_uint256(self.mn_auth_challenge)
         return r
 
     def __repr__(self):