Fix NSE Leak (#2022)

NSE/Sources/NotificationServiceExtension.swift

@@ -18,15 +18,38 @@ import Intents
 import MatrixRustSDK
 import UserNotifications
 
+// The lifecycle of the NSE looks something like the following:
+//  1)  App receives notification
+//  2)  System creates an instance of the extension class
+//      and calls `didReceive` in the background
+//  3)  Extension processes messages / displays whatever
+//      notifications it needs to
+//  4)  Extension notifies its work is complete by calling
+//      the contentHandler
+//  5)  If the extension takes too long to perform its work
+//      (more than 30s), it will be notified and immediately
+//      terminated
+//
+// Note that the NSE does *not* always spawn a new process to
+// handle a new notification and will also try and process notifications
+// in parallel. `didReceive` could be called twice for the same process,
+// but it will always be called on different threads. It may or may not be
+// called on the same instance of `NotificationService` as a previous
+// notification.
+//
+// We keep a global `environment` singleton to ensure that our app context,
+// database, logging, etc. are only ever setup once per *process*
+
+private let settings = NSESettings()
+private let notificationContentBuilder = NotificationContentBuilder(messageEventStringBuilder: RoomMessageEventStringBuilder(attributedStringBuilder: AttributedStringBuilder(permalinkBaseURL: .homeDirectory,
+                                                                                                                                                                              mentionBuilder: PlainMentionBuilder())))
+private let keychainController = KeychainController(service: .sessions,
+                                                    accessGroup: InfoPlistReader.main.keychainAccessGroupIdentifier)
+private var userSessions = [String: NSEUserSession]()
+
 class NotificationServiceExtension: UNNotificationServiceExtension {
-    private let settings = NSESettings()
-    private let notificationContentBuilder = NotificationContentBuilder(messageEventStringBuilder: RoomMessageEventStringBuilder(attributedStringBuilder: AttributedStringBuilder(permalinkBaseURL: .homeDirectory,
-                                                                                                                                                                                  mentionBuilder: PlainMentionBuilder())))
-    private lazy var keychainController = KeychainController(service: .sessions,
-                                                             accessGroup: InfoPlistReader.main.keychainAccessGroupIdentifier)
     private var handler: ((UNNotificationContent) -> Void)?
     private var modifiedContent: UNMutableNotificationContent?
-    private var userSession: NSEUserSession?
 
     override func didReceive(_ request: UNNotificationRequest,
                              withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void) {
@@ -75,8 +98,13 @@ class NotificationServiceExtension: UNNotificationServiceExtension {
         MXLog.info("\(tag) run with roomId: \(roomId), eventId: \(eventId)")
 
         do {
-            let userSession = try NSEUserSession(credentials: credentials, clientSessionDelegate: keychainController)
-            self.userSession = userSession
+            let userSession: NSEUserSession
+            if let existingSession = userSessions[credentials.userID] {
+                userSession = existingSession
+            } else {
+                userSession = try NSEUserSession(credentials: credentials, clientSessionDelegate: keychainController)
+                userSessions[credentials.userID] = userSession
+            }
             
             guard let itemProxy = await userSession.notificationItemProxy(roomID: roomId, eventID: eventId) else {
                 MXLog.info("\(tag) no notification for the event, discard")