async def perform_handshake(
self, network_id: bytes32, protocol_version: str, server_port: int, local_type: NodeType
):
if self.is_outbound:
outbound_handshake = make_msg(
ProtocolMessageTypes.handshake,
Handshake(
network_id,
protocol_version,
chia_full_version_str(),
uint16(server_port),
uint8(local_type.value),
),
)
payload: Optional[Payload] = Payload(outbound_handshake, None)
assert payload is not None
await self._send_message(payload)
payload = await self._read_one_message()
if payload is None:
raise ProtocolError(Err.INVALID_HANDSHAKE)
inbound_handshake = Handshake.from_bytes(payload.msg.data)
if ProtocolMessageTypes(payload.msg.type) != ProtocolMessageTypes.handshake:
raise ProtocolError(Err.INVALID_HANDSHAKE)
if inbound_handshake.protocol_version != protocol_version:
raise ProtocolError(Err.INCOMPATIBLE_PROTOCOL_VERSION)
self.peer_server_port = inbound_handshake.server_port
self.connection_type = NodeType(inbound_handshake.node_type)
else:
try:
payload = await self._read_one_message()
except Exception:
raise ProtocolError(Err.INVALID_HANDSHAKE)
if payload is None:
raise ProtocolError(Err.INVALID_HANDSHAKE)
inbound_handshake = Handshake.from_bytes(payload.msg.data)
if ProtocolMessageTypes(payload.msg.type) != ProtocolMessageTypes.handshake:
raise ProtocolError(Err.INVALID_HANDSHAKE)
if inbound_handshake.protocol_version != protocol_version:
raise ProtocolError(Err.INCOMPATIBLE_PROTOCOL_VERSION)
outbound_handshake = make_msg(
ProtocolMessageTypes.handshake,
Handshake(
network_id,
protocol_version,
chia_full_version_str(),
uint16(server_port),
uint8(local_type.value),
),
)
payload = Payload(outbound_handshake, None)
await self._send_message(payload)
self.peer_server_port = inbound_handshake.server_port
self.connection_type = NodeType(inbound_handshake.node_type)
self.outbound_task = asyncio.create_task(self.outbound_handler())
self.inbound_task = asyncio.create_task(self.inbound_handler())
return True
async def perform_handshake(self, network_id, protocol_version, node_id,
server_port, local_type):
if self.is_outbound:
outbound_handshake = Message(
"handshake",
Handshake(
network_id,
protocol_version,
node_id,
uint16(server_port),
local_type,
),
)
payload = Payload(outbound_handshake, None)
await self._send_message(payload)
payload = await self._read_one_message()
inbound_handshake = Handshake(**payload.msg.data)
if payload.msg.function != "handshake" or not inbound_handshake or not inbound_handshake.node_type:
raise ProtocolError(Err.INVALID_HANDSHAKE)
self.peer_node_id = inbound_handshake.node_id
self.peer_server_port = int(inbound_handshake.server_port)
self.connection_type = inbound_handshake.node_type
else:
payload = await self._read_one_message()
inbound_handshake = Handshake(**payload.msg.data)
if payload.msg.function != "handshake" or not inbound_handshake or not inbound_handshake.node_type:
raise ProtocolError(Err.INVALID_HANDSHAKE)
outbound_handshake = Message(
"handshake",
Handshake(
network_id,
protocol_version,
node_id,
uint16(server_port),
local_type,
),
)
payload = Payload(outbound_handshake, None)
await self._send_message(payload)
self.peer_node_id = inbound_handshake.node_id
self.peer_server_port = int(inbound_handshake.server_port)
self.connection_type = inbound_handshake.node_type
if self.peer_node_id == node_id:
raise ProtocolError(Err.SELF_CONNECTION)
self.outbound_task = asyncio.create_task(self.outbound_handler())
self.inbound_task = asyncio.create_task(self.inbound_handler())
return True
async def perform_handshake(
self, connection: Connection) -> AsyncGenerator[Connection, None]:
"""
Performs handshake with this new connection, and yields the connection. If the handshake
is unsuccessful, or we already have a connection with this peer, the connection is closed,
and nothing is yielded.
"""
# Send handshake message
outbound_handshake = Message(
"handshake",
Handshake(
config["network_id"],
protocol_version,
self._node_id,
uint16(self._port),
self._local_type,
),
)
try:
await connection.send(outbound_handshake)
# Read handshake message
full_message = await connection.read_one_message()
inbound_handshake = Handshake(**full_message.data)
if (full_message.function != "handshake" or not inbound_handshake
or not inbound_handshake.node_type):
raise InvalidHandshake("Invalid handshake")
if inbound_handshake.node_id == self._node_id:
raise InvalidHandshake(
f"Should not connect to ourselves, aborting handshake.")
# Makes sure that we only start one connection with each peer
connection.node_id = inbound_handshake.node_id
connection.peer_server_port = int(inbound_handshake.server_port)
connection.connection_type = inbound_handshake.node_type
if not self.global_connections.add(connection):
raise InvalidHandshake(f"Duplicate connection to {connection}")
# Send Ack message
await connection.send(Message("handshake_ack", HandshakeAck()))
# Read Ack message
full_message = await connection.read_one_message()
if full_message.function != "handshake_ack":
raise InvalidAck("Invalid ack")
if inbound_handshake.version != protocol_version:
raise IncompatibleProtocolVersion(
f"Our node version {protocol_version} is not compatible with peer\
{connection} version {inbound_handshake.version}")
log.info((
f"Handshake with {NodeType(connection.connection_type).name} {connection.get_peername()} "
f"{connection.node_id}"
f" established"))
# Only yield a connection if the handshake is succesful and the connection is not a duplicate.
yield connection
except (
IncompatibleProtocolVersion,
InvalidAck,
InvalidHandshake,
asyncio.IncompleteReadError,
ConnectionResetError,
Exception,
) as e:
log.warning(
f"{e}, handshake not completed. Connection not created.")
# Make sure to close the connection even if it's not in global connections
connection.close()
# Remove the conenction from global connections
self.global_connections.close(connection)
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)
async def perform_handshake(
pair: Tuple[Connection, PeerConnections],
srwt_aiter: push_aiter,
outbound_handshake: Message,
) -> AsyncGenerator[Tuple[Connection, PeerConnections], None]:
"""
Performs handshake with this new connection, and yields the connection. If the handshake
is unsuccessful, or we already have a connection with this peer, the connection is closed,
and nothing is yielded.
"""
connection, global_connections = pair
# Send handshake message
try:
await connection.send(outbound_handshake)
# Read handshake message
full_message = await connection.read_one_message()
inbound_handshake = Handshake(**full_message.data)
if (
full_message.function != "handshake"
or not inbound_handshake
or not inbound_handshake.node_type
):
raise ProtocolError(Err.INVALID_HANDSHAKE)
if inbound_handshake.node_id == outbound_handshake.data.node_id:
raise ProtocolError(Err.SELF_CONNECTION)
# Makes sure that we only start one connection with each peer
connection.node_id = inbound_handshake.node_id
connection.peer_server_port = int(inbound_handshake.server_port)
connection.connection_type = inbound_handshake.node_type
if srwt_aiter.is_stopped():
raise Exception("No longer accepting handshakes, closing.")
if not global_connections.add(connection):
raise ProtocolError(Err.DUPLICATE_CONNECTION, [False])
# Send Ack message
await connection.send(Message("handshake_ack", HandshakeAck()))
# Read Ack message
full_message = await connection.read_one_message()
if full_message.function != "handshake_ack":
raise ProtocolError(Err.INVALID_ACK)
if inbound_handshake.version != protocol_version:
raise ProtocolError(
Err.INCOMPATIBLE_PROTOCOL_VERSION,
[protocol_version, inbound_handshake.version],
)
connection.log.info(
(
f"Handshake with {NodeType(connection.connection_type).name} {connection.get_peername()} "
f"{connection.node_id}"
f" established"
)
)
# Only yield a connection if the handshake is succesful and the connection is not a duplicate.
yield connection, global_connections
except (ProtocolError, asyncio.IncompleteReadError, OSError, Exception,) as e:
connection.log.warning(f"{e}, handshake not completed. Connection not created.")
# Make sure to close the connection even if it's not in global connections
connection.close()
# Remove the conenction from global connections
global_connections.close(connection)
async def perform_handshake(self, network_id: str, protocol_version: str,
server_port: int, local_type: NodeType):
if self.is_outbound:
outbound_handshake = make_msg(
ProtocolMessageTypes.handshake,
Handshake(
network_id,
protocol_version,
chia_full_version_str(),
uint16(server_port),
uint8(local_type.value),
[(uint16(Capability.BASE.value), "1")],
),
)
assert outbound_handshake is not None
await self._send_message(outbound_handshake)
inbound_handshake_msg = await self._read_one_message()
if inbound_handshake_msg is None:
raise ProtocolError(Err.INVALID_HANDSHAKE)
inbound_handshake = Handshake.from_bytes(
inbound_handshake_msg.data)
if ProtocolMessageTypes(inbound_handshake_msg.type
) != ProtocolMessageTypes.handshake:
raise ProtocolError(Err.INVALID_HANDSHAKE)
if inbound_handshake.network_id != network_id:
raise ProtocolError(Err.INCOMPATIBLE_NETWORK_ID)
self.peer_server_port = inbound_handshake.server_port
self.connection_type = NodeType(inbound_handshake.node_type)
else:
try:
message = await self._read_one_message()
except Exception:
raise ProtocolError(Err.INVALID_HANDSHAKE)
if message is None:
raise ProtocolError(Err.INVALID_HANDSHAKE)
inbound_handshake = Handshake.from_bytes(message.data)
if ProtocolMessageTypes(
message.type) != ProtocolMessageTypes.handshake:
raise ProtocolError(Err.INVALID_HANDSHAKE)
if inbound_handshake.network_id != network_id:
raise ProtocolError(Err.INCOMPATIBLE_NETWORK_ID)
outbound_handshake = make_msg(
ProtocolMessageTypes.handshake,
Handshake(
network_id,
protocol_version,
chia_full_version_str(),
uint16(server_port),
uint8(local_type.value),
[(uint16(Capability.BASE.value), "1")],
),
)
await self._send_message(outbound_handshake)
self.peer_server_port = inbound_handshake.server_port
self.connection_type = NodeType(inbound_handshake.node_type)
self.outbound_task = asyncio.create_task(self.outbound_handler())
self.inbound_task = asyncio.create_task(self.inbound_handler())
return True
