Add some convenience functions for ListenConfig (#189)

* Add some convenient functions for ListenConfig * Another helper function * Rename function * Doc comments and fmt * Remove builder-like functions * Fix docs and tests
sigp · Jun 6, 2023 · 6807647 · 6807647
1 parent 70bba1e
commit 6807647
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 10 deletions.
diff --git a/examples/find_nodes.rs b/examples/find_nodes.rs
@@ -22,7 +22,7 @@ use discv5::{
     Discv5, Discv5ConfigBuilder, Discv5Event, ListenConfig,
 };
 use std::{
-    net::{Ipv4Addr, Ipv6Addr},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
     time::Duration,
 };
 use tracing::{info, warn};
@@ -112,8 +112,8 @@ async fn main() {
 
     // the address to listen on.
     let listen_config = match args.socket_kind {
-        SocketKind::Ip4 => ListenConfig::new_ipv4(Ipv4Addr::UNSPECIFIED, port),
-        SocketKind::Ip6 => ListenConfig::new_ipv6(Ipv6Addr::UNSPECIFIED, port6),
+        SocketKind::Ip4 => ListenConfig::from_ip(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port),
+        SocketKind::Ip6 => ListenConfig::from_ip(IpAddr::V6(Ipv6Addr::UNSPECIFIED), port6),
         SocketKind::Ds => ListenConfig::default()
             .with_ipv4(Ipv4Addr::UNSPECIFIED, port)
             .with_ipv6(Ipv6Addr::UNSPECIFIED, port6),

diff --git a/src/socket/mod.rs b/src/socket/mod.rs
@@ -6,7 +6,7 @@ use socket2::{Domain, Protocol, Socket as Socket2, Type};
 use std::{
     collections::HashMap,
     io::Error,
-    net::{Ipv4Addr, Ipv6Addr, SocketAddr},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
     sync::Arc,
     time::Duration,
 };
@@ -157,14 +157,39 @@ impl Socket {
 }
 
 impl ListenConfig {
-    /// Creates a [`ListenConfig`] with an ipv4-only socket.
-    pub fn new_ipv4(ip: Ipv4Addr, port: u16) -> ListenConfig {
-        Self::Ipv4 { ip, port }
+    /// If an [`IpAddr`] is known, a ListenConfig can be created based on the version. This will
+    /// not create a dual stack configuration.
+    pub fn from_ip(ip: IpAddr, port: u16) -> ListenConfig {
+        match ip {
+            IpAddr::V4(ip) => ListenConfig::Ipv4 { ip, port },
+            IpAddr::V6(ip) => ListenConfig::Ipv6 { ip, port },
+        }
     }
 
-    /// Creates a [`ListenConfig`] with an ipv6-only socket.
-    pub fn new_ipv6(ip: Ipv6Addr, port: u16) -> ListenConfig {
-        Self::Ipv6 { ip, port }
+    /// Allows optional ipv4 and ipv6 addresses to be entered to create a [`ListenConfig`]. If both
+    /// are specified a dual-stack configuration will result. This will panic if both parameters
+    /// are None.
+    pub fn from_two_sockets(
+        ipv4: Option<SocketAddrV4>,
+        ipv6: Option<SocketAddrV6>,
+    ) -> ListenConfig {
+        match (ipv4, ipv6) {
+            (Some(ipv4), None) => ListenConfig::Ipv4 {
+                ip: *ipv4.ip(),
+                port: ipv4.port(),
+            },
+            (None, Some(ipv6)) => ListenConfig::Ipv6 {
+                ip: *ipv6.ip(),
+                port: ipv6.port(),
+            },
+            (Some(ipv4), Some(ipv6)) => ListenConfig::DualStack {
+                ipv4: *ipv4.ip(),
+                ipv4_port: ipv4.port(),
+                ipv6: *ipv6.ip(),
+                ipv6_port: ipv6.port(),
+            },
+            (None, None) => panic!("At least one IP address must be entered."),
+        }
     }
 
     // Overrides the ipv4 address and port of ipv4 and dual stack configurations. Ipv6
@@ -229,6 +254,21 @@ impl Default for ListenConfig {
     }
 }
 
+impl From<SocketAddr> for ListenConfig {
+    fn from(socket_addr: SocketAddr) -> Self {
+        match socket_addr {
+            SocketAddr::V4(socket) => ListenConfig::Ipv4 {
+                ip: *socket.ip(),
+                port: socket.port(),
+            },
+            SocketAddr::V6(socket) => ListenConfig::Ipv6 {
+                ip: *socket.ip(),
+                port: socket.port(),
+            },
+        }
+    }
+}
+
 impl Drop for Socket {
     // close the send/recv handlers
     fn drop(&mut self) {