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