Repository: bitcoin/bitcoin
File: test/functional/test_framework/test_framework.py (Python)

Code before:
        if peer_advertises_v2 != from_connection.use_v2transport:
            from_connection.addnode(node=ip_port, command="onetry", v2transport=peer_advertises_v2)
        else:
            # skip the optional third argument if it matches the default, for
            # compatibility with older clients
            from_connection.addnode(ip_port, "onetry")

        if not wait_for_connect:
            return

        # Use subversion as peer id. Test nodes have their node number appended to the user agent string
        from_connection_subver = from_connection.getnetworkinfo()['subversion']
        to_connection_subver = to_connection.getnetworkinfo()['subversion']

        def find_conn(node, peer_subversion, inbound):
            return next(filter(lambda peer: peer['subver'] == peer_subversion and peer['inbound'] == inbound, node.getpeerinfo()), None)

        # poll until version handshake complete to avoid race conditions
        # with transaction relaying
        # See comments in net_processing:
        # * Must have a version message before anything else
        # * Must have a verack message before anything else
        self.wait_until(lambda: find_conn(from_connection, to_connection_subver, inbound=False) is not None)
        self.wait_until(lambda: find_conn(to_connection, from_connection_subver, inbound=True) is not None)

        def check_bytesrecv(peer, msg_type, min_bytes_recv):
            assert peer is not None, "Error: peer disconnected"
            return peer['bytesrecv_per_msg'].pop(msg_type, 0) >= min_bytes_recv

        self.wait_until(lambda: check_bytesrecv(find_conn(from_connection, to_connection_subver, inbound=False), 'verack', 21))
        self.wait_until(lambda: check_bytesrecv(find_conn(to_connection, from_connection_subver, inbound=True), 'verack', 21))

        # The message bytes are counted before processing the message, so make
        # sure it was fully processed by waiting for a ping.
        self.wait_until(lambda: check_bytesrecv(find_conn(from_connection, to_connection_subver, inbound=False), 'pong', 29))
        self.wait_until(lambda: check_bytesrecv(find_conn(to_connection, from_connection_subver, inbound=True), 'pong', 29))

    def disconnect_nodes(self, a, b):
        def disconnect_nodes_helper(node_a, node_b):
            def get_peer_ids(from_connection, node_num):
                result = []
                for peer in from_connection.getpeerinfo():
                    if "testnode{}".format(node_num) in peer['subver']:
                        result.append(peer['id'])
                return result

            peer_ids = get_peer_ids(node_a, node_b.index)
            if not peer_ids:
                self.log.warning("disconnect_nodes: {} and {} were not connected".format(
                    node_a.index,
                    node_b.index,

Review comment (bug):
It would be nice to be explicit about the race conditions rather than keeping this vague. Maybe could add something like: "The node will ignore all messages received during the version handshake window".

Code after:
        if peer_advertises_v2 != from_connection.use_v2transport:
            from_connection.addnode(node=ip_port, command="onetry", v2transport=peer_advertises_v2)
        else:
            # skip the optional third argument if it matches the default, for
            # compatibility with older clients
            from_connection.addnode(ip_port, "onetry")

        if not wait_for_connect:
            return

        # Use subversion as peer id. Test nodes have their node number appended to the user agent string
        from_connection_subver = from_connection.getnetworkinfo()['subversion']
        to_connection_subver = to_connection.getnetworkinfo()['subversion']

        def find_conn(node, peer_subversion, inbound):
            return next(filter(lambda peer: peer['subver'] == peer_subversion and peer['inbound'] == inbound, node.getpeerinfo()), None)

        self.wait_until(lambda: find_conn(from_connection, to_connection_subver, inbound=False) is not None)
        self.wait_until(lambda: find_conn(to_connection, from_connection_subver, inbound=True) is not None)

        def check_bytesrecv(peer, msg_type, min_bytes_recv):
            assert peer is not None, "Error: peer disconnected"
            return peer['bytesrecv_per_msg'].pop(msg_type, 0) >= min_bytes_recv

        # Poll until version handshake (fSuccessfullyConnected) is complete to
        # avoid race conditions, because some message types are blocked from
        # being sent or received before fSuccessfullyConnected.
        #
        # As the flag fSuccessfullyConnected is not exposed, check it by
        # waiting for a pong, which can only happen after the flag was set.
        self.wait_until(lambda: check_bytesrecv(find_conn(from_connection, to_connection_subver, inbound=False), 'pong', 29))
        self.wait_until(lambda: check_bytesrecv(find_conn(to_connection, from_connection_subver, inbound=True), 'pong', 29))

    def disconnect_nodes(self, a, b):
        def disconnect_nodes_helper(node_a, node_b):
            def get_peer_ids(from_connection, node_num):
                result = []
                for peer in from_connection.getpeerinfo():
                    if "testnode{}".format(node_num) in peer['subver']:
                        result.append(peer['id'])
                return result

            peer_ids = get_peer_ids(node_a, node_b.index)
            if not peer_ids:
                self.log.warning("disconnect_nodes: {} and {} were not connected".format(
                    node_a.index,
                    node_b.index,
                ))
                return
            for peer_id in peer_ids:
                try: