1 files changed, 45 insertions, 16 deletions

diff --git a/src/transport/tor.rs b/src/transport/tor.rs
index 643e628..d049453 100644
--- a/src/transport/tor.rs
+++ b/src/transport/tor.rs
@@ -46,7 +46,13 @@ use tracing::{debug, info, warn};
 use crate::identity::NodeKey;
 use crate::locks::MutexExt;
 use crate::transport::challenge::{self, Role};
-use crate::transport::{EncodedMessage, Hex, MessageHandler, ParsedMessage, Transport, wire};
+use crate::transport::{
+    EncodedMessage, HandlerFn, Hex, MessageHandler, ParsedMessage, Transport, hand_on_probe,
+    is_probe, wire,
+};
+
+/// A message read, with the sender's key and its payload length.
+type Inbound = (String, ParsedMessage, usize);
 
 /// The onion service's nickname in arti's state (one per node).
 const NICKNAME: &str = "continuum";
@@ -124,8 +130,10 @@ struct Shared {
     /// One dial at a time per location.
     dialing: Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
     next_connection: AtomicU64,
-    /// Messages read, with the sender's key, for the consumer.
-    inbound: mpsc::Sender<(String, Vec<u8>)>,
+    /// Messages read, for the consumer: in order, and PINGs and PONGs apart, to be
+    /// handed on at once (`is_probe`).
+    inbound: mpsc::Sender<Inbound>,
+    probes: mpsc::Sender<Inbound>,
 }
 
 fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
@@ -139,7 +147,8 @@ pub(crate) struct TorTransport {
     onion: String,
     /// Kept: dropping it stops the onion service.
     _service: Arc<RunningOnionService>,
-    incoming: tokio::sync::Mutex<Option<mpsc::Receiver<(String, Vec<u8>)>>>,
+    /// Messages, and PINGs and PONGs (`Shared::inbound`, `Shared::probes`).
+    incoming: tokio::sync::Mutex<Option<(mpsc::Receiver<Inbound>, mpsc::Receiver<Inbound>)>>,
     stopped: watch::Receiver<bool>,
 }
 
@@ -177,6 +186,7 @@ impl TorTransport {
 
         let local_id = Hex(&config.key.public()).to_string();
         let (inbound, incoming) = mpsc::channel(READS);
+        let (probes, incoming_probes) = mpsc::channel(READS);
         let shared = Arc::new(Shared {
             key: config.key,
             local_id,
@@ -185,6 +195,7 @@ impl TorTransport {
             dialing: Mutex::default(),
             next_connection: AtomicU64::new(0),
             inbound,
+            probes,
         });
         let (stopping, stopped) = watch::channel(false);
         {
@@ -209,7 +220,7 @@ impl TorTransport {
             client,
             onion,
             _service: service,
-            incoming: tokio::sync::Mutex::new(Some(incoming)),
+            incoming: tokio::sync::Mutex::new(Some((incoming, incoming_probes))),
             stopped,
         })
     }
@@ -366,7 +377,23 @@ impl Shared {
                     );
                 }
             }
-            if self.inbound.send((key.to_owned(), body)).await.is_err() {
+            let message = match wire::decode(&body) {
+                Ok(message) => message,
+                Err(error) => {
+                    debug!(peer = %key, error = %error, "Dropping a message that doesn't decode");
+                    continue;
+                }
+            };
+            let channel = if is_probe(&message) {
+                &self.probes
+            } else {
+                &self.inbound
+            };
+            if channel
+                .send((key.to_owned(), message, body.len()))
+                .await
+                .is_err()
+            {
                 return anyhow!("the transport stopped");
             }
         }
@@ -494,19 +521,21 @@ impl Transport for TorTransport {
     /// Handles the messages read from every connection, one at a time, in the order
     /// they were read.
     async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
-        let Some(mut incoming) = self.incoming.lock().await.take() else {
+        let Some((mut incoming, mut probes)) = self.incoming.lock().await.take() else {
             bail!("Tor consumer already started");
         };
-        while let Some((peer, body)) = incoming.recv().await {
-            match wire::decode(&body) {
-                Ok(message) => {
-                    if let Err(error) = handler(peer.clone(), 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");
+        let handler: Arc<HandlerFn> = Arc::from(handler);
+        {
+            let handler = Arc::clone(&handler);
+            tokio::spawn(async move {
+                while let Some((peer, message, payload_len)) = probes.recv().await {
+                    hand_on_probe(&handler, peer, message, payload_len);
                 }
+            });
+        }
+        while let Some((peer, message, payload_len)) = incoming.recv().await {
+            if let Err(error) = handler(peer.clone(), message, payload_len).await {
+                warn!(peer = %peer, error = %format!("{error:#}"), "Failed to process inbound message");
             }
         }
         Ok(())