matrix-org · Anderas · Jan 30, 2023 · Jan 30, 2023
diff --git a/MatrixSDK.xcodeproj/project.pbxproj b/MatrixSDK.xcodeproj/project.pbxproj
@@ -1942,6 +1942,8 @@
 		ED6DAC2228C7A51400ECDCB6 /* MXDateProvider.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED6DAC2028C7A4F000ECDCB6 /* MXDateProvider.swift */; };
 		ED6E87A9294B3BAB00100D9C /* MXAnalyticsDestinationUnitTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED6E87A8294B3BAB00100D9C /* MXAnalyticsDestinationUnitTests.swift */; };
 		ED6E87AA294B3BAB00100D9C /* MXAnalyticsDestinationUnitTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED6E87A8294B3BAB00100D9C /* MXAnalyticsDestinationUnitTests.swift */; };
+		ED6F4EFC2987F0FC007D1191 /* MXEncryptedKeyBackup.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED6F4EFB2987F0FC007D1191 /* MXEncryptedKeyBackup.swift */; };
+		ED6F4EFD2987F0FC007D1191 /* MXEncryptedKeyBackup.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED6F4EFB2987F0FC007D1191 /* MXEncryptedKeyBackup.swift */; };
 		ED7019DD2886C24100FC31B9 /* MXCrossSigningInfoSourceUnitTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED8F1D242885A39800F897E7 /* MXCrossSigningInfoSourceUnitTests.swift */; };
 		ED7019DE2886C24A00FC31B9 /* MXTrustLevelSourceUnitTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED8F1D2F2885AB0300F897E7 /* MXTrustLevelSourceUnitTests.swift */; };
 		ED7019DF2886C25600FC31B9 /* MXDeviceInfoUnitTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = ED8F1D1B2885909E00F897E7 /* MXDeviceInfoUnitTests.swift */; };
@@ -3128,6 +3130,7 @@
 		ED6DAC1D28C79D2000ECDCB6 /* MXUnrequestedForwardedRoomKeyManagerUnitTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MXUnrequestedForwardedRoomKeyManagerUnitTests.swift; sourceTree = "<group>"; };
 		ED6DAC2028C7A4F000ECDCB6 /* MXDateProvider.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MXDateProvider.swift; sourceTree = "<group>"; };
 		ED6E87A8294B3BAB00100D9C /* MXAnalyticsDestinationUnitTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MXAnalyticsDestinationUnitTests.swift; sourceTree = "<group>"; };
+		ED6F4EFB2987F0FC007D1191 /* MXEncryptedKeyBackup.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MXEncryptedKeyBackup.swift; sourceTree = "<group>"; };
 		ED7019E42886C32900FC31B9 /* MXSASTransactionV2.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = MXSASTransactionV2.swift; sourceTree = "<group>"; };
 		ED7019E72886C33100FC31B9 /* MXKeyVerificationRequestV2.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = MXKeyVerificationRequestV2.swift; sourceTree = "<group>"; };
 		ED7019EA2886C33A00FC31B9 /* MXKeyVerificationManagerV2.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = MXKeyVerificationManagerV2.swift; sourceTree = "<group>"; };
@@ -5790,6 +5793,7 @@
 				EDD4197D28DCAA5F007F3757 /* MXNativeKeyBackupEngine.h */,
 				EDD4198028DCAA7B007F3757 /* MXNativeKeyBackupEngine.m */,
 				ED36ED8528DD9E2100C86416 /* MXCryptoKeyBackupEngine.swift */,
+				ED6F4EFB2987F0FC007D1191 /* MXEncryptedKeyBackup.swift */,
 			);
 			path = Engine;
 			sourceTree = "<group>";
@@ -7162,6 +7166,7 @@
 				ECDA764E27BA963D000C48CF /* MXBooleanCapability.m in Sources */,
 				321CFDEB22525DEE004D31DF /* MXIncomingSASTransaction.m in Sources */,
 				EC1165C427107E330089FA56 /* MXRoomListDataFilterOptions.swift in Sources */,
+				ED6F4EFC2987F0FC007D1191 /* MXEncryptedKeyBackup.swift in Sources */,
 				3A108A8025810C96005EEBE9 /* MXKeyData.m in Sources */,
 				3A59A52B25A7B1B000DDA1FC /* MXOutboundSessionInfo.m in Sources */,
 				32A1515C1DB525DA00400192 /* NSObject+sortedKeys.m in Sources */,
