Add enum for AddressFamily

3 files changed, 45 insertions, 18 deletions

diff --git a/talpid-core/src/tunnel/openvpn/mod.rs b/talpid-core/src/tunnel/openvpn/mod.rs
index dd9fdd0e63..9eaa25875f 100644
--- a/talpid-core/src/tunnel/openvpn/mod.rs
+++ b/talpid-core/src/tunnel/openvpn/mod.rs
@@ -41,7 +41,6 @@ use winapi::shared::{
     netioapi::{GetUnicastIpAddressEntry, MIB_UNICASTIPADDRESS_ROW},
     nldef::{IpDadStatePreferred, IpDadStateTentative, NL_DAD_STATE},
     winerror::NO_ERROR,
-    ws2def::AF_UNSPEC,
 };
 #[cfg(windows)]
 use winreg::enums::{KEY_READ, KEY_WRITE};
@@ -1244,7 +1243,7 @@ fn wait_for_ready_device(alias: &str) -> Result<()> {
 
     // Obtain unicast IP addresses
     let mut unicast_rows: Vec<MIB_UNICASTIPADDRESS_ROW> =
-        crate::tunnel::windows::get_unicast_table(AF_UNSPEC as u16)
+        crate::tunnel::windows::get_unicast_table(None)
             .map_err(Error::ObtainUnicastAddress)?
             .into_iter()
             .filter(|row| row.InterfaceLuid.Value == luid.Value)
diff --git a/talpid-core/src/tunnel/openvpn/windows.rs b/talpid-core/src/tunnel/openvpn/windows.rs
index 16869cb20c..6e49bb1f0d 100644
--- a/talpid-core/src/tunnel/openvpn/windows.rs
+++ b/talpid-core/src/tunnel/openvpn/windows.rs
@@ -1,4 +1,4 @@
-use crate::tunnel::windows::{get_ip_interface_entry, set_ip_interface_entry};
+use crate::tunnel::windows::{get_ip_interface_entry, set_ip_interface_entry, AddressFamily};
 use std::{
     ffi::CStr,
     fmt, io, iter, mem,
@@ -18,7 +18,6 @@ use winapi::{
         nldef::RouterDiscoveryDisabled,
         ntdef::FALSE,
         winerror::NO_ERROR,
-        ws2def::{AF_INET, AF_INET6},
     },
     um::{
         libloaderapi::{
@@ -161,8 +160,8 @@ impl WintunAdapter {
     pub fn try_disable_unused_features(&self) {
         // Disable DAD, DHCP, and router discovery
         let luid = self.luid();
-        for family in &[AF_INET, AF_INET6] {
-            if let Ok(mut row) = get_ip_interface_entry(*family as u16, &luid) {
+        for family in &[AddressFamily::Ipv4, AddressFamily::Ipv6] {
+            if let Ok(mut row) = get_ip_interface_entry(*family, &luid) {
                 row.SitePrefixLength = 0;
                 row.RouterDiscoveryBehavior = RouterDiscoveryDisabled;
                 row.DadTransmits = 0;
diff --git a/talpid-core/src/tunnel/windows.rs b/talpid-core/src/tunnel/windows.rs
index bf56b5b168..ea7cb646b7 100644
--- a/talpid-core/src/tunnel/windows.rs
+++ b/talpid-core/src/tunnel/windows.rs
@@ -1,6 +1,6 @@
 use std::{
     ffi::OsStr,
-    io, mem,
+    fmt, io, mem,
     os::windows::{ffi::OsStrExt, io::RawHandle},
     sync::Mutex,
 };
@@ -16,6 +16,22 @@ use winapi::shared::{
     ws2def::{AF_INET, AF_INET6, AF_UNSPEC},
 };
 
+/// Address family. These correspond to the `AF_*` constants.
+#[derive(Debug, Clone, Copy)]
+pub enum AddressFamily {
+    Ipv4 = AF_INET as isize,
+    Ipv6 = AF_INET6 as isize,
+}
+
+impl fmt::Display for AddressFamily {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            AddressFamily::Ipv4 => write!(f, "IPv4 (AF_INET)"),
+            AddressFamily::Ipv6 => write!(f, "IPv6 (AF_INET6)"),
+        }
+    }
+}
+
 /// Context for [`notify_ip_interface_change`]. When it is dropped,
 /// the callback is unregistered.
 pub struct IpNotifierHandle<'a> {
@@ -47,7 +63,7 @@ unsafe extern "system" fn inner_callback(
 /// or changed.
 pub fn notify_ip_interface_change<'a, T: FnMut(&MIB_IPINTERFACE_ROW, u32) + Send + 'a>(
     callback: T,
-    family: u16,
+    family: Option<AddressFamily>,
 ) -> io::Result<Box<IpNotifierHandle<'a>>> {
     let mut context = Box::new(IpNotifierHandle {
         callback: Mutex::new(Box::new(callback)),
@@ -56,7 +72,7 @@ pub fn notify_ip_interface_change<'a, T: FnMut(&MIB_IPINTERFACE_ROW, u32) + Send
 
     let status = unsafe {
         NotifyIpInterfaceChange(
-            family,
+            af_family_from_family(family),
             Some(inner_callback),
             &mut *context as *mut _ as *mut _,
             FALSE,
@@ -72,9 +88,12 @@ pub fn notify_ip_interface_change<'a, T: FnMut(&MIB_IPINTERFACE_ROW, u32) + Send
 }
 
 /// Returns information about a network IP interface.
-pub fn get_ip_interface_entry(family: u16, luid: &NET_LUID) -> io::Result<MIB_IPINTERFACE_ROW> {
+pub fn get_ip_interface_entry(
+    family: AddressFamily,
+    luid: &NET_LUID,
+) -> io::Result<MIB_IPINTERFACE_ROW> {
     let mut row: MIB_IPINTERFACE_ROW = unsafe { mem::zeroed() };
-    row.Family = family;
+    row.Family = family as u16;
     row.InterfaceLuid = *luid;
 
     let result = unsafe { GetIpInterfaceEntry(&mut row) };
@@ -95,7 +114,7 @@ pub fn set_ip_interface_entry(row: &MIB_IPINTERFACE_ROW) -> io::Result<()> {
     }
 }
 
-fn ip_interface_entry_exists(family: u16, luid: &NET_LUID) -> io::Result<bool> {
+fn ip_interface_entry_exists(family: AddressFamily, luid: &NET_LUID) -> io::Result<bool> {
     match get_ip_interface_entry(family, luid) {
         Ok(_) => Ok(true),
         Err(error) if error.raw_os_error() == Some(ERROR_NOT_FOUND as i32) => Ok(false),
@@ -134,12 +153,12 @@ pub async fn wait_for_interfaces(luid: NET_LUID, ipv4: bool, ipv6: bool) -> io::
                 }
             }
         },
-        AF_UNSPEC as u16,
+        None,
     )?;
 
     // Make sure they don't already exist
-    if (!ipv4 || ip_interface_entry_exists(AF_INET as u16, &luid)?)
-        && (!ipv6 || ip_interface_entry_exists(AF_INET6 as u16, &luid)?)
+    if (!ipv4 || ip_interface_entry_exists(AddressFamily::Ipv4, &luid)?)
+        && (!ipv6 || ip_interface_entry_exists(AddressFamily::Ipv6, &luid)?)
     {
         return Ok(());
     }
@@ -148,12 +167,16 @@ pub async fn wait_for_interfaces(luid: NET_LUID, ipv4: bool, ipv6: bool) -> io::
     Ok(())
 }
 
-/// Returns the unicast IP address table.
-pub fn get_unicast_table(family: u16) -> io::Result<Vec<MIB_UNICASTIPADDRESS_ROW>> {
+/// Returns the unicast IP address table. If `family` is `None`, then addresses for all families are
+/// returned.
+pub fn get_unicast_table(
+    family: Option<AddressFamily>,
+) -> io::Result<Vec<MIB_UNICASTIPADDRESS_ROW>> {
     let mut unicast_rows = vec![];
     let mut unicast_table: *mut MIB_UNICASTIPADDRESS_TABLE = std::ptr::null_mut();
 
-    let status = unsafe { GetUnicastIpAddressTable(family, &mut unicast_table) };
+    let status =
+        unsafe { GetUnicastIpAddressTable(af_family_from_family(family), &mut unicast_table) };
     if status != NO_ERROR {
         return Err(io::Error::from_raw_os_error(status as i32));
     }
@@ -180,3 +203,9 @@ pub fn luid_from_alias<T: AsRef<OsStr>>(alias: T) -> io::Result<NET_LUID> {
     }
     Ok(luid)
 }
+
+fn af_family_from_family(family: Option<AddressFamily>) -> u16 {
+    family
+        .map(|family| family as u16)
+        .unwrap_or(AF_UNSPEC as u16)
+}