feat(iroh-net): allow the underlying UdpSockets to be rebound

iroh-net-report/src/reportgen/hairpin.rs

@@ -121,7 +121,7 @@ impl Actor {
             .context("net_report actor gone")?;
         msg_response_rx.await.context("net_report actor died")?;
 
-        if let Err(err) = socket.send_to(&stun::request(txn), dst).await {
+        if let Err(err) = socket.send_to(&stun::request(txn), dst.into()).await {
             warn!(%dst, "failed to send hairpin check");
             return Err(err.into());
         }

iroh-net/src/magicsock.rs

@@ -36,7 +36,7 @@ use futures_util::stream::BoxStream;
 use iroh_base::key::NodeId;
 use iroh_metrics::{inc, inc_by};
 use iroh_relay::protos::stun;
-use netwatch::{interfaces, ip::LocalAddresses, netmon};
+use netwatch::{interfaces, ip::LocalAddresses, netmon, UdpSocket};
 use quinn::AsyncUdpSocket;
 use rand::{seq::SliceRandom, Rng, SeedableRng};
 use smallvec::{smallvec, SmallVec};
@@ -441,11 +441,8 @@ impl MagicSock {
         // Right now however we have one single poller behaving the same for each
         // connection.  It checks all paths and returns Poll::Ready as soon as any path is
         // ready.
-        let ipv4_poller = Arc::new(self.pconn4.clone()).create_io_poller();
-        let ipv6_poller = self
-            .pconn6
-            .as_ref()
-            .map(|sock| Arc::new(sock.clone()).create_io_poller());
+        let ipv4_poller = self.pconn4.create_io_poller();
+        let ipv6_poller = self.pconn6.as_ref().map(|sock| sock.create_io_poller());
         let relay_sender = self.relay_actor_sender.clone();
         Box::pin(IoPoller {
             ipv4_poller,
@@ -1091,10 +1088,9 @@ impl MagicSock {
                     Err(err) if err.kind() == io::ErrorKind::WouldBlock => {
                         // This is the socket .try_send_disco_message_udp used.
                         let sock = self.conn_for_addr(dst)?;
-                        let sock = Arc::new(sock.clone());
-                        let mut poller = sock.create_io_poller();
-                        match poller.as_mut().poll_writable(cx)? {
-                            Poll::Ready(()) => continue,
+                        match sock.as_socket_ref().poll_writable(cx) {
+                            Poll::Ready(Ok(())) => continue,
+                            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
                             Poll::Pending => return Poll::Pending,
                         }
                     }
@@ -1408,6 +1404,9 @@ impl Handle {
         let net_reporter =
             net_report::Client::new(Some(port_mapper.clone()), dns_resolver.clone())?;
 
+        let pconn4_sock = pconn4.as_socket();
+        let pconn6_sock = pconn6.as_ref().map(|p| p.as_socket());
+
         let (actor_sender, actor_receiver) = mpsc::channel(256);
         let (relay_actor_sender, relay_actor_receiver) = mpsc::channel(256);
         let (udp_disco_sender, mut udp_disco_receiver) = mpsc::channel(256);
@@ -1431,9 +1430,9 @@ impl Handle {
             ipv6_reported: Arc::new(AtomicBool::new(false)),
             relay_map,
             my_relay: Default::default(),
-            pconn4: pconn4.clone(),
-            pconn6: pconn6.clone(),
             net_reporter: net_reporter.addr(),
+            pconn4,
+            pconn6,
             disco_secrets: DiscoSecrets::default(),
             node_map,
             relay_actor_sender: relay_actor_sender.clone(),
@@ -1481,8 +1480,8 @@ impl Handle {
                     periodic_re_stun_timer: new_re_stun_timer(false),
                     net_info_last: None,
                     port_mapper,
-                    pconn4,
-                    pconn6,
+                    pconn4: pconn4_sock,
+                    pconn6: pconn6_sock,
                     no_v4_send: false,
                     net_reporter,
                     network_monitor,
@@ -1720,8 +1719,8 @@ struct Actor {
     net_info_last: Option<NetInfo>,
 
     // The underlying UDP sockets used to send/rcv packets.
-    pconn4: UdpConn,
-    pconn6: Option<UdpConn>,
+    pconn4: Arc<UdpSocket>,
+    pconn6: Option<Arc<UdpSocket>>,
 
     /// The NAT-PMP/PCP/UPnP prober/client, for requesting port mappings from NAT devices.
     port_mapper: portmapper::Client,
@@ -1861,6 +1860,14 @@ impl Actor {
         debug!("link change detected: major? {}", is_major);
 
         if is_major {
+            if let Err(err) = self.pconn4.rebind() {
+                warn!("failed to rebind Udp IPv4 socket: {:?}", err);
+            };
+            if let Some(ref pconn6) = self.pconn6 {
+                if let Err(err) = pconn6.rebind() {
+                    warn!("failed to rebind Udp IPv6 socket: {:?}", err);
+                };
+            }
             self.msock.dns_resolver.clear_cache();
             self.msock.re_stun("link-change-major");
             self.close_stale_relay_connections().await;
@@ -1893,14 +1900,6 @@ impl Actor {
                 self.port_mapper.deactivate();
                 self.relay_actor_cancel_token.cancel();
 
-                // Ignore errors from pconnN
-                // They will frequently have been closed already by a call to connBind.Close.
-                debug!("stopping connections");
-                if let Some(ref conn) = self.pconn6 {
-                    conn.close().await.ok();
-                }
-                self.pconn4.close().await.ok();
-
                 debug!("shutdown complete");
                 return true;
             }
@@ -2206,8 +2205,8 @@ impl Actor {
         }
 
         let relay_map = self.msock.relay_map.clone();
-        let pconn4 = Some(self.pconn4.as_socket());
-        let pconn6 = self.pconn6.as_ref().map(|p| p.as_socket());
+        let pconn4 = Some(self.pconn4.clone());
+        let pconn6 = self.pconn6.clone();
 
         debug!("requesting net_report report");
         match self
@@ -3099,6 +3098,45 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_regression_network_change_rebind_wakes_connection_driver(
+    ) -> testresult::TestResult {
+        let _ = iroh_test::logging::setup();
+        let m1 = MagicStack::new(RelayMode::Disabled).await?;
+        let m2 = MagicStack::new(RelayMode::Disabled).await?;
+
+        println!("Net change");
+        m1.endpoint.magic_sock().force_network_change(true).await;
+        tokio::time::sleep(Duration::from_secs(1)).await; // wait for socket rebinding
+
+        let _guard = mesh_stacks(vec![m1.clone(), m2.clone()]).await?;
+
+        let _handle = AbortOnDropHandle::new(tokio::spawn({
+            let endpoint = m2.endpoint.clone();
+            async move {
+                while let Some(incoming) = endpoint.accept().await {
+                    println!("Incoming first conn!");
+                    let conn = incoming.await?;
+                    conn.closed().await;
+                }
+
+                testresult::TestResult::Ok(())
+            }
+        }));
+
+        println!("first conn!");
+        let conn = m1
+            .endpoint
+            .connect(m2.endpoint.node_addr().await?, ALPN)
+            .await?;
+        println!("Closing first conn");
+        conn.close(0u32.into(), b"bye lolz");
+        conn.closed().await;
+        println!("Closed first conn");
+
+        Ok(())
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn test_two_devices_roundtrip_network_change() -> Result<()> {
         time::timeout(

iroh-net/src/magicsock/udp_conn.rs

@@ -1,69 +1,57 @@
 use std::{
     fmt::Debug,
-    future::Future,
     io,
     net::SocketAddr,
     pin::Pin,
     sync::Arc,
-    task::{ready, Context, Poll},
+    task::{Context, Poll},
 };
 
 use anyhow::{bail, Context as _};
 use netwatch::UdpSocket;
 use quinn::AsyncUdpSocket;
-use quinn_udp::{Transmit, UdpSockRef};
-use tokio::io::Interest;
-use tracing::{debug, trace};
+use quinn_udp::Transmit;
+use tracing::debug;
 
 /// A UDP socket implementing Quinn's [`AsyncUdpSocket`].
-#[derive(Clone, Debug)]
+#[derive(Debug, Clone)]
 pub struct UdpConn {
     io: Arc<UdpSocket>,
-    inner: Arc<quinn_udp::UdpSocketState>,
 }
 
 impl UdpConn {
     pub(super) fn as_socket(&self) -> Arc<UdpSocket> {
         self.io.clone()
     }
 
+    pub(super) fn as_socket_ref(&self) -> &UdpSocket {
+        &self.io
+    }
+
     pub(super) fn bind(addr: SocketAddr) -> anyhow::Result<Self> {
         let sock = bind(addr)?;
-        let state = quinn_udp::UdpSocketState::new(quinn_udp::UdpSockRef::from(&sock))?;
-        Ok(Self {
-            io: Arc::new(sock),
-            inner: Arc::new(state),
-        })
+
+        Ok(Self { io: Arc::new(sock) })
     }
 
     pub fn port(&self) -> u16 {
         self.local_addr().map(|p| p.port()).unwrap_or_default()
     }
 
-    #[allow(clippy::unused_async)]
-    pub async fn close(&self) -> Result<(), io::Error> {
-        // Nothing to do atm
-        Ok(())
+    pub(super) fn create_io_poller(&self) -> Pin<Box<dyn quinn::UdpPoller>> {
+        Box::pin(IoPoller {
+            io: self.io.clone(),
+        })
     }
 }
 
 impl AsyncUdpSocket for UdpConn {
     fn create_io_poller(self: Arc<Self>) -> Pin<Box<dyn quinn::UdpPoller>> {
-        let sock = self.io.clone();
-        Box::pin(IoPoller {
-            next_waiter: move || {
-                let sock = sock.clone();
-                async move { sock.writable().await }
-            },
-            waiter: None,
-        })
+        (*self).create_io_poller()
     }
 
     fn try_send(&self, transmit: &Transmit<'_>) -> io::Result<()> {
-        self.io.try_io(Interest::WRITABLE, || {
-            let sock_ref = UdpSockRef::from(&self.io);
-            self.inner.send(sock_ref, transmit)
-        })
+        self.io.try_send_quinn(transmit)
     }
 
     fn poll_recv(
@@ -72,40 +60,23 @@ impl AsyncUdpSocket for UdpConn {
         bufs: &mut [io::IoSliceMut<'_>],
         meta: &mut [quinn_udp::RecvMeta],
     ) -> Poll<io::Result<usize>> {
-        loop {
-            ready!(self.io.poll_recv_ready(cx))?;
-            if let Ok(res) = self.io.try_io(Interest::READABLE, || {
-                self.inner.recv(Arc::as_ref(&self.io).into(), bufs, meta)
-            }) {
-                for meta in meta.iter().take(res) {
-                    trace!(
-                        src = %meta.addr,
-                        len = meta.len,
-                        count = meta.len / meta.stride,
-                        dst = %meta.dst_ip.map(|x| x.to_string()).unwrap_or_default(),
-                        "UDP recv"
-                    );
-                }
-
-                return Poll::Ready(Ok(res));
-            }
-        }
+        self.io.poll_recv_quinn(cx, bufs, meta)
     }
 
     fn local_addr(&self) -> io::Result<SocketAddr> {
         self.io.local_addr()
     }
 
     fn may_fragment(&self) -> bool {
-        self.inner.may_fragment()
+        self.io.may_fragment()
     }
 
     fn max_transmit_segments(&self) -> usize {
-        self.inner.max_gso_segments()
+        self.io.max_gso_segments()
     }
 
     fn max_receive_segments(&self) -> usize {
-        self.inner.gro_segments()
+        self.io.gro_segments()
     }
 }
 
@@ -147,49 +118,14 @@ fn bind(mut addr: SocketAddr) -> anyhow::Result<UdpSocket> {
 }
 
 /// Poller for when the socket is writable.
-///
-/// The tricky part is that we only have `tokio::net::UdpSocket::writable()` to create the
-/// waiter we need, which does not return a named future type.  In order to be able to store
-/// this waiter in a struct without boxing we need to specify the future itself as a type
-/// parameter, which we can only do if we introduce a second type parameter which returns
-/// the future.  So we end up with a function which we do not need, but it makes the types
-/// work.
-#[derive(derive_more::Debug)]
-#[pin_project::pin_project]
-struct IoPoller<F, Fut>
-where
-    F: Fn() -> Fut + Send + Sync + 'static,
-    Fut: Future<Output = io::Result<()>> + Send + Sync + 'static,
-{
-    /// Function which can create a new waiter if there is none.
-    #[debug("next_waiter")]
-    next_waiter: F,
-    /// The waiter which tells us when the socket is writable.
-    #[debug("waiter")]
-    #[pin]
-    waiter: Option<Fut>,
+#[derive(Debug)]
+struct IoPoller {
+    io: Arc<UdpSocket>,
 }
 
-impl<F, Fut> quinn::UdpPoller for IoPoller<F, Fut>
-where
-    F: Fn() -> Fut + Send + Sync + 'static,
-    Fut: Future<Output = io::Result<()>> + Send + Sync + 'static,
-{
+impl quinn::UdpPoller for IoPoller {
     fn poll_writable(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
-        let mut this = self.project();
-        if this.waiter.is_none() {
-            this.waiter.set(Some((this.next_waiter)()));
-        }
-        let result = this
-            .waiter
-            .as_mut()
-            .as_pin_mut()
-            .expect("just set")
-            .poll(cx);
-        if result.is_ready() {
-            this.waiter.set(None);
-        }
-        result
+        self.io.poll_writable(cx)
     }
 }