@@ -7827,6 +7832,7 @@
 				3A858DE227528EEB006322C1 /* MXHomeserverCapabilitiesService.swift in Sources */,
 				ECDA764F27BA963D000C48CF /* MXBooleanCapability.m in Sources */,
 				B14EF24C2397E90400758AF0 /* MXCredentials.m in Sources */,
+				ED6F4EFD2987F0FC007D1191 /* MXEncryptedKeyBackup.swift in Sources */,
 				EC116590270F3ABF0089FA56 /* RLMRealm+MatrixSDK.m in Sources */,
 				B14EF24D2397E90400758AF0 /* MXOutgoingRoomKeyRequest.m in Sources */,
 				B14EF24E2397E90400758AF0 /* MXAllowedCertificates.m in Sources */,

diff --git a/MatrixSDK/Crypto/CryptoMachine/MXCryptoMachine.swift b/MatrixSDK/Crypto/CryptoMachine/MXCryptoMachine.swift
@@ -362,6 +362,15 @@ extension MXCryptoMachine: MXCryptoUserIdentitySource {
 }
 
 extension MXCryptoMachine: MXCryptoRoomEventEncrypting {
+    func isUserTracked(userId: String) -> Bool {
+        do {
+            return try machine.isUserTracked(userId: userId)
+        } catch {
+            log.error("Failed getting tracked status", context: error)
+            return false
+        }
+    }
+
     func addTrackedUsers(_ users: [String]) {
         do {
             try machine.updateTrackedUsers(users: users)

diff --git a/MatrixSDK/Crypto/CryptoMachine/MXCryptoProtocols.swift b/MatrixSDK/Crypto/CryptoMachine/MXCryptoProtocols.swift
@@ -66,6 +66,7 @@ protocol MXCryptoUserIdentitySource: MXCryptoIdentity {
 
 /// Room event encryption
 protocol MXCryptoRoomEventEncrypting: MXCryptoIdentity {
+    func isUserTracked(userId: String) -> Bool
     func addTrackedUsers(_ users: [String])
     func shareRoomKeysIfNecessary(roomId: String, users: [String], settings: EncryptionSettings) async throws
     func encryptRoomEvent(content: [AnyHashable: Any], roomId: String, eventType: String) throws -> [String: Any]

diff --git a/MatrixSDK/Crypto/CryptoMachine/MXCryptoSDKLogger.swift b/MatrixSDK/Crypto/CryptoMachine/MXCryptoSDKLogger.swift
@@ -31,8 +31,15 @@ class MXCryptoSDKLogger: Logger {
     func log(logLine: String) {
         // Excluding some auto-generated logs that are not useful
         // This will be changed in rust-sdk directly
-        guard !logLine.contains("::uniffi_api:") else {
-            return
+        let ignored = [
+            "::uniffi_api:",
+            "::backup_recovery_key: decrypt_v1"
+        ]
+
+        for ignore in ignored {
+            if logLine.contains(ignore) {
+                return
+            }
         }
 
         MXLog.debug("[MXCryptoSDK] \(logLine)")

diff --git a/MatrixSDK/Crypto/Data/MXCryptoConstants.h b/MatrixSDK/Crypto/Data/MXCryptoConstants.h
@@ -89,6 +89,7 @@ typedef enum : NSUInteger
     MXKeyBackupErrorMissingPrivateKeySaltCode,
     MXKeyBackupErrorMissingAuthDataCode,
     MXKeyBackupErrorInvalidOrMissingLocalPrivateKey,
-    MXKeyBackupErrorUnknownAlgorithm
+    MXKeyBackupErrorUnknownAlgorithm,
+    MXKeyBackupErrorAlreadyInProgress,
 
 } MXKeyBackupErrorCode;
diff --git a/MatrixSDK/Crypto/KeyBackup/Engine/MXCryptoKeyBackupEngine.swift b/MatrixSDK/Crypto/KeyBackup/Engine/MXCryptoKeyBackupEngine.swift
@@ -24,12 +24,6 @@ class MXCryptoKeyBackupEngine: NSObject, MXKeyBackupEngine {
     // Batch size chosen arbitrarily, will be moved to CryptoSDK
     private static let ImportBatchSize = 1000
 
-    struct EncryptedSession {
-        let roomId: String
-        let sessionId: String
-        let keyBackup: MXKeyBackupData
-    }
-
     enum Error: Swift.Error {
         case unknownBackupVersion
         case invalidData
@@ -289,19 +283,21 @@ class MXCryptoKeyBackupEngine: NSObject, MXKeyBackupEngine {
 
             let encryptedSessions = keysBackupData.rooms.flatMap { roomId, room in
                 room.sessions.map { sessionId, keyBackup in
-                    EncryptedSession(roomId: roomId, sessionId: sessionId, keyBackup: keyBackup)
+                    MXEncryptedKeyBackup(roomId: roomId, sessionId: sessionId, keyBackup: keyBackup)
                 }
             }
 
-            let count = encryptedSessions.count
-            self.activeImportProgress = Progress(totalUnitCount: Int64(count))
-            self.log.debug("Importing \(count) encrypted sessions")
+            let totalKeysCount = encryptedSessions.count
+            var importedKeysCount: UInt = 0
 
-            let date = Date()
+            self.activeImportProgress = Progress(totalUnitCount: Int64(totalKeysCount))
+            self.log.debug("Importing \(totalKeysCount) encrypted sessions")
 
-            for batchIndex in stride(from: 0, to: count, by: Self.ImportBatchSize) {
+            let startDate = Date()
+
+            for batchIndex in stride(from: 0, to: totalKeysCount, by: Self.ImportBatchSize) {
                 self.log.debug("Decrypting and importing batch \(batchIndex)")
-                let endIndex = min(batchIndex + Self.ImportBatchSize, count)
+                let endIndex = min(batchIndex + Self.ImportBatchSize, totalKeysCount)
                 let batch = encryptedSessions[batchIndex ..< endIndex]
 
                 autoreleasepool {
@@ -317,21 +313,23 @@ class MXCryptoKeyBackupEngine: NSObject, MXKeyBackupEngine {
 
                     do {
                         let result = try self.backup.importDecryptedKeys(roomKeys: sessions, progressListener: self)
-                        self.activeImportProgress?.completedUnitCount += Int64(result.imported)
+                        importedKeysCount += UInt(result.imported)
                     } catch {
                         self.log.error("Failed importing batch of sessions", context: error)
                     }
+
+                    self.activeImportProgress?.completedUnitCount += Int64(Self.ImportBatchSize)
                 }
                 await self.roomEventDecryptor.retryUndecryptedEvents(sessionIds: batch.map(\.sessionId))
             }
 
-            let imported = self.activeImportProgress?.completedUnitCount ?? 0
-            let duration = Date().timeIntervalSince(date) * 1000
-            self.log.debug("Successfully imported \(imported) out of \(count) sessions in \(duration) ms")
+            let imported = importedKeysCount
+            let duration = Date().timeIntervalSince(startDate) * 1000
+            self.log.debug("Successfully imported \(imported) out of \(totalKeysCount) sessions in \(duration) ms")
             self.activeImportProgress = nil
 
             await MainActor.run {
-                success(UInt(count), UInt(imported))
+                success(UInt(totalKeysCount), imported)
             }
         }
     }

diff --git a/MatrixSDK/Crypto/KeyBackup/Engine/MXEncryptedKeyBackup.swift b/MatrixSDK/Crypto/KeyBackup/Engine/MXEncryptedKeyBackup.swift
@@ -0,0 +1,30 @@
+// 
+// Copyright 2023 The Matrix.org Foundation C.I.C
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+import Foundation
+
+/// Helper class for batch importing key backups
+@objcMembers public class MXEncryptedKeyBackup: NSObject {
+    public let roomId: String
+    public let sessionId: String
+    public let keyBackup: MXKeyBackupData
+
+    public init(roomId: String, sessionId: String, keyBackup: MXKeyBackupData) {
+        self.roomId = roomId
+        self.sessionId = sessionId
+        self.keyBackup = keyBackup
+    }
+}
diff --git a/MatrixSDK/Crypto/KeyBackup/Engine/MXNativeKeyBackupEngine.m b/MatrixSDK/Crypto/KeyBackup/Engine/MXNativeKeyBackupEngine.m
@@ -29,6 +29,7 @@
  Maximum number of keys to send at a time to the homeserver.
  */
 NSUInteger const kMXKeyBackupSendKeysMaxCount = 100;
+NSUInteger const kMXKeyBackupImportBatchSize = 1000;
 
 static NSDictionary<NSString*, Class<MXKeyBackupAlgorithm>> *AlgorithmClassesByName;
 static Class DefaultAlgorithmClass;
@@ -38,6 +39,8 @@ @interface MXNativeKeyBackupEngine ()
 @property (nonatomic, weak) MXLegacyCrypto *crypto;
 @property (nonatomic, nullable) MXKeyBackupVersion *keyBackupVersion;
 @property (nonatomic, nullable) id<MXKeyBackupAlgorithm> keyBackupAlgorithm;
+@property (nonatomic, nullable) NSProgress *activeImportProgress;
+@property (nonatomic, nullable) dispatch_queue_t importQueue;
 
 @end
 
@@ -67,6 +70,7 @@ - (instancetype)initWithCrypto:(MXLegacyCrypto *)crypto
     if (self)
     {
         _crypto = crypto;
+        _importQueue = dispatch_queue_create(@"MXNativeKeyBackupEngine".UTF8String, DISPATCH_QUEUE_SERIAL);
     }
     return self;
 }
@@ -533,8 +537,7 @@ - (void)backupKeysWithSuccess:(void (^)(void))success
 
 - (NSProgress *)importProgress
 {
-    // Not implemented for legacy backup
-    return nil;
+    return self.activeImportProgress;
 }
 
 - (void)importKeysWithKeysBackupData:(MXKeysBackupData *)keysBackupData
@@ -543,45 +546,104 @@ - (void)importKeysWithKeysBackupData:(MXKeysBackupData *)keysBackupData
                              success:(void (^)(NSUInteger, NSUInteger))success
                              failure:(void (^)(NSError *))failure
 {
-    id<MXKeyBackupAlgorithm> algorithm = [self getOrCreateKeyBackupAlgorithmFor:keyBackupVersion privateKey:privateKey];
+    // There is no way to cancel import so we may have one ongoing already
+    if (self.activeImportProgress)
+    {
+        MXLogError(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Another import is already ongoing");
+        if (failure)
+        {
+            NSError *error = [NSError errorWithDomain:MXKeyBackupErrorDomain code:MXKeyBackupErrorAlreadyInProgress userInfo:nil];
+            dispatch_async(dispatch_get_main_queue(), ^{
+                failure(error);
+            });
+        }
+        return;
+    }
 
-    NSMutableArray<MXMegolmSessionData*> *sessionDatas = [NSMutableArray array];
+    id<MXKeyBackupAlgorithm> algorithm = [self getOrCreateKeyBackupAlgorithmFor:keyBackupVersion privateKey:privateKey];
 
-    // Restore that data
-    NSUInteger sessionsFromHSCount = 0;
+    // Collect all sessions that we need to decrypt and import
+    NSMutableArray <MXEncryptedKeyBackup *>*encryptedSessions = [[NSMutableArray alloc] init];
     for (NSString *roomId in keysBackupData.rooms)
     {
         for (NSString *sessionId in keysBackupData.rooms[roomId].sessions)
         {
-            sessionsFromHSCount++;
             MXKeyBackupData *keyBackupData = keysBackupData.rooms[roomId].sessions[sessionId];
-            MXMegolmSessionData *sessionData = [algorithm decryptKeyBackupData:keyBackupData forSession:sessionId inRoom:roomId];
-
-            if (sessionData)
-            {
-                [sessionDatas addObject:sessionData];
-            }
+            MXEncryptedKeyBackup *backup = [[MXEncryptedKeyBackup alloc] initWithRoomId:roomId sessionId:sessionId keyBackup:keyBackupData];
+            [encryptedSessions addObject:backup];
         }
     }
 
-    MXLogDebug(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Decrypted %@ keys out of %@ from the backup store on the homeserver", @(sessionDatas.count), @(sessionsFromHSCount));
+    NSUInteger totalKeysCount = encryptedSessions.count;
+    __block NSUInteger importedKeysCount = 0;
 
-    // Do not trigger a backup for them if they come from the backup version we are using
-    BOOL backUp = ![keyBackupVersion.version isEqualToString:self.keyBackupVersion.version];
-    if (backUp)
-    {
-        MXLogDebug(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Those keys will be backed up to backup version: %@", self.keyBackupVersion.version);
-    }
+    self.activeImportProgress = [NSProgress progressWithTotalUnitCount:totalKeysCount];
+    MXLogDebug(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Importing %lu encrypted sessions", totalKeysCount);
 
-    // Import them into the crypto store
-    [self.crypto importMegolmSessionDatas:sessionDatas backUp:backUp success:success failure:^(NSError *error) {
-        if (failure)
+    NSDate *startDate = [NSDate date];
+
+    // Ensure we are on a separate queue so that decrypting and importing can happen in parallel
+    dispatch_async(self.importQueue, ^{
+        dispatch_group_t dispatchGroup = dispatch_group_create();
+
+        // Itterate through the array in memory-isolated batches
+        for (NSInteger batchIndex = 0; batchIndex < totalKeysCount; batchIndex += kMXKeyBackupImportBatchSize)
         {
-            dispatch_async(dispatch_get_main_queue(), ^{
-                failure(error);
-            });
+            MXLogDebug(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Decrypting and importing batch %ld", batchIndex);
+            dispatch_group_enter(dispatchGroup);
+
+            @autoreleasepool {
+
+                // Decrypt batch of sessions
+                NSMutableArray<MXMegolmSessionData*> *sessions = [NSMutableArray array];
+
+                NSInteger endIndex = MIN(batchIndex + kMXKeyBackupImportBatchSize, totalKeysCount);
+                for (NSInteger idx = batchIndex; idx < endIndex; idx++)
+                {
+                    MXEncryptedKeyBackup *session = encryptedSessions[idx];
+                    MXMegolmSessionData *sessionData = [algorithm decryptKeyBackupData:session.keyBackup forSession:session.sessionId inRoom:session.roomId];
+                    if (sessionData)
+                    {
+                        [sessions addObject:sessionData];
+                    }
+                }
+
+                // Do not trigger a backup for them if they come from the backup version we are using
+                BOOL backUp = ![keyBackupVersion.version isEqualToString:self.keyBackupVersion.version];
+                if (backUp)
+                {
+                    MXLogDebug(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Those keys will be backed up to backup version: %@", self.keyBackupVersion.version);
+                }
+
+                // Import them into the crypto store
+                MXWeakify(self);
+                [self.crypto importMegolmSessionDatas:sessions backUp:backUp success:^(NSUInteger total, NSUInteger imported) {
+                    MXStrongifyAndReturnIfNil(self);
+                    MXLogDebug(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Imported batch %ld", batchIndex);
+                    importedKeysCount += imported;
+
+                    self.activeImportProgress.completedUnitCount += kMXKeyBackupImportBatchSize;
+                    dispatch_group_leave(dispatchGroup);
+                } failure:^(NSError *error) {
+                    MXLogErrorDetails(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Failed importing batch of sessions", error);
+
+                    self.activeImportProgress.completedUnitCount += kMXKeyBackupImportBatchSize;
+                    dispatch_group_leave(dispatchGroup);
+                }];
+            }
         }
-    }];
+
+        dispatch_group_notify(dispatchGroup, dispatch_get_main_queue(), ^{
+            NSTimeInterval duration = [[NSDate date] timeIntervalSinceDate:startDate] * 1000;
+
+            MXLogDebug(@"[MXNativeKeyBackupEngine] importKeysWithKeysBackupData: Successfully imported %ld out of %ld sessions in %f ms", importedKeysCount, totalKeysCount, duration);
+            self.activeImportProgress = nil;
+
+            if (success) {
+                success(totalKeysCount, importedKeysCount);
+            }
+        });
+    });
 }
 
 #pragma mark - Private methods -

diff --git a/MatrixSDK/Crypto/MXCryptoV2.swift b/MatrixSDK/Crypto/MXCryptoV2.swift
@@ -656,7 +656,7 @@ class MXCryptoV2: NSObject, MXCrypto {
 
     private func handleRoomMemberEvent(_ event: MXEvent, roomState: MXRoomState?) async {
         guard
-            let userId = event.stateKey,
+            let userId = event.stateKey, !machine.isUserTracked(userId: userId),
             let state = roomState,
             let member = state.members?.member(withUserId: userId)
         else {

diff --git a/MatrixSDKTests/Crypto/Algorithms/RoomEvents/MXRoomEventEncryptionUnitTests.swift b/MatrixSDKTests/Crypto/Algorithms/RoomEvents/MXRoomEventEncryptionUnitTests.swift
@@ -86,6 +86,10 @@ class MXRoomEventEncryptionUnitTests: XCTestCase {
 
     class EncryptorStub: CryptoIdentityStub, MXCryptoRoomEventEncrypting {
         var trackedUsers: Set<String> = []
+        func isUserTracked(userId: String) -> Bool {
+            return trackedUsers.contains(userId)
+        }
+
         func addTrackedUsers(_ users: [String]) {
             trackedUsers = trackedUsers.union(users)
         }

diff --git a/changelog.d/pr-1701.change b/changelog.d/pr-1701.change
@@ -0,0 +1 @@
+Backup: Import legacy backup in batches