From 928d0a592b3495a69a5a4c8a60268afc47d6711f Mon Sep 17 00:00:00 2001
From: Frando <frando@unbiskant.org>
Date: Mon, 16 Sep 2024 15:08:58 +0200
Subject: [PATCH] feat(iroh-net): add Endpoint::watch_node_addr

---
 iroh-net/src/endpoint.rs  | 17 ++++++++++++++++
 iroh-net/src/magicsock.rs | 41 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/iroh-net/src/endpoint.rs b/iroh-net/src/endpoint.rs
index c83181ef639..65c657a8734 100644
--- a/iroh-net/src/endpoint.rs
+++ b/iroh-net/src/endpoint.rs
@@ -647,6 +647,23 @@ impl Endpoint {
         Ok(NodeAddr::from_parts(self.node_id(), relay, addrs))
     }
 
+    /// Watches for changes to the node address.
+    ///
+    /// This stream yields a [`NodeAddr`] whenever either the home relay changed, or when
+    /// the list of discovered direct addresses changed.
+    ///
+    /// When issuing the first call to this method the first direct address discovery might
+    /// still be underway, in this case the first item of the returned stream will become
+    /// ready once the discovery finishes. Later calls to this function return a stream where
+    /// the first item is ready immediately. Subsequent items will become ready whenever
+    /// either the node address or the discovered direct addresses change.
+    ///
+    /// See [`Self::watch_home_relay`] and [`Self::direct_addresses`] for functions to watch
+    /// for changes of only one part of the node address.
+    pub fn watch_node_addr(&self) -> impl Stream<Item = NodeAddr> {
+        self.msock.watch_node_addr()
+    }
+
     /// Returns the [`RelayUrl`] of the Relay server used as home relay.
     ///
     /// Every endpoint has a home Relay server which it chooses as the server with the
diff --git a/iroh-net/src/magicsock.rs b/iroh-net/src/magicsock.rs
index 8972c1bb834..f8c07104667 100644
--- a/iroh-net/src/magicsock.rs
+++ b/iroh-net/src/magicsock.rs
@@ -339,6 +339,47 @@ impl MagicSock {
         current.chain(changes)
     }
 
+    /// Watch for changes to the node address.
+    ///
+    /// Returns a new [`NodeAddr`] whenever either our direct addresses change
+    /// (see [`Self::direct_addresses]) or if our home relay changed
+    /// (see [`Self::watch_home_relay`]).
+    pub(crate) fn watch_node_addr(&self) -> impl Stream<Item = NodeAddr> {
+        enum Item {
+            DirectAddrs(Vec<DirectAddr>),
+            HomeRelay(RelayUrl),
+        }
+        let node_id = self.secret_key.public();
+        let mut current_addrs = self.direct_addrs.get().addrs;
+        let mut home_relay = self.my_relay.get();
+        let addr_stream = self
+            .direct_addrs
+            .watch()
+            .into_stream()
+            .map(|addrs| Item::DirectAddrs(addrs.addrs));
+        let home_relay_stream = self
+            .my_relay
+            .watch()
+            .into_stream()
+            .filter_map(|maybe_relay| maybe_relay)
+            .map(Item::HomeRelay);
+        addr_stream.or(home_relay_stream).map(move |item| {
+            match item {
+                Item::DirectAddrs(addrs) => {
+                    current_addrs = addrs;
+                }
+                Item::HomeRelay(relay_url) => {
+                    home_relay = Some(relay_url);
+                }
+            }
+            NodeAddr::from_parts(
+                node_id,
+                home_relay.clone(),
+                current_addrs.iter().map(|x| x.addr).collect(),
+            )
+        })
+    }
+
     /// Returns a stream that reports the [`ConnectionType`] we have to the
     /// given `node_id`.
     ///