move keepalive and firewall to paths

5 files changed, 231 insertions, 138 deletions

diff --git a/src/transport/firewall.rs b/src/transport/firewall.rs
index 846ae1b..7b0ec0d 100644
--- a/src/transport/firewall.rs
+++ b/src/transport/firewall.rs
@@ -1,13 +1,14 @@
+//! The debug firewall (`/debug/firewall`, with `--debug-api`): drops messages from and
+//! to peers, by node ID, to test partitions on a real swarm. It works at the bottom of
+//! the node, in the paths layer (`transport::paths::Gate`), so it drops everything a
+//! partition would: WHO and ME, PINGs and their PONGs, as well as the node's messages.
+
 use std::collections::HashSet;
-use std::sync::{Arc, RwLock, RwLockReadGuard};
+use std::sync::{RwLock, RwLockReadGuard};
 
-use anyhow::Result;
-use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 
 use crate::locks::RwLockExt;
-use crate::transport::EncodedMessage;
-use crate::transport::{MessageHandler, ParsedMessage, Transport};
 
 #[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub(crate) struct FirewallRules {
@@ -46,99 +47,3 @@ impl Firewall {
         self.read().outbound.contains(peer)
     }
 }
-
-pub(crate) struct FirewallTransport {
-    inner: Arc<dyn Transport>,
-    firewall: Arc<Firewall>,
-}
-
-impl FirewallTransport {
-    pub(crate) fn new(inner: Arc<dyn Transport>, firewall: Arc<Firewall>) -> Self {
-        Self { inner, firewall }
-    }
-}
-
-#[async_trait]
-impl Transport for FirewallTransport {
-    fn local_id(&self) -> &str {
-        self.inner.local_id()
-    }
-
-    fn is_peer_id(&self, id: &str) -> bool {
-        self.inner.is_peer_id(id)
-    }
-
-    fn delivery_limit(&self) -> std::time::Duration {
-        self.inner.delivery_limit()
-    }
-
-    fn local_addresses(&self) -> Vec<String> {
-        self.inner.local_addresses()
-    }
-
-    fn addresses_of(&self, peer: &str) -> Vec<String> {
-        self.inner.addresses_of(peer)
-    }
-
-    fn learn_addresses(&self, peer: &str, addresses: &[String]) {
-        self.inner.learn_addresses(peer, addresses)
-    }
-
-    fn can_reach(&self, peer: &str) -> bool {
-        self.inner.can_reach(peer)
-    }
-
-    fn discovered_seeds(&self) -> Vec<String> {
-        self.inner.discovered_seeds()
-    }
-
-    fn transports(&self) -> Vec<String> {
-        self.inner.transports()
-    }
-
-    fn destination(&self, peer: &str) -> Option<crate::transport::Destination> {
-        self.inner.destination(peer)
-    }
-
-    fn connection(&self, target: &str) -> Option<crate::transport::Connection> {
-        self.inner.connection(target)
-    }
-
-    fn stats(&self) -> crate::transport::stats::TransportStats {
-        self.inner.stats()
-    }
-
-    async fn resolve_seed(&self, seed: &str) -> Result<String> {
-        self.inner.resolve_seed(seed).await
-    }
-
-    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
-        self.inner.encode(peer, message)
-    }
-
-    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
-        if self.firewall.blocks_outbound(&message.peer) {
-            tracing::trace!(peer = %message.peer, "Firewall dropped outbound message");
-            return Ok(());
-        }
-        self.inner.send_encoded(message).await
-    }
-
-    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
-        let firewall = Arc::clone(&self.firewall);
-
-        self.inner
-            .start_consumer(Box::new(move |sender, message, payload_len| {
-                if firewall.blocks_inbound(&sender) {
-                    tracing::trace!(peer = %sender, "Firewall dropped inbound message");
-                    return Box::pin(async { Ok(()) });
-                }
-                handler(sender, message, payload_len)
-            }))
-            .await
-    }
-
-    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
-        self.inner.wait_for_exit().await
-    }
-}
diff --git a/src/transport/identity.rs b/src/transport/identity.rs
index 5857fb0..153dacc 100644
--- a/src/transport/identity.rs
+++ b/src/transport/identity.rs
@@ -37,7 +37,8 @@ use tracing::{debug, info};
 
 use crate::identity::{self, NodeKey};
 use crate::locks::MutexExt;
