[test] Add test framework support to create outbound connections.

In the interest of increasing our P2P test coverage, add support to create
full-relay or block-relay-only connections. To support this, a P2P connection
spins up a listening thread & uses a callback to trigger the node initiating
the connection.

Co-authored-by: Anthony Towns <[email protected]>

Author: amitiuttarwar

test/functional/test_framework/p2p.py

@@ -71,7 +71,11 @@
     NODE_WITNESS,
     sha256,
 )
-from test_framework.util import wait_until_helper
+from test_framework.util import (
+    MAX_NODES,
+    p2p_port,
+    wait_until_helper,
+)
 
 logger = logging.getLogger("TestFramework.p2p")
 
@@ -139,7 +143,7 @@ def __init__(self):
     def is_connected(self):
         return self._transport is not None
 
-    def peer_connect(self, dstaddr, dstport, *, net, timeout_factor):
+    def peer_connect_helper(self, dstaddr, dstport, net, timeout_factor):
         assert not self.is_connected
         self.timeout_factor = timeout_factor
         self.dstaddr = dstaddr
@@ -148,12 +152,20 @@ def peer_connect(self, dstaddr, dstport, *, net, timeout_factor):
         self.on_connection_send_msg = None
         self.recvbuf = b""
         self.magic_bytes = MAGIC_BYTES[net]
-        logger.debug('Connecting to Bitcoin Node: %s:%d' % (self.dstaddr, self.dstport))
+
+    def peer_connect(self, dstaddr, dstport, *, net, timeout_factor):
+        self.peer_connect_helper(dstaddr, dstport, net, timeout_factor)
 
         loop = NetworkThread.network_event_loop
-        conn_gen_unsafe = loop.create_connection(lambda: self, host=self.dstaddr, port=self.dstport)
-        conn_gen = lambda: loop.call_soon_threadsafe(loop.create_task, conn_gen_unsafe)
-        return conn_gen
+        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):
+        self.peer_connect_helper('0', 0, net, timeout_factor)
+
+        logger.debug('Listening for Bitcoin Node with id: {}'.format(connect_id))
+        return lambda: NetworkThread.listen(self, connect_cb, idx=connect_id)
 
     def peer_disconnect(self):
         # Connection could have already been closed by other end.
@@ -312,18 +324,27 @@ def __init__(self, support_addrv2=False, wtxidrelay=True):
         # If the peer supports wtxid-relay
         self.wtxidrelay = wtxidrelay
 
-    def peer_connect(self, *args, services=NODE_NETWORK|NODE_WITNESS, send_version=True, **kwargs):
+    def peer_connect_send_version(self, services):
+        # Send a version msg
+        vt = msg_version()
+        vt.nServices = services
+        vt.addrTo.ip = self.dstaddr
+        vt.addrTo.port = self.dstport
+        vt.addrFrom.ip = "0.0.0.0"
+        vt.addrFrom.port = 0
+        self.on_connection_send_msg = vt  # Will be sent in connection_made callback
+
+    def peer_connect(self, *args, services=NODE_NETWORK | NODE_WITNESS, send_version=True, **kwargs):
         create_conn = super().peer_connect(*args, **kwargs)
 
         if send_version:
-            # Send a version msg
-            vt = msg_version()
-            vt.nServices = services
-            vt.addrTo.ip = self.dstaddr
-            vt.addrTo.port = self.dstport
-            vt.addrFrom.ip = "0.0.0.0"
-            vt.addrFrom.port = 0
-            self.on_connection_send_msg = vt  # Will be sent soon after connection_made
+            self.peer_connect_send_version(services)
+
+        return create_conn
+
+    def peer_accept_connection(self, *args, services=NODE_NETWORK | NODE_WITNESS, **kwargs):
+        create_conn = super().peer_accept_connection(*args, **kwargs)
+        self.peer_connect_send_version(services)
 
         return create_conn
 
@@ -414,6 +435,10 @@ def test_function():
 
         wait_until_helper(test_function, timeout=timeout, lock=p2p_lock, timeout_factor=self.timeout_factor)
 
+    def wait_for_connect(self, timeout=60):
+        test_function = lambda: self.is_connected
+        wait_until_helper(test_function, timeout=timeout, lock=p2p_lock)
+
     def wait_for_disconnect(self, timeout=60):
         test_function = lambda: not self.is_connected
         self.wait_until(test_function, timeout=timeout, check_connected=False)
