Refactor shutdown logic

Co-authored-by: David Lönnhager <david.l@mullvad.net>

1 files changed, 27 insertions, 20 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 33b15bd6eb..9e2cbba0ea 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -567,7 +567,6 @@ pub trait EventListener {
 pub struct Daemon<L: EventListener> {
     tunnel_state: TunnelState,
     target_state: PersistentTargetState,
-    shutting_down: bool,
     #[cfg(target_os = "linux")]
     exclude_pids: split_tunnel::PidManager,
     rx: mpsc::UnboundedReceiver<InternalDaemonEvent>,
@@ -823,7 +822,6 @@ where
                 locked_down: settings.block_when_disconnected,
             },
             target_state,
-            shutting_down: false,
             #[cfg(target_os = "linux")]
             exclude_pids: split_tunnel::PidManager::new().map_err(Error::InitSplitTunneling)?,
             rx: internal_event_rx,
@@ -870,12 +868,26 @@ where
         }
 
         while let Some(event) = self.rx.next().await {
-            self.handle_event(event).await;
-            if self.shutting_down && self.tunnel_state.is_disconnected() {
+            if self.handle_event(event).await {
                 break;
             }
         }
 
+        // Wait for tunnel state machine to disconnect
+        if !self.tunnel_state.is_disconnected() {
+            while let Some(event) = self.rx.next().await {
+                if let InternalDaemonEvent::TunnelStateTransition(transition) = event {
+                    self.handle_tunnel_state_transition(transition).await;
+                } else {
+                    log::trace!("Ignoring event because the daemon is shutting down");
+                }
+
+                if self.tunnel_state.is_disconnected() {
+                    break;
+                }
+            }
+        }
+
         self.finalize().await;
         Ok(())
     }
@@ -925,14 +937,18 @@ where
         )
     }
 
-    async fn handle_event(&mut self, event: InternalDaemonEvent) {
+    async fn handle_event(&mut self, event: InternalDaemonEvent) -> bool {
         use self::InternalDaemonEvent::*;
+        let mut should_stop = false;
         match event {
             TunnelStateTransition(transition) => {
                 self.handle_tunnel_state_transition(transition).await
             }
             Command(command) => self.handle_command(command).await,
-            TriggerShutdown(user_init_shutdown) => self.trigger_shutdown_event(user_init_shutdown),
+            TriggerShutdown(user_init_shutdown) => {
+                self.trigger_shutdown(user_init_shutdown);
+                should_stop = true;
+            }
             NewAppVersionInfo(app_version_info) => {
                 self.handle_new_app_version_info(app_version_info);
             }
@@ -946,6 +962,7 @@ where
             #[cfg(any(windows, target_os = "android", target_os = "macos"))]
             ExcludedPathsEvent(update, tx) => self.handle_new_excluded_paths(update, tx).await,
         }
+        should_stop
     }
 
     async fn handle_tunnel_state_transition(
@@ -1130,11 +1147,6 @@ where
 
     async fn handle_command(&mut self, command: DaemonCommand) {
         use self::DaemonCommand::*;
-        if self.shutting_down {
-            log::trace!("Dropping daemon command because the daemon is shutting down",);
-            return;
-        }
-
         if self.tunnel_state.is_disconnected() {
             self.api_handle.availability.reset_inactivity_timer();
         }
@@ -1409,12 +1421,8 @@ where
         tx: oneshot::Sender<bool>,
         new_target_state: TargetState,
     ) {
-        if !self.shutting_down {
-            let state_change_initated = self.set_target_state(new_target_state).await;
-            Self::oneshot_send(tx, state_change_initated, "state change initiated");
-        } else {
-            log::warn!("Ignoring target state change request due to shutdown");
-        }
+        let state_change_initated = self.set_target_state(new_target_state).await;
+        Self::oneshot_send(tx, state_change_initated, "state change initiated");
     }
 
     fn on_reconnect(&mut self, tx: oneshot::Sender<bool>) {
@@ -1699,7 +1707,7 @@ where
         }
 
         // Shut the daemon down.
-        self.trigger_shutdown_event(false);
+        let _ = self.tx.send(InternalDaemonEvent::TriggerShutdown(false));
 
         self.shutdown_tasks.push(Box::pin(async move {
             if let Err(e) = cleanup::clear_directories().await {
@@ -2628,7 +2636,7 @@ where
         }
     }
 
-    fn trigger_shutdown_event(&mut self, user_init_shutdown: bool) {
+    fn trigger_shutdown(&mut self, user_init_shutdown: bool) {
         // Block all traffic before shutting down to ensure that no traffic can leak on boot or
         // shutdown.
         if !user_init_shutdown
@@ -2639,7 +2647,6 @@ where
             self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(true, tx));
         }
 
-        self.shutting_down = true;
         self.disconnect_tunnel();
     }