1 files changed, 168 insertions, 13 deletions

diff --git a/src/transport/p2p.rs b/src/transport/p2p.rs
index db3628f..3023fb2 100644
--- a/src/transport/p2p.rs
+++ b/src/transport/p2p.rs
@@ -8,9 +8,14 @@
 //! `wire` encodes it. A send is done once the stream is written and closed; a connection
 //! that can't be made, or a stream that can't be written, refuses it.
 //!
-//! What this transport doesn't do yet (`docs/future-work.md`): learn addresses from the
-//! peers it talks to (`identify`), find peers on the LAN (mDNS), or reach nodes behind
-//! NAT (relays, hole punching). It never uses public libp2p infrastructure.
+//! Connected peers tell each other where they listen (libp2p's identify protocol), so a
+//! node learns the addresses of peers that dialled it. A node publishes where it can be
+//! reached (`local_addresses`): the addresses it announces (`--libp2p-announce`, e.g. a
+//! public address forwarded to it), or else those it listens on, loopback left out.
+//!
+//! What this transport doesn't do yet (`docs/future-work.md`): find peers on the LAN
+//! (mDNS), or reach nodes behind NAT (relays, hole punching). It never uses public
+//! libp2p infrastructure.
 
 use std::collections::HashMap;
 use std::path::Path;
@@ -60,6 +65,8 @@ pub(crate) struct Libp2pConfig {
     pub(crate) key: Keypair,
     /// Multiaddrs to listen on.
     pub(crate) listen: Vec<Multiaddr>,
+    /// Multiaddrs to publish instead of the listen addresses, if any.
+    pub(crate) announce: Vec<Multiaddr>,
     /// libp2p's in-memory transport instead of TCP and QUIC: for tests.
     pub(crate) memory: bool,
 }
