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