Allow initializing firewall without args

9 files changed, 103 insertions, 102 deletions

diff --git a/3f3fba4e2066f28a1ad7ac60e86a688a92eb5b5f b/3f3fba4e2066f28a1ad7ac60e86a688a92eb5b5f
new file mode 100644
index 0000000000..cdd25eb781
--- /dev/null
+++ b/3f3fba4e2066f28a1ad7ac60e86a688a92eb5b5f
diff --git a/mullvad-setup/src/main.rs b/mullvad-setup/src/main.rs
index 287da5df6b..947f9bd9ca 100644
--- a/mullvad-setup/src/main.rs
+++ b/mullvad-setup/src/main.rs
@@ -4,7 +4,7 @@ use mullvad_rpc::MullvadRpcRuntime;
 use mullvad_types::version::ParsedAppVersion;
 use std::{path::PathBuf, process, time::Duration};
 use talpid_core::{
-    firewall::{self, Firewall, FirewallArguments, InitialFirewallState},
+    firewall::{self, Firewall},
     future_retry::{constant_interval, retry_future_n},
 };
 use talpid_types::ErrorExt;
@@ -153,15 +153,10 @@ async fn reset_firewall() -> Result<(), Error> {
         return Err(Error::DaemonIsRunning);
     }
 
-    let mut firewall = Firewall::new(FirewallArguments {
-        initial_state: InitialFirewallState::None,
-        allow_lan: true,
-        #[cfg(target_os = "macos")]
-        exclusion_gid: 0,
-    })
-    .map_err(Error::FirewallError)?;
-
-    firewall.reset_policy().map_err(Error::FirewallError)
+    Firewall::new()
+        .map_err(Error::FirewallError)?
+        .reset_policy()
+        .map_err(Error::FirewallError)
 }
 
 async fn remove_wireguard_key() -> Result<(), Error> {
diff --git a/talpid-core/Cargo.toml b/talpid-core/Cargo.toml
index 8b55f3524a..f01682f08c 100644
--- a/talpid-core/Cargo.toml
+++ b/talpid-core/Cargo.toml
@@ -79,7 +79,7 @@ byteorder = "1"
 internet-checksum = "0.2"
 widestring = "0.5"
 winreg = { version = "0.7", features = ["transactions"] }
-winapi = { version = "0.3.6", features = ["combaseapi", "handleapi", "ifdef", "libloaderapi", "netioapi", "psapi", "stringapiset", "synchapi", "tlhelp32", "winbase", "winioctl", "winuser"] }
+winapi = { version = "0.3.6", features = ["combaseapi", "handleapi", "ifdef", "libloaderapi", "netioapi", "psapi", "stringapiset", "synchapi", "tlhelp32", "winbase", "winioctl", "winuser", "dbt"] }
 talpid-platform-metadata = { path = "../talpid-platform-metadata" }
 memoffset = "0.6"
 
diff --git a/talpid-core/src/firewall/android.rs b/talpid-core/src/firewall/android.rs
index d1959d7ae5..fc854bc804 100644
--- a/talpid-core/src/firewall/android.rs
+++ b/talpid-core/src/firewall/android.rs
@@ -1,4 +1,4 @@
-use super::{FirewallArguments, FirewallPolicy, FirewallT};
+use super::{FirewallArguments, FirewallPolicy};
 
 /// Stub error type for Firewall errors on Android.
 #[derive(Debug, err_derive::Error)]
@@ -8,18 +8,20 @@ pub struct Error;
 /// The Android stub implementation for the firewall.
 pub struct Firewall;
 
-impl FirewallT for Firewall {
-    type Error = Error;
+impl Firewall {
+    pub fn from_args(_args: FirewallArguments) -> Result<Self, Error> {
+        Ok(Firewall)
+    }
 
-    fn new(_args: FirewallArguments) -> Result<Self, Self::Error> {
+    pub fn new() -> Result<Self, Error> {
         Ok(Firewall)
     }
 
-    fn apply_policy(&mut self, _policy: FirewallPolicy) -> Result<(), Self::Error> {
+    pub fn apply_policy(&mut self, _policy: FirewallPolicy) -> Result<(), Error> {
         Ok(())
     }
 
-    fn reset_policy(&mut self) -> Result<(), Self::Error> {
+    pub fn reset_policy(&mut self) -> Result<(), Error> {
         Ok(())
     }
 }
diff --git a/talpid-core/src/firewall/linux.rs b/talpid-core/src/firewall/linux.rs
index 480bc3c674..514079036e 100644
--- a/talpid-core/src/firewall/linux.rs
+++ b/talpid-core/src/firewall/linux.rs
@@ -1,4 +1,4 @@
-use super::{FirewallArguments, FirewallPolicy, FirewallT};
+use super::{FirewallArguments, FirewallPolicy};
 use crate::{split_tunnel, tunnel};
 use ipnetwork::IpNetwork;
 use lazy_static::lazy_static;
@@ -104,26 +104,32 @@ struct FirewallTables {
     mangle_v6: Table,
 }
 
-impl FirewallT for Firewall {
-    type Error = Error;
+impl Firewall {
+    pub fn from_args(_args: FirewallArguments) -> Result<Self> {
+        Ok(Firewall(()))
+    }
 
-    fn new(_args: FirewallArguments) -> Result<Self> {
+    pub fn new() -> Result<Self> {
         Ok(Firewall(()))
     }
 
-    fn apply_policy(&mut self, policy: FirewallPolicy) -> Result<()> {
+    pub fn apply_policy(&mut self, policy: FirewallPolicy) -> Result<()> {
         let tables = FirewallTables {
             main: Table::new(&*TABLE_NAME, ProtoFamily::Inet),
             mangle_v4: Table::new(&*MANGLE_TABLE_NAME_V4, ProtoFamily::Ipv4),
             mangle_v6: Table::new(&*MANGLE_TABLE_NAME_V6, ProtoFamily::Ipv6),
         };
         let batch = PolicyBatch::new(&tables).finalize(&policy)?;
-        self.send_and_process(&batch)?;
+        Self::send_and_process(&batch)?;
         Self::apply_kernel_config(&policy);
         self.verify_tables(&[&TABLE_NAME, &MANGLE_TABLE_NAME_V4, &MANGLE_TABLE_NAME_V6])
     }
 
-    fn reset_policy(&mut self) -> Result<()> {
+    pub fn reset_policy(&mut self) -> Result<()> {
+        Self::clear_policy()
+    }
+
+    pub fn clear_policy() -> Result<()> {
         let tables = [
             Table::new(&*TABLE_NAME, ProtoFamily::Inet),
             Table::new(&*MANGLE_TABLE_NAME_V4, ProtoFamily::Ipv4),
@@ -139,12 +145,10 @@ impl FirewallT for Firewall {
         }
         let batch = batch.finalize();
         log::debug!("Removing table and chain from netfilter");
-        self.send_and_process(&batch)?;
+        Self::send_and_process(&batch)?;
         Ok(())
     }
-}
 
-impl Firewall {
     fn apply_kernel_config(policy: &FirewallPolicy) {
         if *DONT_SET_SRC_VALID_MARK {
             log::debug!("Not setting src_valid_mark");
@@ -158,7 +162,7 @@ impl Firewall {
         }
     }
 
-    fn send_and_process(&self, batch: &FinalizedBatch) -> Result<()> {
+    fn send_and_process(batch: &FinalizedBatch) -> Result<()> {
         let socket = mnl::Socket::new(mnl::Bus::Netfilter).map_err(Error::NetlinkOpenError)?;
         socket.send_all(batch).map_err(Error::NetlinkSendError)?;
 
diff --git a/talpid-core/src/firewall/macos.rs b/talpid-core/src/firewall/macos.rs
index fc81e47db4..9ae293b868 100644
--- a/talpid-core/src/firewall/macos.rs
+++ b/talpid-core/src/firewall/macos.rs
@@ -1,4 +1,4 @@
-use super::{FirewallArguments, FirewallPolicy, FirewallT};
+use super::{FirewallArguments, FirewallPolicy};
 use ipnetwork::IpNetwork;
 use pfctl::{DropAction, FilterRuleAction, Uid};
 use std::{
@@ -24,10 +24,12 @@ pub struct Firewall {
     _exclusion_gid: u32,
 }
 
-impl FirewallT for Firewall {
-    type Error = Error;
+impl Firewall {
+    pub fn from_args(args: FirewallArguments) -> Result<Self> {
+        Self::new(args.exclusion_gid)
+    }
 
-    fn new(args: FirewallArguments) -> Result<Self> {
+    fn new(exclusion_gid: u32) -> Result<Self> {
         // Allows controlling whether firewall rules should log to pflog0. Useful for debugging the
         // rules.
         let firewall_debugging = env::var("TALPID_FIREWALL_DEBUG");
@@ -43,17 +45,17 @@ impl FirewallT for Firewall {
             pf: pfctl::PfCtl::new()?,
             pf_was_enabled: None,
             rule_logging,
-            _exclusion_gid: args.exclusion_gid,
+            _exclusion_gid: exclusion_gid,
         })
     }
 
-    fn apply_policy(&mut self, policy: FirewallPolicy) -> Result<()> {
+    pub fn apply_policy(&mut self, policy: FirewallPolicy) -> Result<()> {
         self.enable()?;
         self.add_anchor()?;
         self.set_rules(policy)
     }
 
-    fn reset_policy(&mut self) -> Result<()> {
+    pub fn reset_policy(&mut self) -> Result<()> {
         // Implemented this way to not early return on an error.
         // We always want all three methods to run, and then return
         // the first error it encounterd, if any.
@@ -61,9 +63,7 @@ impl FirewallT for Firewall {
             .and(self.remove_anchor())
             .and(self.restore_state())
     }
-}
 
-impl Firewall {
     fn set_rules(&mut self, policy: FirewallPolicy) -> Result<()> {
         let mut new_filter_rules = vec![];
 
diff --git a/talpid-core/src/firewall/mod.rs b/talpid-core/src/firewall/mod.rs
index 7231f8f20c..f2ac129ba2 100644
--- a/talpid-core/src/firewall/mod.rs
+++ b/talpid-core/src/firewall/mod.rs
@@ -225,10 +225,6 @@ pub struct FirewallArguments {
     pub initial_state: InitialFirewallState,
     /// This argument is required for the blocked state to configure the firewall correctly.
     pub allow_lan: bool,
-    #[cfg(target_os = "macos")]
-    /// This argument is required on macOS to know which group's traffic should be excluded, if at
-    /// all.
-    pub exclusion_gid: u32,
 }
 
 /// State to enter during firewall init.
@@ -240,10 +236,17 @@ pub enum InitialFirewallState {
 }
 
 impl Firewall {
-    /// Returns a new `Firewall`, ready to apply policies.
-    pub fn new(args: FirewallArguments) -> Result<Self, Error> {
+    /// Creates a firewall instance with the given arguments.
+    pub fn from_args(args: FirewallArguments) -> Result<Self, Error> {
         Ok(Firewall {
-            inner: imp::Firewall::new(args)?,
+            inner: imp::Firewall::from_args(args)?,
+        })
+    }
+
+    /// Createsa new firewall instance.
+    pub fn new() -> Result<Self, Error> {
+        Ok(Firewall {
+            inner: imp::Firewall::new()?,
         })
     }
 
@@ -261,19 +264,3 @@ impl Firewall {
         self.inner.reset_policy()
     }
 }
-
-/// Abstract firewall interaction trait. Used by the OS specific implementations.
-trait FirewallT: Sized {
-    /// The error type thrown by the implementer of this trait
-    type Error: std::error::Error;
-
-    /// Create new instance
-    fn new(args: FirewallArguments) -> Result<Self, Self::Error>;
-
-    /// Enable the given FirewallPolicy
-    fn apply_policy(&mut self, policy: FirewallPolicy) -> Result<(), Self::Error>;
-
-    /// Revert the system firewall state to what it was before this instance started
-    /// modifying the system.
-    fn reset_policy(&mut self) -> Result<(), Self::Error>;
-}
diff --git a/talpid-core/src/firewall/windows.rs b/talpid-core/src/firewall/windows.rs
index 81640038b6..dca86d2257 100644
--- a/talpid-core/src/firewall/windows.rs
+++ b/talpid-core/src/firewall/windows.rs
@@ -3,7 +3,7 @@ use crate::{logging::windows::log_sink, tunnel::TunnelMetadata};
 use std::{net::IpAddr, path::Path, ptr};
 
 use self::winfw::*;
-use super::{FirewallArguments, FirewallPolicy, FirewallT, InitialFirewallState};
+use super::{FirewallArguments, FirewallPolicy, InitialFirewallState};
 use crate::winnet;
 use talpid_types::{
     net::{AllowedEndpoint, Endpoint},
@@ -47,40 +47,55 @@ pub enum Error {
 /// Timeout for acquiring the WFP transaction lock
 const WINFW_TIMEOUT_SECONDS: u32 = 5;
 
+const LOGGING_CONTEXT: &[u8] = b"WinFw\0";
+
 /// The Windows implementation for the firewall and DNS.
 pub struct Firewall(());
 
-impl FirewallT for Firewall {
-    type Error = Error;
-
-    fn new(args: FirewallArguments) -> Result<Self, Self::Error> {
-        let logging_context = b"WinFw\0".as_ptr();
-
+impl Firewall {
+    pub fn from_args(args: FirewallArguments) -> Result<Self, Error> {
         if let InitialFirewallState::Blocked(allowed_endpoint) = args.initial_state {
-            let cfg = &WinFwSettings::new(args.allow_lan);
-            let allowed_endpoint = WinFwAllowedEndpointContainer::from(allowed_endpoint);
-            unsafe {
-                WinFw_InitializeBlocked(
-                    WINFW_TIMEOUT_SECONDS,
-                    &cfg,
-                    &allowed_endpoint.as_endpoint(),
-                    Some(log_sink),
-                    logging_context,
-                )
-                .into_result()?
-            };
+            Self::initialize_blocked(allowed_endpoint, args.allow_lan)
         } else {
-            unsafe {
-                WinFw_Initialize(WINFW_TIMEOUT_SECONDS, Some(log_sink), logging_context)
-                    .into_result()?
-            };
+            Self::new()
         }
+    }
+
+    pub fn new() -> Result<Self, Error> {
+        unsafe {
+            WinFw_Initialize(
+                WINFW_TIMEOUT_SECONDS,
+                Some(log_sink),
+                LOGGING_CONTEXT.as_ptr(),
+            )
+            .into_result()?
+        };
 
         log::trace!("Successfully initialized windows firewall module");
         Ok(Firewall(()))
     }
 
-    fn apply_policy(&mut self, policy: FirewallPolicy) -> Result<(), Self::Error> {
+    fn initialize_blocked(
+        allowed_endpoint: AllowedEndpoint,
+        allow_lan: bool,
+    ) -> Result<Self, Error> {
+        let cfg = &WinFwSettings::new(allow_lan);
+        let allowed_endpoint = WinFwAllowedEndpointContainer::from(allowed_endpoint);
+        unsafe {
+            WinFw_InitializeBlocked(
+                WINFW_TIMEOUT_SECONDS,
+                &cfg,
+                &allowed_endpoint.as_endpoint(),
+                Some(log_sink),
+                LOGGING_CONTEXT.as_ptr(),
+            )
+            .into_result()?
+        };
+        log::trace!("Successfully initialized windows firewall module to a blocking state");
+        Ok(Firewall(()))
+    }
+
+    pub fn apply_policy(&mut self, policy: FirewallPolicy) -> Result<(), Error> {
         match policy {
             FirewallPolicy::Connecting {
                 peer_endpoint,
@@ -122,27 +137,11 @@ impl FirewallT for Firewall {
         }
     }
 
-    fn reset_policy(&mut self) -> Result<(), Self::Error> {
+    pub fn reset_policy(&mut self) -> Result<(), Error> {
         unsafe { WinFw_Reset().into_result().map_err(Error::ResettingPolicy) }?;
         Ok(())
     }
-}
 
-impl Drop for Firewall {
-    fn drop(&mut self) {
-        if unsafe {
-            WinFw_Deinitialize(WinFwCleanupPolicy::ContinueBlocking)
-                .into_result()
-                .is_ok()
-        } {
-            log::trace!("Successfully deinitialized windows firewall module");
-        } else {
-            log::error!("Failed to deinitialize windows firewall module");
-        };
-    }
-}
-
-impl Firewall {
     fn set_connecting_state(
         &mut self,
         endpoint: &Endpoint,
@@ -257,6 +256,20 @@ impl Firewall {
     }
 }
 
+impl Drop for Firewall {
+    fn drop(&mut self) {
+        if unsafe {
+            WinFw_Deinitialize(WinFwCleanupPolicy::ContinueBlocking)
+                .into_result()
+                .is_ok()
+        } {
+            log::trace!("Successfully deinitialized windows firewall module");
+        } else {
+            log::error!("Failed to deinitialize windows firewall module");
+        };
+    }
+}
+
 fn widestring_ip(ip: IpAddr) -> WideCString {
     WideCString::from_str_truncate(ip.to_string())
 }
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index 8c4446f3f3..154864228b 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -241,7 +241,7 @@ impl TunnelStateMachine {
             exclusion_gid,
         };
 
-        let firewall = Firewall::new(args).map_err(Error::InitFirewallError)?;
+        let firewall = Firewall::from_args(args).map_err(Error::InitFirewallError)?;
         let route_manager = RouteManager::new(HashSet::new())
             .await
             .map_err(Error::InitRouteManagerError)?;