Davids refactor

Author: Serock3

mullvad-daemon/src/lib.rs

@@ -33,7 +33,7 @@ use api::AccessMethodEvent;
 use device::{AccountEvent, PrivateAccountAndDevice, PrivateDeviceEvent};
 use futures::{
     channel::{mpsc, oneshot},
-    future::{abortable, AbortHandle, Future, LocalBoxFuture},
+    future::{abortable, AbortHandle, Future},
     StreamExt,
 };
 use geoip::GeoIpHandler;
@@ -569,7 +569,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>,
@@ -825,7 +824,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,
@@ -872,8 +870,7 @@ 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;
             }
         }
@@ -883,8 +880,27 @@ where
     }
 
     async fn finalize(self) {
-        let (event_listener, shutdown_tasks, api_runtime, tunnel_state_machine_handle) =
-            self.shutdown();
+        let (event_listener, shutdown_tasks, api_runtime, tunnel_state_machine_handle) = {
+            let Daemon {
+                event_listener,
+                mut shutdown_tasks,
+                api_runtime,
+                tunnel_state_machine_handle,
+                target_state,
+                account_manager,
+                ..
+            } = self;
+
+            shutdown_tasks.push(Box::pin(target_state.finalize()));
+            shutdown_tasks.push(Box::pin(account_manager.shutdown()));
+
+            (
+                event_listener,
+                shutdown_tasks,
+                api_runtime,
+                tunnel_state_machine_handle,
+            )
+        };
         for future in shutdown_tasks {
             future.await;
         }
@@ -902,45 +918,18 @@ where
         }
     }
 
-    /// Shuts down the daemon without shutting down the underlying event listener and the shutdown
-    /// callbacks
-    fn shutdown<'a>(
-        self,
-    ) -> (
-        L,
-        Vec<LocalBoxFuture<'a, ()>>,
-        mullvad_api::Runtime,
-        TunnelStateMachineHandle,
-    ) {
-        let Daemon {
-            event_listener,
-            mut shutdown_tasks,
-            api_runtime,
-            tunnel_state_machine_handle,
-            target_state,
-            account_manager,
-            ..
-        } = self;
-
-        shutdown_tasks.push(Box::pin(target_state.finalize()));
-        shutdown_tasks.push(Box::pin(account_manager.shutdown()));
-
-        (
-            event_listener,
-            shutdown_tasks,
-            api_runtime,
-            tunnel_state_machine_handle,
-        )
-    }
-
-    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.shutdown(user_init_shutdown).await;
+                should_stop = true;
+            }
             NewAppVersionInfo(app_version_info) => {
                 self.handle_new_app_version_info(app_version_info);
             }
@@ -954,6 +943,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(
@@ -1138,11 +1128,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();
         }
@@ -1417,12 +1402,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>) {
@@ -1707,7 +1688,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 {
@@ -2636,7 +2617,7 @@ where
         }
     }
 
-    fn trigger_shutdown_event(&mut self, user_init_shutdown: bool) {
+    async fn 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
@@ -2647,8 +2628,21 @@ where
             self.send_tunnel_command(TunnelCommand::BlockWhenDisconnected(true, tx));
         }
 
-        self.shutting_down = true;
         self.disconnect_tunnel();
+
+        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;
+                }
+            }
+        }
     }
 
     fn on_prepare_restart(&mut self) {

wireguard-go-rs/libwg/wireguard-go

@@ -0,0 +1 @@
+Subproject commit c14351e8550c8110df236a2a4354599747552c15