feat: support user data in local swarm discovery

Author: Frando

Cargo.lock

@@ -41,3 +41,6 @@ unexpected_cfgs = { level = "warn", check-cfg = ["cfg(iroh_docsrs)", "cfg(iroh_l
 
 [workspace.lints.clippy]
 unused-async = "warn"
+
+[patch.crates-io]
+swarm-discovery = { git = "https://github.com/Frando/swarm-discovery.git", branch = "feat/txt" }

Cargo.toml

@@ -195,7 +195,7 @@ pub trait Discovery: std::fmt::Debug + Send + Sync {
 impl<T: Discovery> Discovery for Arc<T> {}
 
 /// The information that can be published about a node in discovery services.
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone, Default, Eq, PartialEq)]
 pub struct DiscoveryData {
     relay_url: Option<RelayUrl>,
     direct_addresses: BTreeSet<SocketAddr>,
@@ -248,6 +248,11 @@ impl DiscoveryData {
             direct_addresses: self.direct_addresses,
         }
     }
+
+    /// Returns the optional user-defined data.
+    pub fn user_data(&self) -> Option<&UserData> {
+        self.user_data.as_ref()
+    }
 }
 
 impl From<NodeAddr> for DiscoveryData {

iroh/src/discovery.rs

@@ -37,7 +37,7 @@ use std::{
 
 use anyhow::Result;
 use derive_more::FromStr;
-use iroh_base::{NodeAddr, NodeId, PublicKey, RelayUrl};
+use iroh_base::{NodeAddr, NodeId, PublicKey};
 use n0_future::{
     boxed::BoxStream,
     task::{self, AbortOnDropHandle, JoinSet},
@@ -65,6 +65,8 @@ const N0_LOCAL_SWARM: &str = "iroh.local.swarm";
 /// Used in the [`crate::endpoint::Source::Discovery`] enum variant as the `name`.
 pub const NAME: &str = "local.swarm.discovery";
 
+const USER_DATA_ATTRIBUTE: &str = "user-data";
+
 /// How long we will wait before we stop sending discovery items
 const DISCOVERY_DURATION: Duration = Duration::from_secs(10);
 
@@ -75,7 +77,7 @@ pub struct LocalSwarmDiscovery {
     handle: AbortOnDropHandle<()>,
     sender: mpsc::Sender<Message>,
     /// When `local_addrs` changes, we re-publish our info.
-    local_addrs: Watchable<Option<(Option<RelayUrl>, BTreeSet<SocketAddr>)>>,
+    local_addrs: Watchable<Option<DiscoveryData>>,
 }
 
 #[derive(Debug)]
@@ -148,8 +150,7 @@ impl LocalSwarmDiscovery {
             &rt,
         )?;
 
-        let local_addrs: Watchable<Option<(Option<RelayUrl>, BTreeSet<SocketAddr>)>> =
-            Watchable::default();
+        let local_addrs: Watchable<Option<DiscoveryData>> = Watchable::default();
         let mut addrs_change = local_addrs.watch();
         let discovery_fut = async move {
             let mut node_addrs: HashMap<PublicKey, Peer> = HashMap::default();
@@ -166,14 +167,20 @@ impl LocalSwarmDiscovery {
                     msg = recv.recv() => {
                         msg
                     }
-                    Ok(Some((_url, addrs))) = addrs_change.updated() => {
-                        tracing::trace!(?addrs, "LocalSwarmDiscovery address changed");
+                    Ok(Some(data)) = addrs_change.updated() => {
+                        println!("PUBLISH {data:?}");
+                        tracing::trace!(?data, "LocalSwarmDiscovery address changed");
                         discovery.remove_all();
                         let addrs =
-                            LocalSwarmDiscovery::socketaddrs_to_addrs(addrs);
+                            LocalSwarmDiscovery::socketaddrs_to_addrs(data.direct_addrs());
                         for addr in addrs {
                             discovery.add(addr.0, addr.1)
                         }
+                        if let Some(user_data) = data.user_data() {
+                            if let Err(err) = discovery.set_txt_attribute(USER_DATA_ATTRIBUTE.to_string(), Some(user_data.to_string())) {
+                                warn!("Failed to set user data in local swarm discovery: {err:?}");
+                            }
+                        }
                         continue;
                     }
                 };
@@ -320,7 +327,7 @@ impl LocalSwarmDiscovery {
                 sender.send(Message::Discovery(node_id, peer)).await.ok();
             });
         };
-        let addrs = LocalSwarmDiscovery::socketaddrs_to_addrs(socketaddrs);
+        let addrs = LocalSwarmDiscovery::socketaddrs_to_addrs(&socketaddrs);
         let node_id_str = data_encoding::BASE32_NOPAD
             .encode(node_id.as_bytes())
             .to_ascii_lowercase();
@@ -333,7 +340,7 @@ impl LocalSwarmDiscovery {
         discoverer.spawn(rt)
     }
 
-    fn socketaddrs_to_addrs(socketaddrs: BTreeSet<SocketAddr>) -> HashMap<u16, Vec<IpAddr>> {
+    fn socketaddrs_to_addrs(socketaddrs: &BTreeSet<SocketAddr>) -> HashMap<u16, Vec<IpAddr>> {
         let mut addrs: HashMap<u16, Vec<IpAddr>> = HashMap::default();
         for socketaddr in socketaddrs {
             addrs
@@ -351,6 +358,10 @@ fn peer_to_discovery_item(peer: &Peer, node_id: &NodeId) -> DiscoveryItem {
         .iter()
         .map(|(ip, port)| SocketAddr::new(*ip, *port))
         .collect();
+    let user_data = match peer.txt_attribute(USER_DATA_ATTRIBUTE) {
+        Some(Some(user_data)) => user_data.parse().ok(),
+        _ => None,
+    };
     DiscoveryItem {
         node_addr: NodeAddr {
             node_id: *node_id,
@@ -359,7 +370,7 @@ fn peer_to_discovery_item(peer: &Peer, node_id: &NodeId) -> DiscoveryItem {
         },
         provenance: NAME,
         last_updated: None,
-        user_data: None,
+        user_data,
     }
 }
 
@@ -380,12 +391,7 @@ impl Discovery for LocalSwarmDiscovery {
     }
 
     fn publish(&self, data: &DiscoveryData) {
-        self.local_addrs
-            .set(Some((
-                data.relay_url().cloned(),
-                data.direct_addrs().clone(),
-            )))
-            .ok();
+        self.local_addrs.set(Some(data.clone())).ok();
     }
 
     fn subscribe(&self) -> Option<BoxStream<DiscoveryItem>> {
@@ -413,6 +419,7 @@ mod tests {
         use tracing_test::traced_test;
 
         use super::super::*;
+        use crate::discovery::UserData;
 
         #[tokio::test]
         #[traced_test]
@@ -421,7 +428,10 @@ mod tests {
             let (node_id_b, discovery_b) = make_discoverer()?;
 
             // make addr info for discoverer b
-            let addr_info = DiscoveryData::new(None, BTreeSet::from(["0.0.0.0:11111".parse()?]));
+            let user_data: UserData = "foobar".parse()?;
+            let addr_info = DiscoveryData::new(None, BTreeSet::from(["0.0.0.0:11111".parse()?]))
+                .with_user_data(user_data.clone());
+            println!("info {addr_info:?}");
 
             // pass in endpoint, this is never used
             let ep = crate::endpoint::Builder::default().bind().await?;
@@ -441,8 +451,10 @@ mod tests {
                 .unwrap()?;
             assert_eq!(s1_res.node_addr.relay_url.as_ref(), addr_info.relay_url());
             assert_eq!(&s1_res.node_addr.direct_addresses, addr_info.direct_addrs());
+            assert_eq!(s1_res.user_data.as_ref(), Some(&user_data));
             assert_eq!(s2_res.node_addr.relay_url.as_ref(), addr_info.relay_url());
             assert_eq!(&s2_res.node_addr.direct_addresses, addr_info.direct_addrs());
+            assert_eq!(s2_res.user_data.as_ref(), Some(&user_data));
 
             Ok(())
         }
@@ -457,9 +469,11 @@ mod tests {
             let (_, discovery) = make_discoverer()?;
             let addr_info = DiscoveryData::new(None, BTreeSet::from(["0.0.0.0:11111".parse()?]));
 
-            for _ in 0..num_nodes {
+            for i in 0..num_nodes {
                 let (node_id, discovery) = make_discoverer()?;
-                node_ids.insert(node_id);
+                let user_data: UserData = format!("node{i}").parse()?;
+                let addr_info = addr_info.clone().with_user_data(user_data.clone());
+                node_ids.insert((node_id, Some(user_data)));
                 discovery.publish(&addr_info);
                 discoverers.push(discovery);
             }
@@ -470,8 +484,8 @@ mod tests {
                 let mut got_ids = BTreeSet::new();
                 while got_ids.len() != num_nodes {
                     if let Some(item) = events.next().await {
-                        if node_ids.contains(&item.node_addr.node_id) {
-                            got_ids.insert(item.node_addr.node_id);
+                        if node_ids.contains(&(item.node_addr.node_id, item.user_data.clone())) {
+                            got_ids.insert((item.node_addr.node_id, item.user_data));
                         }
                     } else {
                         anyhow::bail!(