Merge pull request #315 from ralexstokes/add-ping-protocol

Add `ping` protocol.

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

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()