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 create_request(self, message: Message, timeout: int = 15):
""" Sends a message and waits for a response. """
if self.closed:
return None
event = asyncio.Event()
payload = Payload(message, token_bytes(8))
self.pending_requests[payload.id] = event
await self.outgoing_queue.put(payload)
async def time_out(req_id, req_timeout):
await asyncio.sleep(req_timeout)
if req_id in self.pending_requests:
self.pending_requests[req_id].set()
asyncio.create_task(time_out(payload.id, timeout))
await event.wait()
self.pending_requests.pop(payload.id)
result: Optional[Message] = None
if payload.id in self.request_results:
result_payload: Payload = self.request_results[payload.id]
result = result_payload.msg
self.log.info(
f"<- {result_payload.msg.function} from: {self.peer_host}:{self.peer_port}"
)
self.request_results.pop(payload.id)
return result
async def api_call(payload: Payload, connection: WSChiaConnection,
task_id):
start_time = time.time()
try:
if self.received_message_callback is not None:
await self.received_message_callback(connection)
full_message = payload.msg
connection.log.info(
f"<- {ProtocolMessageTypes(full_message.type).name} from peer "
f"{connection.peer_node_id} {connection.peer_host}")
message_type: str = ProtocolMessageTypes(
full_message.type).name
f = getattr(self.api, message_type, None)
if f is None:
self.log.error(
f"Non existing function: {message_type}")
raise ProtocolError(Err.INVALID_PROTOCOL_MESSAGE,
[message_type])
if not hasattr(f, "api_function"):
self.log.error(
f"Peer trying to call non api function {message_type}"
)
raise ProtocolError(Err.INVALID_PROTOCOL_MESSAGE,
[message_type])
if hasattr(f, "peer_required"):
coroutine = f(full_message.data, connection)
else:
coroutine = f(full_message.data)
response: Optional[Message] = await asyncio.wait_for(
coroutine, timeout=300)
connection.log.debug(
f"Time taken to process {message_type} from {connection.peer_node_id} is "
f"{time.time() - start_time} seconds")
if response is not None:
payload_id = payload.id
response_payload = Payload(response, payload_id)
await connection.reply_to_request(response_payload)
except Exception as e:
if self.connection_close_task is None:
tb = traceback.format_exc()
connection.log.error(
f"Exception: {e}, closing connection {connection.get_peer_info()}. {tb}"
)
else:
connection.log.debug(
f"Exception: {e} while closing connection")
pass
await connection.close()
finally:
if task_id in self.api_tasks:
self.api_tasks.pop(task_id)
if task_id in self.tasks_from_peer[
connection.peer_node_id]:
self.tasks_from_peer[connection.peer_node_id].remove(
task_id)
async def _read_one_message(self) -> Optional[Payload]:
try:
message: WSMessage = await self.ws.receive(30)
except asyncio.TimeoutError:
# self.ws._closed if we didn't receive a ping / pong
if self.ws._closed:
asyncio.create_task(self.close())
await asyncio.sleep(3)
return None
return None
if self.connection_type is not None:
connection_type_str = NodeType(self.connection_type).name.lower()
else:
connection_type_str = ""
if message.type == WSMsgType.CLOSING:
self.log.debug(
f"Closing connection to {connection_type_str} {self.peer_host}:"
f"{self.peer_server_port}/"
f"{self.peer_port}"
)
asyncio.create_task(self.close())
await asyncio.sleep(3)
elif message.type == WSMsgType.CLOSE:
self.log.debug(
f"Peer closed connection {connection_type_str} {self.peer_host}:"
f"{self.peer_server_port}/"
f"{self.peer_port}"
)
asyncio.create_task(self.close())
await asyncio.sleep(3)
elif message.type == WSMsgType.CLOSED:
if not self.closed:
asyncio.create_task(self.close())
await asyncio.sleep(3)
return None
elif message.type == WSMsgType.BINARY:
data = message.data
full_message_loaded: Payload = Payload.from_bytes(data)
self.bytes_read += len(data)
self.last_message_time = time.time()
return full_message_loaded
else:
self.log.error(f"Unexpected WebSocket message type: {message}")
asyncio.create_task(self.close())
await asyncio.sleep(3)
return None
async def create_request(self, message: Message, timeout: int):
""" Sends a message and waits for a response. """
if self.closed:
return None
# We will wait for this event, it will be set either by the response, or the timeout
event = asyncio.Event()
# The request nonce is an integer between 0 and 2**16 - 1, which is used to match requests to responses
request_id = self.request_nonce
self.request_nonce = uint16(self.request_nonce + 1) if self.request_nonce != (2 ** 16 - 1) else uint16(0)
payload = Payload(message, request_id)
self.pending_requests[payload.id] = event
await self.outgoing_queue.put(payload)
# If the timeout passes, we set the event
async def time_out(req_id, req_timeout):
await asyncio.sleep(req_timeout)
if req_id in self.pending_requests:
self.pending_requests[req_id].set()
asyncio.create_task(time_out(payload.id, timeout))
await event.wait()
self.pending_requests.pop(payload.id)
result: Optional[Message] = None
if payload.id in self.request_results:
result_payload: Payload = self.request_results[payload.id]
result = result_payload.msg
self.log.info(
f"<- {ProtocolMessageTypes(result_payload.msg.type).name} from: {self.peer_host}:{self.peer_port}"
)
self.request_results.pop(payload.id)
return result
async def send_messages(self, messages: List[Message]):
if self.closed:
return
for message in messages:
payload = Payload(message, None)
await self.outgoing_queue.put(payload)
async def send_message(self, message: Message):
""" Send message sends a message with no tracking / callback. """
if self.closed:
return
payload = Payload(message, None)
await self.outgoing_queue.put(payload)