Merge branch 'fix-connectivity-check'

5 files changed, 230 insertions, 195 deletions

diff --git a/talpid-core/src/tunnel/wireguard/mod.rs b/talpid-core/src/tunnel/wireguard/mod.rs
index 4edcc93e94..30919c1b39 100644
--- a/talpid-core/src/tunnel/wireguard/mod.rs
+++ b/talpid-core/src/tunnel/wireguard/mod.rs
@@ -23,9 +23,6 @@ mod wireguard_go;
 #[cfg(target_os = "linux")]
 mod wireguard_kernel;
 
-#[cfg(target_os = "linux")]
-mod network_manager;
-
 use self::wireguard_go::WgGoTunnel;
 
 type Result<T> = std::result::Result<T, Error>;
@@ -152,14 +149,17 @@ impl WireguardMonitor {
     ) -> Result<Box<dyn Tunnel>> {
         #[cfg(target_os = "linux")]
         if !*FORCE_USERSPACE_WIREGUARD {
-            match network_manager::NetworkManager::new(config) {
+            match wireguard_kernel::NetworkManagerTunnel::new(
+                route_manager.runtime_handle(),
+                config,
+            ) {
                 Ok(tunnel) => {
                     log::debug!("Using NetworkManager to use kernel WireGuard implementation");
                     return Ok(Box::new(tunnel));
                 }
                 Err(err) => {
                     if !err.should_use_userspace() {
-                        match wireguard_kernel::KernelTunnel::new(
+                        match wireguard_kernel::NetlinkTunnel::new(
                             route_manager.runtime_handle(),
                             config,
                         ) {
diff --git a/talpid-core/src/tunnel/wireguard/stats.rs b/talpid-core/src/tunnel/wireguard/stats.rs
index be45d1af47..dbfd9fcc0a 100644
--- a/talpid-core/src/tunnel/wireguard/stats.rs
+++ b/talpid-core/src/tunnel/wireguard/stats.rs
@@ -1,3 +1,7 @@
+#[cfg(target_os = "linux")]
+use super::wireguard_kernel::wg_message::{DeviceMessage, DeviceNla, PeerNla};
+
+
 #[derive(err_derive::Error, Debug, PartialEq)]
 pub enum Error {
     #[error(display = "Failed to parse integer from string \"_0\"")]
@@ -58,6 +62,29 @@ impl Stats {
             _ => Err(Error::KeyNotFoundError),
         }
     }
+
+    #[cfg(target_os = "linux")]
+    pub fn parse_device_message(message: &DeviceMessage) -> Self {
+        // iterate over device attributes
+        let mut tx_bytes = 0;
+        let mut rx_bytes = 0;
+        for nla in &message.nlas {
+            if let DeviceNla::Peers(peers) = nla {
+                // iterate over all peer attributes
+                let peer_iter = peers.iter().map(|peer| peer.0.as_slice()).flatten();
+
+                for peer_nla in peer_iter {
+                    match peer_nla {
+                        PeerNla::TxBytes(bytes) => tx_bytes += *bytes,
+                        PeerNla::RxBytes(bytes) => rx_bytes += *bytes,
+                        _ => continue,
+                    };
+                }
+            }
+        }
+
+        Self { tx_bytes, rx_bytes }
+    }
 }
 
 
diff --git a/talpid-core/src/tunnel/wireguard/wireguard_kernel/mod.rs b/talpid-core/src/tunnel/wireguard/wireguard_kernel/mod.rs
index b1fa3a0f65..284c93be13 100644
--- a/talpid-core/src/tunnel/wireguard/wireguard_kernel/mod.rs
+++ b/talpid-core/src/tunnel/wireguard/wireguard_kernel/mod.rs
@@ -1,4 +1,4 @@
-use super::{stats::Stats, Config, Tunnel, TunnelError};
+use super::{Config, Tunnel, TunnelError};
 use futures::future::{abortable, AbortHandle};
 use netlink_packet_core::{constants::*, NetlinkDeserializable};
 use netlink_packet_route::{
@@ -19,11 +19,16 @@ use tokio::stream::StreamExt;
 
 mod parsers;
 
-mod wg_message;
-use wg_message::{DeviceMessage, DeviceNla, PeerNla};
-mod nl_message;
+pub mod wg_message;
+use wg_message::{DeviceMessage, DeviceNla};
+pub mod nl_message;
 use nl_message::{ControlNla, NetlinkControlMessage};
 
+pub mod netlink_tunnel;
+pub use netlink_tunnel::NetlinkTunnel;
+pub mod nm_tunnel;
+pub use nm_tunnel::NetworkManagerTunnel;
+
 
 #[derive(err_derive::Error, Debug)]
 #[error(no_from)]
@@ -75,141 +80,26 @@ pub enum Error {
 
     #[error(display = "Failed to delete device")]
     DeleteDeviceError(#[error(source)] rtnetlink::Error),
-}
 
-pub struct KernelTunnel {
-    interface_index: u32,
-    netlink_connections: Handle,
-    tokio_handle: tokio::runtime::Handle,
+    #[error(display = "NetworkManager error")]
+    NetworkManager(#[error(source)] nm_tunnel::Error),
 }
 
 const MULLVAD_INTERFACE_NAME: &str = "wg-mullvad";
 
-impl KernelTunnel {
-    pub fn new(tokio_handle: tokio::runtime::Handle, config: &Config) -> Result<Self, Error> {
-        tokio_handle.clone().block_on(async {
-            let mut netlink_connections = Handle::connect().await?;
-            let interface_index = netlink_connections
-                .create_device(MULLVAD_INTERFACE_NAME.to_string(), config.mtu as u32)
-                .await?;
-
-            let mut tunnel = Self {
-                interface_index,
-                netlink_connections,
-                tokio_handle,
-            };
-
-            if let Err(err) = tunnel.setup(config).await {
-                if let Err(teardown_err) = tunnel
-                    .netlink_connections
-                    .delete_device(interface_index)
-                    .await
-                {
-                    log::error!(
-                        "Failed to tear down WireGuard interface after failing to apply config: {}",
-                        teardown_err
-                    );
-                }
-                return Err(err);
-            }
-
-
-            Ok(tunnel)
-        })
-    }
-
-    async fn setup(&mut self, config: &Config) -> Result<(), Error> {
-        self.netlink_connections
-            .wg_handle
-            .set_config(self.interface_index, config)
-            .await?;
-
-        for tunnel_ip in config.tunnel.addresses.iter() {
-            self.netlink_connections
-                .set_ip_address(self.interface_index, *tunnel_ip)
-                .await?;
-        }
-
-        Ok(())
-    }
-}
-
-impl Tunnel for KernelTunnel {
-    fn get_interface_name(&self) -> String {
-        let mut wg = self.netlink_connections.wg_handle.clone();
-        let result = self.tokio_handle.block_on(async move {
-            let device = wg.get_by_index(self.interface_index).await?;
-            for nla in device.nlas {
-                if let DeviceNla::IfName(name) = nla {
-                    return Ok(name);
-                }
-            }
-            return Err(Error::Truncated);
-        });
 
-        match result {
-            Ok(name) => name.to_string_lossy().to_string(),
-            Err(err) => {
-                log::error!("Failed to deduce interface name at runtime, will attempt to use the default name. {}", err);
-                MULLVAD_INTERFACE_NAME.to_string()
-            }
+impl Error {
+    pub fn should_use_userspace(&self) -> bool {
+        match self {
+            Error::NetworkManager(nm_tunnel::Error::NMTooOld(_)) => true,
+            _ => false,
         }
     }
-
-    fn stop(self: Box<Self>) -> std::result::Result<(), TunnelError> {
-        let Self {
-            mut netlink_connections,
-            interface_index,
-            tokio_handle,
-        } = *self;
-        tokio_handle.block_on(async move {
-            if let Err(err) = netlink_connections.delete_device(interface_index).await {
-                log::error!("Failed to remove WireGuard device - {}", err);
-                Err(TunnelError::FatalStartWireguardError)
-            } else {
-                Ok(())
-            }
-        })
-    }
-
-    fn get_tunnel_stats(&self) -> std::result::Result<Stats, TunnelError> {
-        let mut wg = self.netlink_connections.wg_handle.clone();
-        let interface_index = self.interface_index;
-        let result = self.tokio_handle.block_on(async move {
-            let device = wg.get_by_index(interface_index).await.map_err(|err| {
-                log::error!("Failed to fetch WireGuard device config: {}", err);
-                TunnelError::GetConfigError
-            })?;
-
-            // iterate over device attributes
-            let mut tx_bytes = 0;
-            let mut rx_bytes = 0;
-            for nla in device.nlas {
-                if let DeviceNla::Peers(peers) = nla {
-                    // iterate over all peer attributes
-                    let peer_iter = peers.iter().map(|peer| peer.0.as_slice()).flatten();
-
-                    for peer_nla in peer_iter {
-                        match peer_nla {
-                            PeerNla::TxBytes(bytes) => tx_bytes += *bytes,
-                            PeerNla::RxBytes(bytes) => rx_bytes += *bytes,
-                            _ => continue,
-                        };
-                    }
-                }
-            }
-
-            Ok(Stats { tx_bytes, rx_bytes })
-        });
-
-        result
-    }
 }
 
-
 #[derive(Debug)]
 pub struct Handle {
-    wg_handle: WireguardConnection,
+    pub wg_handle: WireguardConnection,
     route_handle: rtnetlink::Handle,
     wg_abort_handle: AbortHandle,
     route_abort_handle: AbortHandle,
@@ -370,7 +260,7 @@ impl Drop for Handle {
 }
 
 #[derive(Debug, Clone)]
-struct WireguardConnection {
+pub struct WireguardConnection {
     connection: ConnectionHandle<DeviceMessage>,
     message_type: u16,
 }
diff --git a/talpid-core/src/tunnel/wireguard/wireguard_kernel/netlink_tunnel.rs b/talpid-core/src/tunnel/wireguard/wireguard_kernel/netlink_tunnel.rs
new file mode 100644
index 0000000000..b5a4f1429a
--- /dev/null
+++ b/talpid-core/src/tunnel/wireguard/wireguard_kernel/netlink_tunnel.rs
@@ -0,0 +1,113 @@
+use super::{
+    super::stats::Stats, wg_message::DeviceNla, Config, Error, Handle, Tunnel, TunnelError,
+    MULLVAD_INTERFACE_NAME,
+};
+
+
+pub struct NetlinkTunnel {
+    interface_index: u32,
+    netlink_connections: Handle,
+    tokio_handle: tokio::runtime::Handle,
+}
+
+impl NetlinkTunnel {
+    pub fn new(tokio_handle: tokio::runtime::Handle, config: &Config) -> Result<Self, Error> {
+        tokio_handle.clone().block_on(async {
+            let mut netlink_connections = Handle::connect().await?;
+            let interface_index = netlink_connections
+                .create_device(MULLVAD_INTERFACE_NAME.to_string(), config.mtu as u32)
+                .await?;
+
+            let mut tunnel = Self {
+                interface_index,
+                netlink_connections,
+                tokio_handle,
+            };
+
+            if let Err(err) = tunnel.setup(config).await {
+                if let Err(teardown_err) = tunnel
+                    .netlink_connections
+                    .delete_device(interface_index)
+                    .await
+                {
+                    log::error!(
+                        "Failed to tear down WireGuard interface after failing to apply config: {}",
+                        teardown_err
+                    );
+                }
+                return Err(err);
+            }
+
+
+            Ok(tunnel)
+        })
+    }
+
+    async fn setup(&mut self, config: &Config) -> Result<(), Error> {
+        self.netlink_connections
+            .wg_handle
+            .set_config(self.interface_index, config)
+            .await?;
+
+        for tunnel_ip in config.tunnel.addresses.iter() {
+            self.netlink_connections
+                .set_ip_address(self.interface_index, *tunnel_ip)
+                .await?;
+        }
+
+        Ok(())
+    }
+}
+
+impl Tunnel for NetlinkTunnel {
+    fn get_interface_name(&self) -> String {
+        let mut wg = self.netlink_connections.wg_handle.clone();
+        let result = self.tokio_handle.block_on(async move {
+            let device = wg.get_by_index(self.interface_index).await?;
+            for nla in device.nlas {
+                if let DeviceNla::IfName(name) = nla {
+                    return Ok(name);
+                }
+            }
+            return Err(Error::Truncated);
+        });
+
+        match result {
+            Ok(name) => name.to_string_lossy().to_string(),
+            Err(err) => {
+                log::error!("Failed to deduce interface name at runtime, will attempt to use the default name. {}", err);
+                MULLVAD_INTERFACE_NAME.to_string()
+            }
+        }
+    }
+
+    fn stop(self: Box<Self>) -> std::result::Result<(), TunnelError> {
+        let Self {
+            mut netlink_connections,
+            interface_index,
+            tokio_handle,
+        } = *self;
+        tokio_handle.block_on(async move {
+            if let Err(err) = netlink_connections.delete_device(interface_index).await {
+                log::error!("Failed to remove WireGuard device - {}", err);
+                Err(TunnelError::FatalStartWireguardError)
+            } else {
+                Ok(())
+            }
+        })
+    }
+
+    fn get_tunnel_stats(&self) -> std::result::Result<Stats, TunnelError> {
+        let mut wg = self.netlink_connections.wg_handle.clone();
+        let interface_index = self.interface_index;
+        let result = self.tokio_handle.block_on(async move {
+            let device = wg.get_by_index(interface_index).await.map_err(|err| {
+                log::error!("Failed to fetch WireGuard device config: {}", err);
+                TunnelError::GetConfigError
+            })?;
+            Ok(Stats::parse_device_message(&device))
+        });
+
+        result
+    }
+}
diff --git a/talpid-core/src/tunnel/wireguard/network_manager.rs b/talpid-core/src/tunnel/wireguard/wireguard_kernel/nm_tunnel.rs
index 978128b1c6..a0ec39bddd 100644
--- a/talpid-core/src/tunnel/wireguard/network_manager.rs
+++ b/talpid-core/src/tunnel/wireguard/wireguard_kernel/nm_tunnel.rs
@@ -1,4 +1,7 @@
-use super::{config::Config, stats::Stats, Tunnel, TunnelError};
+use super::{
+    super::stats::Stats, wg_message::DeviceNla, Config, Error as WgKernelError, Handle, Tunnel,
+    TunnelError, MULLVAD_INTERFACE_NAME,
+};
 use dbus::{
     arg::{RefArg, Variant},
     blocking::{stdintf::org_freedesktop_dbus::Properties, BlockingSender, Connection, Proxy},
@@ -12,14 +15,12 @@ const NM_INTERFACE_SETTINGS: &str = "org.freedesktop.NetworkManager.Settings";
 const NM_INTERFACE_SETTINGS_CONNECTION: &str = "org.freedesktop.NetworkManager.Settings.Connection";
 const NM_SETTINGS_PATH: &str = "/org/freedesktop/NetworkManager/Settings";
 const NM_DEVICE: &str = "org.freedesktop.NetworkManager.Device";
-const NM_DEVICE_STATISTICS: &str = "org.freedesktop.NetworkManager.Device.Statistics";
 const NM_CONNECTION_ACTIVE: &str = "org.freedesktop.NetworkManager.Connection.Active";
 const NM_MANAGER: &str = "org.freedesktop.NetworkManager";
 const NM_MANAGER_PATH: &str = "/org/freedesktop/NetworkManager";
 
 const NM_ADD_CONNECTION_VOLATILE: u32 = 0x2;
 
-const TRAFFIC_STATS_REFRESH_RATE_MS: u32 = 1000;
 const RPC_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3);
 
 const DBUS_UNKNOWN_METHOD: &str = "org.freedesktop.DBus.Error.UnknownMethod";
@@ -42,9 +43,6 @@ pub enum Error {
     #[error(display = "Failed to match the returned D-Bus object with expected type")]
     MatchDBusTypeError(#[error(source)] dbus::arg::TypeMismatchError),
 
-    #[error(display = "Failed to set statistics refresh rate")]
-    SetStatsRefreshError(#[error(source)] dbus::Error),
-
     #[error(display = "Error while removing ")]
     DeviceRemovalError(#[error(source)] dbus::Error),
 
@@ -53,20 +51,17 @@ pub enum Error {
 
     #[error(display = "NetworkManager is too old - {}", _0)]
     NMTooOld(String),
-}
 
-impl Error {
-    pub fn should_use_userspace(&self) -> bool {
-        match self {
-            Error::NMTooOld(_) => true,
-            _ => false,
-        }
-    }
+    #[error(display = "Cannot obtain the tunnel interface index")]
+    FindInterfaceIndex,
 }
 
-pub struct NetworkManager {
+pub struct NetworkManagerTunnel {
     dbus_connection: Connection,
     tunnel: Option<WireguardTunnel>,
+    interface_index: u32,
+    netlink_connections: Handle,
+    tokio_handle: tokio::runtime::Handle,
 }
 
 
@@ -75,18 +70,47 @@ type VariantMap = HashMap<String, VariantRefArg>;
 // settings are a{sa{sv}}
 type DbusSettings = HashMap<String, VariantMap>;
 
-impl NetworkManager {
-    pub fn new(config: &Config) -> Result<Self> {
-        let mut dbus_connection = Connection::new_system().map_err(Error::Dbus)?;
-        Self::ensure_nm_is_new_enough(&dbus_connection)?;
-        let tunnel = Some(Self::create_wg_tunnel(&mut dbus_connection, config)?);
-        let manager = NetworkManager {
-            dbus_connection,
-            tunnel,
-        };
+impl NetworkManagerTunnel {
+    pub fn new(
+        tokio_handle: tokio::runtime::Handle,
+        config: &Config,
+    ) -> std::result::Result<Self, WgKernelError> {
+        let mut dbus_connection = Connection::new_system()
+            .map_err(|error| WgKernelError::NetworkManager(Error::Dbus(error)))?;
+        Self::ensure_nm_is_new_enough(&dbus_connection).map_err(WgKernelError::NetworkManager)?;
+        let tunnel = Some(
+            Self::create_wg_tunnel(&mut dbus_connection, config)
+                .map_err(WgKernelError::NetworkManager)?,
+        );
+
+        tokio_handle.clone().block_on(async move {
+            let netlink_connections = Handle::connect().await?;
+            let interface_index = Self::find_device_index(&netlink_connections)
+                .await?
+                .ok_or(WgKernelError::NetworkManager(Error::FindInterfaceIndex))?;
 
-        manager.set_stats_refresh_rate()?;
-        Ok(manager)
+            Ok(NetworkManagerTunnel {
+                dbus_connection,
+                tunnel,
+                interface_index,
+                tokio_handle,
+                netlink_connections,
+            })
+        })
+    }
+
+    async fn find_device_index(
+        netlink_connections: &Handle,
+    ) -> std::result::Result<Option<u32>, WgKernelError> {
+        let mut wg = netlink_connections.wg_handle.clone();
+        let device = wg.get_by_name(MULLVAD_INTERFACE_NAME.to_string()).await?;
+
+        for nla in &device.nlas {
+            if let DeviceNla::IfIndex(index) = nla {
+                return Ok(Some(*index));
+            }
+        }
+        Ok(None)
     }
 
     fn ensure_nm_is_new_enough(connection: &Connection) -> Result<()> {
@@ -235,10 +259,13 @@ impl NetworkManager {
         wireguard_config.insert("peers".into(), Variant(Box::new(peer_configs)));
 
         connection_config.insert("type".into(), Variant(Box::new("wireguard".to_string())));
-        connection_config.insert("id".into(), Variant(Box::new("wg-mullvad".to_string())));
+        connection_config.insert(
+            "id".into(),
+            Variant(Box::new(MULLVAD_INTERFACE_NAME.to_string())),
+        );
         connection_config.insert(
             "interface-name".into(),
-            Variant(Box::new("wg-mullvad".to_string())),
+            Variant(Box::new(MULLVAD_INTERFACE_NAME.to_string())),
         );
         connection_config.insert("autoconnect".into(), Variant(Box::new(true)));
 
@@ -285,20 +312,6 @@ impl NetworkManager {
         settings
     }
 
-    fn set_stats_refresh_rate(&self) -> Result<()> {
-        if let Some(tunnel) = self.tunnel.as_ref() {
-            tunnel
-                .device_proxy(&self.dbus_connection)
-                .set(
-                    NM_DEVICE_STATISTICS,
-                    "RefreshRateMs",
-                    TRAFFIC_STATS_REFRESH_RATE_MS,
-                )
-                .map_err(Error::SetStatsRefreshError)?;
-        }
-        Ok(())
-    }
-
     fn convert_address_to_dbus(address: &IpAddr) -> VariantMap {
         let mut map: VariantMap = HashMap::new();
         map.insert(
@@ -339,7 +352,7 @@ impl NetworkManager {
     }
 }
 
-impl Tunnel for NetworkManager {
+impl Tunnel for NetworkManagerTunnel {
     fn get_interface_name(&self) -> String {
         if let Some(tunnel) = self.tunnel.as_ref() {
             let interface_name = tunnel
@@ -353,7 +366,7 @@ impl Tunnel for NetworkManager {
                 Err(error) => log::error!("Failed to fetch interface name from NM: {}", error),
             }
         }
-        "wg-mullvad".to_string()
+        MULLVAD_INTERFACE_NAME.to_string()
     }
 
     fn stop(mut self: Box<Self>) -> std::result::Result<(), TunnelError> {
@@ -366,25 +379,17 @@ impl Tunnel for NetworkManager {
     }
 
     fn get_tunnel_stats(&self) -> std::result::Result<Stats, TunnelError> {
-        if let Some(tunnel) = self.tunnel.as_ref() {
-            let device = tunnel.device_proxy(&self.dbus_connection);
-            let get_device_stats = || -> std::result::Result<Stats, dbus::Error> {
-                let rx_bytes = device.get(NM_DEVICE_STATISTICS, "RxBytes")?;
-                let tx_bytes = device.get(NM_DEVICE_STATISTICS, "TxBytes")?;
-
-                Ok(Stats { rx_bytes, tx_bytes })
-            };
+        let mut wg = self.netlink_connections.wg_handle.clone();
+        let interface_index = self.interface_index;
+        let result = self.tokio_handle.block_on(async move {
+            let device = wg.get_by_index(interface_index).await.map_err(|err| {
+                log::error!("Failed to fetch WireGuard device config: {}", err);
+                TunnelError::GetConfigError
+            })?;
+            Ok(Stats::parse_device_message(&device))
+        });
 
-            match get_device_stats() {
-                Ok(stats) => Ok(stats),
-                Err(err) => {
-                    log::error!("Failed to read tunnel stats from NM: {}", err);
-                    Err(TunnelError::StatsError(super::stats::Error::NoTunnelDevice))
-                }
-            }
-        } else {
-            Err(TunnelError::StatsError(super::stats::Error::NoTunnelDevice))
-        }
+        result
     }
 }