feat(client): allow both IPv6 and IPv4 local addresses

Cargo.toml

@@ -55,6 +55,9 @@ tokio-util = { version = "0.3", features = ["codec"] }
 tower-util = "0.3"
 url = "1.0"
 
+[target.'cfg(any(target_os = "linux", target_os = "macos"))'.dev-dependencies]
+pnet = "0.25.0"
+
 [features]
 default = [
     "runtime",

src/client/connect/dns.rs

@@ -200,27 +200,33 @@ impl IpAddrs {
         None
     }
 
-    pub(super) fn split_by_preference(self, local_addr: Option<IpAddr>) -> (IpAddrs, IpAddrs) {
-        if let Some(local_addr) = local_addr {
-            let preferred = self
-                .iter
-                .filter(|addr| addr.is_ipv6() == local_addr.is_ipv6())
-                .collect();
-
-            (IpAddrs::new(preferred), IpAddrs::new(vec![]))
-        } else {
-            let preferring_v6 = self
-                .iter
-                .as_slice()
-                .first()
-                .map(SocketAddr::is_ipv6)
-                .unwrap_or(false);
-
-            let (preferred, fallback) = self
-                .iter
-                .partition::<Vec<_>, _>(|addr| addr.is_ipv6() == preferring_v6);
-
-            (IpAddrs::new(preferred), IpAddrs::new(fallback))
+    #[inline]
+    fn filter(self, predicate: impl FnMut(&SocketAddr) -> bool) -> IpAddrs {
+        IpAddrs::new(self.iter.filter(predicate).collect())
+    }
+
+    pub(super) fn split_by_preference(
+        self,
+        local_addr_ipv4: Option<Ipv4Addr>,
+        local_addr_ipv6: Option<Ipv6Addr>,
+    ) -> (IpAddrs, IpAddrs) {
+        match (local_addr_ipv4, local_addr_ipv6) {
+            (Some(_), None) => (self.filter(SocketAddr::is_ipv4), IpAddrs::new(vec![])),
+            (None, Some(_)) => (self.filter(SocketAddr::is_ipv6), IpAddrs::new(vec![])),
+            _ => {
+                let preferring_v6 = self
+                    .iter
+                    .as_slice()
+                    .first()
+                    .map(SocketAddr::is_ipv6)
+                    .unwrap_or(false);
+
+                let (preferred, fallback) = self
+                    .iter
+                    .partition::<Vec<_>, _>(|addr| addr.is_ipv6() == preferring_v6);
+
+                (IpAddrs::new(preferred), IpAddrs::new(fallback))
+            }
         }
     }
 
@@ -355,34 +361,50 @@ mod tests {
 
     #[test]
     fn test_ip_addrs_split_by_preference() {
-        let v4_addr = (Ipv4Addr::new(127, 0, 0, 1), 80).into();
-        let v6_addr = (Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), 80).into();
+        let ip_v4 = Ipv4Addr::new(127, 0, 0, 1);
+        let ip_v6 = Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1);
+        let v4_addr = (ip_v4, 80).into();
+        let v6_addr = (ip_v6, 80).into();
+
+        let (mut preferred, mut fallback) = IpAddrs {
+            iter: vec![v4_addr, v6_addr].into_iter(),
+        }
+        .split_by_preference(None, None);
+        assert!(preferred.next().unwrap().is_ipv4());
+        assert!(fallback.next().unwrap().is_ipv6());
+
+        let (mut preferred, mut fallback) = IpAddrs {
+            iter: vec![v6_addr, v4_addr].into_iter(),
+        }
+        .split_by_preference(None, None);
+        assert!(preferred.next().unwrap().is_ipv6());
+        assert!(fallback.next().unwrap().is_ipv4());
 
         let (mut preferred, mut fallback) = IpAddrs {
             iter: vec![v4_addr, v6_addr].into_iter(),
         }
-        .split_by_preference(None);
+        .split_by_preference(Some(ip_v4), Some(ip_v6));
         assert!(preferred.next().unwrap().is_ipv4());
         assert!(fallback.next().unwrap().is_ipv6());
 
         let (mut preferred, mut fallback) = IpAddrs {
             iter: vec![v6_addr, v4_addr].into_iter(),
         }
-        .split_by_preference(None);
+        .split_by_preference(Some(ip_v4), Some(ip_v6));
         assert!(preferred.next().unwrap().is_ipv6());
         assert!(fallback.next().unwrap().is_ipv4());
 
         let (mut preferred, fallback) = IpAddrs {
             iter: vec![v4_addr, v6_addr].into_iter(),
         }
-        .split_by_preference(Some(v4_addr.ip()));
+        .split_by_preference(Some(ip_v4), None);
         assert!(preferred.next().unwrap().is_ipv4());
         assert!(fallback.is_empty());
 
         let (mut preferred, fallback) = IpAddrs {
             iter: vec![v4_addr, v6_addr].into_iter(),
         }
-        .split_by_preference(Some(v6_addr.ip()));
+        .split_by_preference(None, Some(ip_v6));
         assert!(preferred.next().unwrap().is_ipv6());
         assert!(fallback.is_empty());
     }