Inside Bitcoin

751        NetworkThread.protos_accept_done = []
752
753    def run(self):
754        """Start the network thread."""
755        NetworkThread.network_event_loop = asyncio.SelectorEventLoop() if platform.system() == "Windows" else asyncio.new_event_loop()
756        self.network_event_loop.run_forever()
757
758    def close(self, *, timeout):
759        """Close the connections and network event loop."""
760        for p in NetworkThread.protos_accept_done:
761            p.peer_disconnect()
762        NetworkThread.protos_accept_done.clear()
763
764        listeners = list(NetworkThread.listeners.values())
765        NetworkThread.listeners.clear()
766
767        async def close_listeners():
768            for listener in listeners:
769                listener.close()
770            for listener in listeners:
771                await listener.wait_closed()
772        future = asyncio.run_coroutine_threadsafe(close_listeners(), self.network_event_loop)
773        future.result(timeout=timeout)
774
775        self.network_event_loop.call_soon_threadsafe(self.network_event_loop.stop)
776        wait_until_helper_internal(lambda: not self.network_event_loop.is_running(), timeout=timeout)
777        self.join(timeout)
778        self.network_event_loop.close()
779        # Safe to remove event loop.
780        NetworkThread.network_event_loop = None
781
782    @classmethod
783    def listen(cls, p2p, callback, port=None, addr=None, idx=1):
784        """ Ensure a listening server is running on the given port, and run the
785        protocol specified by `p2p` on the next connection to it. Once ready
786        for connections, call `callback`."""
787
788        if port is None:
789            assert 0 < idx <= MAX_NODES
790            port = p2p_port(MAX_NODES - idx)
791        if addr is None:
792            addr = '127.0.0.1'
793
794        def exception_handler(loop, context):
795            if not p2p.reconnect:
796                loop.default_exception_handler(context)
797
798        cls.network_event_loop.set_exception_handler(exception_handler)
799        coroutine = cls.create_listen_server(addr, port, callback, p2p)
800        cls.network_event_loop.call_soon_threadsafe(cls.network_event_loop.create_task, coroutine)
801
802    @classmethod
803    async def create_listen_server(cls, addr, port, callback, proto):
804        def peer_protocol():
805            """Returns a function that does the protocol handling for a new
806            connection. To allow different connections to have different
807            behaviors, the protocol function is first put in the cls.protos
808            dict. When the connection is made, the function removes the
809            protocol function from that dict, and returns it so the event loop
810            can start executing it."""
811            response = cls.protos.get((addr, port))
812            # remove protocol function from dict only when reconnection doesn't need to happen/already happened
813            if not proto.reconnect:
814                cls.protos_accept_done.append(response)
815                cls.protos[(addr, port)] = None
816            return response
817
818        if port == 0 or (addr, port) not in cls.listeners:
819            # When creating a listener on a given (addr, port) we only need to
820            # do it once. If we want different behaviors for different
821            # connections, we can accomplish this by providing different
822            # `proto` functions
823
824            if port == 0:
825                # Manually create the socket in order to set the range to be
826                # used for the port before the bind() call.
827                if ipaddress.ip_address(addr).version == 4:
828                    address_family = socket.AF_INET
829                else:
830                    address_family = socket.AF_INET6
831                s = socket.socket(address_family)
832                set_ephemeral_port_range(s)
833                s.bind((addr, 0))
834                s.listen()
835                listener = await cls.network_event_loop.create_server(peer_protocol, sock=s)
836                port = listener.sockets[0].getsockname()[1]
837            else:
838                listener = await cls.network_event_loop.create_server(peer_protocol, addr, port)
839
840            logger.debug("Listening server on %s:%d should be started" % (addr, port))
841            cls.listeners[(addr, port)] = listener
842
843        cls.protos[(addr, port)] = proto
844        callback(addr, port)
845
846
847class P2PDataStore(P2PInterface):
848    """A P2P data store class.
849
850    Keeps a block and transaction store and responds correctly to getdata and getheaders requests."""
851
852    def __init__(self):
853        super().__init__()
854        # store of blocks. key is block hash, value is a CBlock object
855        self.block_store = {}
856        self.last_block_hash = ''
857        # store of txs. key is txid, value is a CTransaction object
858        self.tx_store = {}
859        self.getdata_requests = []
860
861    def on_getdata(self, message):
862        """Check for the tx/block in our stores and if found, reply with MSG_TX or MSG_BLOCK."""
863        for inv in message.inv:
864            self.getdata_requests.append(inv.hash)
865            invtype = inv.type & MSG_TYPE_MASK
866            if (invtype == MSG_TX or invtype == MSG_WTX) and inv.hash in self.tx_store.keys():
867                self.send_without_ping(msg_tx(self.tx_store[inv.hash]))
868            elif invtype == MSG_BLOCK and inv.hash in self.block_store.keys():
869                self.send_without_ping(msg_block(self.block_store[inv.hash]))
870            else:
871                logger.debug('getdata message type {} received.'.format(hex(inv.type)))
872
873    def on_getheaders(self, message):
874        """Search back through our block store for the locator, and reply with a headers message if found."""
875
876        locator, hash_stop = message.locator, message.hashstop
877
878        # Assume that the most recent block added is the tip
879        if not self.block_store:
880            return
881
882        headers_list = [self.block_store[self.last_block_hash]]
883        while headers_list[-1].hash_int not in locator.vHave:
884            # Walk back through the block store, adding headers to headers_list
885            # as we go.
886            prev_block_hash = headers_list[-1].hashPrevBlock
887            if prev_block_hash in self.block_store:
888                prev_block_header = CBlockHeader(self.block_store[prev_block_hash])
889                headers_list.append(prev_block_header)
890                if prev_block_header.hash_int == hash_stop:
891                    # if this is the hashstop header, stop here
892                    break
893            else:
894                logger.debug('block hash {} not found in block store'.format(hex(prev_block_hash)))
895                break
896
897        # Truncate the list if there are too many headers
898        headers_list = headers_list[:-MAX_HEADERS_RESULTS - 1:-1]
899        response = msg_headers(headers_list)
900
901        if response is not None:
902            self.send_without_ping(response)
903
904    def send_blocks_and_test(self, blocks, node, *, success=True, force_send=False, reject_reason=None, expect_disconnect=False, timeout=60, is_decoy=False):
905        """Send blocks to test node and test whether the tip advances.
906
907         - add all blocks to our block_store
908         - send a headers message for the final block
909         - the on_getheaders handler will ensure that any getheaders are responded to
910         - if force_send is False: wait for getdata for each of the blocks. The on_getdata handler will
911           ensure that any getdata messages are responded to. Otherwise send the full block unsolicited.
912         - if success is True: assert that the node's tip is the last block in blocks at the end of the operation.
913         - if success is False: assert that the node's tip isn't the last block in blocks at the end of the operation
914         - if reject_reason is set: assert that the correct reject message is logged"""
915
916        with p2p_lock:
917            for block in blocks:
918                self.block_store[block.hash_int] = block
919                self.last_block_hash = block.hash_int
920
921        reject_reason = [reject_reason] if reject_reason else []
922        with node.assert_debug_log(expected_msgs=reject_reason):
923            if is_decoy:  # since decoy messages are ignored by the recipient - no need to wait for response
924                force_send = True
925            if force_send:
926                for b in blocks:
927                    self.send_without_ping(msg_block(block=b), is_decoy)
928            else:
929                self.send_without_ping(msg_headers([CBlockHeader(block) for block in blocks]))
930                self.wait_until(
931                    lambda: blocks[-1].hash_int in self.getdata_requests,
932                    timeout=timeout,
933                    check_connected=success,
934                )
935
936            if expect_disconnect:
937                self.wait_for_disconnect(timeout=timeout)
938            else:
939                self.sync_with_ping(timeout=timeout)
940
941            if success:
942                self.wait_until(lambda: node.getbestblockhash() == blocks[-1].hash_hex, timeout=timeout)
943            else:
944                assert_not_equal(node.getbestblockhash(), blocks[-1].hash_hex)
945
946    def send_txs_and_test(self, txs, node, *, success=True, reject_reason=None):
947        """Send txs to test node and test whether they're accepted to the mempool.
948
949         - add all txs to our tx_store
950         - send tx messages for all txs
951         - if success is True/False: assert that the txs are/are not accepted to the mempool
952         - if reject_reason is set: assert that the correct reject message is logged."""
953
954        with p2p_lock:
955            for tx in txs:
956                self.tx_store[tx.txid_int] = tx
957
958        reject_reason = [reject_reason] if reject_reason else []
959        with node.assert_debug_log(expected_msgs=reject_reason):
960            for tx in txs:
961                self.send_without_ping(msg_tx(tx))
962
963            self.sync_with_ping()
964
965            raw_mempool = node.getrawmempool()
966            if success:
967                # Check that all txs are now in the mempool
968                for tx in txs:
969                    assert tx.txid_hex in raw_mempool, "{} not found in mempool".format(tx.txid_hex)
970            else:
971                # Check that none of the txs are now in the mempool
972                for tx in txs:
973                    assert tx.txid_hex not in raw_mempool, "{} tx found in mempool".format(tx.txid_hex)
974
975class P2PTxInvStore(P2PInterface):
976    """A P2PInterface which stores a count of how many times each txid has been announced."""
977    def __init__(self, **kwargs):
978        super().__init__(**kwargs)
979        self.tx_invs_received = defaultdict(int)
980
981    def on_inv(self, message):
982        super().on_inv(message) # Send getdata in response.
983        # Store how many times invs have been received for each tx.
984        for i in message.inv:
985            if (i.type == MSG_TX) or (i.type == MSG_WTX):
986                # save txid
987                self.tx_invs_received[i.hash] += 1
988
989    def get_invs(self):
990        with p2p_lock:
991            return list(self.tx_invs_received.keys())
992
993    def wait_for_broadcast(self, txns, *, timeout=60):
994        """Waits for the txns (list of txids) to complete initial broadcast.
995        The mempool should mark unbroadcast=False for these transactions.
996        """
997        # Wait until invs have been received (and getdatas sent) for each txid.
998        self.wait_until(lambda: set(self.tx_invs_received.keys()) == set([int(tx, 16) for tx in txns]), timeout=timeout)
999        # Flush messages and wait for the getdatas to be processed
1000        self.sync_with_ping()