From 049194f5d2978f10ac2b56e81a0ffe9dffbf10fb Mon Sep 17 00:00:00 2001
From: Max Goedjen <max.goedjen@gmail.com>
Date: Thu, 12 Nov 2020 22:54:21 -0800
Subject: [PATCH] Relaunch

---
 Secretive/App.swift                           | 44 ++++++++++++++-----
 .../Controllers/JustUpdatedChecker.swift      |  4 +-
 .../Controllers/LaunchAgentController.swift   | 11 ++++-
 Secretive/Views/SetupView.swift               |  2 +-
 4 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/Secretive/App.swift b/Secretive/App.swift
index c26e9bf..e12712d 100644
--- a/Secretive/App.swift
+++ b/Secretive/App.swift
@@ -28,21 +28,16 @@ struct Secretive: App {
                 .onAppear {
                     if !hasRunSetup {
                         showingSetup = true
-                    } else {
-                        if agentStatusChecker.running && justUpdatedChecker.justUpdated {
-                            // Relaunch the agent, since it'll be running from earlier update still
-                            _ = LaunchAgentController().install()
-                        }
                     }
                 }
-                .onReceive(NotificationCenter.default.publisher(for: NSApplication.willBecomeActiveNotification)) { _ in
+                .onReceive(NotificationCenter.default.publisher(for: NSApplication.didBecomeActiveNotification)) { _ in
+                    guard hasRunSetup else { return }
                     agentStatusChecker.check()
-                    if hasRunSetup && !agentStatusChecker.running {
-                        // We've run setup, we didn't just update, launchd is just not doing it's thing.
-                        // Force a launch directly.
-                        LaunchAgentController().forceLaunch { _ in
-                            agentStatusChecker.check()
-                        }
+                    if agentStatusChecker.running && justUpdatedChecker.justUpdated {
+                        // Relaunch the agent, since it'll be running from earlier update still
+                        reinstallAgent()
+                    } else if !agentStatusChecker.running {
+                        forceLaunchAgent()
                     }
                 }
         }
@@ -69,6 +64,31 @@ struct Secretive: App {
 
 }
 
+extension Secretive {
+
+    private func reinstallAgent() {
+        justUpdatedChecker.check()
+        LaunchAgentController().install {
+            // Wait a second for launchd to kick in (next runloop isn't enough).
+            DispatchQueue.main.asyncAfter(deadline: .now() + 1) {
+                agentStatusChecker.check()
+                if !agentStatusChecker.running {
+                    forceLaunchAgent()
+                }
+            }
+        }
+    }
+
+    private func forceLaunchAgent() {
+        // We've run setup, we didn't just update, launchd is just not doing it's thing.
+        // Force a launch directly.
+        LaunchAgentController().forceLaunch { _ in
+            agentStatusChecker.check()
+        }
+    }
+
+}
+
 
 private enum Constants {
     static let helpURL = URL(string: "https://github.com/maxgoedjen/secretive/blob/main/FAQ.md")!
diff --git a/Secretive/Controllers/JustUpdatedChecker.swift b/Secretive/Controllers/JustUpdatedChecker.swift
index 95a7ed3..4c86f68 100644
--- a/Secretive/Controllers/JustUpdatedChecker.swift
+++ b/Secretive/Controllers/JustUpdatedChecker.swift
@@ -18,9 +18,7 @@ class JustUpdatedChecker: ObservableObject, JustUpdatedCheckerProtocol {
         let lastBuild = UserDefaults.standard.object(forKey: Constants.previousVersionUserDefaultsKey) as? String ?? "None"
         let currentBuild = Bundle.main.infoDictionary!["CFBundleShortVersionString"] as! String
         UserDefaults.standard.set(currentBuild, forKey: Constants.previousVersionUserDefaultsKey)
-        if lastBuild != currentBuild {
-            justUpdated = true
-        }
+        justUpdated = lastBuild != currentBuild
     }
 
 
diff --git a/Secretive/Controllers/LaunchAgentController.swift b/Secretive/Controllers/LaunchAgentController.swift
index 80df05d..97915f8 100644
--- a/Secretive/Controllers/LaunchAgentController.swift
+++ b/Secretive/Controllers/LaunchAgentController.swift
@@ -5,10 +5,17 @@ import OSLog
 
 struct LaunchAgentController {
 
-    func install() -> Bool {
+    func install(completion: (() -> Void)? = nil) {
         Logger().debug("Installing agent")
         _ = setEnabled(false)
-        return setEnabled(true)
+        // This is definitely a bit of a "seems to work better" thing but:
+        // Seems to more reliably hit if these are on separate runloops, otherwise it seems like it sometimes doesn't kill old
+        // and start new?
+        DispatchQueue.main.asyncAfter(deadline: .now() + 0.1) {
+            _  = setEnabled(true)
+            completion?()
+        }
+
     }
 
     func forceLaunch(completion: ((Bool) -> Void)?) {
diff --git a/Secretive/Views/SetupView.swift b/Secretive/Views/SetupView.swift
index b935ee7..ad358f2 100644
--- a/Secretive/Views/SetupView.swift
+++ b/Secretive/Views/SetupView.swift
@@ -157,7 +157,7 @@ struct SecretAgentSetupView: View {
     }
 
     func install() {
-        _ = LaunchAgentController().install()
+        LaunchAgentController().install()
         buttonAction()
     }