mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2026-02-10 07:00:54 +00:00
Convert message IO to work w/ a RawConnection.
This commit is contained in:
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
|
||||||
|
|||||||
Reference in New Issue
Block a user