Convert message IO to work w/ a RawConnection
.
This commit is contained in:
parent
7c004a4e14
commit
eb5ef39399
10
libp2p/io/exceptions.py
Normal file
10
libp2p/io/exceptions.py
Normal file
|
@ -0,0 +1,10 @@
|
||||||
|
class MsgioException(Exception):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class MissingLengthException(MsgioException):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class MissingMessageException(MsgioException):
|
||||||
|
pass
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue
Block a user