Refactor daemon event loop

1 files changed, 32 insertions, 24 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 224d7e4b2c..468f4fca38 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -853,30 +853,8 @@ impl Daemon {
     /// shutdown event is received.
     pub async fn run(mut self) -> Result<(), Error> {
         self.handle_initial_target_state();
-
-        while let Some(event) = self.rx.next().await {
-            if self.handle_event(event).await {
-                break;
-            }
-        }
-
-        // Wait for tunnel state machine to disconnect
-        if !self.tunnel_state.is_disconnected() {
-            self.disconnect_tunnel();
-
-            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.handle_events().await;
+        self.disconnect_tunnel_and_wait().await;
         self.finalize().await;
         Ok(())
     }
@@ -900,6 +878,27 @@ impl Daemon {
         TunnelCommand::Connect
     }
 
+    /// Begin disconnecting and wait for the tunnel state machine to be disconnected
+    async fn disconnect_tunnel_and_wait(&mut self) {
+        if self.tunnel_state.is_disconnected() {
+            return;
+        }
+
+        self.disconnect_tunnel();
+
+        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;
+            }
+        }
+    }
+
     /// Destroy daemon safely, by dropping all objects in the correct order, waiting for them to
     /// be destroyed, and executing shutdown tasks
     async fn finalize(self) {
@@ -927,6 +926,15 @@ impl Daemon {
         drop(api_runtime);
     }
 
+    /// Handle internal daemon events until a shutdown event is received
+    async fn handle_events(&mut self) {
+        while let Some(event) = self.rx.next().await {
+            if self.handle_event(event).await {
+                break;
+            }
+        }
+    }
+
     async fn handle_event(&mut self, event: InternalDaemonEvent) -> bool {
         use self::InternalDaemonEvent::*;
         let mut should_stop = false;