Fix notification service

Telegram/NotificationService/Sources/NotificationService.swift

@@ -598,17 +598,40 @@ private final class NotificationServiceHandler {
             return nil
         }
 
-        let _ = (self.accountManager.currentAccountRecord(allocateIfNotExists: false)
+        let _ = (self.accountManager.accountRecords()
         |> take(1)
         |> deliverOn(self.queue)).start(next: { [weak self] records in
-            guard let strongSelf = self, let record = records else {
+            var recordId: AccountRecordId?
+            var isCurrentAccount: Bool = false
+
+            if let keyId = notificationPayloadKeyId(data: payloadData) {
+                outer: for listRecord in records.records {
+                    for attribute in listRecord.attributes {
+                        if case let .backupData(backupData) = attribute {
+                            if let notificationEncryptionKeyId = backupData.data?.notificationEncryptionKeyId {
+                                if keyId == notificationEncryptionKeyId {
+                                    recordId = listRecord.id
+                                    isCurrentAccount = records.currentRecord?.id == listRecord.id
+                                    break outer
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+
+            guard let strongSelf = self, let recordId = recordId else {
+                let content = NotificationContent()
+                updateCurrentContent(content)
+                completed()
+
                 return
             }
 
             let _ = (standaloneStateManager(
                 accountManager: strongSelf.accountManager,
                 networkArguments: networkArguments,
-                id: record.0,
+                id: recordId,
                 encryptionParameters: strongSelf.encryptionParameters,
                 rootPath: rootPath,
                 auxiliaryMethods: accountAuxiliaryMethods
@@ -618,6 +641,8 @@ private final class NotificationServiceHandler {
                     return
                 }
                 guard let stateManager = stateManager else {
+                    let content = NotificationContent()
+                    updateCurrentContent(content)
                     completed()
                     return
                 }
@@ -626,18 +651,31 @@ private final class NotificationServiceHandler {
                 strongSelf.notificationKeyDisposable.set((existingMasterNotificationsKey(postbox: stateManager.postbox)
                 |> deliverOn(strongSelf.queue)).start(next: { notificationsKey in
                     guard let strongSelf = self else {
+                        let content = NotificationContent()
+                        updateCurrentContent(content)
+                        completed()
+
                         return
                     }
                     guard let notificationsKey = notificationsKey else {
+                        let content = NotificationContent()
+                        updateCurrentContent(content)
                         completed()
+
                         return
                     }
                     guard let decryptedPayload = decryptedNotificationPayload(key: notificationsKey, data: payloadData) else {
+                        let content = NotificationContent()
+                        updateCurrentContent(content)
                         completed()
+
                         return
                     }
                     guard let payloadJson = try? JSONSerialization.jsonObject(with: decryptedPayload, options: []) as? [String: Any] else {
+                        let content = NotificationContent()
+                        updateCurrentContent(content)
                         completed()
+
                         return
                     }
 
@@ -663,6 +701,10 @@ private final class NotificationServiceHandler {
                         if let channelIdValue = Int64(channelIdString) {
                             peerId = PeerId(namespace: Namespaces.Peer.CloudChannel, id: PeerId.Id._internalFromInt64Value(channelIdValue))
                         }
+                    } else if let encryptionIdString = payloadJson["encryption_id"] as? String {
+                        if let encryptionIdValue = Int64(encryptionIdString) {
+                            peerId = PeerId(namespace: Namespaces.Peer.SecretChat, id: PeerId.Id._internalFromInt64Value(encryptionIdValue))
+                        }
                     }
 
                     enum Action {
@@ -735,7 +777,7 @@ private final class NotificationServiceHandler {
                             }
 
                             content.userInfo["peerId"] = "\(peerId.toInt64())"
-                            content.userInfo["accountId"] = "\(record.0.int64)"
+                            content.userInfo["accountId"] = "\(recordId.int64)"
 
                             if let silentString = payloadJson["silent"] as? String {
                                 if let silentValue = Int(silentString), silentValue != 0 {
@@ -888,7 +930,9 @@ private final class NotificationServiceHandler {
                                                     return
                                                 }
 
+                                                if isCurrentAccount {
                                                     content.badge = Int(value.0)
+                                                }
 
                                                 if let image = mediaAttachment as? TelegramMediaImage, let resource = largestImageRepresentation(image.representations)?.resource {
                                                     if let mediaData = mediaData {
@@ -1037,7 +1081,9 @@ private final class NotificationServiceHandler {
                                         )
                                         |> deliverOn(strongSelf.queue)).start(next: { value in
                                             var content = NotificationContent()
+                                            if isCurrentAccount {
                                                 content.badge = Int(value.0)
+                                            }
 
                                             updateCurrentContent(content)
 
@@ -1080,7 +1126,9 @@ private final class NotificationServiceHandler {
                                         )
                                         |> deliverOn(strongSelf.queue)).start(next: { value in
                                             var content = NotificationContent()
+                                            if isCurrentAccount {
                                                 content.badge = Int(value.0)
+                                            }
 
                                             updateCurrentContent(content)
 

submodules/TelegramCore/Sources/Account/Account.swift

@@ -695,11 +695,17 @@ private func masterNotificationsKey(masterNotificationKeyValue: Atomic<MasterNot
     }
 
     return postbox.transaction(ignoreDisabled: ignoreDisabled, { transaction -> MasterNotificationKey? in
+        let result = masterNotificationsKey(transaction: transaction, createIfNotExists: createIfNotExists)
+        let _ = masterNotificationKeyValue.swap(result)
+        return result
+    })
+}
+
+func masterNotificationsKey(transaction: Transaction, createIfNotExists: Bool) -> MasterNotificationKey? {
     if let value = transaction.keychainEntryForKey("master-notification-secret"), !value.isEmpty {
         let authKeyHash = sha1Digest(value)
         let authKeyId = authKeyHash.subdata(in: authKeyHash.count - 8 ..< authKeyHash.count)
         let keyData = MasterNotificationKey(id: authKeyId, data: value)
-            let _ = masterNotificationKeyValue.swap(keyData)
         return keyData
     } else if createIfNotExists {
         var secretData = Data(count: 256)
@@ -716,12 +722,18 @@ private func masterNotificationsKey(masterNotificationKeyValue: Atomic<MasterNot
         let authKeyHash = sha1Digest(secretData)
         let authKeyId = authKeyHash.subdata(in: authKeyHash.count - 8 ..< authKeyHash.count)
         let keyData = MasterNotificationKey(id: authKeyId, data: secretData)
-            let _ = masterNotificationKeyValue.swap(keyData)
         return keyData
     } else {
         return nil
     }
-    })
+}
+
+public func notificationPayloadKeyId(data: Data) -> Data? {
+    if data.count < 8 {
+        return nil
+    }
+
+    return data.subdata(in: 0 ..< 8)
 }
 
 public func decryptedNotificationPayload(key: MasterNotificationKey, data: Data) -> Data? {
@@ -792,7 +804,14 @@ public func accountBackupData(postbox: Postbox) -> Signal<AccountBackupData?, No
         guard let authKey = datacenterAuthInfo.authKey else {
             return nil
         }
-        return AccountBackupData(masterDatacenterId: state.masterDatacenterId, peerId: state.peerId.toInt64(), masterDatacenterKey: authKey, masterDatacenterKeyId: datacenterAuthInfo.authKeyId)
+        let notificationsKey = masterNotificationsKey(transaction: transaction, createIfNotExists: true)
+        return AccountBackupData(
+            masterDatacenterId: state.masterDatacenterId,
+            peerId: state.peerId.toInt64(),
+            masterDatacenterKey: authKey,
+            masterDatacenterKeyId: datacenterAuthInfo.authKeyId,
+            notificationEncryptionKeyId: notificationsKey?.id
+        )
     }
 }
 

submodules/TelegramCore/Sources/SyncCore/SyncCore_AccountBackupDataAttribute.swift

@@ -6,12 +6,20 @@ public struct AccountBackupData: Codable, Equatable {
     public var peerId: Int64
     public var masterDatacenterKey: Data
     public var masterDatacenterKeyId: Int64
+    public var notificationEncryptionKeyId: Data?
 
-    public init(masterDatacenterId: Int32, peerId: Int64, masterDatacenterKey: Data, masterDatacenterKeyId: Int64) {
+    public init(
+        masterDatacenterId: Int32,
+        peerId: Int64,
+        masterDatacenterKey: Data,
+        masterDatacenterKeyId: Int64,
+        notificationEncryptionKeyId: Data?
+    ) {
     	self.masterDatacenterId = masterDatacenterId
     	self.peerId = peerId
     	self.masterDatacenterKey = masterDatacenterKey
     	self.masterDatacenterKeyId = masterDatacenterKeyId
+        self.notificationEncryptionKeyId = notificationEncryptionKeyId
     }
 }
 

submodules/TelegramUI/Sources/AppDelegate.swift

@@ -2277,6 +2277,8 @@ private func notificationPayloadKey(data: Data) -> Data? {
 private func accountIdFromNotification(_ notification: UNNotification, sharedContext: Signal<SharedApplicationContext, NoError>) -> Signal<AccountRecordId?, NoError> {
     if let id = notification.request.content.userInfo["accountId"] as? Int64 {
         return .single(AccountRecordId(rawValue: id))
+    } else if let idString = notification.request.content.userInfo["accountId"] as? String, let id = Int64(idString) {
+        return .single(AccountRecordId(rawValue: id))
     } else {
         var encryptedData: Data?
         if var encryptedPayload = notification.request.content.userInfo["p"] as? String {
@@ -2336,6 +2338,8 @@ private func accountIdFromNotification(_ notification: UNNotification, sharedCon
 private func peerIdFromNotification(_ notification: UNNotification) -> PeerId? {
     if let peerId = notification.request.content.userInfo["peerId"] as? Int64 {
         return PeerId(peerId)
+    } else if let peerIdString = notification.request.content.userInfo["peerId"] as? String, let peerId = Int64(peerIdString) {
+        return PeerId(peerId)
     } else {
         let payload = notification.request.content.userInfo
         var peerId: PeerId?