feat: add connection monitor

Frando · Frando · commit 50c9724791d0 · 2025-11-06T12:09:49.000+01:00
diff --git a/iroh/examples/monitor-connections.rs b/iroh/examples/monitor-connections.rs
@@ -0,0 +1,115 @@
+use std::{sync::Arc, time::Duration};
+
+use iroh::{
+    Endpoint, RelayMode,
+    endpoint::{ConnectionInfo, ConnectionMonitor},
+};
+use n0_error::{Result, StackResultExt, StdResultExt, ensure_any};
+use n0_future::task::AbortOnDropHandle;
+use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
+use tracing::{Instrument, info, info_span};
+
+const ALPN: &[u8] = b"iroh/test";
+
+#[tokio::main]
+async fn main() -> Result {
+    tracing_subscriber::fmt()
+        .with_env_filter(
+            tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()),
+        )
+        .init();
+
+    let monitor = Monitor::new();
+    let server = Endpoint::empty_builder(RelayMode::Disabled)
+        .alpns(vec![ALPN.to_vec()])
+        .monitor_connections(monitor.clone())
+        .bind()
+        .instrument(info_span!("server"))
+        .await?;
+    let server_addr = server.addr();
+
+    let count = 2;
+
+    let client_task = tokio::spawn(
+        async move {
+            let client = Endpoint::empty_builder(RelayMode::Disabled)
+                .bind()
+                .instrument(info_span!("client"))
+                .await?;
+            for _i in 0..count {
+                let conn = client.connect(server_addr.clone(), ALPN).await?;
+                let mut s = conn.accept_uni().await.anyerr()?;
+                let data = s.read_to_end(2).await.anyerr()?;
+                ensure_any!(data == b"hi", "unexpected data");
+                conn.close(23u32.into(), b"bye");
+            }
+            client.close().await;
+            n0_error::Ok(client)
+        }
+        .instrument(info_span!("client")),
+    );
+
+    let server_task = tokio::spawn(
+        async move {
+            for _i in 0..count {
+                let conn = server
+                    .accept()
+                    .await
+                    .context("server endpoint closed")?
+                    .await?;
+                let mut s = conn.open_uni().await.anyerr()?;
+                s.write_all(b"hi").await.anyerr()?;
+                conn.closed().await;
+            }
+            info!("done");
+            server.close().await;
+            n0_error::Ok(())
+        }
+        .instrument(info_span!("server")),
+    );
+    client_task.await.std_context("client")?.context("client")?;
+    server_task.await.std_context("server")?.context("server")?;
+    tokio::time::sleep(Duration::from_secs(1)).await;
+    drop(monitor);
+    Ok(())
+}
+
+/// Our connection monitor impl.
+///
+/// This here only logs connection open and close events via tracing.
+/// It could also maintain a datastructure of all connections, or send the stats to some metrics service.
+#[derive(Clone)]
+struct Monitor {
+    tx: UnboundedSender<ConnectionInfo>,
+    _task: Arc<AbortOnDropHandle<()>>,
+}
+
+impl ConnectionMonitor for Monitor {
+    fn on_connection(&self, connection: ConnectionInfo) {
+        self.tx.send(connection).ok();
+    }
+}
+
+impl Monitor {
+    fn new() -> Self {
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        let task = tokio::spawn(Self::run(rx).instrument(info_span!("watcher")));
+        Self {
+            tx,
+            _task: Arc::new(AbortOnDropHandle::new(task)),
+        }
+    }
+
+    async fn run(mut rx: UnboundedReceiver<ConnectionInfo>) {
+        loop {
+            tokio::select! {
+                Some(conn) = rx.recv() => {
+                    let alpn = String::from_utf8_lossy(conn.alpn()).to_string();
+                    let remote = conn.remote_id().fmt_short();
+                    info!(%remote, %alpn, rtt=?conn.rtt(), "new connection");
+                }
+                else => break,
+            }
+        }
+    }
+}
diff --git a/iroh/src/endpoint.rs b/iroh/src/endpoint.rs
@@ -117,6 +117,8 @@ pub struct Builder {
     #[cfg(any(test, feature = "test-utils"))]
     path_selection: PathSelection,
     max_tls_tickets: usize,
+    #[debug("{}", connection_monitor.as_ref().map(|_| "Some(Box<dyn ConnectionMonitor>)").unwrap_or("None"))]
+    connection_monitor: Option<Box<dyn ConnectionMonitor>>,
 }
 
 impl Builder {
@@ -158,6 +160,7 @@ impl Builder {
             #[cfg(any(test, feature = "test-utils"))]
             path_selection: PathSelection::default(),
             max_tls_tickets: DEFAULT_MAX_TLS_TICKETS,
+            connection_monitor: None,
         }
     }
 
@@ -208,6 +211,7 @@ impl Builder {
             // #[cfg(any(test, feature = "test-utils"))]
             // path_selection: self.path_selection,
             metrics,
+            connection_monitor: self.connection_monitor,
         };
 
         let msock = magicsock::MagicSock::spawn(msock_opts).await?;