-use crate::transport::paths::{PathHandler, Paths, scheme_of};
+use crate::transport::firewall::Firewall;
+use crate::transport::paths::{Gate, PathHandler, Paths, scheme_of};
 use crate::transport::stats::{IdentityCounts, IdentityStats, TransportStats};
 use crate::transport::traffic::Traffic;
 use crate::transport::{
@@ -96,6 +97,8 @@ pub(crate) struct IdentityTransport {
 
 struct Shared {
     paths: Arc<Paths>,
+    /// The debug firewall, by node ID; paths applies it (`Gate`).
+    firewall: Option<Arc<Firewall>>,
     key: Arc<NodeKey>,
     node_id: String,
     /// WHO and ME are counted here, under the address (`traffic`).
@@ -208,27 +211,54 @@ impl Book {
 
 impl IdentityTransport {
     /// Over `transport`, whose addresses have `config.scheme`, and `config.more`.
+    /// `firewall`: the debug firewall, if any, which paths applies by address.
     pub(crate) fn new(
         transport: Arc<dyn Transport>,
         config: IdentityConfig,
         traffic: Arc<Traffic>,
+        firewall: Option<Arc<Firewall>>,
     ) -> Self {
         let node_id = config.key.id();
-        let paths = Arc::new(Paths::new(
-            std::iter::once((config.scheme, transport))
-                .chain(config.more)
-                .collect(),
-            Some(Arc::clone(&traffic)),
-        ));
-        Self {
-            shared: Arc::new(Shared {
-                paths,
+        let shared = Arc::new_cyclic(|shared: &std::sync::Weak<Shared>| {
+            let gate: std::sync::Weak<dyn Gate> = shared.clone();
+            Shared {
+                paths: Arc::new(Paths::new(
+                    std::iter::once((config.scheme, transport))
+                        .chain(config.more)
+                        .collect(),
+                    Some(Arc::clone(&traffic)),
+                    firewall.is_some().then_some(gate),
+                )),
+                firewall,
                 key: config.key,
                 node_id,
                 traffic,
                 book: Mutex::default(),
-            }),
-        }
+            }
+        });
+        Self { shared }
+    }
+}
+
+/// The firewall's rules are by node ID: an address is blocked if the node proven at it
+/// is. Addresses no node proved itself at aren't (nothing could name them).
+impl Gate for Shared {
+    fn blocks_inbound(&self, address: &str) -> bool {
+        self.blocks(address, Firewall::blocks_inbound)
+    }
+
+    fn blocks_outbound(&self, address: &str) -> bool {
+        self.blocks(address, Firewall::blocks_outbound)
+    }
+}
+
+impl Shared {
+    fn blocks(&self, address: &str, rule: fn(&Firewall, &str) -> bool) -> bool {
+        let Some(firewall) = &self.firewall else {
+            return false;
+        };
+        let node = self.lock().at.get(address).map(|(node, _)| node.clone());
+        node.is_some_and(|node| rule(firewall, &node))
     }
 }
 
@@ -461,6 +491,11 @@ impl Transport for IdentityTransport {
         })
     }
 
+    /// The paths layer answers them (`Paths::receive`).
+    fn answers_pings(&self) -> bool {
+        true
+    }
+
     /// The identity layer's, and each transport's.
     fn stats(&self) -> TransportStats {
         let mut stats = self.shared.paths.stats();
diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index 147e560..f60a9f6 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -649,6 +649,13 @@ pub(crate) trait Transport: Send + Sync {
         None
     }
 
+    /// Whether the transport answers PINGs with PONGs itself (the paths layer,
+    /// `transport::paths`); if not (a bare transport, in tests), the node does. By
+    /// default not.
+    fn answers_pings(&self) -> bool {
+        false
+    }
+
     /// Figures about the transport itself (`stats`). By default none.
     fn stats(&self) -> stats::TransportStats {
         stats::TransportStats::default()
diff --git a/src/transport/paths.rs b/src/transport/paths.rs
index 30aab8d..dda2d87 100644
--- a/src/transport/paths.rs
+++ b/src/transport/paths.rs
@@ -17,16 +17,24 @@
 //!
 //! Paths measures the addresses of such nodes itself: every `PROBE_EVERY`, each address
 //! it had a choice of within `ACTIVE_FOR` gets a PING with a random nonce, which the
-//! other side answers with a PONG at once, over the same path. Answers give each
+//! other side's paths answers at once with a PONG, over the same path. Answers give each
 //! path's round-trip time (smoothed); `UNANSWERED_DOWN` PINGs in a row unanswered, or a
 //! send that failed within `DOWN_FOR`, make a path down. `rank` puts paths that are up
 //! first, the one in use first among them, and the rest by round-trip time; it changes
 //! the one in use when another is faster by `SWITCH_GAIN` (a quarter) and at least
 //! `SWITCH_MIN` (both measured), so it doesn't flap. Until both are measured, the
-//! configured order decides. PING and PONG end here: the identity layer never sees them.
+//! configured order decides.
+//!
+//! PINGs and PONGs then go up as other messages do: they tell the node the peer is
+//! alive. The node's keepalive is a PING too (`swarm::peer::keepalive`), which paths
+//! sends like any message, and notes as the probe of the path it went over.
+//!
+//! The debug firewall (`transport::firewall`) works here, below everything else, through
+//! a `Gate`: messages from an address it blocks are dropped as they arrive, and sends to
+//! one are dropped silently, as by a lossy network.
 
 use std::collections::HashMap;
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, Weak};
 use std::time::Duration;
 
 use anyhow::{Context, Result, anyhow};
@@ -36,7 +44,14 @@ use tracing::{debug, info};
 use crate::locks::MutexExt;
 use crate::transport::stats::{PathCounts, PathInfo, PathsStats, TransportStats};
 use crate::transport::traffic::Traffic;
-use crate::transport::{Connection, EncodedMessage, MessageKind, ParsedMessage, Transport};
+use crate::transport::{Connection, EncodedMessage, ParsedMessage, Transport};
+
+/// Whether to drop what comes from or goes to an address (the debug firewall,
+/// `transport::firewall`; the identity layer knows which node is where).
+pub(crate) trait Gate: Send + Sync {
+    fn blocks_inbound(&self, address: &str) -> bool;
+    fn blocks_outbound(&self, address: &str) -> bool;
+}
 
 /// Called with the sender's address, the message and its `payload_len`.
 pub(crate) type PathHandler = dyn Fn(String, ParsedMessage, usize) -> std::pin::Pin<Box<dyn Future<Output = Result<()>> + Send>>
@@ -250,6 +265,16 @@ impl Table {
         due
     }
 
+    /// The node sent a PING with `nonce` to `address` (its keepalive): the probe of that
+    /// path, if paths measures it.
+    fn pinged(&mut self, address: &str, nonce: u64, now: Instant) {
+        if let Some(path) = self.paths.get_mut(address) {
+            path.probe = Some((nonce, now));
+            path.last_probe = Some(now);
+            self.counts.probes_sent += 1;
+        }
+    }
+
     /// A PONG with `nonce` from `address`: the round trip, if it answers our PING out.
     fn pong(&mut self, address: &str, nonce: u64, now: Instant) -> Option<Duration> {
         let path = self.paths.get_mut(address)?;
@@ -298,18 +323,21 @@ impl Table {
 pub(crate) struct Paths {
     /// In order of preference.
     links: Vec<Link>,
-    /// PINGs and PONGs are counted here, under the address (`traffic`).
+    /// Its own PINGs and PONGs are counted here, under the address (`traffic`).
     traffic: Option<Arc<Traffic>>,
+    /// The debug firewall, if any.
+    gate: Option<Weak<dyn Gate>>,
     /// A std mutex: never held across an await.
     table: Mutex<Table>,
 }
 
 impl Paths {
-    /// Over `transports`, by scheme, in order of preference; PINGs and PONGs counted in
-    /// `traffic`.
+    /// Over `transports`, by scheme, in order of preference; its own PINGs and PONGs
+    /// counted in `traffic`, the firewall through `gate`.
     pub(crate) fn new(
         transports: Vec<(&'static str, Arc<dyn Transport>)>,
         traffic: Option<Arc<Traffic>>,
+        gate: Option<Weak<dyn Gate>>,
     ) -> Self {
         Self {
             links: transports
@@ -317,10 +345,15 @@ impl Paths {
                 .map(|(scheme, inner)| Link { scheme, inner })
                 .collect(),
             traffic,
+            gate,
             table: Mutex::default(),
         }
     }
 
+    fn gate(&self) -> Option<Arc<dyn Gate>> {
+        self.gate.as_ref()?.upgrade()
+    }
+
     /// The transport `address` is on, and the transport identity at it
     /// (`Transport::locate`, which notes where that is if the address says); None if no
     /// transport of ours has its scheme, or it isn't one of that transport's.
@@ -519,6 +552,9 @@ impl Paths {
                 Ok(()) => {
                     let mut table = self.table.lock_or_recover();
                     table.sent(address);
+                    if let Some(ParsedMessage::Ping(nonce)) = message.message.as_deref() {
+                        table.pinged(address, *nonce, Instant::now());
+                    }
                     if !failures.is_empty() {
                         table.counts.failovers += 1;
                         drop(table);
@@ -545,8 +581,15 @@ impl Paths {
         }
     }
 
-    /// Sends `message` to its `address`.
+    /// Sends `message` to its `address` (dropped silently if the firewall blocks it).
     pub(crate) async fn send(&self, message: &EncodedMessage) -> Result<()> {
+        if self
+            .gate()
+            .is_some_and(|gate| gate.blocks_outbound(&message.address))
+        {
+            tracing::trace!(to = %message.address, "Firewall dropped outbound message");
+            return Ok(());
+        }
         let (link, target) = self
             .split(&message.address)
             .with_context(|| format!("no transport for {}", message.address))?;
@@ -583,7 +626,8 @@ impl Paths {
         stats
     }
 
-    /// Sends a PING or PONG with `nonce` to `address`, counted (`traffic`).
+    /// Sends a PING (or with `pong`, a PONG) with `nonce` to `address`, counted
+    /// (`traffic`).
     async fn send_probe(&self, address: &str, pong: bool, nonce: u64) -> Result<()> {
         let message = if pong {
             ParsedMessage::Pong(nonce)
@@ -622,8 +666,8 @@ impl Paths {
         }
     }
 
-    /// A message from `address`: a PING is answered, a PONG noted; anything else goes to
-    /// `handler`.
+    /// A message from `address`, handed on to `handler`: a PING answered and a PONG
+    /// noted first; dropped if the firewall blocks it.
     async fn receive(
         self: Arc<Self>,
         address: String,
@@ -631,29 +675,31 @@ impl Paths {
         payload_len: usize,
         handler: &PathHandler,
     ) -> Result<()> {
+        if self
+            .gate()
+            .is_some_and(|gate| gate.blocks_inbound(&address))
+        {
+            tracing::trace!(from = %address, "Firewall dropped inbound message");
+            return Ok(());
+        }
         match message {
             ParsedMessage::Ping(nonce) => {
-                if let Some(traffic) = &self.traffic {
-                    traffic.received(&address, MessageKind::Ping, payload_len);
-                }
+                let paths = Arc::clone(&self);
+                let to = address.clone();
                 tokio::spawn(async move {
-                    if let Err(error) = self.send_probe(&address, true, nonce).await {
-                        debug!(to = %address, error = %format!("{error:#}"), "Couldn't answer a PING");
+                    if let Err(error) = paths.send_probe(&to, true, nonce).await {
+                        debug!(to = %to, error = %format!("{error:#}"), "Couldn't answer a PING");
                     }
                 });
-                Ok(())
             }
             ParsedMessage::Pong(nonce) => {
-                if let Some(traffic) = &self.traffic {
-                    traffic.received(&address, MessageKind::Pong, payload_len);
-                }
                 self.table
                     .lock_or_recover()
                     .pong(&address, nonce, Instant::now());
-                Ok(())
             }
-            message => handler(address, message, payload_len).await,
+            _ => {}
         }
+        handler(address, message, payload_len).await
     }
 
     /// Every transport's consumer, each message handed to `handler` under the sender's
@@ -712,7 +758,7 @@ impl Paths {
 }
 
 /// A PING's nonce: random, so a PONG can't be made up for a PING not seen.
-fn random_nonce() -> u64 {
+pub(crate) fn random_nonce() -> u64 {
     let mut bytes = [0u8; 8];
     if getrandom::fill(&mut bytes).is_err() {
         // No randomness: a nonce still tells answers apart, if not unguessably.
@@ -764,6 +810,7 @@ mod tests {
                     ("other", second.transport(id) as Arc<dyn Transport>),
                 ],
                 None,
+                None,
             ))
         };
         (on_both("a"), on_both("b"))
@@ -896,6 +943,99 @@ mod tests {
         assert!(table.paths.is_empty() && table.chosen.is_empty());
     }
 
+    /// The node's keepalive PING on a path is its probe: paths doesn't ping it too, and
+    /// the PONG measures it.
+    #[test]
+    fn keepalive_pings_are_probes() {
+        let mut table = Table::default();
+        let now = Instant::now();
+        let mut nonce = 100;
+        let mut next = || {
+            nonce += 1;
+            nonce
+        };
+        table.rank(paths(&["libp2p:x", "mycelium:x"]), now);
+        table.pinged("libp2p:x", 7, now);
+        let due = table.due(now, &mut next);
+        assert_eq!(due.len(), 1);
+        assert_eq!(due[0].0, "mycelium:x");
+        assert!(
+            table
+                .pong("libp2p:x", 7, now + Duration::from_millis(20))
+                .is_some()
+        );
+        assert_eq!(table.counts.probes_sent, 2);
+        // A peer with one path isn't measured.
+        table.pinged("libp2p:z", 8, now);
+        assert!(!table.paths.contains_key("libp2p:z"));
+    }
+
+    /// Blocks one address, both ways.
+    struct Blocking(&'static str);
+
+    impl Gate for Blocking {
+        fn blocks_inbound(&self, address: &str) -> bool {
+            address == self.0
+        }
+        fn blocks_outbound(&self, address: &str) -> bool {
+            address == self.0
+        }
+    }
+
+    /// The firewall drops what comes from a blocked address before anything sees it
+    /// (a PING isn't answered), and sends to one silently.
+    #[tokio::test]
+    async fn the_firewall_drops_both_ways() {
+        let net = SimNet::new(4);
+        let gate: Arc<dyn Gate> = Arc::new(Blocking("sim:a"));
+        let b = Arc::new(Paths::new(
+            vec![("sim", net.transport("b") as Arc<dyn Transport>)],
+            None,
+            Some(Arc::downgrade(&gate)),
+        ));
+        let a = Arc::new(Paths::new(
+            vec![("sim", net.transport("a") as Arc<dyn Transport>)],
+            None,
+            None,
+        ));
+        let received = |paths: &Arc<Paths>| {
+            let (sender, received) = tokio::sync::mpsc::unbounded_channel();
+            let handler: Arc<PathHandler> = Arc::new(move |from, message, _| {
+                let sender = sender.clone();
+                Box::pin(async move {
+                    let _ = sender.send((from, message.kind()));
+                    Ok(())
+                })
+            });
+            let paths = Arc::clone(paths);
+            tokio::spawn(async move { paths.start_consumer(handler).await });
+            received
+        };
+        let (mut at_a, mut at_b) = (received(&a), received(&b));
+
+        // a's PING never reaches b, so nothing answers it.
+        let ping = a.encode("sim:b", &ParsedMessage::Ping(5)).unwrap();
+        a.send(&ping).await.unwrap();
+        // b's send to a "works", and goes nowhere.
+        let status = b.encode("sim:a", &some_status()).unwrap();
+        b.send(&status).await.unwrap();
+        tokio::time::sleep(Duration::from_millis(200)).await;
+        assert!(at_b.try_recv().is_err());
+        assert!(at_a.try_recv().is_err());
+
+        // Without the gate: a PING is answered by b's paths, and goes on up there.
+        drop(gate);
+        a.send(&ping).await.unwrap();
+        assert_eq!(
+            at_b.recv().await.unwrap(),
+            ("sim:a".to_owned(), crate::transport::MessageKind::Ping)
+        );
+        assert_eq!(
+            at_a.recv().await.unwrap(),
+            ("sim:b".to_owned(), crate::transport::MessageKind::Pong)
+        );
+    }
+
     /// A transport whose sends all fail.
     struct Broken;
 
@@ -948,10 +1088,12 @@ mod tests {
                 ("sim", net.transport("a") as Arc<dyn Transport>),
             ],
             None,
+            None,
         );
         let b = Arc::new(Paths::new(
             vec![("sim", net.transport("b") as Arc<dyn Transport>)],
             None,
+            None,
         ));
         let (sender, mut received) = tokio::sync::mpsc::unbounded_channel();
         let handler: Arc<PathHandler> = Arc::new(move |from, _, _| {
diff --git a/src/transport/traffic.rs b/src/transport/traffic.rs
index 9b5f9c9..b070023 100644
--- a/src/transport/traffic.rs
+++ b/src/transport/traffic.rs
@@ -5,8 +5,8 @@
 //! (for Mycelium: the topic, or the base64 of its HTTP API). Counters only grow (from
 //! the start of the process); rates are for whoever reads them.
 //!
-//! A message dropped by the debug firewall still counts as sent (the firewall plays a
-//! lossy network) but not as received.
+//! A message dropped by the debug firewall (in the paths layer, below) still counts as
+//! sent (the firewall plays a lossy network) but not as received.
 
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex, MutexGuard};
@@ -188,6 +188,10 @@ impl Transport for CountingTransport {
         self.inner.stats()
     }
 
+    fn answers_pings(&self) -> bool {
+        self.inner.answers_pings()
+    }
+
     async fn resolve_seed(&self, seed: &str) -> Result<String> {
         self.inner.resolve_seed(seed).await
     }