Inside Bitcoin

501    def peer_connect(self, *, services=P2P_SERVICES, send_version, **kwargs):
502        create_conn = super().peer_connect(**kwargs)
503
504        if send_version:
505            self.peer_connect_send_version(services)
506
507        return create_conn
508
509    def peer_accept_connection(self, *args, services=P2P_SERVICES, **kwargs):
510        create_conn = super().peer_accept_connection(*args, **kwargs)
511        self.peer_connect_send_version(services)
512
513        return create_conn
514
515    # Message receiving methods
516
517    def on_message(self, message):
518        """Receive message and dispatch message to appropriate callback.
519
520        We keep a count of how many of each message type has been received
521        and the most recent message of each type."""
522        with p2p_lock:
523            try:
524                msgtype = message.msgtype.decode('ascii')
525                self.message_count[msgtype] += 1
526                self.last_message[msgtype] = message
527                getattr(self, 'on_' + msgtype)(message)
528            except Exception:
529                print("ERROR delivering %s (%s)" % (repr(message), sys.exc_info()[0]))
530                raise
531
532    # Callback methods. Can be overridden by subclasses in individual test
533    # cases to provide custom message handling behaviour.
534
535    def on_open(self):
536        pass
537
538    def on_close(self):
539        pass
540
541    def on_addr(self, message): pass
542    def on_addrv2(self, message): pass
543    def on_block(self, message): pass
544    def on_blocktxn(self, message): pass
545    def on_cfcheckpt(self, message): pass
546    def on_cfheaders(self, message): pass
547    def on_cfilter(self, message): pass
548    def on_cmpctblock(self, message): pass
549    def on_feature(self, message): pass
550    def on_feefilter(self, message): pass
551    def on_filteradd(self, message): pass
552    def on_filterclear(self, message): pass
553    def on_filterload(self, message): pass
554    def on_getaddr(self, message): pass
555    def on_getblocks(self, message): pass
556    def on_getblocktxn(self, message): pass
557    def on_getdata(self, message): pass
558    def on_getheaders(self, message): pass
559    def on_headers(self, message): pass
560    def on_mempool(self, message): pass
561    def on_merkleblock(self, message): pass
562    def on_notfound(self, message): pass
563    def on_pong(self, message): pass
564    def on_sendaddrv2(self, message): pass
565    def on_sendcmpct(self, message): pass
566    def on_sendheaders(self, message): pass
567    def on_sendtxrcncl(self, message): pass
568    def on_tx(self, message): pass
569    def on_wtxidrelay(self, message): pass
570
571    def on_inv(self, message):
572        want = msg_getdata()
573        for i in message.inv:
574            if i.type != 0:
575                want.inv.append(i)
576        if len(want.inv):
577            self.send_without_ping(want)
578
579    def on_ping(self, message):
580        self.send_without_ping(msg_pong(message.nonce))
581
582    def on_verack(self, message):
583        pass
584
585    def on_version(self, message):
586        assert message.nVersion >= MIN_P2P_VERSION_SUPPORTED, "Version {} received. Test framework only supports versions greater than {}".format(message.nVersion, MIN_P2P_VERSION_SUPPORTED)
587        # for inbound connections, reply to version with own version message
588        # (could be due to v1 reconnect after a failed v2 handshake)
589        if not self.p2p_connected_to_node:
590            self.send_version()
591            self.reconnect = False
592        if message.nVersion >= 70016 and self.wtxidrelay:
593            self.send_without_ping(msg_wtxidrelay())
594        if self.support_addrv2:
595            self.send_without_ping(msg_sendaddrv2())
596        self.send_without_ping(msg_verack())
597        self.nServices = message.nServices
598        self.relay = message.relay
599        if self.p2p_connected_to_node:
600            self.send_without_ping(msg_getaddr())
601
602    # Connection helper methods
603
604    def wait_until(self, test_function_in, *, timeout=60, check_connected=True, check_interval=0.05):
605        def test_function():
606            if check_connected:
607                assert self.is_connected
608            return test_function_in()
609
610        wait_until_helper_internal(test_function, timeout=timeout, lock=p2p_lock, timeout_factor=self.timeout_factor, check_interval=check_interval)
611
612    def wait_for_connect(self, *, timeout=60):
613        def test_function():
614            return self.is_connected
615
616        self.wait_until(test_function, timeout=timeout, check_connected=False)
617
618    def wait_for_disconnect(self, *, timeout=60):
619        def test_function():
620            return not self.is_connected
621
622        self.wait_until(test_function, timeout=timeout, check_connected=False)
623
624    def wait_for_reconnect(self, *, timeout=60):
625        def test_function():
626            return self.is_connected and self.last_message.get('version') and not self.supports_v2_p2p
627        self.wait_until(test_function, timeout=timeout, check_connected=False)
628
629    # Message receiving helper methods
630
631    def wait_for_tx(self, txid, *, timeout=60):
632        def test_function():
633            if not self.last_message.get('tx'):
634                return False
635            return self.last_message['tx'].tx.txid_hex == txid
636
637        self.wait_until(test_function, timeout=timeout)
638
639    def wait_for_block(self, blockhash, *, timeout=60):
640        def test_function():
641            return self.last_message.get("block") and self.last_message["block"].block.hash_int == blockhash
642
643        self.wait_until(test_function, timeout=timeout)
644
645    def wait_for_header(self, blockhash, *, timeout=60):
646        def test_function():
647            last_headers = self.last_message.get('headers')
648            if not last_headers:
649                return False
650            return last_headers.headers[0].hash_int == int(blockhash, 16)
651
652        self.wait_until(test_function, timeout=timeout)
653
654    def wait_for_merkleblock(self, blockhash, *, timeout=60):
655        def test_function():
656            last_filtered_block = self.last_message.get('merkleblock')
657            if not last_filtered_block:
658                return False
659            return last_filtered_block.merkleblock.header.hash_int == int(blockhash, 16)
660
661        self.wait_until(test_function, timeout=timeout)
662
663    def wait_for_getdata(self, hash_list, *, timeout=60):
664        """Waits for a getdata message.
665
666        The object hashes in the inventory vector must match the provided hash_list."""
667        def test_function():
668            last_data = self.last_message.get("getdata")
669            if not last_data:
670                return False
671            return [x.hash for x in last_data.inv] == hash_list
672
673        self.wait_until(test_function, timeout=timeout)
674
675    def wait_for_getheaders(self, block_hash=None, *, timeout=60):
676        """Waits for a getheaders message containing a specific block hash.
677
678        If no block hash is provided, checks whether any getheaders message has been received by the node."""
679        def test_function():
680            last_getheaders = self.last_message.pop("getheaders", None)
681            if block_hash is None:
682                return last_getheaders
683            if last_getheaders is None:
684                return False
685            return block_hash == last_getheaders.locator.vHave[0]
686
687        self.wait_until(test_function, timeout=timeout)
688
689    def wait_for_inv(self, expected_inv, *, timeout=60):
690        """Waits for an INV message and checks that the first inv object in the message was as expected."""
691        if len(expected_inv) > 1:
692            raise NotImplementedError("wait_for_inv() will only verify the first inv object")
693
694        def test_function():
695            return self.last_message.get("inv") and \
696                                self.last_message["inv"].inv[0].type == expected_inv[0].type and \
697                                self.last_message["inv"].inv[0].hash == expected_inv[0].hash
698
699        self.wait_until(test_function, timeout=timeout)
700
701    def wait_for_verack(self, *, timeout=60):
702        def test_function():
703            return "verack" in self.last_message
704
705        self.wait_until(test_function, timeout=timeout)
706
707    # Message sending helper functions
708
709    def send_version(self):
710        if self.on_connection_send_msg:
711            self.send_without_ping(self.on_connection_send_msg)
712            self.on_connection_send_msg = None  # Never used again
713
714    def send_and_ping(self, message, *, timeout=60):
715        self.send_without_ping(message)
716        self.sync_with_ping(timeout=timeout)
717
718    def sync_with_ping(self, *, timeout=60):
719        """Ensure ProcessMessages and SendMessages is called on this connection"""
720        # Sending two pings back-to-back, requires that the node calls
721        # `ProcessMessage` twice, and thus ensures `SendMessages` must have
722        # been called at least once
723        self.send_without_ping(msg_ping(nonce=0))
724        self.send_without_ping(msg_ping(nonce=self.ping_counter))
725
726        def test_function():
727            return self.last_message.get("pong") and self.last_message["pong"].nonce == self.ping_counter
728
729        self.wait_until(test_function, timeout=timeout)
730        self.ping_counter += 1
731
732
733# One lock for synchronizing all data access between the network event loop (see
734# NetworkThread below) and the thread running the test logic.  For simplicity,
735# P2PConnection acquires this lock whenever delivering a message to a P2PInterface.
736# This lock should be acquired in the thread running the test logic to synchronize
737# access to any data shared with the P2PInterface or P2PConnection.
738p2p_lock = threading.Lock()
739
740
741class NetworkThread(threading.Thread):
742    network_event_loop = None
743
744    def __init__(self):
745        super().__init__(name="NetworkThread")
746        # There is only one event loop and no more than one thread must be created
747        assert not self.network_event_loop
748
749        NetworkThread.listeners = {}
750        NetworkThread.protos = {}