From a049d1bd08c8cdb3b693520f24f8a82572dcaab1 Mon Sep 17 00:00:00 2001 From: stratospher <44024636+stratospher@users.noreply.github.com> Date: Sat, 5 Feb 2022 22:11:02 +0530 Subject: [PATCH] [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). --- test/functional/test_framework/p2p.py | 24 ++++++++++++++++++--- test/functional/test_framework/test_node.py | 9 ++++---- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/test/functional/test_framework/p2p.py b/test/functional/test_framework/p2p.py index eef55b62e45..004a9edc1c6 100755 --- a/test/functional/test_framework/p2p.py +++ b/test/functional/test_framework/p2p.py @@ -80,6 +80,9 @@ from test_framework.util import ( p2p_port, wait_until_helper_internal, ) +from test_framework.v2_p2p import ( + EncryptedP2PState, +) logger = logging.getLogger("TestFramework.p2p") @@ -159,11 +162,16 @@ class P2PConnection(asyncio.Protocol): # The underlying transport of the connection. # Should only call methods on this from the NetworkThread, c.f. call_soon_threadsafe self._transport = None + self.v2_state = None # EncryptedP2PState object needed for v2 p2p connections @property def is_connected(self): return self._transport is not None + @property + def supports_v2_p2p(self): + return self.v2_state is not None + def peer_connect_helper(self, dstaddr, dstport, net, timeout_factor): assert not self.is_connected self.timeout_factor = timeout_factor @@ -174,16 +182,20 @@ class P2PConnection(asyncio.Protocol): self.recvbuf = b"" self.magic_bytes = MAGIC_BYTES[net] - def peer_connect(self, dstaddr, dstport, *, net, timeout_factor): + def peer_connect(self, dstaddr, dstport, *, net, timeout_factor, supports_v2_p2p): self.peer_connect_helper(dstaddr, dstport, net, timeout_factor) + if supports_v2_p2p: + self.v2_state = EncryptedP2PState(initiating=True, net=net) loop = NetworkThread.network_event_loop logger.debug('Connecting to Bitcoin Node: %s:%d' % (self.dstaddr, self.dstport)) coroutine = loop.create_connection(lambda: self, host=self.dstaddr, port=self.dstport) return lambda: loop.call_soon_threadsafe(loop.create_task, coroutine) - def peer_accept_connection(self, connect_id, connect_cb=lambda: None, *, net, timeout_factor): + def peer_accept_connection(self, connect_id, connect_cb=lambda: None, *, net, timeout_factor, supports_v2_p2p): self.peer_connect_helper('0', 0, net, timeout_factor) + if supports_v2_p2p: + self.v2_state = EncryptedP2PState(initiating=False, net=net) logger.debug('Listening for Bitcoin Node with id: {}'.format(connect_id)) return lambda: NetworkThread.listen(self, connect_cb, idx=connect_id) @@ -199,7 +211,13 @@ class P2PConnection(asyncio.Protocol): assert not self._transport logger.debug("Connected & Listening: %s:%d" % (self.dstaddr, self.dstport)) self._transport = transport - if self.on_connection_send_msg: + # in an inbound connection to the TestNode with P2PConnection as the initiator, [TestNode <---- P2PConnection] + # send the initial handshake immediately + if self.supports_v2_p2p and self.v2_state.initiating and not self.v2_state.tried_v2_handshake: + send_handshake_bytes = self.v2_state.initiate_v2_handshake() + self.send_raw_message(send_handshake_bytes) + # if v2 connection, send `on_connection_send_msg` after initial v2 handshake. + if self.on_connection_send_msg and not self.supports_v2_p2p: self.send_message(self.on_connection_send_msg) self.on_connection_send_msg = None # Never used again self.on_open() diff --git a/test/functional/test_framework/test_node.py b/test/functional/test_framework/test_node.py index 850aa20db2b..444976e54f7 100755 --- a/test/functional/test_framework/test_node.py +++ b/test/functional/test_framework/test_node.py @@ -642,7 +642,7 @@ class TestNode(): assert_msg += "with expected error " + expected_msg self._raise_assertion_error(assert_msg) - def add_p2p_connection(self, p2p_conn, *, wait_for_verack=True, send_version=True, **kwargs): + def add_p2p_connection(self, p2p_conn, *, wait_for_verack=True, send_version=True, supports_v2_p2p=False, **kwargs): """Add an inbound p2p connection to the node. This method adds the p2p connection to the self.p2ps list and also @@ -653,7 +653,8 @@ class TestNode(): kwargs['dstaddr'] = '127.0.0.1' p2p_conn.p2p_connected_to_node = True - p2p_conn.peer_connect(**kwargs, send_version=send_version, net=self.chain, timeout_factor=self.timeout_factor)() + p2p_conn.peer_connect(**kwargs, send_version=send_version, net=self.chain, timeout_factor=self.timeout_factor, supports_v2_p2p=supports_v2_p2p)() + self.p2ps.append(p2p_conn) p2p_conn.wait_until(lambda: p2p_conn.is_connected, check_connected=False) if send_version: @@ -684,7 +685,7 @@ class TestNode(): return p2p_conn - def add_outbound_p2p_connection(self, p2p_conn, *, wait_for_verack=True, p2p_idx, connection_type="outbound-full-relay", **kwargs): + def add_outbound_p2p_connection(self, p2p_conn, *, wait_for_verack=True, p2p_idx, connection_type="outbound-full-relay", supports_v2_p2p=False, **kwargs): """Add an outbound p2p connection from node. Must be an "outbound-full-relay", "block-relay-only", "addr-fetch" or "feeler" connection. @@ -701,7 +702,7 @@ class TestNode(): self.addconnection('%s:%d' % (address, port), connection_type) p2p_conn.p2p_connected_to_node = False - p2p_conn.peer_accept_connection(connect_cb=addconnection_callback, connect_id=p2p_idx + 1, net=self.chain, timeout_factor=self.timeout_factor, **kwargs)() + 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)() if connection_type == "feeler": # feeler connections are closed as soon as the node receives a `version` message