@@ -106,6 +113,35 @@ enum Command {
 struct Book {
     addresses: HashMap<PeerId, Vec<Multiaddr>>,
     listening: Vec<Multiaddr>,
+    /// `Libp2pConfig::announce`.
+    announced: Vec<Multiaddr>,
+}
+
+/// Whether `address` only reaches this host (or this link): no use to other hosts, and
+/// dialled by one, it reaches that host itself.
+fn is_local(address: &Multiaddr) -> bool {
+    match address.iter().next() {
+        Some(Protocol::Ip4(ip)) => ip.is_loopback() || ip.is_link_local() || ip.is_unspecified(),
+        Some(Protocol::Ip6(ip)) => {
+            ip.is_loopback() || ip.is_unicast_link_local() || ip.is_unspecified()
+        }
+        _ => false,
+    }
+}
+
+/// The addresses of `all` worth publishing: all but local ones, unless they're all
+/// local (a test on one host).
+fn published(all: &[Multiaddr]) -> Vec<Multiaddr> {
+    let reachable: Vec<Multiaddr> = all
+        .iter()
+        .filter(|address| !is_local(address))
+        .cloned()
+        .collect();
+    if reachable.is_empty() {
+        all.to_vec()
+    } else {
+        reachable
+    }
 }
 
 impl Book {
@@ -160,11 +196,18 @@ impl Libp2pTransport {
                 .listen_on(address.clone())
                 .with_context(|| format!("Listening on {address}"))?;
         }
-        let mut control = swarm.behaviour().new_control();
+        for address in &config.announce {
+            // Identify tells peers these.
+            swarm.add_external_address(address.clone());
+        }
+        let mut control = swarm.behaviour().stream.new_control();
         let incoming = control
             .accept(PROTOCOL)
             .map_err(|error| anyhow!("Accepting {PROTOCOL}: {error}"))?;
-        let book = Arc::new(Mutex::new(Book::default()));
+        let book = Arc::new(Mutex::new(Book {
+            announced: config.announce,
+            ..Book::default()
+        }));
         let (commands, receiver) = mpsc::unbounded_channel();
         let (stopping, stopped) = watch::channel(false);
         tokio::spawn(run(swarm, receiver, Arc::clone(&book), stopping));
@@ -204,7 +247,25 @@ impl Libp2pTransport {
     }
 }
 
-type Behaviour = libp2p_stream::Behaviour;
+/// The protocol version identify says (informational: nothing checks it).
+const IDENTIFY_VERSION: &str = "/continuum/1";
+
+#[derive(::libp2p::swarm::NetworkBehaviour)]
+struct Behaviour {
+    stream: libp2p_stream::Behaviour,
+    identify: ::libp2p::identify::Behaviour,
+}
+
+impl Behaviour {
+    fn new(key: &Keypair) -> Self {
+        let config = ::libp2p::identify::Config::new(IDENTIFY_VERSION.to_owned(), key.public())
+            .with_agent_version(format!("continuum/{}", env!("CARGO_PKG_VERSION")));
+        Self {
+            stream: libp2p_stream::Behaviour::new(),
+            identify: ::libp2p::identify::Behaviour::new(config),
+        }
+    }
+}
 
 fn network_swarm(key: Keypair) -> Result<Swarm<Behaviour>> {
     Ok(::libp2p::SwarmBuilder::with_existing_identity(key)
@@ -216,7 +277,7 @@ fn network_swarm(key: Keypair) -> Result<Swarm<Behaviour>> {
         )?
         .with_quic()
         .with_dns()?
-        .with_behaviour(|_| Behaviour::new())?
+        .with_behaviour(Behaviour::new)?
         .with_swarm_config(|config| config.with_idle_connection_timeout(IDLE_TIMEOUT))
         .build())
 }
@@ -235,7 +296,7 @@ fn memory_swarm(key: Keypair) -> Result<Swarm<Behaviour>> {
         .boxed();
     Ok(Swarm::new(
         transport,
-        Behaviour::new(),
+        Behaviour::new(&key),
         peer_id,
         ::libp2p::swarm::Config::with_tokio_executor().with_idle_connection_timeout(IDLE_TIMEOUT),
     ))
@@ -327,6 +388,22 @@ async fn run(
                 SwarmEvent::ListenerError { error, .. } => {
                     warn!(error = %error, "libp2p listener failed");
                 }
+                SwarmEvent::Behaviour(BehaviourEvent::Identify(
+                    ::libp2p::identify::Event::Received { peer_id, info, .. },
+                )) => {
+                    debug!(
+                        peer = %peer_id,
+                        listen = ?info.listen_addrs,
+                        observed = %info.observed_addr,
+                        agent = %info.agent_version,
+                        "libp2p peer identified"
+                    );
+                    let mut book = lock(&book);
+                    // Last first: `remember` puts each in front.
+                    for address in published(&info.listen_addrs).into_iter().rev() {
+                        book.remember(peer_id, address);
+                    }
+                }
                 _ => {}
             },
         }
@@ -371,12 +448,18 @@ impl Transport for Libp2pTransport {
         stream_timeout(MAX_MESSAGE)
     }
 
-    /// Where we listen, each with our peer ID.
+    /// Where we can be reached, each with our peer ID: the announced addresses, or else
+    /// those we listen on (`published`).
     fn local_addresses(&self) -> Vec<String> {
-        lock(&self.book)
-            .listening
-            .iter()
-            .filter_map(|address| address.clone().with_p2p(self.peer_id).ok())
+        let book = lock(&self.book);
+        let addresses = if book.announced.is_empty() {
+            published(&book.listening)
+        } else {
+            book.announced.clone()
+        };
+        addresses
+            .into_iter()
+            .filter_map(|address| address.with_p2p(self.peer_id).ok())
             .map(|address| address.to_string())
             .collect()
     }
@@ -534,6 +617,7 @@ mod tests {
         let transport = Libp2pTransport::start(Libp2pConfig {
             key,
             listen: vec![address.clone()],
+            announce: Vec::new(),
             memory: true,
         })
         .unwrap();
@@ -620,6 +704,77 @@ mod tests {
         );
     }
 
+    /// Waits up to 5 s for `condition`.
+    async fn eventually(mut condition: impl FnMut() -> bool) -> bool {
+        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
+        while !condition() {
+            if tokio::time::Instant::now() >= deadline {
+                return false;
+            }
+            tokio::time::sleep(Duration::from_millis(10)).await;
+        }
+        true
+    }
+
+    /// A node learns where a peer that dialled it listens (identify), and can then
+    /// pass it on.
+    #[tokio::test]
+    async fn a_dialled_node_learns_where_the_dialler_listens() {
+        let (a, a_address) = memory_node();
+        let (b, b_address) = memory_node();
+        let (a, b) = (Arc::new(a), Arc::new(b));
+        let _at_a = received(Arc::clone(&a));
+        b.resolve_seed(&a_address.to_string()).await.unwrap();
+        b.send_message(a.local_id(), &ping()).await.unwrap();
+        let b_id = b.local_id().to_owned();
+        assert!(
+            eventually(|| a.locations(&b_id) == [b_address.to_string()]).await,
+            "a knows {:?} of b",
+            a.locations(&b_id)
+        );
+    }
+
+    #[tokio::test]
+    async fn announced_addresses_are_published_instead() {
+        let key = Keypair::generate_ed25519();
+        let peer = key.public().to_peer_id();
+        let port = NEXT_PORT.fetch_add(1, Ordering::Relaxed);
+        let announced: Multiaddr = "/ip4/203.0.113.7/tcp/9661".parse().unwrap();
+        let transport = Libp2pTransport::start(Libp2pConfig {
+            key,
+            listen: vec![format!("/memory/{port}").parse().unwrap()],
+            announce: vec![announced.clone()],
+            memory: true,
+        })
+        .unwrap();
+        assert_eq!(
+            transport.local_addresses(),
+            [announced.with_p2p(peer).unwrap().to_string()]
+        );
+    }
+
+    #[test]
+    fn local_addresses_are_left_out_unless_there_are_only_local_ones() {
+        let parse = |text: &str| text.parse::<Multiaddr>().unwrap();
+        let loopback = parse("/ip4/127.0.0.1/tcp/1");
+        let link = parse("/ip6/fe80::1/tcp/1");
+        let lan = parse("/ip4/192.168.1.2/tcp/1");
+        let wireguard = parse("/ip6/fd00::2/tcp/1");
+        assert_eq!(
+            published(&[
+                loopback.clone(),
+                link.clone(),
+                lan.clone(),
+                wireguard.clone()
+            ]),
+            [lan, wireguard]
+        );
+        assert_eq!(
+            published(&[loopback.clone(), link.clone()]),
+            [loopback, link]
+        );
+    }
+
     #[tokio::test]
     async fn unknown_peers_are_refused() {
         let (b, _) = memory_node();