diff --git a/libp2p/network/connection/raw_connection.py b/libp2p/network/connection/raw_connection.py index fe09c6f..2bdb12a 100644 --- a/libp2p/network/connection/raw_connection.py +++ b/libp2p/network/connection/raw_connection.py @@ -1,5 +1,6 @@ import asyncio +from .exceptions import RawConnError from .raw_connection_interface import IRawConnection @@ -23,19 +24,28 @@ class RawConnection(IRawConnection): self._drain_lock = asyncio.Lock() async def write(self, data: bytes) -> None: - self.writer.write(data) + try: + self.writer.write(data) + except ConnectionResetError: + raise RawConnError() # Reference: https://github.com/ethereum/lahja/blob/93610b2eb46969ff1797e0748c7ac2595e130aef/lahja/asyncio/endpoint.py#L99-L102 # noqa: E501 # Use a lock to serialize drain() calls. Circumvents this bug: # https://bugs.python.org/issue29930 async with self._drain_lock: - await self.writer.drain() + try: + await self.writer.drain() + except ConnectionResetError: + raise RawConnError() async def read(self, n: int = -1) -> bytes: """ Read up to ``n`` bytes from the underlying stream. This call is delegated directly to the underlying ``self.reader``. """ - return await self.reader.read(n) + try: + return await self.reader.read(n) + except ConnectionResetError: + raise RawConnError() async def close(self) -> None: self.writer.close()