1 files changed, 21 insertions, 16 deletions
diff --git a/src/transport/libp2p.rs b/src/transport/libp2p.rs
index 8655f23..76d07ff 100644
--- a/src/transport/libp2p.rs
+++ b/src/transport/libp2p.rs
@@ -62,7 +62,8 @@ use crate::transport::stats::{
     Connections, DialFailures, HolePunches, Libp2pStats, Reachability, RelaySlots, TransportStats,
 };
 use crate::transport::{
-    Connection, EncodedMessage, MessageHandler, ParsedMessage, Transport, wire,
+    Connection, EncodedMessage, HandlerFn, MessageHandler, ParsedMessage, Transport, hand_on_probe,
+    is_probe, wire,
 };
 
 const PROTOCOL: StreamProtocol = StreamProtocol::new("/continuum/1");
@@ -1385,24 +1386,35 @@ impl Transport for Libp2pTransport {
     }
 
     /// Reads inbound streams (`READS` at once) and handles their messages one at a
-    /// time, in the order they were read.
+    /// time, in the order they were read; PINGs and PONGs at once (`is_probe`).
     async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
         let Some(mut incoming) = self.incoming.lock().await.take() else {
             bail!("libp2p consumer already started");
         };
-        let (sender, mut received) = mpsc::channel::<(PeerId, Vec<u8>)>(READS);
+        let handler: Arc<HandlerFn> = Arc::from(handler);
+        let (sender, mut received) = mpsc::channel::<(PeerId, ParsedMessage, usize)>(READS);
         let reads = Arc::new(Semaphore::new(READS));
+        let probes = Arc::clone(&handler);
         tokio::spawn(async move {
             while let Some((peer, stream)) = incoming.next().await {
                 let Ok(permit) = Arc::clone(&reads).acquire_owned().await else {
                     break;
                 };
                 let sender = sender.clone();
+                let probes = Arc::clone(&probes);
                 tokio::spawn(async move {
                     match read_message(stream).await {
-                        Ok(body) => {
-                            let _ = sender.send((peer, body)).await;
-                        }
+                        Ok(body) => match wire::decode(&body) {
+                            Ok(message) if is_probe(&message) => {
+                                hand_on_probe(&probes, peer.to_string(), message, body.len());
+                            }
+                            Ok(message) => {
+                                let _ = sender.send((peer, message, body.len())).await;
+                            }
+                            Err(error) => {
+                                debug!(peer = %peer, error = %error, "Dropping a message that doesn't decode");
+                            }
+                        },
                         Err(error) => {
                             debug!(peer = %peer, error = %format!("{error:#}"), "Dropping an unreadable stream");
                         }
@@ -1411,16 +1423,9 @@ impl Transport for Libp2pTransport {
                 });
             }
         });
-        while let Some((peer, body)) = received.recv().await {
-            match wire::decode(&body) {
-                Ok(message) => {
-                    if let Err(error) = handler(peer.to_string(), message, body.len()).await {
-                        warn!(peer = %peer, error = %format!("{error:#}"), "Failed to process inbound message");
-                    }
-                }
-                Err(error) => {
-                    debug!(peer = %peer, error = %error, "Dropping a message that doesn't decode");
-                }
+        while let Some((peer, message, payload_len)) = received.recv().await {
+            if let Err(error) = handler(peer.to_string(), message, payload_len).await {
+                warn!(peer = %peer, error = %format!("{error:#}"), "Failed to process inbound message");
             }
         }
         Ok(())