Convert message IO to work w/ a RawConnection.

This commit is contained in:
Alex Stokes 2019-08-24 22:47:56 +02:00
parent 7c004a4e14
commit eb5ef39399
No known key found for this signature in database
GPG Key ID: 51CE1721B245C086
2 changed files with 22 additions and 4 deletions

10
libp2p/io/exceptions.py Normal file
View File

@ -0,0 +1,10 @@
class MsgioException(Exception):
pass
class MissingLengthException(MsgioException):
pass
class MissingMessageException(MsgioException):
pass

View File

@ -1,4 +1,6 @@
import asyncio from libp2p.network.connection.raw_connection_interface import IRawConnection
from .exceptions import MissingLengthException, MissingMessageException
SIZE_LEN_BYTES = 4 SIZE_LEN_BYTES = 4
@ -10,7 +12,13 @@ def encode(msg_bytes: bytes) -> bytes:
return len_prefix + msg_bytes return len_prefix + msg_bytes
async def read_next_message(reader: asyncio.StreamReader) -> bytes: async def read_next_message(reader: IRawConnection) -> bytes:
len_bytes = await reader.readexactly(SIZE_LEN_BYTES) len_bytes = await reader.read(SIZE_LEN_BYTES)
if len(len_bytes) != SIZE_LEN_BYTES:
raise MissingLengthException()
len_int = int.from_bytes(len_bytes, "big") len_int = int.from_bytes(len_bytes, "big")
return await reader.readexactly(len_int) next_msg = await reader.read(len_int)
if len(next_msg) != len_int:
# TODO makes sense to keep reading until this condition is true?
raise MissingMessageException()
return next_msg