Skip to content

Commit a049d1b

Browse files
committed
[test] Introduce EncryptedP2PState object in P2PConnection
Instantiate this object when the connection supports v2 P2P transport protocol. - When a P2PConnection is opened, perform initiate_v2_handshake() if the connection is an initiator. application layer messages are only sent after the initial v2 handshake is over (for both initiator and responder).
1 parent b89fa59 commit a049d1b

File tree

2 files changed

+26
-7
lines changed

2 files changed

+26
-7
lines changed

test/functional/test_framework/p2p.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,9 @@
8080
p2p_port,
8181
wait_until_helper_internal,
8282
)
83+
from test_framework.v2_p2p import (
84+
EncryptedP2PState,
85+
)
8386

8487
logger = logging.getLogger("TestFramework.p2p")
8588

@@ -159,11 +162,16 @@ def __init__(self):
159162
# The underlying transport of the connection.
160163
# Should only call methods on this from the NetworkThread, c.f. call_soon_threadsafe
161164
self._transport = None
165+
self.v2_state = None # EncryptedP2PState object needed for v2 p2p connections
162166

163167
@property
164168
def is_connected(self):
165169
return self._transport is not None
166170

171+
@property
172+
def supports_v2_p2p(self):
173+
return self.v2_state is not None
174+
167175
def peer_connect_helper(self, dstaddr, dstport, net, timeout_factor):
168176
assert not self.is_connected
169177
self.timeout_factor = timeout_factor
@@ -174,16 +182,20 @@ def peer_connect_helper(self, dstaddr, dstport, net, timeout_factor):
174182
self.recvbuf = b""
175183
self.magic_bytes = MAGIC_BYTES[net]
176184

177-
def peer_connect(self, dstaddr, dstport, *, net, timeout_factor):
185+
def peer_connect(self, dstaddr, dstport, *, net, timeout_factor, supports_v2_p2p):
178186
self.peer_connect_helper(dstaddr, dstport, net, timeout_factor)
187+
if supports_v2_p2p:
188+
self.v2_state = EncryptedP2PState(initiating=True, net=net)
179189

180190
loop = NetworkThread.network_event_loop
181191
logger.debug('Connecting to Bitcoin Node: %s:%d' % (self.dstaddr, self.dstport))
182192
coroutine = loop.create_connection(lambda: self, host=self.dstaddr, port=self.dstport)
183193
return lambda: loop.call_soon_threadsafe(loop.create_task, coroutine)
184194

185-
def peer_accept_connection(self, connect_id, connect_cb=lambda: None, *, net, timeout_factor):
195+
def peer_accept_connection(self, connect_id, connect_cb=lambda: None, *, net, timeout_factor, supports_v2_p2p):
186196
self.peer_connect_helper('0', 0, net, timeout_factor)
197+
if supports_v2_p2p:
198+
self.v2_state = EncryptedP2PState(initiating=False, net=net)
187199

188200
logger.debug('Listening for Bitcoin Node with id: {}'.format(connect_id))
189201
return lambda: NetworkThread.listen(self, connect_cb, idx=connect_id)
@@ -199,7 +211,13 @@ def connection_made(self, transport):
199211
assert not self._transport
200212
logger.debug("Connected & Listening: %s:%d" % (self.dstaddr, self.dstport))
201213
self._transport = transport
202-
if self.on_connection_send_msg:
214+
# in an inbound connection to the TestNode with P2PConnection as the initiator, [TestNode <---- P2PConnection]
215+
# send the initial handshake immediately
216+
if self.supports_v2_p2p and self.v2_state.initiating and not self.v2_state.tried_v2_handshake:
217+
send_handshake_bytes = self.v2_state.initiate_v2_handshake()
218+
self.send_raw_message(send_handshake_bytes)
219+
# if v2 connection, send `on_connection_send_msg` after initial v2 handshake.
220+
if self.on_connection_send_msg and not self.supports_v2_p2p:
203221
self.send_message(self.on_connection_send_msg)
204222
self.on_connection_send_msg = None # Never used again
205223
self.on_open()

test/functional/test_framework/test_node.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -642,7 +642,7 @@ def assert_start_raises_init_error(self, extra_args=None, expected_msg=None, mat
642642
assert_msg += "with expected error " + expected_msg
643643
self._raise_assertion_error(assert_msg)
644644

645-
def add_p2p_connection(self, p2p_conn, *, wait_for_verack=True, send_version=True, **kwargs):
645+
def add_p2p_connection(self, p2p_conn, *, wait_for_verack=True, send_version=True, supports_v2_p2p=False, **kwargs):
646646
"""Add an inbound p2p connection to the node.
647647
648648
This method adds the p2p connection to the self.p2ps list and also
@@ -653,7 +653,8 @@ def add_p2p_connection(self, p2p_conn, *, wait_for_verack=True, send_version=Tru
653653
kwargs['dstaddr'] = '127.0.0.1'
654654

655655
p2p_conn.p2p_connected_to_node = True
656-
p2p_conn.peer_connect(**kwargs, send_version=send_version, net=self.chain, timeout_factor=self.timeout_factor)()
656+
p2p_conn.peer_connect(**kwargs, send_version=send_version, net=self.chain, timeout_factor=self.timeout_factor, supports_v2_p2p=supports_v2_p2p)()
657+
657658
self.p2ps.append(p2p_conn)
658659
p2p_conn.wait_until(lambda: p2p_conn.is_connected, check_connected=False)
659660
if send_version:
@@ -684,7 +685,7 @@ def add_p2p_connection(self, p2p_conn, *, wait_for_verack=True, send_version=Tru
684685

685686
return p2p_conn
686687

687-
def add_outbound_p2p_connection(self, p2p_conn, *, wait_for_verack=True, p2p_idx, connection_type="outbound-full-relay", **kwargs):
688+
def add_outbound_p2p_connection(self, p2p_conn, *, wait_for_verack=True, p2p_idx, connection_type="outbound-full-relay", supports_v2_p2p=False, **kwargs):
688689
"""Add an outbound p2p connection from node. Must be an
689690
"outbound-full-relay", "block-relay-only", "addr-fetch" or "feeler" connection.
690691
@@ -701,7 +702,7 @@ def addconnection_callback(address, port):
701702
self.addconnection('%s:%d' % (address, port), connection_type)
702703

703704
p2p_conn.p2p_connected_to_node = False
704-
p2p_conn.peer_accept_connection(connect_cb=addconnection_callback, connect_id=p2p_idx + 1, net=self.chain, timeout_factor=self.timeout_factor, **kwargs)()
705+
p2p_conn.peer_accept_connection(connect_cb=addconnection_callback, connect_id=p2p_idx + 1, net=self.chain, timeout_factor=self.timeout_factor, supports_v2_p2p=supports_v2_p2p, **kwargs)()
705706

706707
if connection_type == "feeler":
707708
# feeler connections are closed as soon as the node receives a `version` message

0 commit comments

Comments
 (0)