diff --git a/libp2p/host/ping.py b/libp2p/host/ping.py new file mode 100644 index 00000000..f47e83d8 --- /dev/null +++ b/libp2p/host/ping.py @@ -0,0 +1,56 @@ +import asyncio +import logging + +from libp2p.network.stream.exceptions import StreamEOF, StreamReset +from libp2p.network.stream.net_stream_interface import INetStream +from libp2p.peer.id import ID as PeerID + +ID = "/ipfs/ping/1.0.0" +PING_LENGTH = 32 +RESP_TIMEOUT = 60 + +logger = logging.getLogger("libp2p.host.ping") + + +async def _handle_ping(stream: INetStream, peer_id: PeerID) -> bool: + """ + Return a boolean indicating if we expect more pings from the peer at ``peer_id``. + """ + try: + payload = await asyncio.wait_for(stream.read(PING_LENGTH), RESP_TIMEOUT) + except asyncio.TimeoutError as error: + logger.debug("Timed out waiting for ping from %s: %s", peer_id, error) + raise + except StreamEOF: + logger.debug("Other side closed while waiting for ping from %s", peer_id) + return False + except StreamReset as error: + logger.debug( + "Other side reset while waiting for ping from %s: %s", peer_id, error + ) + raise + except Exception as error: + 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()) + + await stream.write(payload) + return True + + +async def handle_ping(stream: INetStream) -> None: + """ + ``handle_ping`` responds to incoming ping requests until one side + errors or closes the ``stream``. + """ + peer_id = stream.muxed_conn.peer_id + + while True: + try: + should_continue = await _handle_ping(stream, peer_id) + if not should_continue: + return + except Exception: + await stream.reset() + return diff --git a/tests/host/test_ping.py b/tests/host/test_ping.py new file mode 100644 index 00000000..d6f02449 --- /dev/null +++ b/tests/host/test_ping.py @@ -0,0 +1,70 @@ +import asyncio +import secrets + +import pytest + +from libp2p.host.ping import ID, PING_LENGTH, handle_ping +from libp2p.peer.peerinfo import info_from_p2p_addr +from tests.utils import set_up_nodes_by_transport_opt + + +def _add_ping_to(host): + host.set_stream_handler(ID, handle_ping) + + +@pytest.mark.asyncio +async def test_ping_once(): + transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]] + (host_a, host_b) = await set_up_nodes_by_transport_opt(transport_opt_list) + + # NOTE: this will fail after we add ping as a default handler + # as a forced reminder to fix this test by removing the calls to + # `_add_ping_to` + assert host_a.get_mux().handlers == {} + assert host_b.get_mux().handlers == {} + _add_ping_to(host_a) + _add_ping_to(host_b) + + addr = host_a.get_addrs()[0] + info = info_from_p2p_addr(addr) + await host_b.connect(info) + + stream = await host_b.new_stream(host_a.get_id(), (ID,)) + some_ping = secrets.token_bytes(PING_LENGTH) + await stream.write(some_ping) + some_pong = await stream.read(PING_LENGTH) + assert some_ping == some_pong + await stream.close() + + +SOME_PING_COUNT = 3 + + +@pytest.mark.asyncio +async def test_ping_several(): + transport_opt_list = [["/ip4/127.0.0.1/tcp/0"], ["/ip4/127.0.0.1/tcp/0"]] + (host_a, host_b) = await set_up_nodes_by_transport_opt(transport_opt_list) + + # NOTE: this will fail after we add ping as a default handler + # as a forced reminder to fix this test by removing the calls to + # `_add_ping_to` + assert host_a.get_mux().handlers == {} + assert host_b.get_mux().handlers == {} + _add_ping_to(host_a) + _add_ping_to(host_b) + + addr = host_a.get_addrs()[0] + info = info_from_p2p_addr(addr) + await host_b.connect(info) + + stream = await host_b.new_stream(host_a.get_id(), (ID,)) + for _ in range(SOME_PING_COUNT): + some_ping = secrets.token_bytes(PING_LENGTH) + await stream.write(some_ping) + some_pong = await stream.read(PING_LENGTH) + assert some_ping == some_pong + # NOTE: simulate some time to sleep to mirror a real + # world usage where a peer sends pings on some periodic interval + # NOTE: this interval can be `0` for this test. + await asyncio.sleep(0) + await stream.close()