def test_azip(self): i1 = ("abcdefgh") i2 = list(range(20)) i3 = list(str(_) for _ in range(20)) ai1 = iter_to_aiter(i1) ai2 = iter_to_aiter(i2) ai3 = iter_to_aiter(i3) ai = azip(ai1, ai2, ai3) r = run(get_n(ai)) self.assertEqual(r, list(zip(i1, i2, i3)))
def test_join_aiters(self): int_vals = [1, 2, 3, 4] str_vals = "abcdefg" list_of_lists = [int_vals, str_vals] iter_of_aiters = [iter_to_aiter(_) for _ in list_of_lists] aiter_of_aiters = iter_to_aiter(iter_of_aiters) r = run(get_n(join_aiters(aiter_of_aiters))) r1 = [_ for _ in r if isinstance(_, int)] r2 = [_ for _ in r if isinstance(_, str)] self.assertEqual(r1, int_vals) self.assertEqual(r2, list(str_vals))
async def start_client( self, target_node: PeerInfo, on_connect: OnConnectFunc = None, ) -> bool: """ Tries to connect to the target node, adding one connection into the pipeline, if successful. An on connect method can also be specified, and this will be saved into the instance variables. """ if self._server is not None: if (self._host == target_node.host or target_node.host == "127.0.0.1") and self._port == target_node.port: self.global_connections.peers.remove(target_node) return False if self._pipeline_task.done(): return False try: reader, writer = await asyncio.open_connection( target_node.host, int(target_node.port)) except ( ConnectionRefusedError, TimeoutError, OSError, asyncio.TimeoutError, ) as e: log.warning( f"Could not connect to {target_node}. {type(e)}{str(e)}. Aborting and removing peer." ) self.global_connections.peers.remove(target_node) return False asyncio.create_task( self._add_to_srwt_aiter( iter_to_aiter([(reader, writer, on_connect)]))) return True
async def start_client( self, target_node: PeerInfo, on_connect: OnConnectFunc = None, auth: bool = False, ) -> bool: """ Tries to connect to the target node, adding one connection into the pipeline, if successful. An on connect method can also be specified, and this will be saved into the instance variables. """ if self._pipeline_task.done(): self.log.warning("Starting client after server closed") return False ssl_context = ssl_context_for_client(self.root_path, self.config, auth=auth) try: reader, writer = await asyncio.open_connection( target_node.host, int(target_node.port), ssl=ssl_context ) except Exception as e: self.log.warning( f"Could not connect to {target_node}. {type(e)}{str(e)}. Aborting and removing peer." ) self.global_connections.peers.remove(target_node) return False if not self._srwt_aiter.is_stopped(): self._srwt_aiter.push(iter_to_aiter([(reader, writer, on_connect)])) ssl_object = writer.get_extra_info(name="ssl_object") peer_cert = ssl_object.getpeercert() self.log.info(f"Server authed as {peer_cert}") return True
async def start_client( self, target_node: PeerInfo, on_connect: OnConnectFunc = None, config: Dict = None, ) -> bool: """ Tries to connect to the target node, adding one connection into the pipeline, if successful. An on connect method can also be specified, and this will be saved into the instance variables. """ if self._server is not None: if ( target_node.host == "127.0.0.1" or target_node.host == "0.0.0.0" or target_node.host == "::1" or target_node.host == "0:0:0:0:0:0:0:1" ) and self._port == target_node.port: self.global_connections.peers.remove(target_node) return False if self._pipeline_task.done(): return False ssl_context = ssl._create_unverified_context(purpose=ssl.Purpose.SERVER_AUTH) certfile, keyfile, password, cafile = self.loadSSLConfig("client_ssl", config) ssl_context.load_cert_chain( certfile=certfile, keyfile=keyfile, password=password ) if cafile is None: ssl_context.verify_mode = ssl.CERT_NONE else: ssl_context.load_verify_locations(cafile=cafile) ssl_context.verify_mode = ssl.CERT_REQUIRED try: reader, writer = await asyncio.open_connection( target_node.host, int(target_node.port), ssl=ssl_context ) except ( ConnectionRefusedError, TimeoutError, OSError, asyncio.TimeoutError, ) as e: self.log.warning( f"Could not connect to {target_node}. {type(e)}{str(e)}. Aborting and removing peer." ) self.global_connections.peers.remove(target_node) return False self._tasks.append( asyncio.create_task( self._add_to_srwt_aiter(iter_to_aiter([(reader, writer, on_connect)])) ) ) ssl_object = writer.get_extra_info(name="ssl_object") peer_cert = ssl_object.getpeercert() self.log.info(f"Server authed as {peer_cert}") return True
def test_gated_aiter(self): ai = iter_to_aiter(range(3000000000)) aiter = gated_aiter(ai) aiter.push(9) r = run(get_n(aiter, 3)) r.extend(run(get_n(aiter, 4))) aiter.push(11) aiter.stop() r.extend(run(get_n(aiter))) self.assertEqual(r, list(range(20)))
async def start_client( self, target_node: PeerInfo, on_connect: OnConnectFunc = None, auth: bool = False, is_feeler: bool = False, ) -> bool: """ Tries to connect to the target node, adding one connection into the pipeline, if successful. An on connect method can also be specified, and this will be saved into the instance variables. """ self.log.info(f"Trying to connect with {target_node}.") if self._pipeline_task.done(): self.log.error("Starting client after server closed") return False ssl_context = ssl_context_for_client(self.root_path, self.config, auth=auth) try: # Sometimes open_connection takes a long time, so we add it as a task, and cancel # the task in the event of closing the node. peer_info = (target_node.host, target_node.port) if peer_info in self._pending_connections: return False self._pending_connections.append(peer_info) oc_task: asyncio.Task = asyncio.create_task( asyncio.open_connection( target_node.host, int(target_node.port), ssl=ssl_context ) ) self._oc_tasks.append(oc_task) reader, writer = await oc_task self._pending_connections.remove(peer_info) self._oc_tasks.remove(oc_task) except Exception as e: self.log.warning( f"Could not connect to {target_node}. {type(e)}{str(e)}. Aborting and removing peer." ) self.global_connections.failed_connection(target_node) if self.global_connections.introducer_peers is not None: self.global_connections.introducer_peers.remove(target_node) if peer_info in self._pending_connections: self._pending_connections.remove(peer_info) return False if not self._srwt_aiter.is_stopped(): self._srwt_aiter.push( iter_to_aiter([(reader, writer, on_connect, True, is_feeler)]) ) ssl_object = writer.get_extra_info(name="ssl_object") peer_cert = ssl_object.getpeercert() self.log.info(f"Server authed as {peer_cert}") return True
def initialize_pipeline(self, aiter, api: Any, server_port: int) -> asyncio.Task: """ A pipeline that starts with (StreamReader, StreamWriter), maps it though to connections, messages, executes a local API call, and returns responses. """ # Maps a stream reader, writer and NodeType to a Connection object connections_aiter = map_aiter( partial_func.partial_async(self.stream_reader_writer_to_connection, server_port), aiter, ) # Performs a handshake with the peer handshaked_connections_aiter = join_aiters( map_aiter(self.perform_handshake, connections_aiter)) forker = aiter_forker(handshaked_connections_aiter) handshake_finished_1 = forker.fork(is_active=True) handshake_finished_2 = forker.fork(is_active=True) # Reads messages one at a time from the TCP connection messages_aiter = join_aiters( map_aiter(self.connection_to_message, handshake_finished_1, 100)) # Handles each message one at a time, and yields responses to send back or broadcast responses_aiter = join_aiters( map_aiter( partial_func.partial_async_gen(self.handle_message, api), messages_aiter, 100, )) # Uses a forked aiter, and calls the on_connect function to send some initial messages # as soon as the connection is established on_connect_outbound_aiter = join_aiters( map_aiter(self.connection_to_outbound, handshake_finished_2, 100)) # Also uses the instance variable _outbound_aiter, which clients can use to send messages # at any time, not just on_connect. outbound_aiter_mapped = map_aiter(lambda x: (None, x), self._outbound_aiter) responses_aiter = join_aiters( iter_to_aiter([ responses_aiter, on_connect_outbound_aiter, outbound_aiter_mapped ])) # For each outbound message, replicate for each peer that we need to send to expanded_messages_aiter = join_aiters( map_aiter(self.expand_outbound_messages, responses_aiter, 100)) # This will run forever. Sends each message through the TCP connection, using the # length encoding and CBOR serialization async def serve_forever(): async for connection, message in expanded_messages_aiter: log.info( f"-> {message.function} to peer {connection.get_peername()}" ) try: await connection.send(message) except ( ConnectionResetError, BrokenPipeError, RuntimeError, TimeoutError, ) as e: log.error( f"Cannot write to {connection}, already closed. Error {e}." ) self.global_connections.close(connection, True) # We will return a task for this, so user of start_chia_server or start_chia_client can wait until # the server is closed. return asyncio.get_running_loop().create_task(serve_forever())
async def initialize_pipeline( srwt_aiter, api: Any, server_port: int, outbound_aiter: push_aiter, global_connections: PeerConnections, local_type: NodeType, node_id: bytes32, network_id: bytes32, log: logging.Logger, ): """ A pipeline that starts with (StreamReader, StreamWriter), maps it though to connections, messages, executes a local API call, and returns responses. """ # Maps a stream reader, writer and NodeType to a Connection object connections_aiter = map_aiter( partial_func.partial_async( stream_reader_writer_to_connection, server_port, local_type, log, ), join_aiters(srwt_aiter), ) def add_global_connections(connection): return connection, global_connections connections_with_global_connections_aiter = map_aiter( add_global_connections, connections_aiter ) # Performs a handshake with the peer outbound_handshake = Message( "handshake", Handshake( network_id, protocol_version, node_id, uint16(server_port), local_type, ), ) handshaked_connections_aiter = join_aiters( map_aiter( lambda _: perform_handshake(_, srwt_aiter, outbound_handshake), connections_with_global_connections_aiter, ) ) forker = aiter_forker(handshaked_connections_aiter) handshake_finished_1 = forker.fork(is_active=True) handshake_finished_2 = forker.fork(is_active=True) # Reads messages one at a time from the TCP connection messages_aiter = join_aiters( map_aiter(connection_to_message, handshake_finished_1, 100) ) # Handles each message one at a time, and yields responses to send back or broadcast responses_aiter = join_aiters( map_aiter( partial_func.partial_async_gen(handle_message, api), messages_aiter, 100, ) ) # Uses a forked aiter, and calls the on_connect function to send some initial messages # as soon as the connection is established on_connect_outbound_aiter = join_aiters( map_aiter(connection_to_outbound, handshake_finished_2, 100) ) # Also uses the instance variable _outbound_aiter, which clients can use to send messages # at any time, not just on_connect. outbound_aiter_mapped = map_aiter( lambda x: (None, x, global_connections), outbound_aiter ) responses_aiter = join_aiters( iter_to_aiter( [responses_aiter, on_connect_outbound_aiter, outbound_aiter_mapped] ) ) # For each outbound message, replicate for each peer that we need to send to expanded_messages_aiter = join_aiters( map_aiter(expand_outbound_messages, responses_aiter, 100) ) # This will run forever. Sends each message through the TCP connection, using the # length encoding and CBOR serialization async for connection, message in expanded_messages_aiter: if message is None: # Does not ban the peer, this is just a graceful close of connection. global_connections.close(connection, True) continue if connection.is_closing(): connection.log.info( f"Closing, so will not send {message.function} to peer {connection.get_peername()}" ) continue connection.log.info( f"-> {message.function} to peer {connection.get_peername()}" ) try: await connection.send(message) except (RuntimeError, TimeoutError, OSError,) as e: connection.log.warning( f"Cannot write to {connection}, already closed. Error {e}." ) global_connections.close(connection, True)