@@ -527,6 +552,8 @@ def __init__(self):
         # There is only one event loop and no more than one thread must be created
         assert not self.network_event_loop
 
+        NetworkThread.listeners = {}
+        NetworkThread.protos = {}
         NetworkThread.network_event_loop = asyncio.new_event_loop()
 
     def run(self):
@@ -542,6 +569,48 @@ def close(self, timeout=10):
         # Safe to remove event loop.
         NetworkThread.network_event_loop = None
 
+    @classmethod
+    def listen(cls, p2p, callback, port=None, addr=None, idx=1):
+        """ Ensure a listening server is running on the given port, and run the
+        protocol specified by `p2p` on the next connection to it. Once ready
+        for connections, call `callback`."""
+
+        if port is None:
+            assert 0 < idx <= MAX_NODES
+            port = p2p_port(MAX_NODES - idx)
+        if addr is None:
+            addr = '127.0.0.1'
+
+        coroutine = cls.create_listen_server(addr, port, callback, p2p)
+        cls.network_event_loop.call_soon_threadsafe(cls.network_event_loop.create_task, coroutine)
+
+    @classmethod
+    async def create_listen_server(cls, addr, port, callback, proto):
+        def peer_protocol():
+            """Returns a function that does the protocol handling for a new
+            connection. To allow different connections to have different
+            behaviors, the protocol function is first put in the cls.protos
+            dict. When the connection is made, the function removes the
+            protocol function from that dict, and returns it so the event loop
+            can start executing it."""
+            response = cls.protos.get((addr, port))
+            cls.protos[(addr, port)] = None
+            return response
+
+        if (addr, port) not in cls.listeners:
+            # When creating a listener on a given (addr, port) we only need to
+            # do it once. If we want different behaviors for different
+            # connections, we can accomplish this by providing different
+            # `proto` functions
+
+            listener = await cls.network_event_loop.create_server(peer_protocol, addr, port)
+            logger.debug("Listening server on %s:%d should be started" % (addr, port))
+            cls.listeners[(addr, port)] = listener
+
+        cls.protos[(addr, port)] = proto
+        callback(addr, port)
+
+
 class P2PDataStore(P2PInterface):
     """A P2P data store class.
 

test/functional/test_framework/test_node.py

@@ -71,6 +71,7 @@ def __init__(self, i, datadir, *, chain, rpchost, timewait, timeout_factor, bitc
         """
 
         self.index = i
+        self.p2p_conn_index = 1
         self.datadir = datadir
         self.bitcoinconf = os.path.join(self.datadir, "bitcoin.conf")
         self.stdout_dir = os.path.join(self.datadir, "stdout")
@@ -517,7 +518,7 @@ def assert_start_raises_init_error(self, extra_args=None, expected_msg=None, mat
                 self._raise_assertion_error(assert_msg)
 
     def add_p2p_connection(self, p2p_conn, *, wait_for_verack=True, **kwargs):
-        """Add a 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
         returns the connection to the caller."""
@@ -546,6 +547,29 @@ def add_p2p_connection(self, p2p_conn, *, wait_for_verack=True, **kwargs):
 
         return p2p_conn
 
+    def add_outbound_p2p_connection(self, p2p_conn, *, p2p_idx, connection_type="outbound-full-relay", **kwargs):
+        """Add an outbound p2p connection from node. Either
+        full-relay("outbound-full-relay") or
+        block-relay-only("block-relay-only") connection.
+
+        This method adds the p2p connection to the self.p2ps list and returns
+        the connection to the caller.
+        """
+
+        def addconnection_callback(address, port):
+            self.log.debug("Connecting to %s:%d %s" % (address, port, connection_type))
+            self.addconnection('%s:%d' % (address, port), connection_type)
+
+        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.wait_for_connect()
+        self.p2ps.append(p2p_conn)
+
+        p2p_conn.wait_for_verack()
+        p2p_conn.sync_with_ping()
+
+        return p2p_conn
+
     def num_test_p2p_connections(self):
         """Return number of test framework p2p connections to the node."""
         return len([peer for peer in self.getpeerinfo() if peer['subver'] == MY_SUBVERSION.decode("utf-8")])
@@ -555,6 +579,7 @@ def disconnect_p2ps(self):
         for p in self.p2ps:
             p.peer_disconnect()
         del self.p2ps[:]
+
         wait_until_helper(lambda: self.num_test_p2p_connections() == 0, timeout_factor=self.timeout_factor)