Make cached target state start Firewall correctly

8 files changed, 41 insertions, 35 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index e159e33be7..0c5f11ab9d 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -558,6 +558,21 @@ where
         let tunnel_parameters_generator = MullvadTunnelParametersGenerator {
             tx: internal_event_tx.clone(),
         };
+
+
+        let initial_target_state = if settings.get_account_token().is_some() {
+            if settings.auto_connect {
+                // Note: Auto-connect overrides the cached target state
+                info!("Automatically connecting since auto-connect is turned on");
+                TargetState::Secured
+            } else {
+                cached_target_state.unwrap_or(TargetState::Unsecured)
+            }
+        } else {
+            TargetState::Unsecured
+        };
+
+
         let tunnel_command_tx = tunnel_state_machine::spawn(
             settings.allow_lan,
             settings.block_when_disconnected,
@@ -567,6 +582,7 @@ where
             cache_dir.clone(),
             internal_event_tx.to_specialized_sender(),
             tunnel_state_machine_shutdown_tx,
+            initial_target_state != TargetState::Secured,
             #[cfg(target_os = "android")]
             android_context,
         )
@@ -579,18 +595,6 @@ where
         // Attempt to download a fresh relay list
         relay_selector.update().await;
 
-        let initial_target_state = if settings.get_account_token().is_some() {
-            if settings.auto_connect {
-                // Note: Auto-connect overrides the cached target state
-                info!("Automatically connecting since auto-connect is turned on");
-                TargetState::Secured
-            } else {
-                cached_target_state.unwrap_or(TargetState::Unsecured)
-            }
-        } else {
-            TargetState::Unsecured
-        };
-
         let mut daemon = Daemon {
             tunnel_command_tx,
             tunnel_state: TunnelState::Disconnected,
diff --git a/mullvad-setup/src/main.rs b/mullvad-setup/src/main.rs
index c53394a473..7c587cbe72 100644
--- a/mullvad-setup/src/main.rs
+++ b/mullvad-setup/src/main.rs
@@ -72,7 +72,7 @@ async fn reset_firewall() -> Result<(), Error> {
 
     let mut firewall = Firewall::new(FirewallArguments {
         initialize_blocked: false,
-        allow_lan: None,
+        allow_lan: true,
     })
     .map_err(Error::FirewallError)?;
 
diff --git a/talpid-core/src/firewall/mod.rs b/talpid-core/src/firewall/mod.rs
index 4d8d3f0459..658fba72ae 100644
--- a/talpid-core/src/firewall/mod.rs
+++ b/talpid-core/src/firewall/mod.rs
@@ -186,7 +186,7 @@ pub struct FirewallArguments {
     /// Determines whether the firewall should atomically enter the blocked state during init.
     pub initialize_blocked: bool,
     /// This argument is required for the blocked state to configure the firewall correctly.
-    pub allow_lan: Option<bool>,
+    pub allow_lan: bool,
 }
 
 impl Firewall {
diff --git a/talpid-core/src/firewall/windows.rs b/talpid-core/src/firewall/windows.rs
index 0aa34a5f06..9da2c611b2 100644
--- a/talpid-core/src/firewall/windows.rs
+++ b/talpid-core/src/firewall/windows.rs
@@ -56,7 +56,7 @@ impl FirewallT for Firewall {
         let logging_context = b"WinFw\0".as_ptr();
 
         if args.initialize_blocked {
-            let cfg = &WinFwSettings::new(args.allow_lan.unwrap());
+            let cfg = &WinFwSettings::new(args.allow_lan);
             unsafe {
                 WinFw_InitializeBlocked(
                     WINFW_TIMEOUT_SECONDS,
diff --git a/talpid-core/src/tunnel_state_machine/disconnected_state.rs b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
index 0b1005a355..b9a13c7b16 100644
--- a/talpid-core/src/tunnel_state_machine/disconnected_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnected_state.rs
@@ -10,7 +10,10 @@ use talpid_types::ErrorExt;
 pub struct DisconnectedState;
 
 impl DisconnectedState {
-    fn set_firewall_policy(shared_values: &mut SharedTunnelStateValues) {
+    fn set_firewall_policy(
+        shared_values: &mut SharedTunnelStateValues,
+        should_reset_firewall: bool,
+    ) {
         let result = if shared_values.block_when_disconnected {
             let policy = FirewallPolicy::Blocked {
                 allow_lan: shared_values.allow_lan,
@@ -20,11 +23,13 @@ impl DisconnectedState {
                     "Failed to apply blocking firewall policy for disconnected state",
                 )
             })
-        } else {
+        } else if should_reset_firewall {
             shared_values
                 .firewall
                 .reset_policy()
                 .map_err(|e| e.display_chain_with_msg("Failed to reset firewall policy"))
+        } else {
+            Ok(())
         };
         if let Err(error_chain) = result {
             log::error!("{}", error_chain);
@@ -33,11 +38,11 @@ impl DisconnectedState {
 }
 
 impl TunnelState for DisconnectedState {
-    type Bootstrap = ();
+    type Bootstrap = bool;
 
     fn enter(
         shared_values: &mut SharedTunnelStateValues,
-        _: Self::Bootstrap,
+        should_reset_firewall: Self::Bootstrap,
     ) -> (TunnelStateWrapper, TunnelStateTransition) {
         #[cfg(target_os = "linux")]
         if let Err(error) = shared_values.route_manager.disable_exclusions_routes() {
@@ -46,7 +51,7 @@ impl TunnelState for DisconnectedState {
                 error.display_chain_with_msg("Failed to disable exclusions routes")
             );
         }
-        Self::set_firewall_policy(shared_values);
+        Self::set_firewall_policy(shared_values, should_reset_firewall);
         #[cfg(target_os = "android")]
         shared_values.tun_provider.close_tun();
 
@@ -72,14 +77,14 @@ impl TunnelState for DisconnectedState {
                         .set_allow_lan(allow_lan)
                         .expect("Failed to set allow LAN parameter");
 
-                    Self::set_firewall_policy(shared_values);
+                    Self::set_firewall_policy(shared_values, true);
                 }
                 SameState(self)
             }
             Ok(TunnelCommand::BlockWhenDisconnected(block_when_disconnected)) => {
                 if shared_values.block_when_disconnected != block_when_disconnected {
                     shared_values.block_when_disconnected = block_when_disconnected;
-                    Self::set_firewall_policy(shared_values);
+                    Self::set_firewall_policy(shared_values, true);
                 }
                 SameState(self)
             }
diff --git a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
index 21e3198978..634cae45f2 100644
--- a/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
+++ b/talpid-core/src/tunnel_state_machine/disconnecting_state.rs
@@ -125,7 +125,7 @@ impl DisconnectingState {
         }
 
         match self.after_disconnect {
-            AfterDisconnect::Nothing => DisconnectedState::enter(shared_values, ()),
+            AfterDisconnect::Nothing => DisconnectedState::enter(shared_values, true),
             AfterDisconnect::Block(cause) => ErrorState::enter(shared_values, cause),
             AfterDisconnect::Reconnect(retry_attempt) => {
                 ConnectingState::enter(shared_values, retry_attempt)
diff --git a/talpid-core/src/tunnel_state_machine/error_state.rs b/talpid-core/src/tunnel_state_machine/error_state.rs
index 48ffc39e04..d6a434b055 100644
--- a/talpid-core/src/tunnel_state_machine/error_state.rs
+++ b/talpid-core/src/tunnel_state_machine/error_state.rs
@@ -115,7 +115,7 @@ impl TunnelState for ErrorState {
             }
             Ok(TunnelCommand::Connect) => NewState(ConnectingState::enter(shared_values, 0)),
             Ok(TunnelCommand::Disconnect) | Err(_) => {
-                NewState(DisconnectedState::enter(shared_values, ()))
+                NewState(DisconnectedState::enter(shared_values, true))
             }
             Ok(TunnelCommand::Block(reason)) => NewState(ErrorState::enter(shared_values, reason)),
         }
diff --git a/talpid-core/src/tunnel_state_machine/mod.rs b/talpid-core/src/tunnel_state_machine/mod.rs
index 4d96a2fe3c..8a3b588927 100644
--- a/talpid-core/src/tunnel_state_machine/mod.rs
+++ b/talpid-core/src/tunnel_state_machine/mod.rs
@@ -84,6 +84,7 @@ pub async fn spawn(
     cache_dir: impl AsRef<Path> + Send + 'static,
     state_change_listener: impl Sender<TunnelStateTransition> + Send + 'static,
     shutdown_tx: oneshot::Sender<()>,
+    reset_firewall: bool,
     #[cfg(target_os = "android")] android_context: AndroidContext,
 ) -> Result<Arc<mpsc::UnboundedSender<TunnelCommand>>, Error> {
     let (command_tx, mut command_rx) = mpsc::unbounded();
@@ -126,6 +127,7 @@ pub async fn spawn(
             resource_dir,
             cache_dir,
             command_adapter_rx,
+            reset_firewall,
         );
         let state_machine = match state_machine {
             Ok(state_machine) => {
@@ -200,17 +202,11 @@ impl TunnelStateMachine {
         resource_dir: PathBuf,
         cache_dir: impl AsRef<Path>,
         commands: old_mpsc::UnboundedReceiver<TunnelCommand>,
+        reset_firewall: bool,
     ) -> Result<Self, Error> {
-        let args = if block_when_disconnected {
-            FirewallArguments {
-                initialize_blocked: true,
-                allow_lan: Some(allow_lan),
-            }
-        } else {
-            FirewallArguments {
-                initialize_blocked: false,
-                allow_lan: None,
-            }
+        let args = FirewallArguments {
+            initialize_blocked: block_when_disconnected || !reset_firewall,
+            allow_lan,
         };
 
         let firewall = Firewall::new(args).map_err(Error::InitFirewallError)?;
@@ -230,7 +226,8 @@ impl TunnelStateMachine {
             resource_dir,
         };
 
-        let (initial_state, _) = DisconnectedState::enter(&mut shared_values, ());
+        let (initial_state, _) = DisconnectedState::enter(&mut shared_values, reset_firewall);
+
         Ok(TunnelStateMachine {
             current_state: Some(initial_state),
             commands,