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