Persist blocking firewall rules across a reboot conditionally

5 files changed, 131 insertions, 15 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index efdbea5cc6..c33ecf542b 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -83,6 +83,8 @@ use std::{
     sync::{Arc, Weak},
     time::Duration,
 };
+#[cfg(not(target_os = "android"))]
+use talpid_core::tunnel_state_machine::BlockWhenDisconnected;
 use talpid_core::{
     mpsc::Sender,
     split_tunnel,
@@ -875,7 +877,9 @@ impl Daemon {
             tunnel_state_machine::InitialTunnelState {
                 allow_lan: settings.allow_lan,
                 #[cfg(not(target_os = "android"))]
-                block_when_disconnected: settings.block_when_disconnected,
+                block_when_disconnected: BlockWhenDisconnected::from(
+                    settings.block_when_disconnected,
+                ),
                 dns_config: dns::addresses_from_options(&settings.tunnel_options.dns_options),
                 allowed_endpoint: access_mode_handler
                     .get_current()
@@ -2385,7 +2389,7 @@ impl Daemon {
             Ok(settings_changed) => {
                 if settings_changed {
                     self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(
-                        block_when_disconnected,
+                        BlockWhenDisconnected::from(block_when_disconnected),
                         oneshot_map(tx, |tx, ()| {
                             Self::oneshot_send(tx, Ok(()), "set_block_when_disconnected response");
                         }),
@@ -3087,7 +3091,7 @@ impl Daemon {
         {
             let (tx, _rx) = oneshot::channel();
             self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(
-                self.settings.block_when_disconnected,
+                BlockWhenDisconnected::from(self.settings.block_when_disconnected),
                 tx,
             ));
         }
@@ -3149,7 +3153,10 @@ impl Daemon {
         {
             log::debug!("Blocking firewall during shutdown");
             let (tx, _rx) = oneshot::channel();
-            self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(true, tx));
+            self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(
+                BlockWhenDisconnected::yes(),
+                tx,
+            ));
         }
 
         self.disconnect_tunnel();
@@ -3164,8 +3171,21 @@ impl Daemon {
         //       without causing the service to be restarted.
         #[cfg(not(target_os = "android"))]
         if *self.target_state == TargetState::Secured {
+            let persist = if cfg!(target_os = "windows") {
+                // During app upgrades, as a safety measure, we make the firewall filters
+                // non-persistent. If the installation of the new version fails and
+                // the user is left in blocked state with no app, they can reboot
+                // to regain internet access.
+                self.settings.settings().block_when_disconnected
+                    || self.settings.settings().auto_connect
+            } else {
+                true
+            };
             let (tx, _rx) = oneshot::channel();
-            self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(true, tx));
+            self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(
+                BlockWhenDisconnected::yes().persist(persist),
+                tx,
+            ));
         }
         self.target_state.lock();
 
diff --git a/talpid-core/src/firewall/mod.rs b/talpid-core/src/firewall/mod.rs
index 053317f2a5..4255854624 100644
--- a/talpid-core/src/firewall/mod.rs
+++ b/talpid-core/src/firewall/mod.rs
@@ -326,4 +326,10 @@ impl Firewall {
         log::info!("Resetting firewall policy");
         self.inner.reset_policy()
     }
+
+    /// Sets whether the firewall should persist the blocking rules across a reboot.
+    #[cfg(target_os = "windows")]
+    pub fn persist(&mut self, persist: bool) {
+        self.inner.persist(persist);
+    }
 }
diff --git a/talpid-core/src/firewall/windows/mod.rs b/talpid-core/src/firewall/windows/mod.rs
index 7b273ed6c1..13d3f5db90 100644
--- a/talpid-core/src/firewall/windows/mod.rs
+++ b/talpid-core/src/firewall/windows/mod.rs
@@ -185,6 +185,10 @@ impl Firewall {
         Ok(())
     }
 
+    pub fn persist(&mut self, persist: bool) {
+        self.persist = persist;
+    }
+
     fn set_connecting_state(
         &mut self,
         peer_endpoint: &AllowedEndpoint,
@@ -231,7 +235,8 @@ impl Firewall {
 
 impl Drop for Firewall {
     fn drop(&mut self) {
-        // TODO: Comment mee
+        // Deinitialize WinFW with or without persistent filters.
+        // All other filters should still remain intact.
         let cleanup_policy = if self.persist {
             WinFwCleanupPolicy::ContinueBlocking
         } else {
diff --git a/talpid-core/src/tunnel_state_machine/disconnected_state.rs b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
index 8f96ff7b90..427a6e5cb3 100644
--- a/talpid-core/src/tunnel_state_machine/disconnected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
@@ -30,7 +30,7 @@ impl DisconnectedState {
             );
         }
         #[cfg(target_os = "macos")]
-        if shared_values.block_when_disconnected {
+        if shared_values.block_when_disconnected.bool() {
             if let Err(err) = Self::setup_local_dns_config(shared_values) {
                 log::error!(
                     "{}",
@@ -64,7 +64,7 @@ impl DisconnectedState {
                 // Being disconnected and having lockdown mode enabled implies that your internet
                 // access is locked down
                 #[cfg(not(target_os = "android"))]
-                locked_down: shared_values.block_when_disconnected,
+                locked_down: shared_values.block_when_disconnected.bool(),
             },
         )
     }
@@ -74,7 +74,15 @@ impl DisconnectedState {
         shared_values: &mut SharedTunnelStateValues,
         should_reset_firewall: bool,
     ) {
-        let result = if shared_values.block_when_disconnected {
+        let result = if shared_values.block_when_disconnected.bool() {
+            #[cfg(target_os = "windows")]
+            {
+                // Respect the persist flag of BlockWhenDisconnected.
+                shared_values
+                    .firewall
+                    .persist(shared_values.block_when_disconnected.should_persist());
+            }
+
             let policy = FirewallPolicy::Blocked {
                 allow_lan: shared_values.allow_lan,
                 allowed_endpoint: Some(shared_values.allowed_endpoint.clone()),
@@ -110,7 +118,7 @@ impl DisconnectedState {
         shared_values: &mut SharedTunnelStateValues,
         should_reset_firewall: bool,
     ) {
-        if should_reset_firewall && !shared_values.block_when_disconnected {
+        if should_reset_firewall && !shared_values.block_when_disconnected.bool() {
             if let Err(error) = shared_values.split_tunnel.clear_tunnel_addresses() {
                 log::error!(
                     "{}",
@@ -198,7 +206,7 @@ impl TunnelState for DisconnectedState {
                     #[cfg(windows)]
                     Self::register_split_tunnel_addresses(shared_values, true);
                     #[cfg(target_os = "macos")]
-                    if block_when_disconnected {
+                    if block_when_disconnected.bool() {
                         if let Err(err) = Self::setup_local_dns_config(shared_values) {
                             log::error!(
                                 "{}",
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index 059edd417d..2deb7c8f7e 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -94,7 +94,7 @@ pub struct InitialTunnelState {
     pub allow_lan: bool,
     /// Block traffic unless connected to the VPN.
     #[cfg(not(target_os = "android"))]
-    pub block_when_disconnected: bool,
+    pub block_when_disconnected: BlockWhenDisconnected,
     /// DNS configuration to use
     pub dns_config: DnsConfig,
     /// A single endpoint that is allowed to communicate outside the tunnel, i.e.
@@ -200,7 +200,7 @@ pub enum TunnelCommand {
     Dns(crate::dns::DnsConfig, oneshot::Sender<()>),
     /// Enable or disable the block_when_disconnected feature.
     #[cfg(not(target_os = "android"))]
-    BlockWhenDisconnected(bool, oneshot::Sender<()>),
+    BlockWhenDisconnected(BlockWhenDisconnected, oneshot::Sender<()>),
     /// Notify the state machine of the connectivity of the device.
     Connectivity(Connectivity),
     /// Open tunnel connection.
@@ -234,6 +234,82 @@ enum EventResult {
     Close(Result<Option<ErrorStateCause>, oneshot::Canceled>),
 }
 
+/// If firewall should apply blocking rules in the disconnected state.
+/// Argument of TunnelCommand::BlockWhenDisconnected message.
+///
+/// Semantically equivalent to a boolean value, but is grouped togetether with the persist
+/// parameter on Windows for cohesiveness.
+#[derive(Clone, Copy, Debug)]
+pub enum BlockWhenDisconnected {
+    /// Firewall should *not* apply blocking rules.
+    Disabled,
+    /// Firewall should apply blocking rules.
+    Enabled {
+        /// If blocked state should be persisted across a reboot (restart of BFE)
+        persist: bool,
+    },
+}
+
+impl BlockWhenDisconnected {
+    /// `true`. Apply blocking firewall rules in the disconnected state.
+    pub const fn yes() -> Self {
+        BlockWhenDisconnected::Enabled { persist: true }
+    }
+
+    /// `false`. Do *not* apply blocking firewall rules in the disconnected state.
+    pub const fn no() -> Self {
+        BlockWhenDisconnected::Disabled
+    }
+
+    /// [self] as a boolean value.
+    pub const fn bool(&self) -> bool {
+        matches!(self, BlockWhenDisconnected::Enabled { .. })
+    }
+
+    /// If [BlockWhenDisconnected] should persist across reboots.
+    ///
+    /// Semantically meaningless on non-Windows platforms, will always return true.
+    pub const fn should_persist(&self) -> bool {
+        if cfg!(target_os = "windows") {
+            matches!(&self, BlockWhenDisconnected::Enabled { persist: true })
+        } else {
+            true
+        }
+    }
+
+    /// Semantically meaningless on non-Windows platforms
+    #[cfg(not(target_os = "windows"))]
+    pub fn persist(self, _persist: bool) -> Self {
+        self
+    }
+
+    /// Semantically meaningless on non-Windows platforms
+    #[cfg(target_os = "windows")]
+    pub fn persist(self, persist: bool) -> Self {
+        match self {
+            BlockWhenDisconnected::Disabled => BlockWhenDisconnected::Disabled,
+            // Forget previous value of persist
+            BlockWhenDisconnected::Enabled { .. } => BlockWhenDisconnected::Enabled { persist },
+        }
+    }
+}
+
+impl From<bool> for BlockWhenDisconnected {
+    fn from(block: bool) -> Self {
+        if block {
+            BlockWhenDisconnected::yes()
+        } else {
+            BlockWhenDisconnected::no()
+        }
+    }
+}
+
+impl PartialEq for BlockWhenDisconnected {
+    fn eq(&self, other: &Self) -> bool {
+        self.bool() == other.bool()
+    }
+}
+
 /// Asynchronous handling of the tunnel state machine.
 ///
 /// This type implements `Stream`, and attempts to advance the state machine based on the events
@@ -295,7 +371,8 @@ impl TunnelStateMachine {
 
         let fw_args = FirewallArguments {
             #[cfg(not(target_os = "android"))]
-            initial_state: if args.settings.block_when_disconnected || !args.settings.reset_firewall
+            initial_state: if args.settings.block_when_disconnected.bool()
+                || !args.settings.reset_firewall
             {
                 InitialFirewallState::Blocked(args.settings.allowed_endpoint.clone())
             } else {
@@ -474,7 +551,7 @@ struct SharedTunnelStateValues {
     allow_lan: bool,
     /// Should network access be allowed when in the disconnected state.
     #[cfg(not(target_os = "android"))]
-    block_when_disconnected: bool,
+    block_when_disconnected: BlockWhenDisconnected,
     /// True when the computer is known to be offline.
     connectivity: Connectivity,
     /// DNS configuration to use.