Merge branch 'prune-daemon-execution-state'

1 files changed, 27 insertions, 65 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 3affb7915d..9e2cbba0ea 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -70,7 +70,6 @@ use std::collections::HashSet;
 use std::os::unix::io::RawFd;
 use std::{
     marker::PhantomData,
-    mem,
     path::PathBuf,
     pin::Pin,
     sync::{Arc, Weak},
@@ -433,49 +432,6 @@ impl From<(AccessMethodEvent, oneshot::Sender<()>)> for InternalDaemonEvent {
     }
 }
 
-#[derive(Clone, Debug, Eq, PartialEq)]
-enum DaemonExecutionState {
-    Running,
-    Exiting,
-    Finished,
-}
-
-impl DaemonExecutionState {
-    pub fn shutdown(&mut self, tunnel_state: &TunnelState) {
-        use self::DaemonExecutionState::*;
-
-        match self {
-            Running => {
-                match tunnel_state {
-                    TunnelState::Disconnected { .. } => mem::replace(self, Finished),
-                    _ => mem::replace(self, Exiting),
-                };
-            }
-            Exiting | Finished => {}
-        };
-    }
-
-    pub fn disconnected(&mut self) {
-        use self::DaemonExecutionState::*;
-
-        match self {
-            Exiting => {
-                let _ = mem::replace(self, Finished);
-            }
-            Running | Finished => {}
-        };
-    }
-
-    pub fn is_running(&self) -> bool {
-        use self::DaemonExecutionState::*;
-
-        match self {
-            Running => true,
-            Exiting | Finished => false,
-        }
-    }
-}
-
 pub struct DaemonCommandChannel {
     sender: DaemonCommandSender,
     receiver: mpsc::UnboundedReceiver<InternalDaemonEvent>,
@@ -611,7 +567,6 @@ pub trait EventListener {
 pub struct Daemon<L: EventListener> {
     tunnel_state: TunnelState,
     target_state: PersistentTargetState,
-    state: DaemonExecutionState,
     #[cfg(target_os = "linux")]
     exclude_pids: split_tunnel::PidManager,
     rx: mpsc::UnboundedReceiver<InternalDaemonEvent>,
@@ -867,7 +822,6 @@ where
                 locked_down: settings.block_when_disconnected,
             },
             target_state,
-            state: DaemonExecutionState::Running,
             #[cfg(target_os = "linux")]
             exclude_pids: split_tunnel::PidManager::new().map_err(Error::InitSplitTunneling)?,
             rx: internal_event_rx,
@@ -914,12 +868,26 @@ where
         }
 
         while let Some(event) = self.rx.next().await {
-            self.handle_event(event).await;
-            if self.state == DaemonExecutionState::Finished {
+            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(())
     }
@@ -969,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);
             }
@@ -990,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(
@@ -1038,7 +1011,6 @@ where
         }
 
         match &tunnel_state {
-            TunnelState::Disconnected { .. } => self.state.disconnected(),
             TunnelState::Connecting { .. } => {
                 log::debug!("Settings: {}", self.settings.summary());
             }
@@ -1175,11 +1147,6 @@ where
 
     async fn handle_command(&mut self, command: DaemonCommand) {
         use self::DaemonCommand::*;
-        if !self.state.is_running() {
-            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();
         }
@@ -1454,12 +1421,8 @@ where
         tx: oneshot::Sender<bool>,
         new_target_state: TargetState,
     ) {
-        if self.state.is_running() {
-            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>) {
@@ -1744,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 {
@@ -2673,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
@@ -2684,7 +2647,6 @@ where
             self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(true, tx));
         }
 
-        self.state.shutdown(&self.tunnel_state);
         self.disconnect_tunnel();
     }