varun-r-mallya/py-libp2p
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/py-libp2p.git synced 2026-02-12 16:10:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix ping protocol

Browse Source
This commit is contained in:
mhchia
2019-12-01 19:17:32 +08:00
parent 6149aacc01
commit eb494e8682
4 changed files with 29 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
libp2p/host/ping.py
View File

@ -1,4 +1,4 @@
import asyncio import trio
import logging import logging
from libp2p.network.stream.exceptions import StreamClosed, StreamEOF, StreamReset from libp2p.network.stream.exceptions import StreamClosed, StreamEOF, StreamReset
@ -16,22 +16,23 @@ logger = logging.getLogger("libp2p.host.ping")
async def _handle_ping(stream: INetStream, peer_id: PeerID) -> bool: async def _handle_ping(stream: INetStream, peer_id: PeerID) -> bool:
"""Return a boolean indicating if we expect more pings from the peer at """Return a boolean indicating if we expect more pings from the peer at
``peer_id``.""" ``peer_id``."""
try: with trio.fail_after(RESP_TIMEOUT):
payload = await asyncio.wait_for(stream.read(PING_LENGTH), RESP_TIMEOUT) try:
except asyncio.TimeoutError as error: payload = await stream.read(PING_LENGTH)
logger.debug("Timed out waiting for ping from %s: %s", peer_id, error) except trio.TooSlowError as error:
raise logger.debug("Timed out waiting for ping from %s: %s", peer_id, error)
except StreamEOF: raise
logger.debug("Other side closed while waiting for ping from %s", peer_id) except StreamEOF:
return False logger.debug("Other side closed while waiting for ping from %s", peer_id)
except StreamReset as error: return False
logger.debug( except StreamReset as error:
"Other side reset while waiting for ping from %s: %s", peer_id, error logger.debug(
) "Other side reset while waiting for ping from %s: %s", peer_id, error
raise )
except Exception as error: raise
logger.debug("Error while waiting to read ping for %s: %s", peer_id, error) except Exception as error:
raise logger.debug("Error while waiting to read ping for %s: %s", peer_id, error)
raise
logger.debug("Received ping from %s with data: 0x%s", peer_id, payload.hex()) logger.debug("Received ping from %s with data: 0x%s", peer_id, payload.hex())

1
libp2p/stream_muxer/mplex/mplex.py
View File

@ -1,4 +1,3 @@
import asyncio
import logging import logging
import math import math
from typing import Any # noqa: F401 from typing import Any # noqa: F401

3
libp2p/stream_muxer/mplex/mplex_stream.py
View File

@ -104,6 +104,7 @@ class MplexStream(IMuxedStream):
# and then return. # and then return.
try: try:
data = self.incoming_data_channel.receive_nowait() data = self.incoming_data_channel.receive_nowait()
self._buf.extend(data)
except trio.EndOfChannel: except trio.EndOfChannel:
raise MplexStreamEOF raise MplexStreamEOF
except trio.WouldBlock: except trio.WouldBlock:
@ -111,6 +112,7 @@ class MplexStream(IMuxedStream):
# catch all kinds of errors here. # catch all kinds of errors here.
try: try:
data = await self.incoming_data_channel.receive() data = await self.incoming_data_channel.receive()
self._buf.extend(data)
except trio.EndOfChannel: except trio.EndOfChannel:
if self.event_reset.is_set(): if self.event_reset.is_set():
raise MplexStreamReset raise MplexStreamReset
@ -125,7 +127,6 @@ class MplexStream(IMuxedStream):
"`incoming_data_channel` is closed but stream is not reset. " "`incoming_data_channel` is closed but stream is not reset. "
"This should never happen." "This should never happen."
) from error ) from error
self._buf.extend(data)
self._buf.extend(self._read_return_when_blocked()) self._buf.extend(self._read_return_when_blocked())
payload = self._buf[:n] payload = self._buf[:n]
self._buf = self._buf[len(payload) :] self._buf = self._buf[len(payload) :]

17
tests/host/test_ping.py
View File

@ -1,4 +1,4 @@
import asyncio import trio
import secrets import secrets
import pytest import pytest
@ -7,12 +7,13 @@ from libp2p.host.ping import ID, PING_LENGTH
from libp2p.tools.factories import host_pair_factory from libp2p.tools.factories import host_pair_factory
@pytest.mark.asyncio @pytest.mark.trio
async def test_ping_once(): async def test_ping_once(is_host_secure):
async with host_pair_factory() as (host_a, host_b): async with host_pair_factory(is_host_secure) as (host_a, host_b):
stream = await host_b.new_stream(host_a.get_id(), (ID,)) stream = await host_b.new_stream(host_a.get_id(), (ID,))
some_ping = secrets.token_bytes(PING_LENGTH) some_ping = secrets.token_bytes(PING_LENGTH)
await stream.write(some_ping) await stream.write(some_ping)
await trio.sleep(0.01)
some_pong = await stream.read(PING_LENGTH) some_pong = await stream.read(PING_LENGTH)
assert some_ping == some_pong assert some_ping == some_pong
await stream.close() await stream.close()
@ -21,9 +22,9 @@ async def test_ping_once():
SOME_PING_COUNT = 3 SOME_PING_COUNT = 3
@pytest.mark.asyncio @pytest.mark.trio
async def test_ping_several(): async def test_ping_several(is_host_secure):
async with host_pair_factory() as (host_a, host_b): async with host_pair_factory(is_host_secure) as (host_a, host_b):
stream = await host_b.new_stream(host_a.get_id(), (ID,)) stream = await host_b.new_stream(host_a.get_id(), (ID,))
for _ in range(SOME_PING_COUNT): for _ in range(SOME_PING_COUNT):
some_ping = secrets.token_bytes(PING_LENGTH) some_ping = secrets.token_bytes(PING_LENGTH)
@ -33,5 +34,5 @@ async def test_ping_several():
# NOTE: simulate some time to sleep to mirror a real # NOTE: simulate some time to sleep to mirror a real
# world usage where a peer sends pings on some periodic interval # world usage where a peer sends pings on some periodic interval
# NOTE: this interval can be `0` for this test. # NOTE: this interval can be `0` for this test.
await asyncio.sleep(0) await trio.sleep(0)
await stream.close() await stream.close()