Fix NSE Leak (#2022)

2026-11-02 21:33:26 +00:00 · 2023-11-03 21:30:58 +01:00
parent 4985d410b3
commit ce72a9b5f2
11 changed files with 77 additions and 22 deletions
@@ -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")
@@ -24,6 +24,7 @@ final class NSEUserSession {
    private(set) lazy var mediaProvider: MediaProviderProtocol = MediaProvider(mediaLoader: MediaLoader(client: baseClient),
                                                                               imageCache: .onlyOnDisk,
                                                                               backgroundTaskService: nil)
+    private let delegateHandle: TaskHandle?

    init(credentials: KeychainCredentials, clientSessionDelegate: ClientSessionDelegate) throws {
        userID = credentials.userID
@@ -35,7 +36,7 @@ final class NSEUserSession {
                                           sessionDelegate: clientSessionDelegate)
            .build()
        
-        baseClient.setDelegate(delegate: ClientDelegateWrapper())
+        delegateHandle = baseClient.setDelegate(delegate: ClientDelegateWrapper())
        try baseClient.restoreSession(session: credentials.restorationToken.session)
        
        notificationClient = try baseClient
@@ -62,6 +63,10 @@ final class NSEUserSession {
            }
        }
    }
+    
+    deinit {
+        delegateHandle?.cancel()
+    }
 }

 private class ClientDelegateWrapper: ClientDelegate {