0
0
Fork 0
mirror of https://github.com/bitcoin/bitcoin.git synced 2025-02-10 10:52:31 -05:00

[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).
This commit is contained in:
stratospher 2022-02-05 22:11:02 +05:30
parent b89fa59e71
commit a049d1bd08
2 changed files with 26 additions and 7 deletions

View file

@ -80,6 +80,9 @@ from test_framework.util import (
p2p_port, p2p_port,
wait_until_helper_internal, wait_until_helper_internal,
) )
from test_framework.v2_p2p import (
EncryptedP2PState,
)
logger = logging.getLogger("TestFramework.p2p") logger = logging.getLogger("TestFramework.p2p")
@ -159,11 +162,16 @@ class P2PConnection(asyncio.Protocol):
# The underlying transport of the connection. # The underlying transport of the connection.
# Should only call methods on this from the NetworkThread, c.f. call_soon_threadsafe # Should only call methods on this from the NetworkThread, c.f. call_soon_threadsafe
self._transport = None self._transport = None
self.v2_state = None # EncryptedP2PState object needed for v2 p2p connections
@property @property
def is_connected(self): def is_connected(self):
return self._transport is not None 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): def peer_connect_helper(self, dstaddr, dstport, net, timeout_factor):
assert not self.is_connected assert not self.is_connected
self.timeout_factor = timeout_factor self.timeout_factor = timeout_factor
@ -174,16 +182,20 @@ class P2PConnection(asyncio.Protocol):
self.recvbuf = b"" self.recvbuf = b""
self.magic_bytes = MAGIC_BYTES[net] 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) 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 loop = NetworkThread.network_event_loop
logger.debug('Connecting to Bitcoin Node: %s:%d' % (self.dstaddr, self.dstport)) logger.debug('Connecting to Bitcoin Node: %s:%d' % (self.dstaddr, self.dstport))
coroutine = loop.create_connection(lambda: self, host=self.dstaddr, port=self.dstport) coroutine = loop.create_connection(lambda: self, host=self.dstaddr, port=self.dstport)
return lambda: loop.call_soon_threadsafe(loop.create_task, coroutine) 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) 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)) logger.debug('Listening for Bitcoin Node with id: {}'.format(connect_id))
return lambda: NetworkThread.listen(self, connect_cb, idx=connect_id) return lambda: NetworkThread.listen(self, connect_cb, idx=connect_id)
@ -199,7 +211,13 @@ class P2PConnection(asyncio.Protocol):
assert not self._transport assert not self._transport
logger.debug("Connected & Listening: %s:%d" % (self.dstaddr, self.dstport)) logger.debug("Connected & Listening: %s:%d" % (self.dstaddr, self.dstport))
self._transport = transport 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.send_message(self.on_connection_send_msg)
self.on_connection_send_msg = None # Never used again self.on_connection_send_msg = None # Never used again
self.on_open() self.on_open()

View file

@ -642,7 +642,7 @@ class TestNode():
assert_msg += "with expected error " + expected_msg assert_msg += "with expected error " + expected_msg
self._raise_assertion_error(assert_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. """Add an inbound p2p connection to the node.
This method adds the p2p connection to the self.p2ps list and also 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' kwargs['dstaddr'] = '127.0.0.1'
p2p_conn.p2p_connected_to_node = True 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) self.p2ps.append(p2p_conn)
p2p_conn.wait_until(lambda: p2p_conn.is_connected, check_connected=False) p2p_conn.wait_until(lambda: p2p_conn.is_connected, check_connected=False)
if send_version: if send_version:
@ -684,7 +685,7 @@ class TestNode():
return p2p_conn 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 """Add an outbound p2p connection from node. Must be an
"outbound-full-relay", "block-relay-only", "addr-fetch" or "feeler" connection. "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) self.addconnection('%s:%d' % (address, port), connection_type)
p2p_conn.p2p_connected_to_node = False 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": if connection_type == "feeler":
# feeler connections are closed as soon as the node receives a `version` message # feeler connections are closed as soon as the node receives a `version` message