diff options
| -rw-r--r-- | src/transport/libp2p.rs | 37 | ||||
| -rw-r--r-- | src/transport/mod.rs | 22 | ||||
| -rw-r--r-- | src/transport/mycelium/mod.rs | 10 | ||||
| -rw-r--r-- | src/transport/tor.rs | 61 |
4 files changed, 96 insertions, 34 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(()) diff --git a/src/transport/mod.rs b/src/transport/mod.rs index 50841e0..7c31dc6 100644 --- a/src/transport/mod.rs +++ b/src/transport/mod.rs @@ -569,6 +569,28 @@ pub(crate) type HandlerFn = dyn Fn(String, ParsedMessage, usize) -> Pin<Box<dyn + 'static; pub(crate) type MessageHandler = Box<HandlerFn>; +/// Whether `message` is a PING or PONG. The paths layer times them, so a transport's +/// consumer hands them on as soon as they are read (`hand_on_probe`), not after the +/// messages read before them are handled (which can take a while: sync). +pub(crate) fn is_probe(message: &ParsedMessage) -> bool { + matches!(message, ParsedMessage::Ping(_) | ParsedMessage::Pong(_)) +} + +/// Hands a PING or PONG (`is_probe`) to `handler` in a task of its own. +pub(crate) fn hand_on_probe( + handler: &std::sync::Arc<HandlerFn>, + sender: String, + message: ParsedMessage, + payload_len: usize, +) { + let handler = std::sync::Arc::clone(handler); + tokio::spawn(async move { + if let Err(error) = handler(sender.clone(), message, payload_len).await { + tracing::warn!(peer = %sender, error = %format!("{error:#}"), "Failed to process inbound message"); + } + }); +} + #[async_trait] pub(crate) trait Transport: Send + Sync { /// This node's identity on the transport (for Mycelium: its public key). diff --git a/src/transport/mycelium/mod.rs b/src/transport/mycelium/mod.rs index 61aab85..2bb399c 100644 --- a/src/transport/mycelium/mod.rs +++ b/src/transport/mycelium/mod.rs @@ -8,7 +8,9 @@ use thiserror::Error; use url::Url; use crate::transport::wire; -use crate::transport::{EncodedMessage, MessageHandler, ParsedMessage, Transport}; +use crate::transport::{ + EncodedMessage, HandlerFn, MessageHandler, ParsedMessage, Transport, hand_on_probe, is_probe, +}; pub(crate) mod daemon; @@ -289,6 +291,8 @@ impl Transport for MyceliumTransport { } drop(tx); + // PINGs and PONGs aren't handled after the message before them (`is_probe`). + let handler: std::sync::Arc<HandlerFn> = std::sync::Arc::from(handler); while let Some(inbound_message) = rx.recv().await { match parse_message(&inbound_message) { Ok((sender, parsed)) => { @@ -300,7 +304,9 @@ impl Transport for MyceliumTransport { "Received message" ); let payload_len = base64_decoded_len(&inbound_message.payload); - if let Err(error) = handler(sender, parsed, payload_len).await { + if is_probe(&parsed) { + hand_on_probe(&handler, sender, parsed, payload_len); + } else if let Err(error) = handler(sender, parsed, payload_len).await { tracing::warn!(error = %format!("{error:#}"), "Failed to process inbound message"); } } 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(()) |