@@ -432,6 +436,44 @@ impl Builder {
         self.max_tls_tickets = n;
         self
     }
+
+    // TODO docs
+    /// Register a handler that is invoked for each connection the endpoint accepts or initiates.
+    ///
+    /// The [`ConnectionMonitor::on_connection`] method is invoked synchronosuly, from within a tokio
+    /// context. So you can spawn tasks if needed.
+    /// Make sure that whatever you do with the connection info here is non-blocking.
+    /// Usually you'd want to send the info over a broadcast or unbounded channel,
+    /// or insert it into some persistent datastructure.
+    ///
+    /// The `ConnectionInfo` internally contains a weak reference to the connection,
+    /// so keeping the struct alive does not keep the connection alive.
+    /// Note however that `ConnectionInfo` does keep an allocation per connection alive
+    /// so to not leak memory you should drop the `ConnectionInfo` eventually
+    ///
+    /// [`ConnectionMonitor`] is implemented for `Fn(ConnectionInfo)`, so you can
+    /// also pass a closure that takes [`ConnectionInfo`] to this function.
+    pub fn monitor_connections(mut self, monitor: impl ConnectionMonitor) -> Self {
+        self.connection_monitor = Some(Box::new(monitor));
+        self
+    }
+}
+
+/// Monitor each connection accepted or initiated by the endpoint.
+pub trait ConnectionMonitor: Send + Sync + 'static {
+    /// Called for each new connection the endpoint accepts or initiates.
+    ///
+    /// This is only called when a connection is fully established.
+    fn on_connection(&self, connection: ConnectionInfo);
+}
+
+impl<T> ConnectionMonitor for T
+where
+    T: Fn(ConnectionInfo) + Send + Sync + 'static,
+{
+    fn on_connection(&self, connection: ConnectionInfo) {
+        (self)(connection)
+    }
 }
 
 /// Configuration for a [`quinn::Endpoint`] that cannot be changed at runtime.
diff --git a/iroh/src/magicsock.rs b/iroh/src/magicsock.rs
@@ -63,7 +63,7 @@ use crate::{
     defaults::timeouts::NET_REPORT_TIMEOUT,
     disco::{self, SendAddr},
     discovery::{ConcurrentDiscovery, Discovery, EndpointData, UserData},
-    endpoint::ConnectionInfo,
+    endpoint::{ConnectionInfo, ConnectionMonitor},
     key::{DecryptionError, SharedSecret, public_ed_box, secret_ed_box},
     metrics::EndpointMetrics,
     net_report::{self, IfStateDetails, Report},
@@ -148,6 +148,8 @@ pub(crate) struct Options {
     // #[cfg(any(test, feature = "test-utils"))]
     // pub(crate) path_selection: PathSelection,
     pub(crate) metrics: EndpointMetrics,
+    #[debug("{}", connection_monitor.as_ref().map(|_| "Some(Box<dyn ConnectionMonitor>)").unwrap_or("None"))]
+    pub(crate) connection_monitor: Option<Box<dyn ConnectionMonitor>>,
 }
 
 /// Handle for [`MagicSock`].
@@ -219,6 +221,9 @@ pub(crate) struct MagicSock {
 
     /// Metrics
     pub(crate) metrics: EndpointMetrics,
+
+    #[debug("{}", connection_monitor.as_ref().map(|_| "Some(Box<dyn ConnectionMonitor>)").unwrap_or("None"))]
+    connection_monitor: Option<Box<dyn ConnectionMonitor>>,
 }
 
 #[allow(missing_docs)]
@@ -274,6 +279,11 @@ impl MagicSock {
         conn: ConnectionInfo,
         paths_info: n0_watcher::Watchable<PathsInfo>,
     ) {
+        // Inform the monitor about the new connection.
+        if let Some(monitor) = self.connection_monitor.as_ref() {
+            monitor.on_connection(conn.clone());
+        }
+
         // TODO: Spawning tasks like this is obviously bad.  But it is solvable:
         //   - This is only called from inside Connection::new.
         //   - Connection::new is called from:
@@ -942,6 +952,7 @@ impl Handle {
             // #[cfg(any(test, feature = "test-utils"))]
             // path_selection,
             metrics,
+            connection_monitor,
         } = opts;
 
         let discovery = ConcurrentDiscovery::default();
@@ -1018,6 +1029,7 @@ impl Handle {
             local_addrs_watch: transports.local_addrs_watch(),
             #[cfg(not(wasm_browser))]
             ip_bind_addrs: transports.ip_bind_addrs(),
+            connection_monitor,
         });
 
         let mut endpoint_config = quinn::EndpointConfig::default();
@@ -1919,6 +1931,7 @@ mod tests {
             // path_selection: PathSelection::default(),
             discovery_user_data: None,
             metrics: Default::default(),
+            connection_monitor: None,
         }
     }
 
@@ -2351,6 +2364,7 @@ mod tests {
             insecure_skip_relay_cert_verify: false,
             // path_selection: PathSelection::default(),
             metrics: Default::default(),
+            connection_monitor: None,
         };
         let msock = MagicSock::spawn(opts).await?;
         Ok(msock)
