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)
