async def receive_handshake(reader: asyncio.StreamReader,
stats: PeerConnectionStats = None) -> Optional[Handshake]:
"""
Receives and decodes the handshake message sent by the peer.
:param reader: `asyncio.StreamReader` to read handshake from.
:param stats: `PeerConnectionStats` to populate data into.
:return: Decoded `Handshake` if successfully read, otherwise None.
:raises `PeerError`: if disconnected.
"""
assert reader is not None
if reader.at_eof() or reader.exception():
raise PeerError("Cannot receive message on disconnected reader.")
data = await reader.readexactly(Handshake.msg_len)
if stats:
stats.bytes_downloaded += Handshake.msg_len
return Handshake.decode(data)
async def _receive_from_peer(reader: asyncio.StreamReader,
stats: PeerConnectionStats) -> ProtocolMessage:
"""
Receives and decodes the next message sent by the peer.
:param reader: `StreamReader` to read from.
:param stats: `PeerConnectionStats` to update.
:return: The specific instance of the `ProtocolMessage` received.
:raises `PeerError`: on exception or if disconnected.
"""
assert reader is not None
if reader.at_eof() or reader.exception():
raise PeerError("Cannot receive message on disconnected reader.")
try:
msg_len = struct.unpack(">I", await reader.readexactly(4))[0]
stats.bytes_downloaded += 4
if msg_len == 0:
return KeepAlive()
msg_id = struct.unpack(">B", await reader.readexactly(1))[0]
if msg_id is None or (not (0 <= msg_id <= 8)):
raise PeerError("Unknown message received: %s" % msg_id)
# the msg_len includes 1 byte for the id
msg_len -= 1
if msg_len == 0:
return MESSAGE_TYPES[msg_id].decode()
msg_data = await reader.readexactly(msg_len)
if stats:
stats.bytes_downloaded += msg_len
return MESSAGE_TYPES[msg_id].decode(msg_data)
except Exception as e:
raise PeerError from e
