From 80354a2a7224a55f89d1c95d1996dbf1e93e9f61 Mon Sep 17 00:00:00 2001
From: Gui Rambo <insidegui@gmail.com>
Date: Mon, 7 Jun 2021 20:08:07 -0300
Subject: [PATCH] Restarting WWDCAgent when a new version is detected (i.e.
 after app updates)

---
 WWDC/AppDelegate.swift         |  4 +++
 WWDC/WWDCAgentController.swift | 50 ++++++++++++++++++++++++++++++++--
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/WWDC/AppDelegate.swift b/WWDC/AppDelegate.swift
index a9c6b2f2..1268a891 100644
--- a/WWDC/AppDelegate.swift
+++ b/WWDC/AppDelegate.swift
@@ -21,6 +21,8 @@ extension Notification.Name {
 
 class AppDelegate: NSObject, NSApplicationDelegate {
     
+    private lazy var agentController = WWDCAgentController()
+    
     private let log = OSLog(subsystem: "io.wwdc.app", category: String(describing: AppDelegate.self))
 
     private lazy var commandsReceiver = AppCommandsReceiver()
@@ -62,6 +64,8 @@ class AppDelegate: NSObject, NSApplicationDelegate {
 
     func applicationDidFinishLaunching(_ notification: Notification) {
         UserDefaults.standard.register(defaults: ["NSApplicationCrashOnExceptions": false])
+        
+        agentController.registerAgentVersion()
 
         NSApp.registerForRemoteNotifications(matching: [])
 
diff --git a/WWDC/WWDCAgentController.swift b/WWDC/WWDCAgentController.swift
index f055bc17..c3986181 100644
--- a/WWDC/WWDCAgentController.swift
+++ b/WWDC/WWDCAgentController.swift
@@ -8,6 +8,7 @@
 
 import Cocoa
 import ServiceManagement
+import os.log
 
 extension Notification.Name {
     static let WWDCAgentEnabledPreferenceChanged = Notification.Name("io.wwdc.app.AgentEnabledPreferenceChanged")
@@ -15,6 +16,8 @@ extension Notification.Name {
 
 final class WWDCAgentController: NSObject {
     
+    private let log = OSLog(subsystem: "io.wwdc.app", category: String(describing: WWDCAgentController.self))
+    
     static private(set) var isAgentEnabled: Bool {
         get { UserDefaults.standard.bool(forKey: #function) }
         set {
@@ -31,7 +34,12 @@ final class WWDCAgentController: NSObject {
         }
     }
     
-    private lazy var agentBundleIdentifier: String = {
+    private(set) var lastRunAgentBuild: String? {
+        get { UserDefaults.standard.string(forKey: #function) }
+        set { UserDefaults.standard.set(newValue, forKey: #function) }
+    }
+    
+    private lazy var agentBundle: Bundle = {
         let agentURL = Bundle.main.bundleURL
             .appendingPathComponent("Contents/Library/LoginItems", isDirectory: true)
             .appendingPathComponent("WWDCAgent.app")
@@ -40,13 +48,25 @@ final class WWDCAgentController: NSObject {
             preconditionFailure("Couldn't instatiate agent bundle")
         }
         
-        guard let identifier = bundle.bundleIdentifier else {
+        return bundle
+    }()
+    
+    private lazy var agentBundleIdentifier: String = {
+        guard let identifier = agentBundle.bundleIdentifier else {
             preconditionFailure("Failed to read identifier for agent bundle")
         }
         
         return identifier
     }()
     
+    private lazy var currentAgentBundleBuild: String = {
+        guard let build = agentBundle.infoDictionary?["CFBundleVersion"] as? String else {
+            preconditionFailure("Failed to read CFBundleVersion for agent bundle")
+        }
+        
+        return build
+    }()
+    
     func enableAgent() -> Bool {
         Self.isAgentEnabled = SMLoginItemSetEnabled(agentBundleIdentifier as CFString, true)
         
@@ -58,5 +78,31 @@ final class WWDCAgentController: NSObject {
         
         Self.isAgentEnabled = false
     }
+    
+    func registerAgentVersion() {
+        guard currentAgentBundleBuild != lastRunAgentBuild else {
+            os_log("Agent bundle build has not changed", log: self.log, type: .debug)
+            return
+        }
+        
+        if let build = lastRunAgentBuild {
+            os_log("Registering new agent build: %{public}@", log: self.log, type: .debug, build)
+        } else {
+            os_log("Registering agent build for the first time: %{public}@", log: self.log, type: .debug, currentAgentBundleBuild)
+        }
+        
+        lastRunAgentBuild = currentAgentBundleBuild
+        
+        guard let agentApp = NSRunningApplication.runningApplications(withBundleIdentifier: agentBundleIdentifier).first else {
+            os_log("Couldn't find agent app running, ignoring", log: self.log, type: .debug)
+            return
+        }
+        
+        os_log("Found old agent running with pid %{public}d, restarting", log: self.log, type: .debug, agentApp.processIdentifier)
+        
+        if !agentApp.forceTerminate() {
+            os_log("Failed to terminate agent", log: self.log, type: .fault)
+        }
+    }
 
 }