fix: use more specific poll_writable

Author: dignifiedquire

Cargo.lock

@@ -40,3 +40,9 @@ unexpected_cfgs = { level = "warn", check-cfg = ["cfg(iroh_docsrs)", "cfg(iroh_l
 
 [workspace.lints.clippy]
 unused-async = "warn"
+
+
+[patch.crates-io]
+iroh-quinn = { git = "https://github.com/n0-computer/quinn", branch = "iroh-fix-poll-writable" }
+iroh-quinn-udp = { git = "https://github.com/n0-computer/quinn", branch = "iroh-fix-poll-writable" }
+iroh-quinn-proto = { git = "https://github.com/n0-computer/quinn", branch = "iroh-fix-poll-writable" }

Cargo.toml

@@ -468,6 +468,17 @@ impl MagicSock {
         Err(io::Error::other("no valid socket available"))
     }
 
+    pub(super) fn addr_for_send(
+        &self,
+        dest: NodeIdMappedAddr,
+    ) -> Option<(PublicKey, Option<SocketAddr>, Option<RelayUrl>)> {
+        self.node_map.addr_for_send(
+            dest,
+            self.ipv6_reported.load(Ordering::Relaxed),
+            &self.metrics.magicsock,
+        )
+    }
+
     /// Implementation for AsyncUdpSocket::try_send
     #[instrument(skip_all)]
     fn try_send(&self, transmit: &quinn_udp::Transmit) -> io::Result<()> {
@@ -1621,7 +1632,8 @@ enum DiscoBoxError {
 
 impl AsyncUdpSocket for MagicSock {
     fn create_io_poller(self: Arc<Self>) -> Pin<Box<dyn quinn::UdpPoller>> {
-        self.transports.create_io_poller()
+        self.transports
+            .create_io_poller(self.clone(), self.ip_mapped_addrs.clone())
     }
 
     fn try_send(&self, transmit: &quinn_udp::Transmit) -> io::Result<()> {

iroh/src/magicsock.rs

@@ -265,6 +265,22 @@ impl NodeMap {
         Some((public_key, udp_addr, relay_url, msgs))
     }
 
+    #[allow(clippy::type_complexity)]
+    pub(super) fn addr_for_send(
+        &self,
+        addr: NodeIdMappedAddr,
+        have_ipv6: bool,
+        metrics: &Metrics,
+    ) -> Option<(PublicKey, Option<SocketAddr>, Option<RelayUrl>)> {
+        let mut inner = self.inner.lock().expect("poisoned");
+        let ep = inner.get_mut(NodeStateKey::NodeIdMappedAddr(addr))?;
+        let public_key = *ep.public_key();
+        trace!(dest = %addr, node_id = %public_key.fmt_short(), "dst mapped to NodeId");
+        let now = Instant::now();
+        let (udp_addr, relay_url) = ep.addr_for_send(&now, have_ipv6, metrics);
+        Some((public_key, udp_addr, relay_url))
+    }
+
     pub(super) fn notify_shutdown(&self) {
         let mut inner = self.inner.lock().expect("poisoned");
         for (_, ep) in inner.node_states_mut() {

iroh/src/magicsock/node_map.rs

@@ -277,7 +277,7 @@ impl NodeState {
     /// Returns the address(es) that should be used for sending the next packet.
     ///
     /// This may return to send on one, both or no paths.
-    fn addr_for_send(
+    pub(super) fn addr_for_send(
         &mut self,
         now: &Instant,
         have_ipv6: bool,

iroh/src/magicsock/node_map/node_state.rs

@@ -8,19 +8,20 @@ use std::{
 
 use iroh_base::{NodeId, RelayUrl};
 use n0_watcher::Watcher;
-#[cfg(not(wasm_browser))]
-use netwatch::UdpSocket;
 use relay::RelayDatagramSendChannelSender;
 use tracing::{trace, warn};
 
 #[cfg(not(wasm_browser))]
 mod ip;
 mod relay;
 
+#[cfg(not(wasm_browser))]
+use self::ip::IpIoPoller;
 #[cfg(not(wasm_browser))]
 pub(crate) use self::ip::IpTransport;
 pub(crate) use self::relay::{RelayActorConfig, RelayTransport};
-use super::NetInfo;
+use super::{MagicSock, MappedAddr, NetInfo};
+use crate::net_report::IpMappedAddresses;
 
 /// Manages the different underlying data transports that the magicsock
 /// can support.
@@ -250,7 +251,11 @@ impl Transports {
         }
     }
 
-    pub(crate) fn create_io_poller(&self) -> Pin<Box<dyn quinn::UdpPoller>> {
+    pub(crate) fn create_io_poller(
+        &self,
+        msock: Arc<MagicSock>,
+        ip_mapped_addrs: IpMappedAddresses,
+    ) -> Pin<Box<dyn quinn::UdpPoller>> {
         #[cfg(not(wasm_browser))]
         let ip_pollers = self.ip.iter().map(|t| t.create_io_poller()).collect();
 
@@ -260,6 +265,8 @@ impl Transports {
             #[cfg(not(wasm_browser))]
             ip_pollers,
             relay_pollers,
+            ip_mapped_addrs,
+            msock,
         })
     }
 
@@ -380,29 +387,74 @@ impl Addr {
 #[derive(Debug)]
 pub struct IoPoller {
     #[cfg(not(wasm_browser))]
-    ip_pollers: Vec<Arc<UdpSocket>>,
+    ip_pollers: Vec<IpIoPoller>,
     relay_pollers: Vec<RelayDatagramSendChannelSender>,
+    ip_mapped_addrs: IpMappedAddresses,
+    msock: Arc<MagicSock>, // :(
 }
 
 impl quinn::UdpPoller for IoPoller {
-    fn poll_writable(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
-        // This version returns Ready as soon as any of them are ready.
-
+    fn poll_writable(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context,
+        transmit: &quinn_proto::Transmit,
+    ) -> Poll<io::Result<()>> {
         let this = &mut *self;
-        #[cfg(not(wasm_browser))]
-        for poller in &mut this.ip_pollers {
-            match poller.poll_writable(cx) {
-                Poll::Ready(_) => return Poll::Ready(Ok(())),
-                Poll::Pending => (),
-            }
-        }
 
-        for poller in &mut this.relay_pollers {
-            match poller.poll_writable(cx) {
-                Poll::Ready(_) => return Poll::Ready(Ok(())),
-                Poll::Pending => (),
+        match MappedAddr::from(transmit.destination) {
+            MappedAddr::None(_dest) => {
+                // return Poll::Ready(Err(io::Error::other("Cannot convert to a mapped address.")));
+            }
+            MappedAddr::NodeId(dest) => {
+                // Get the node's relay address and best direct address, as well
+                // as any pings that need to be sent for hole-punching purposes.
+                match this.msock.addr_for_send(dest) {
+                    Some((_node_id, udp_addr, relay_url)) => {
+                        #[cfg(not(wasm_browser))]
+                        if let Some(addr) = udp_addr {
+                            for poller in &mut this.ip_pollers {
+                                if poller.is_valid_send_addr(&addr) {
+                                    match poller.poll_writable(cx) {
+                                        Poll::Ready(_) => return Poll::Ready(Ok(())),
+                                        Poll::Pending => (),
+                                    }
+                                }
+                            }
+                        }
+                        if let Some(_url) = relay_url {
+                            for poller in &mut this.relay_pollers {
+                                match poller.poll_writable(cx) {
+                                    Poll::Ready(_) => return Poll::Ready(Ok(())),
+                                    Poll::Pending => (),
+                                }
+                            }
+                        }
+                    }
+                    None => {
+                        // return Poll::Ready(Err(io::Error::other(
+                        //     "no NodeState for mapped address",
+                        // )));
+                    }
+                }
             }
+            #[cfg(not(wasm_browser))]
+            MappedAddr::Ip(addr) => match this.ip_mapped_addrs.get_ip_addr(&addr) {
+                Some(addr) => {
+                    for poller in &mut this.ip_pollers {
+                        if poller.is_valid_send_addr(&addr) {
+                            match poller.poll_writable(cx) {
+                                Poll::Ready(_) => return Poll::Ready(Ok(())),
+                                Poll::Pending => (),
+                            }
+                        }
+                    }
+                }
+                None => {
+                    // return Poll::Ready(Err(io::Error::other("unknown mapped address")));
+                }
+            },
         }
+
         Poll::Pending
     }
 }

iroh/src/magicsock/transports.rs

@@ -38,8 +38,11 @@ impl IpTransport {
         }
     }
 
-    pub(super) fn create_io_poller(&self) -> Arc<UdpSocket> {
-        self.socket.clone()
+    pub(super) fn create_io_poller(&self) -> IpIoPoller {
+        IpIoPoller {
+            bind_addr: self.bind_addr,
+            socket: self.socket.clone(),
+        }
     }
 
     pub(super) fn poll_send(
@@ -57,6 +60,7 @@ impl IpTransport {
             segment_size: transmit.segment_size,
             src_ip: src,
         });
+        trace!("send res: {:?}", res);
 
         match res {
             Ok(res) => {
@@ -153,3 +157,24 @@ impl IpTransport {
         self.socket.clone()
     }
 }
+
+#[derive(Debug)]
+pub(super) struct IpIoPoller {
+    socket: Arc<UdpSocket>,
+    bind_addr: SocketAddr,
+}
+
+impl IpIoPoller {
+    pub(super) fn is_valid_send_addr(&self, addr: &SocketAddr) -> bool {
+        #[allow(clippy::match_like_matches_macro)]
+        match (self.bind_addr, addr) {
+            (SocketAddr::V4(_), SocketAddr::V4(..)) => true,
+            (SocketAddr::V6(_), SocketAddr::V6(..)) => true,
+            _ => false,
+        }
+    }
+
+    pub(super) fn poll_writable(&self, cx: &mut std::task::Context<'_>) -> Poll<io::Result<()>> {
+        self.socket.poll_writable(cx)
+    }
+}

iroh/src/magicsock/transports/ip.rs

@@ -446,7 +446,6 @@ mod tests {
     use n0_watcher::Watcher;
     use quinn::ApplicationClose;
     use testresult::TestResult;
-    use tracing_test::traced_test;
 
     use super::*;
     use crate::{endpoint::ConnectionError, RelayMode};
@@ -490,17 +489,24 @@ mod tests {
         }
     }
 
+    // #[tokio::test(flavor = "multi_thread")]
     #[tokio::test]
-    #[traced_test]
     async fn test_limiter() -> Result<()> {
-        let e1 = Endpoint::builder().bind().await?;
+        tracing_subscriber::fmt::try_init().ok();
+        let e1 = Endpoint::builder()
+            .relay_mode(RelayMode::Disabled)
+            .bind()
+            .await?;
         // deny all access
         let proto = AccessLimit::new(Echo, |_node_id| false);
         let r1 = Router::builder(e1.clone()).accept(ECHO_ALPN, proto).spawn();
 
         let addr1 = r1.endpoint().node_addr().initialized().await?;
-
-        let e2 = Endpoint::builder().bind().await?;
+        dbg!(&addr1);
+        let e2 = Endpoint::builder()
+            .relay_mode(RelayMode::Disabled)
+            .bind()
+            .await?;
 
         println!("connecting");
         let conn = e2.connect(addr1, ECHO_ALPN).await?;
@@ -515,8 +521,10 @@ mod tests {
         Ok(())
     }
 
+    // #[tokio::test(flavor = "multi_thread")]
     #[tokio::test]
     async fn test_graceful_shutdown() -> TestResult {
+        tracing_subscriber::fmt::try_init().ok();
         #[derive(Debug, Clone, Default)]
         struct TestProtocol {
             connections: Arc<Mutex<Vec<Connection>>>,
@@ -534,34 +542,44 @@ mod tests {
             }
 
             fn shutdown(&self) -> BoxFuture<()> {
+                eprintln!("shutting down");
                 let this = self.clone();
                 Box::pin(async move {
                     tokio::time::sleep(Duration::from_millis(100)).await;
+                    eprintln!("waiting for lock");
                     let mut connections = this.connections.lock().expect("poisoned");
+                    eprintln!("got lock");
                     for conn in connections.drain(..) {
+                        eprintln!("closing");
                         conn.close(42u32.into(), b"shutdown");
                     }
                 })
             }
         }
 
+        eprintln!("creating ep1");
         let endpoint = Endpoint::builder()
             .relay_mode(RelayMode::Disabled)
             .bind()
             .await?;
         let router = Router::builder(endpoint)
             .accept(TEST_ALPN, TestProtocol::default())
             .spawn();
+        eprintln!("waiting for node addr");
         let addr = router.endpoint().node_addr().initialized().await?;
 
+        eprintln!("creating ep2");
         let endpoint2 = Endpoint::builder()
             .relay_mode(RelayMode::Disabled)
             .bind()
             .await?;
+        eprintln!("connecting to {:?}", addr);
         let conn = endpoint2.connect(addr, TEST_ALPN).await?;
 
+        eprintln!("starting shutdown");
         router.shutdown().await?;
 
+        eprintln!("waiting for closed conn");
         let reason = conn.closed().await;
         assert_eq!(
             reason,