Fix ping protocol

libp2p/host/ping.py

@@ -1,4 +1,4 @@
-import asyncio
+import trio
 import logging
 
 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:
     """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
+    with trio.fail_after(RESP_TIMEOUT):
+        try:
+            payload = await stream.read(PING_LENGTH)
+        except trio.TooSlowError 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())
 

libp2p/stream_muxer/mplex/mplex.py

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

libp2p/stream_muxer/mplex/mplex_stream.py

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