p2p.py raw

   1  #!/usr/bin/env python3
   2  # Copyright (c) 2010 ArtForz -- public domain half-a-node
   3  # Copyright (c) 2012 Jeff Garzik
   4  # Copyright (c) 2010-present The Bitcoin Core developers
   5  # Distributed under the MIT software license, see the accompanying
   6  # file COPYING or http://www.opensource.org/licenses/mit-license.php.
   7  """Test objects for interacting with a bitcoind node over the p2p protocol.
   8  
   9  The P2PInterface objects interact with the bitcoind nodes under test using the
  10  node's p2p interface. They can be used to send messages to the node, and
  11  callbacks can be registered that execute when messages are received from the
  12  node. Messages are sent to/received from the node on an asyncio event loop.
  13  State held inside the objects must be guarded by the p2p_lock to avoid data
  14  races between the main testing thread and the event loop.
  15  
  16  P2PConnection: A low-level connection object to a node's P2P interface
  17  P2PInterface: A high-level interface object for communicating to a node over P2P
  18  P2PDataStore: A p2p interface class that keeps a store of transactions and blocks
  19                and can respond correctly to getdata and getheaders messages
  20  P2PTxInvStore: A p2p interface class that inherits from P2PDataStore, and keeps
  21                a count of how many times each txid has been announced."""
  22  
  23  import asyncio
  24  from collections import defaultdict
  25  import ipaddress
  26  from io import BytesIO
  27  import logging
  28  import platform
  29  import socket
  30  import struct
  31  import sys
  32  import threading
  33  
  34  from test_framework.messages import (
  35      CBlockHeader,
  36      MAX_HEADERS_RESULTS,
  37      msg_addr,
  38      msg_addrv2,
  39      msg_block,
  40      MSG_BLOCK,
  41      msg_blocktxn,
  42      msg_cfcheckpt,
  43      msg_cfheaders,
  44      msg_cfilter,
  45      msg_cmpctblock,
  46      msg_feature,
  47      msg_feefilter,
  48      msg_filteradd,
  49      msg_filterclear,
  50      msg_filterload,
  51      msg_getaddr,
  52      msg_getblocks,
  53      msg_getblocktxn,
  54      msg_getcfcheckpt,
  55      msg_getcfheaders,
  56      msg_getcfilters,
  57      msg_getdata,
  58      msg_getheaders,
  59      msg_headers,
  60      msg_inv,
  61      msg_mempool,
  62      msg_merkleblock,
  63      msg_notfound,
  64      msg_ping,
  65      msg_pong,
  66      msg_sendaddrv2,
  67      msg_sendcmpct,
  68      msg_sendheaders,
  69      msg_sendtxrcncl,
  70      msg_tx,
  71      MSG_TX,
  72      MSG_TYPE_MASK,
  73      msg_verack,
  74      msg_version,
  75      MSG_WTX,
  76      msg_wtxidrelay,
  77      NODE_NETWORK,
  78      NODE_WITNESS,
  79      MAGIC_BYTES,
  80      sha256,
  81  )
  82  from test_framework.netutil import (
  83      set_ephemeral_port_range,
  84  )
  85  from test_framework.util import (
  86      assert_not_equal,
  87      MAX_NODES,
  88      p2p_port,
  89      wait_until_helper_internal,
  90  )
  91  from test_framework.v2_p2p import (
  92      EncryptedP2PState,
  93      MSGTYPE_TO_SHORTID,
  94      SHORTID,
  95  )
  96  
  97  logger = logging.getLogger("TestFramework.p2p")
  98  
  99  # The minimum P2P version that this test framework supports
 100  MIN_P2P_VERSION_SUPPORTED = 60001
 101  # The P2P version that this test framework implements and sends in its `version` message
 102  # Version 70016 supports wtxid relay
 103  # Version 70017 supports feature
 104  P2P_VERSION = 70017
 105  # The services that this test framework offers in its `version` message
 106  P2P_SERVICES = NODE_NETWORK | NODE_WITNESS
 107  # The P2P user agent string that this test framework sends in its `version` message
 108  P2P_SUBVERSION = "/python-p2p-tester:0.0.3/"
 109  # Value for relay that this test framework sends in its `version` message
 110  P2P_VERSION_RELAY = 1
 111  # Delay after receiving a tx inv before requesting transactions from non-preferred peers, in seconds
 112  NONPREF_PEER_TX_DELAY = 2
 113  # Delay for requesting transactions via txids if we have wtxid-relaying peers, in seconds
 114  TXID_RELAY_DELAY = 2
 115  # Delay for requesting transactions if the peer has MAX_PEER_TX_REQUEST_IN_FLIGHT or more requests
 116  OVERLOADED_PEER_TX_DELAY = 2
 117  # How long to wait before downloading a transaction from an additional peer
 118  GETDATA_TX_INTERVAL = 60
 119  
 120  MESSAGEMAP = {
 121      b"addr": msg_addr,
 122      b"addrv2": msg_addrv2,
 123      b"block": msg_block,
 124      b"blocktxn": msg_blocktxn,
 125      b"cfcheckpt": msg_cfcheckpt,
 126      b"cfheaders": msg_cfheaders,
 127      b"cfilter": msg_cfilter,
 128      b"cmpctblock": msg_cmpctblock,
 129      b"feature": msg_feature,
 130      b"feefilter": msg_feefilter,
 131      b"filteradd": msg_filteradd,
 132      b"filterclear": msg_filterclear,
 133      b"filterload": msg_filterload,
 134      b"getaddr": msg_getaddr,
 135      b"getblocks": msg_getblocks,
 136      b"getblocktxn": msg_getblocktxn,
 137      b"getcfcheckpt": msg_getcfcheckpt,
 138      b"getcfheaders": msg_getcfheaders,
 139      b"getcfilters": msg_getcfilters,
 140      b"getdata": msg_getdata,
 141      b"getheaders": msg_getheaders,
 142      b"headers": msg_headers,
 143      b"inv": msg_inv,
 144      b"mempool": msg_mempool,
 145      b"merkleblock": msg_merkleblock,
 146      b"notfound": msg_notfound,
 147      b"ping": msg_ping,
 148      b"pong": msg_pong,
 149      b"sendaddrv2": msg_sendaddrv2,
 150      b"sendcmpct": msg_sendcmpct,
 151      b"sendheaders": msg_sendheaders,
 152      b"sendtxrcncl": msg_sendtxrcncl,
 153      b"tx": msg_tx,
 154      b"verack": msg_verack,
 155      b"version": msg_version,
 156      b"wtxidrelay": msg_wtxidrelay,
 157  }
 158  
 159  
 160  class P2PConnection(asyncio.Protocol):
 161      """A low-level connection object to a node's P2P interface.
 162  
 163      This class is responsible for:
 164  
 165      - opening and closing the TCP connection to the node
 166      - reading bytes from and writing bytes to the socket
 167      - deserializing and serializing the P2P message header
 168      - logging messages as they are sent and received
 169  
 170      This class contains no logic for handing the P2P message payloads. It must be
 171      sub-classed and the on_message() callback overridden."""
 172  
 173      def __init__(self):
 174          # The underlying transport of the connection.
 175          # Should only call methods on this from the NetworkThread, c.f. call_soon_threadsafe
 176          self._transport = None
 177          # This lock is acquired before sending messages over the socket. There's an implied lock order and
 178          # p2p_lock must not be acquired after _send_lock as it could result in deadlocks.
 179          self._send_lock = threading.Lock()
 180          self.v2_state = None  # EncryptedP2PState object needed for v2 p2p connections
 181          self.reconnect = False  # set if reconnection needs to happen
 182  
 183      @property
 184      def is_connected(self):
 185          return self._transport is not None
 186  
 187      @property
 188      def supports_v2_p2p(self):
 189          return self.v2_state is not None
 190  
 191      def peer_connect_helper(self, dstaddr, dstport, net, timeout_factor):
 192          assert not self.is_connected
 193          self.timeout_factor = timeout_factor
 194          self.dstaddr = dstaddr
 195          self.dstport = dstport
 196          # The initial message to send after the connection was made:
 197          self.on_connection_send_msg = None
 198          self.recvbuf = b""
 199          self.magic_bytes = MAGIC_BYTES[net]
 200          self.p2p_connected_to_node = dstport != 0
 201  
 202      def peer_connect(self, dstaddr, dstport, *, net, timeout_factor, supports_v2_p2p):
 203          self.peer_connect_helper(dstaddr, dstport, net, timeout_factor)
 204          if supports_v2_p2p:
 205              self.v2_state = EncryptedP2PState(initiating=True, net=net)
 206  
 207          loop = NetworkThread.network_event_loop
 208          logger.debug('Connecting to Bitcoin Node: %s:%d' % (self.dstaddr, self.dstport))
 209          coroutine = loop.create_connection(lambda: self, host=self.dstaddr, port=self.dstport)
 210          return lambda: loop.call_soon_threadsafe(loop.create_task, coroutine)
 211  
 212      def peer_accept_connection(self, connect_id, connect_cb=lambda: None, *, net, timeout_factor, supports_v2_p2p, reconnect):
 213          self.peer_connect_helper('0', 0, net, timeout_factor)
 214          self.reconnect = reconnect
 215          if supports_v2_p2p:
 216              self.v2_state = EncryptedP2PState(initiating=False, net=net)
 217  
 218          logger.debug('Listening for Bitcoin Node with id: {}'.format(connect_id))
 219          return lambda: NetworkThread.listen(self, connect_cb, idx=connect_id)
 220  
 221      def peer_disconnect(self):
 222          # Connection could have already been closed by other end.
 223          NetworkThread.network_event_loop.call_soon_threadsafe(lambda: self._transport and self._transport.abort())
 224  
 225      # Connection and disconnection methods
 226  
 227      def connection_made(self, transport):
 228          """asyncio callback when a connection is opened."""
 229          assert not self._transport
 230          info = transport.get_extra_info("socket")
 231          us = info.getsockname()
 232          them = info.getpeername()
 233          logger.debug(f"Connected: us={us[0]}:{us[1]}, them={them[0]}:{them[1]}")
 234          self.dstaddr = them[0]
 235          self.dstport = them[1]
 236          self._transport = transport
 237          # in an inbound connection to the TestNode with P2PConnection as the initiator, [TestNode <---- P2PConnection]
 238          # send the initial handshake immediately
 239          if self.supports_v2_p2p and self.v2_state.initiating and not self.v2_state.tried_v2_handshake:
 240              send_handshake_bytes = self.v2_state.initiate_v2_handshake()
 241              logger.debug(f"sending {len(self.v2_state.sent_garbage)} bytes of garbage data")
 242              self.send_raw_message(send_handshake_bytes)
 243          # for v1 outbound connections, send version message immediately after opening
 244          # (for v2 outbound connections, send it after the initial v2 handshake)
 245          if self.p2p_connected_to_node and not self.supports_v2_p2p:
 246              self.send_version()
 247          self.on_open()
 248  
 249      def connection_lost(self, exc):
 250          """asyncio callback when a connection is closed."""
 251          # don't display warning if reconnection needs to be attempted using v1 P2P
 252          if exc and not self.reconnect:
 253              logger.warning("Connection lost to {}:{} due to {}".format(self.dstaddr, self.dstport, exc))
 254          else:
 255              logger.debug("Closed connection to: %s:%d" % (self.dstaddr, self.dstport))
 256          self._transport = None
 257          self.recvbuf = b""
 258          self.on_close()
 259  
 260      # v2 handshake method
 261      def _on_data_v2_handshake(self):
 262          """v2 handshake performed before P2P messages are exchanged (see BIP324). P2PConnection is the initiator
 263          (in inbound connections to TestNode) and the responder (in outbound connections from TestNode).
 264          Performed by:
 265              * initiator using `initiate_v2_handshake()`, `complete_handshake()` and `authenticate_handshake()`
 266              * responder using `respond_v2_handshake()`, `complete_handshake()` and `authenticate_handshake()`
 267  
 268          `initiate_v2_handshake()` is immediately done by the initiator when the connection is established in
 269          `connection_made()`. The rest of the initial v2 handshake functions are handled here.
 270          """
 271          if not self.v2_state.peer:
 272              if not self.v2_state.initiating and not self.v2_state.sent_garbage:
 273                  # if the responder hasn't sent garbage yet, the responder is still reading ellswift bytes
 274                  # reads ellswift bytes till the first mismatch from 12 bytes V1_PREFIX
 275                  length, send_handshake_bytes = self.v2_state.respond_v2_handshake(BytesIO(self.recvbuf))
 276                  self.recvbuf = self.recvbuf[length:]
 277                  if send_handshake_bytes == -1:
 278                      self.v2_state = None
 279                      return
 280                  elif send_handshake_bytes:
 281                      logger.debug(f"sending {len(self.v2_state.sent_garbage)} bytes of garbage data")
 282                      self.send_raw_message(send_handshake_bytes)
 283                  elif send_handshake_bytes == b"":
 284                      return  # only after send_handshake_bytes are sent can `complete_handshake()` be done
 285  
 286              # `complete_handshake()` reads the remaining ellswift bytes from recvbuf
 287              # and sends response after deriving shared ECDH secret using received ellswift bytes
 288              length, response = self.v2_state.complete_handshake(BytesIO(self.recvbuf))
 289              self.recvbuf = self.recvbuf[length:]
 290              if response:
 291                  self.send_raw_message(response)
 292              else:
 293                  return  # only after response is sent can `authenticate_handshake()` be done
 294  
 295          # `self.v2_state.peer` is instantiated only after shared ECDH secret/BIP324 derived keys and ciphers
 296          # is derived in `complete_handshake()`.
 297          # so `authenticate_handshake()` which uses the BIP324 derived ciphers gets called after `complete_handshake()`.
 298          assert self.v2_state.peer
 299          length, is_mac_auth = self.v2_state.authenticate_handshake(self.recvbuf)
 300          if not is_mac_auth:
 301              raise ValueError("invalid v2 mac tag in handshake authentication")
 302          self.recvbuf = self.recvbuf[length:]
 303          if self.v2_state.tried_v2_handshake:
 304              # for v2 outbound connections, send version message immediately after v2 handshake
 305              if self.p2p_connected_to_node:
 306                  self.send_version()
 307              # process post-v2-handshake data immediately, if available
 308              if len(self.recvbuf) > 0:
 309                  self._on_data()
 310  
 311      # Socket read methods
 312  
 313      def data_received(self, t):
 314          """asyncio callback when data is read from the socket."""
 315          if len(t) > 0:
 316              self.recvbuf += t
 317              if self.supports_v2_p2p and not self.v2_state.tried_v2_handshake:
 318                  self._on_data_v2_handshake()
 319              else:
 320                  self._on_data()
 321  
 322      def _on_data(self):
 323          """Try to read P2P messages from the recv buffer.
 324  
 325          This method reads data from the buffer in a loop. It deserializes,
 326          parses and verifies the P2P header, then passes the P2P payload to
 327          the on_message callback for processing."""
 328          try:
 329              while True:
 330                  if self.supports_v2_p2p:
 331                      # v2 P2P messages are read
 332                      msglen, msg = self.v2_state.v2_receive_packet(self.recvbuf)
 333                      if msglen == -1:
 334                          raise ValueError("invalid v2 mac tag " + repr(self.recvbuf))
 335                      elif msglen == 0:  # need to receive more bytes in recvbuf
 336                          return
 337                      self.recvbuf = self.recvbuf[msglen:]
 338  
 339                      if msg is None:  # ignore decoy messages
 340                          return
 341                      assert msg  # application layer messages (which aren't decoy messages) are non-empty
 342                      shortid = msg[0]  # 1-byte short message type ID
 343                      if shortid == 0:
 344                          # next 12 bytes are interpreted as ASCII message type if shortid is b'\x00'
 345                          if len(msg) < 13:
 346                              raise IndexError("msg needs minimum required length of 13 bytes")
 347                          msgtype = msg[1:13].rstrip(b'\x00')
 348                          msg = msg[13:]  # msg is set to be payload
 349                      else:
 350                          # a 1-byte short message type ID
 351                          msgtype = SHORTID.get(shortid, f"unknown-{shortid}")
 352                          msg = msg[1:]
 353                  else:
 354                      # v1 P2P messages are read
 355                      if len(self.recvbuf) < 4:
 356                          return
 357                      if self.recvbuf[:4] != self.magic_bytes:
 358                          raise ValueError("magic bytes mismatch: {} != {}".format(repr(self.magic_bytes), repr(self.recvbuf)))
 359                      if len(self.recvbuf) < 4 + 12 + 4 + 4:
 360                          return
 361                      msgtype = self.recvbuf[4:4+12].split(b"\x00", 1)[0]
 362                      msglen = struct.unpack("<i", self.recvbuf[4+12:4+12+4])[0]
 363                      checksum = self.recvbuf[4+12+4:4+12+4+4]
 364                      if len(self.recvbuf) < 4 + 12 + 4 + 4 + msglen:
 365                          return
 366                      msg = self.recvbuf[4+12+4+4:4+12+4+4+msglen]
 367                      th = sha256(msg)
 368                      h = sha256(th)
 369                      if checksum != h[:4]:
 370                          raise ValueError("got bad checksum " + repr(self.recvbuf))
 371                      self.recvbuf = self.recvbuf[4+12+4+4+msglen:]
 372                  if msgtype not in MESSAGEMAP:
 373                      raise ValueError("Received unknown msgtype from %s:%d: '%s' %s" % (self.dstaddr, self.dstport, msgtype, repr(msg)))
 374                  f = BytesIO(msg)
 375                  t = MESSAGEMAP[msgtype]()
 376                  t.deserialize(f)
 377                  self._log_message("receive", t)
 378                  self.on_message(t)
 379          except Exception as e:
 380              if not self.reconnect:
 381                  logger.exception(f"Error reading message: {repr(e)}")
 382              raise
 383  
 384      def on_message(self, message):
 385          """Callback for processing a P2P payload. Must be overridden by derived class."""
 386          raise NotImplementedError
 387  
 388      # Socket write methods
 389  
 390      def send_without_ping(self, message, is_decoy=False):
 391          """Send a P2P message over the socket.
 392  
 393          This method takes a P2P payload, builds the P2P header and adds
 394          the message to the send buffer to be sent over the socket.
 395  
 396          When a message does not lead to a disconnect, send_and_ping is usually
 397          preferred to send a message. This can help to reduce intermittent test
 398          failures due to a missing sync. Also, it includes a call to
 399          sync_with_ping, allowing for concise test code.
 400          """
 401          with self._send_lock:
 402              tmsg = self.build_message(message, is_decoy)
 403              self._log_message("send", message)
 404              return self.send_raw_message(tmsg)
 405  
 406      def send_raw_message(self, raw_message_bytes):
 407          if not self.is_connected:
 408              raise IOError('Not connected')
 409  
 410          def maybe_write():
 411              if not self._transport:
 412                  return
 413              if self._transport.is_closing():
 414                  return
 415              self._transport.write(raw_message_bytes)
 416          NetworkThread.network_event_loop.call_soon_threadsafe(maybe_write)
 417  
 418      # Class utility methods
 419  
 420      def build_message(self, message, is_decoy=False):
 421          """Build a serialized P2P message"""
 422          msgtype = message.msgtype
 423          data = message.serialize()
 424          if self.supports_v2_p2p:
 425              if msgtype in SHORTID.values():
 426                  tmsg = MSGTYPE_TO_SHORTID.get(msgtype).to_bytes(1, 'big')
 427              else:
 428                  tmsg = b"\x00"
 429                  tmsg += msgtype
 430                  tmsg += b"\x00" * (12 - len(msgtype))
 431              tmsg += data
 432              return self.v2_state.v2_enc_packet(tmsg, ignore=is_decoy)
 433          else:
 434              tmsg = self.magic_bytes
 435              tmsg += msgtype
 436              tmsg += b"\x00" * (12 - len(msgtype))
 437              tmsg += len(data).to_bytes(4, "little")
 438              th = sha256(data)
 439              h = sha256(th)
 440              tmsg += h[:4]
 441              tmsg += data
 442              return tmsg
 443  
 444      def _log_message(self, direction, msg):
 445          """Logs a message being sent or received over the connection."""
 446          if direction == "send":
 447              log_message = "Send message to "
 448          elif direction == "receive":
 449              log_message = "Received message from "
 450          log_message += "%s:%d: %s" % (self.dstaddr, self.dstport, repr(msg)[:500])
 451          if len(log_message) > 500:
 452              log_message += "... (msg truncated)"
 453          logger.debug(log_message)
 454  
 455  
 456  class P2PInterface(P2PConnection):
 457      """A high-level P2P interface class for communicating with a Bitcoin node.
 458  
 459      This class provides high-level callbacks for processing P2P message
 460      payloads, as well as convenience methods for interacting with the
 461      node over P2P.
 462  
 463      Individual testcases should subclass this and override the on_* methods
 464      if they want to alter message handling behaviour."""
 465      def __init__(self, support_addrv2=False, wtxidrelay=True):
 466          super().__init__()
 467  
 468          # Track number of messages of each type received.
 469          # Should be read-only in a test.
 470          self.message_count = defaultdict(int)
 471  
 472          # Track the most recent message of each type.
 473          # To wait for a message to be received, pop that message from
 474          # this and use self.wait_until.
 475          self.last_message = {}
 476  
 477          # A count of the number of ping messages we've sent to the node
 478          self.ping_counter = 1
 479  
 480          # The network services received from the peer
 481          self.nServices = 0
 482  
 483          self.support_addrv2 = support_addrv2
 484  
 485          # If the peer supports wtxid-relay
 486          self.wtxidrelay = wtxidrelay
 487  
 488      def peer_connect_send_version(self, services):
 489          # Send a version msg
 490          vt = msg_version()
 491          vt.nVersion = P2P_VERSION
 492          vt.strSubVer = P2P_SUBVERSION
 493          vt.relay = P2P_VERSION_RELAY
 494          vt.nServices = services
 495          vt.addrTo.ip = self.dstaddr
 496          vt.addrTo.port = self.dstport
 497          vt.addrFrom.ip = "0.0.0.0"
 498          vt.addrFrom.port = 0
 499          self.on_connection_send_msg = vt  # Will be sent in connection_made callback
 500  
 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          test_function = lambda: self.is_connected
 614          self.wait_until(test_function, timeout=timeout, check_connected=False)
 615  
 616      def wait_for_disconnect(self, *, timeout=60):
 617          test_function = lambda: not self.is_connected
 618          self.wait_until(test_function, timeout=timeout, check_connected=False)
 619  
 620      def wait_for_reconnect(self, *, timeout=60):
 621          def test_function():
 622              return self.is_connected and self.last_message.get('version') and not self.supports_v2_p2p
 623          self.wait_until(test_function, timeout=timeout, check_connected=False)
 624  
 625      # Message receiving helper methods
 626  
 627      def wait_for_tx(self, txid, *, timeout=60):
 628          def test_function():
 629              if not self.last_message.get('tx'):
 630                  return False
 631              return self.last_message['tx'].tx.txid_hex == txid
 632  
 633          self.wait_until(test_function, timeout=timeout)
 634  
 635      def wait_for_block(self, blockhash, *, timeout=60):
 636          def test_function():
 637              return self.last_message.get("block") and self.last_message["block"].block.hash_int == blockhash
 638  
 639          self.wait_until(test_function, timeout=timeout)
 640  
 641      def wait_for_header(self, blockhash, *, timeout=60):
 642          def test_function():
 643              last_headers = self.last_message.get('headers')
 644              if not last_headers:
 645                  return False
 646              return last_headers.headers[0].hash_int == int(blockhash, 16)
 647  
 648          self.wait_until(test_function, timeout=timeout)
 649  
 650      def wait_for_merkleblock(self, blockhash, *, timeout=60):
 651          def test_function():
 652              last_filtered_block = self.last_message.get('merkleblock')
 653              if not last_filtered_block:
 654                  return False
 655              return last_filtered_block.merkleblock.header.hash_int == int(blockhash, 16)
 656  
 657          self.wait_until(test_function, timeout=timeout)
 658  
 659      def wait_for_getdata(self, hash_list, *, timeout=60):
 660          """Waits for a getdata message.
 661  
 662          The object hashes in the inventory vector must match the provided hash_list."""
 663          def test_function():
 664              last_data = self.last_message.get("getdata")
 665              if not last_data:
 666                  return False
 667              return [x.hash for x in last_data.inv] == hash_list
 668  
 669          self.wait_until(test_function, timeout=timeout)
 670  
 671      def wait_for_getheaders(self, block_hash=None, *, timeout=60):
 672          """Waits for a getheaders message containing a specific block hash.
 673  
 674          If no block hash is provided, checks whether any getheaders message has been received by the node."""
 675          def test_function():
 676              last_getheaders = self.last_message.pop("getheaders", None)
 677              if block_hash is None:
 678                  return last_getheaders
 679              if last_getheaders is None:
 680                  return False
 681              return block_hash == last_getheaders.locator.vHave[0]
 682  
 683          self.wait_until(test_function, timeout=timeout)
 684  
 685      def wait_for_inv(self, expected_inv, *, timeout=60):
 686          """Waits for an INV message and checks that the first inv object in the message was as expected."""
 687          if len(expected_inv) > 1:
 688              raise NotImplementedError("wait_for_inv() will only verify the first inv object")
 689  
 690          def test_function():
 691              return self.last_message.get("inv") and \
 692                                  self.last_message["inv"].inv[0].type == expected_inv[0].type and \
 693                                  self.last_message["inv"].inv[0].hash == expected_inv[0].hash
 694  
 695          self.wait_until(test_function, timeout=timeout)
 696  
 697      def wait_for_verack(self, *, timeout=60):
 698          def test_function():
 699              return "verack" in self.last_message
 700  
 701          self.wait_until(test_function, timeout=timeout)
 702  
 703      # Message sending helper functions
 704  
 705      def send_version(self):
 706          if self.on_connection_send_msg:
 707              self.send_without_ping(self.on_connection_send_msg)
 708              self.on_connection_send_msg = None  # Never used again
 709  
 710      def send_and_ping(self, message, *, timeout=60):
 711          self.send_without_ping(message)
 712          self.sync_with_ping(timeout=timeout)
 713  
 714      def sync_with_ping(self, *, timeout=60):
 715          """Ensure ProcessMessages and SendMessages is called on this connection"""
 716          # Sending two pings back-to-back, requires that the node calls
 717          # `ProcessMessage` twice, and thus ensures `SendMessages` must have
 718          # been called at least once
 719          self.send_without_ping(msg_ping(nonce=0))
 720          self.send_without_ping(msg_ping(nonce=self.ping_counter))
 721  
 722          def test_function():
 723              return self.last_message.get("pong") and self.last_message["pong"].nonce == self.ping_counter
 724  
 725          self.wait_until(test_function, timeout=timeout)
 726          self.ping_counter += 1
 727  
 728  
 729  # One lock for synchronizing all data access between the network event loop (see
 730  # NetworkThread below) and the thread running the test logic.  For simplicity,
 731  # P2PConnection acquires this lock whenever delivering a message to a P2PInterface.
 732  # This lock should be acquired in the thread running the test logic to synchronize
 733  # access to any data shared with the P2PInterface or P2PConnection.
 734  p2p_lock = threading.Lock()
 735  
 736  
 737  class NetworkThread(threading.Thread):
 738      network_event_loop = None
 739  
 740      def __init__(self):
 741          super().__init__(name="NetworkThread")
 742          # There is only one event loop and no more than one thread must be created
 743          assert not self.network_event_loop
 744  
 745          NetworkThread.listeners = {}
 746          NetworkThread.protos = {}
 747  
 748      def run(self):
 749          """Start the network thread."""
 750          NetworkThread.network_event_loop = asyncio.SelectorEventLoop() if platform.system() == "Windows" else asyncio.new_event_loop()
 751          self.network_event_loop.run_forever()
 752  
 753      def close(self, *, timeout):
 754          """Close the connections and network event loop."""
 755          self.network_event_loop.call_soon_threadsafe(self.network_event_loop.stop)
 756          wait_until_helper_internal(lambda: not self.network_event_loop.is_running(), timeout=timeout)
 757          self.network_event_loop.close()
 758          self.join(timeout)
 759          # Safe to remove event loop.
 760          NetworkThread.network_event_loop = None
 761  
 762      @classmethod
 763      def listen(cls, p2p, callback, port=None, addr=None, idx=1):
 764          """ Ensure a listening server is running on the given port, and run the
 765          protocol specified by `p2p` on the next connection to it. Once ready
 766          for connections, call `callback`."""
 767  
 768          if port is None:
 769              assert 0 < idx <= MAX_NODES
 770              port = p2p_port(MAX_NODES - idx)
 771          if addr is None:
 772              addr = '127.0.0.1'
 773  
 774          def exception_handler(loop, context):
 775              if not p2p.reconnect:
 776                  loop.default_exception_handler(context)
 777  
 778          cls.network_event_loop.set_exception_handler(exception_handler)
 779          coroutine = cls.create_listen_server(addr, port, callback, p2p)
 780          cls.network_event_loop.call_soon_threadsafe(cls.network_event_loop.create_task, coroutine)
 781  
 782      @classmethod
 783      async def create_listen_server(cls, addr, port, callback, proto):
 784          def peer_protocol():
 785              """Returns a function that does the protocol handling for a new
 786              connection. To allow different connections to have different
 787              behaviors, the protocol function is first put in the cls.protos
 788              dict. When the connection is made, the function removes the
 789              protocol function from that dict, and returns it so the event loop
 790              can start executing it."""
 791              response = cls.protos.get((addr, port))
 792              # remove protocol function from dict only when reconnection doesn't need to happen/already happened
 793              if not proto.reconnect:
 794                  cls.protos[(addr, port)] = None
 795              return response
 796  
 797          if port == 0 or (addr, port) not in cls.listeners:
 798              # When creating a listener on a given (addr, port) we only need to
 799              # do it once. If we want different behaviors for different
 800              # connections, we can accomplish this by providing different
 801              # `proto` functions
 802  
 803              if port == 0:
 804                  # Manually create the socket in order to set the range to be
 805                  # used for the port before the bind() call.
 806                  if ipaddress.ip_address(addr).version == 4:
 807                      address_family = socket.AF_INET
 808                  else:
 809                      address_family = socket.AF_INET6
 810                  s = socket.socket(address_family)
 811                  set_ephemeral_port_range(s)
 812                  s.bind((addr, 0))
 813                  s.listen()
 814                  listener = await cls.network_event_loop.create_server(peer_protocol, sock=s)
 815                  port = listener.sockets[0].getsockname()[1]
 816              else:
 817                  listener = await cls.network_event_loop.create_server(peer_protocol, addr, port)
 818  
 819              logger.debug("Listening server on %s:%d should be started" % (addr, port))
 820              cls.listeners[(addr, port)] = listener
 821  
 822          cls.protos[(addr, port)] = proto
 823          callback(addr, port)
 824  
 825  
 826  class P2PDataStore(P2PInterface):
 827      """A P2P data store class.
 828  
 829      Keeps a block and transaction store and responds correctly to getdata and getheaders requests."""
 830  
 831      def __init__(self):
 832          super().__init__()
 833          # store of blocks. key is block hash, value is a CBlock object
 834          self.block_store = {}
 835          self.last_block_hash = ''
 836          # store of txs. key is txid, value is a CTransaction object
 837          self.tx_store = {}
 838          self.getdata_requests = []
 839  
 840      def on_getdata(self, message):
 841          """Check for the tx/block in our stores and if found, reply with MSG_TX or MSG_BLOCK."""
 842          for inv in message.inv:
 843              self.getdata_requests.append(inv.hash)
 844              invtype = inv.type & MSG_TYPE_MASK
 845              if (invtype == MSG_TX or invtype == MSG_WTX) and inv.hash in self.tx_store.keys():
 846                  self.send_without_ping(msg_tx(self.tx_store[inv.hash]))
 847              elif invtype == MSG_BLOCK and inv.hash in self.block_store.keys():
 848                  self.send_without_ping(msg_block(self.block_store[inv.hash]))
 849              else:
 850                  logger.debug('getdata message type {} received.'.format(hex(inv.type)))
 851  
 852      def on_getheaders(self, message):
 853          """Search back through our block store for the locator, and reply with a headers message if found."""
 854  
 855          locator, hash_stop = message.locator, message.hashstop
 856  
 857          # Assume that the most recent block added is the tip
 858          if not self.block_store:
 859              return
 860  
 861          headers_list = [self.block_store[self.last_block_hash]]
 862          while headers_list[-1].hash_int not in locator.vHave:
 863              # Walk back through the block store, adding headers to headers_list
 864              # as we go.
 865              prev_block_hash = headers_list[-1].hashPrevBlock
 866              if prev_block_hash in self.block_store:
 867                  prev_block_header = CBlockHeader(self.block_store[prev_block_hash])
 868                  headers_list.append(prev_block_header)
 869                  if prev_block_header.hash_int == hash_stop:
 870                      # if this is the hashstop header, stop here
 871                      break
 872              else:
 873                  logger.debug('block hash {} not found in block store'.format(hex(prev_block_hash)))
 874                  break
 875  
 876          # Truncate the list if there are too many headers
 877          headers_list = headers_list[:-MAX_HEADERS_RESULTS - 1:-1]
 878          response = msg_headers(headers_list)
 879  
 880          if response is not None:
 881              self.send_without_ping(response)
 882  
 883      def send_blocks_and_test(self, blocks, node, *, success=True, force_send=False, reject_reason=None, expect_disconnect=False, timeout=60, is_decoy=False):
 884          """Send blocks to test node and test whether the tip advances.
 885  
 886           - add all blocks to our block_store
 887           - send a headers message for the final block
 888           - the on_getheaders handler will ensure that any getheaders are responded to
 889           - if force_send is False: wait for getdata for each of the blocks. The on_getdata handler will
 890             ensure that any getdata messages are responded to. Otherwise send the full block unsolicited.
 891           - if success is True: assert that the node's tip is the last block in blocks at the end of the operation.
 892           - if success is False: assert that the node's tip isn't the last block in blocks at the end of the operation
 893           - if reject_reason is set: assert that the correct reject message is logged"""
 894  
 895          with p2p_lock:
 896              for block in blocks:
 897                  self.block_store[block.hash_int] = block
 898                  self.last_block_hash = block.hash_int
 899  
 900          reject_reason = [reject_reason] if reject_reason else []
 901          with node.assert_debug_log(expected_msgs=reject_reason):
 902              if is_decoy:  # since decoy messages are ignored by the recipient - no need to wait for response
 903                  force_send = True
 904              if force_send:
 905                  for b in blocks:
 906                      self.send_without_ping(msg_block(block=b), is_decoy)
 907              else:
 908                  self.send_without_ping(msg_headers([CBlockHeader(block) for block in blocks]))
 909                  self.wait_until(
 910                      lambda: blocks[-1].hash_int in self.getdata_requests,
 911                      timeout=timeout,
 912                      check_connected=success,
 913                  )
 914  
 915              if expect_disconnect:
 916                  self.wait_for_disconnect(timeout=timeout)
 917              else:
 918                  self.sync_with_ping(timeout=timeout)
 919  
 920              if success:
 921                  self.wait_until(lambda: node.getbestblockhash() == blocks[-1].hash_hex, timeout=timeout)
 922              else:
 923                  assert_not_equal(node.getbestblockhash(), blocks[-1].hash_hex)
 924  
 925      def send_txs_and_test(self, txs, node, *, success=True, reject_reason=None):
 926          """Send txs to test node and test whether they're accepted to the mempool.
 927  
 928           - add all txs to our tx_store
 929           - send tx messages for all txs
 930           - if success is True/False: assert that the txs are/are not accepted to the mempool
 931           - if reject_reason is set: assert that the correct reject message is logged."""
 932  
 933          with p2p_lock:
 934              for tx in txs:
 935                  self.tx_store[tx.txid_int] = tx
 936  
 937          reject_reason = [reject_reason] if reject_reason else []
 938          with node.assert_debug_log(expected_msgs=reject_reason):
 939              for tx in txs:
 940                  self.send_without_ping(msg_tx(tx))
 941  
 942              self.sync_with_ping()
 943  
 944              raw_mempool = node.getrawmempool()
 945              if success:
 946                  # Check that all txs are now in the mempool
 947                  for tx in txs:
 948                      assert tx.txid_hex in raw_mempool, "{} not found in mempool".format(tx.txid_hex)
 949              else:
 950                  # Check that none of the txs are now in the mempool
 951                  for tx in txs:
 952                      assert tx.txid_hex not in raw_mempool, "{} tx found in mempool".format(tx.txid_hex)
 953  
 954  class P2PTxInvStore(P2PInterface):
 955      """A P2PInterface which stores a count of how many times each txid has been announced."""
 956      def __init__(self, **kwargs):
 957          super().__init__(**kwargs)
 958          self.tx_invs_received = defaultdict(int)
 959  
 960      def on_inv(self, message):
 961          super().on_inv(message) # Send getdata in response.
 962          # Store how many times invs have been received for each tx.
 963          for i in message.inv:
 964              if (i.type == MSG_TX) or (i.type == MSG_WTX):
 965                  # save txid
 966                  self.tx_invs_received[i.hash] += 1
 967  
 968      def get_invs(self):
 969          with p2p_lock:
 970              return list(self.tx_invs_received.keys())
 971  
 972      def wait_for_broadcast(self, txns, *, timeout=60):
 973          """Waits for the txns (list of txids) to complete initial broadcast.
 974          The mempool should mark unbroadcast=False for these transactions.
 975          """
 976          # Wait until invs have been received (and getdatas sent) for each txid.
 977          self.wait_until(lambda: set(self.tx_invs_received.keys()) == set([int(tx, 16) for tx in txns]), timeout=timeout)
 978          # Flush messages and wait for the getdatas to be processed
 979          self.sync_with_ping()
 980  
 981  def start_p2p_listener(network_thread, listener):
 982      listen_addr = ""
 983      listen_port = 0
 984  
 985      def on_listen_done(addr, port):
 986          nonlocal listen_addr
 987          nonlocal listen_port
 988          listen_addr = addr
 989          listen_port = port
 990  
 991      # Use port=0 to let the OS assign an available port. This
 992      # avoids "address already in use" errors when tests run
 993      # concurrently or ports are still in TIME_WAIT state.
 994      network_thread.listen(
 995          addr="127.0.0.1",
 996          port=0,
 997          p2p=listener,
 998          callback=on_listen_done)
 999  
1000      # Wait until the callback has been called.
1001      wait_until_helper_internal(lambda: listen_port != 0)
1002  
1003      return listen_addr, listen_port
1004