Reorganize WG kernel tunnel

5 files changed, 192 insertions, 188 deletions

diff --git a/talpid-core/src/tunnel/wireguard/mod.rs b/talpid-core/src/tunnel/wireguard/mod.rs
index 22dc1c77f8..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(route_manager.runtime_handle(), 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 528d79199d..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::{
@@ -20,10 +20,15 @@ use tokio::stream::StreamExt;
 mod parsers;
 
 pub mod wg_message;
-use wg_message::{DeviceMessage, DeviceNla, PeerNla};
-mod nl_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,138 +80,23 @@ 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 {
     pub wg_handle: WireguardConnection,
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 748f6df7a0..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},
@@ -53,21 +56,7 @@ pub enum Error {
     FindInterfaceIndex,
 }
 
-impl Error {
-    pub fn should_use_userspace(&self) -> bool {
-        match self {
-            Error::NMTooOld(_) => true,
-            _ => false,
-        }
-    }
-}
-
-use super::wireguard_kernel::{
-    wg_message::{DeviceNla, PeerNla},
-    Handle,
-};
-
-pub struct NetworkManager {
+pub struct NetworkManagerTunnel {
     dbus_connection: Connection,
     tunnel: Option<WireguardTunnel>,
     interface_index: u32,
@@ -81,21 +70,26 @@ type VariantMap = HashMap<String, VariantRefArg>;
 // settings are a{sa{sv}}
 type DbusSettings = HashMap<String, VariantMap>;
 
-impl NetworkManager {
-    pub fn new(tokio_handle: tokio::runtime::Handle, 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)?);
+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
-                .map_err(|_| Error::FindInterfaceIndex)?;
+            let netlink_connections = Handle::connect().await?;
             let interface_index = Self::find_device_index(&netlink_connections)
                 .await?
-                .ok_or(Error::FindInterfaceIndex)?;
+                .ok_or(WgKernelError::NetworkManager(Error::FindInterfaceIndex))?;
 
-            Ok(NetworkManager {
+            Ok(NetworkManagerTunnel {
                 dbus_connection,
                 tunnel,
                 interface_index,
@@ -105,15 +99,11 @@ impl NetworkManager {
         })
     }
 
-    async fn find_device_index(netlink_connections: &Handle) -> Result<Option<u32>> {
+    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("wg-mullvad".to_string())
-            .await
-            .map_err(|err| {
-                log::error!("Failed to fetch WireGuard device config: {}", err);
-                Error::FindInterfaceIndex
-            })?;
+        let device = wg.get_by_name(MULLVAD_INTERFACE_NAME.to_string()).await?;
 
         for nla in &device.nlas {
             if let DeviceNla::IfIndex(index) = nla {
@@ -269,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)));
 
@@ -359,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
@@ -373,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> {
@@ -393,26 +386,7 @@ impl Tunnel for NetworkManager {
                 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 })
+            Ok(Stats::parse_device_message(&device))
         });
 
         result