From d0c691bb7fdaaff9cea6ce689bad5eee995ec03d Mon Sep 17 00:00:00 2001 From: Matthew Chen Date: Tue, 20 Mar 2018 17:22:19 -0400 Subject: [PATCH] Lazy attachment restores. --- Signal.xcodeproj/project.pbxproj | 4 + Signal/src/AppDelegate.m | 6 +- Signal/src/Signal-Bridging-Header.h | 2 + Signal/src/util/OWSBackup.h | 14 +- Signal/src/util/OWSBackup.m | 144 ++++++++++++++++++ Signal/src/util/OWSBackupAPI.swift | 9 ++ Signal/src/util/OWSBackupIO.h | 2 + Signal/src/util/OWSBackupIO.m | 15 +- Signal/src/util/OWSBackupJob.m | 6 +- Signal/src/util/OWSBackupLazyRestoreJob.swift | 90 +++++++++++ .../Messages/Attachments/TSAttachmentStream.h | 12 ++ .../Messages/Attachments/TSAttachmentStream.m | 30 ++++ .../src/Storage/OWSPrimaryStorage.m | 1 + SignalServiceKit/src/Storage/TSDatabaseView.h | 5 + SignalServiceKit/src/Storage/TSDatabaseView.m | 95 +++++++++--- SignalServiceKit/src/Util/OWSFileSystem.h | 2 + SignalServiceKit/src/Util/OWSFileSystem.m | 15 ++ 17 files changed, 422 insertions(+), 30 deletions(-) create mode 100644 Signal/src/util/OWSBackupLazyRestoreJob.swift diff --git a/Signal.xcodeproj/project.pbxproj b/Signal.xcodeproj/project.pbxproj index c35a62037..6ee946dd6 100644 --- a/Signal.xcodeproj/project.pbxproj +++ b/Signal.xcodeproj/project.pbxproj @@ -205,6 +205,7 @@ 34D1F0C01F8EC1760066283D /* MessageRecipientStatusUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 34D1F0BF1F8EC1760066283D /* MessageRecipientStatusUtils.swift */; }; 34D2CCD4206294B900CB1A14 /* OWSScreenLock.swift in Sources */ = {isa = PBXBuildFile; fileRef = 34D2CCD3206294B900CB1A14 /* OWSScreenLock.swift */; }; 34D2CCDA2062E7D000CB1A14 /* OWSScreenLockUI.m in Sources */ = {isa = PBXBuildFile; fileRef = 34D2CCD92062E7D000CB1A14 /* OWSScreenLockUI.m */; }; + 34D2CCD220618B3000CB1A14 /* OWSBackupLazyRestoreJob.swift in Sources */ = {isa = PBXBuildFile; fileRef = 34D2CCD120618B2F00CB1A14 /* OWSBackupLazyRestoreJob.swift */; }; 34D5CCA91EAE3D30005515DB /* AvatarViewHelper.m in Sources */ = {isa = PBXBuildFile; fileRef = 34D5CCA81EAE3D30005515DB /* AvatarViewHelper.m */; }; 34D8C0271ED3673300188D7C /* DebugUIMessages.m in Sources */ = {isa = PBXBuildFile; fileRef = 34D8C0241ED3673300188D7C /* DebugUIMessages.m */; }; 34D8C0281ED3673300188D7C /* DebugUITableViewController.m in Sources */ = {isa = PBXBuildFile; fileRef = 34D8C0261ED3673300188D7C /* DebugUITableViewController.m */; }; @@ -812,6 +813,7 @@ 34D2CCD3206294B900CB1A14 /* OWSScreenLock.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = OWSScreenLock.swift; sourceTree = ""; }; 34D2CCD82062E7D000CB1A14 /* OWSScreenLockUI.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = OWSScreenLockUI.h; sourceTree = ""; }; 34D2CCD92062E7D000CB1A14 /* OWSScreenLockUI.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = OWSScreenLockUI.m; sourceTree = ""; }; + 34D2CCD120618B2F00CB1A14 /* OWSBackupLazyRestoreJob.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = OWSBackupLazyRestoreJob.swift; sourceTree = ""; }; 34D5CCA71EAE3D30005515DB /* AvatarViewHelper.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = AvatarViewHelper.h; sourceTree = ""; }; 34D5CCA81EAE3D30005515DB /* AvatarViewHelper.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = AvatarViewHelper.m; sourceTree = ""; }; 34D8C0231ED3673300188D7C /* DebugUIMessages.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = DebugUIMessages.h; sourceTree = ""; }; @@ -1974,6 +1976,7 @@ 34D2CCD3206294B900CB1A14 /* OWSScreenLock.swift */, 34D2CCD82062E7D000CB1A14 /* OWSScreenLockUI.h */, 34D2CCD92062E7D000CB1A14 /* OWSScreenLockUI.m */, + 34D2CCD120618B2F00CB1A14 /* OWSBackupLazyRestoreJob.swift */, 4579431C1E7C8CE9008ED0C0 /* Pastelog.h */, 4579431D1E7C8CE9008ED0C0 /* Pastelog.m */, 450DF2041E0D74AC003D14BE /* Platform.swift */, @@ -3143,6 +3146,7 @@ 458DE9D91DEE7B360071BB03 /* OWSWebRTCDataProtos.pb.m in Sources */, 76EB063C18170B33006006FC /* NumberUtil.m in Sources */, 451166C01FD86B98000739BA /* AccountManager.swift in Sources */, + 34D2CCD220618B3000CB1A14 /* OWSBackupLazyRestoreJob.swift in Sources */, 3430FE181F7751D4000EC51B /* GiphyAPI.swift in Sources */, 34A55F3720485465002CC6DE /* OWS2FARegistrationViewController.m in Sources */, 340FC8AD204DAC8D007AEB0F /* OWSLinkedDevicesTableViewController.m in Sources */, diff --git a/Signal/src/AppDelegate.m b/Signal/src/AppDelegate.m index 6d2b0d1cd..3ae51edc6 100644 --- a/Signal/src/AppDelegate.m +++ b/Signal/src/AppDelegate.m @@ -418,8 +418,8 @@ static NSString *const kURLHostVerifyPrefix = @"verify"; // Every time we change or add a database view in such a way that // might cause a delay on launch, we need to bump this constant. // - // We upgraded YapDatabase in v2.20.0 and need to regenerate all database views. - NSString *kLastVersionWithDatabaseViewChange = @"2.20.0"; + // We added a database view in v2.23.0. + NSString *kLastVersionWithDatabaseViewChange = @"2.23.0"; BOOL mayNeedUpgrade = ([TSAccountManager isRegistered] && lastLaunchedAppVersion && (!lastCompletedLaunchAppVersion || [VersionMigrations isVersion:lastCompletedLaunchAppVersion @@ -1134,6 +1134,8 @@ static NSString *const kURLHostVerifyPrefix = @"verify"; [self ensureRootViewController]; [OWSBackup.sharedManager setup]; + + [OWSBackupLazyRestoreJob run]; } - (void)registrationStateDidChange diff --git a/Signal/src/Signal-Bridging-Header.h b/Signal/src/Signal-Bridging-Header.h index 0c93f10df..a950f6f37 100644 --- a/Signal/src/Signal-Bridging-Header.h +++ b/Signal/src/Signal-Bridging-Header.h @@ -19,6 +19,8 @@ #import "NotificationsManager.h" #import "OWSAnyTouchGestureRecognizer.h" #import "OWSAudioPlayer.h" +#import "OWSBackup.h" +#import "OWSBackupIO.h" #import "OWSBezierPathView.h" #import "OWSCallNotificationsAdaptee.h" #import "OWSDatabaseMigration.h" diff --git a/Signal/src/util/OWSBackup.h b/Signal/src/util/OWSBackup.h index bcae8148f..924ab6672 100644 --- a/Signal/src/util/OWSBackup.h +++ b/Signal/src/util/OWSBackup.h @@ -20,13 +20,15 @@ typedef NS_ENUM(NSUInteger, OWSBackupState) { OWSBackupState_Succeeded, }; +@class OWSBackupIO; +@class TSAttachmentStream; @class TSThread; @interface OWSBackup : NSObject - (instancetype)init NS_UNAVAILABLE; -+ (instancetype)sharedManager; ++ (instancetype)sharedManager NS_SWIFT_NAME(shared()); - (void)setup; @@ -71,6 +73,16 @@ typedef NS_ENUM(NSUInteger, OWSBackupState) { - (void)logBackupRecords; - (void)clearAllCloudKitRecords; +#pragma mark - Lazy Restore + +- (NSArray *)attachmentsForLazyRestore; + +- (NSArray *)attachmentIdsForLazyRestore; + +- (void)lazyRestoreAttachment:(TSAttachmentStream *)attachment + backupIO:(OWSBackupIO *)backupIO + completion:(OWSBackupBoolBlock)completion; + @end NS_ASSUME_NONNULL_END diff --git a/Signal/src/util/OWSBackup.m b/Signal/src/util/OWSBackup.m index 420cfcaf6..197602cf5 100644 --- a/Signal/src/util/OWSBackup.m +++ b/Signal/src/util/OWSBackup.m @@ -4,6 +4,7 @@ #import "OWSBackup.h" #import "OWSBackupExportJob.h" +#import "OWSBackupIO.h" #import "OWSBackupImportJob.h" #import "Signal-Swift.h" #import @@ -483,6 +484,149 @@ NS_ASSUME_NONNULL_BEGIN }]; } +#pragma mark - Lazy Restore + +- (NSArray *)attachmentsForLazyRestore +{ + NSMutableArray *attachments = [NSMutableArray new]; + [self.dbConnection readWithBlock:^(YapDatabaseReadTransaction *transaction) { + id ext = [transaction ext:TSMessageDatabaseViewExtensionName]; + if (!ext) { + OWSProdLogAndFail(@"%@ Could not load database view.", self.logTag); + return; + } + + [ext enumerateKeysAndObjectsInGroup:TSLazyRestoreAttachmentsGroup + usingBlock:^( + NSString *collection, NSString *key, id object, NSUInteger index, BOOL *stop) { + if (![object isKindOfClass:[TSAttachmentStream class]]) { + OWSProdLogAndFail(@"%@ Unexpected object: %@ in collection:%@", + self.logTag, + [object class], + collection); + return; + } + [attachments addObject:object]; + }]; + }]; + return attachments; +} + +- (NSArray *)attachmentIdsForLazyRestore +{ + NSMutableArray *attachmentIds = [NSMutableArray new]; + [self.dbConnection readWithBlock:^(YapDatabaseReadTransaction *transaction) { + id ext = [transaction ext:TSMessageDatabaseViewExtensionName]; + if (!ext) { + OWSProdLogAndFail(@"%@ Could not load database view.", self.logTag); + return; + } + + [ext enumerateKeysInGroup:TSLazyRestoreAttachmentsGroup + usingBlock:^(NSString *collection, NSString *key, NSUInteger index, BOOL *stop) { + [attachmentIds addObject:key]; + }]; + }]; + return attachmentIds; +} + +- (void)lazyRestoreAttachment:(TSAttachmentStream *)attachment + backupIO:(OWSBackupIO *)backupIO + completion:(OWSBackupBoolBlock)completion +{ + OWSAssert(attachment); + OWSAssert(backupIO); + OWSAssert(completion); + + NSString *_Nullable attachmentFilePath = [attachment filePath]; + if (attachmentFilePath.length < 1) { + DDLogError(@"%@ Attachment has invalid file path.", self.logTag); + return completion(NO); + } + if ([NSFileManager.defaultManager fileExistsAtPath:attachmentFilePath]) { + DDLogError(@"%@ Attachment already has file.", self.logTag); + return completion(NO); + } + + NSString *_Nullable recordName = attachment.backupRestoreRecordName; + NSData *_Nullable encryptionKey = attachment.backupRestoreEncryptionKey; + if (recordName.length < 1 || encryptionKey.length < 1) { + return completion(NO); + } + + // Use a predictable file path so that multiple "import backup" attempts + // will leverage successful file downloads from previous attempts. + // + // TODO: This will also require imports using a predictable jobTempDirPath. + NSString *_Nullable tempFilePath = [backupIO createTempFile]; + if (!tempFilePath) { + return completion(NO); + } + + [OWSBackupAPI downloadFileFromCloudWithRecordName:recordName + toFileUrl:[NSURL fileURLWithPath:tempFilePath] + success:^{ + dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ + [self lazyRestoreAttachment:attachment + backupIO:backupIO + encryptedFilePath:tempFilePath + encryptionKey:encryptionKey + completion:completion]; + }); + } + failure:^(NSError *error) { + dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ + completion(NO); + }); + }]; +} + +- (void)lazyRestoreAttachment:(TSAttachmentStream *)attachment + backupIO:(OWSBackupIO *)backupIO + encryptedFilePath:(NSString *)encryptedFilePath + encryptionKey:(NSData *)encryptionKey + completion:(OWSBackupBoolBlock)completion +{ + OWSAssert(attachment); + OWSAssert(backupIO); + OWSAssert(encryptedFilePath.length > 0); + OWSAssert(encryptionKey.length > 0); + OWSAssert(completion); + + NSData *_Nullable data = [NSData dataWithContentsOfFile:encryptedFilePath]; + if (!data) { + DDLogError(@"%@ Could not load encrypted file.", self.logTag); + return completion(NO); + } + + NSString *_Nullable decryptedFilePath = [backupIO createTempFile]; + if (!decryptedFilePath) { + return completion(NO); + } + + if (![backupIO decryptFileAsFile:encryptedFilePath dstFilePath:decryptedFilePath encryptionKey:encryptionKey]) { + DDLogError(@"%@ Could not load decrypt file.", self.logTag); + return completion(NO); + } + + NSString *_Nullable attachmentFilePath = [attachment filePath]; + if (attachmentFilePath.length < 1) { + DDLogError(@"%@ Attachment has invalid file path.", self.logTag); + return completion(NO); + } + NSError *error; + BOOL success = + [NSFileManager.defaultManager moveItemAtPath:decryptedFilePath toPath:attachmentFilePath error:&error]; + if (!success || error) { + DDLogError(@"%@ Attachment file could not be restored: %@.", self.logTag, error); + return completion(NO); + } + + [attachment updateWithBackupRestoreComplete]; + + completion(YES); +} + #pragma mark - Notifications - (void)postDidChangeNotification diff --git a/Signal/src/util/OWSBackupAPI.swift b/Signal/src/util/OWSBackupAPI.swift index a3a53e0f0..3cdc1955f 100644 --- a/Signal/src/util/OWSBackupAPI.swift +++ b/Signal/src/util/OWSBackupAPI.swift @@ -6,6 +6,15 @@ import Foundation import SignalServiceKit import CloudKit +// We don't worry about atomic writes. Each backup export +// will diff against last successful backup. +// +// Note that all of our CloudKit records are immutable. +// "Persistent" records are only uploaded once. +// "Ephemeral" records are always uploaded to a new record name. +// +// TODO: We could store known encryption ids locally to +// facilitate "resume" of failed backup exports. @objc public class OWSBackupAPI: NSObject { // If we change the record types, we need to ensure indices diff --git a/Signal/src/util/OWSBackupIO.h b/Signal/src/util/OWSBackupIO.h index fabb0088d..eb25f0277 100644 --- a/Signal/src/util/OWSBackupIO.h +++ b/Signal/src/util/OWSBackupIO.h @@ -20,6 +20,8 @@ NS_ASSUME_NONNULL_BEGIN - (instancetype)initWithJobTempDirPath:(NSString *)jobTempDirPath; +- (nullable NSString *)createTempFile; + #pragma mark - Encrypt - (nullable OWSBackupEncryptedItem *)encryptFileAsTempFile:(NSString *)srcFilePath; diff --git a/Signal/src/util/OWSBackupIO.m b/Signal/src/util/OWSBackupIO.m index 02b50d9fb..75c91bb5d 100644 --- a/Signal/src/util/OWSBackupIO.m +++ b/Signal/src/util/OWSBackupIO.m @@ -44,6 +44,16 @@ static const compression_algorithm SignalCompressionAlgorithm = COMPRESSION_LZMA return self; } +- (nullable NSString *)createTempFile +{ + NSString *filePath = [self.jobTempDirPath stringByAppendingPathComponent:[NSUUID UUID].UUIDString]; + if (![OWSFileSystem ensureFileExists:filePath]) { + OWSProdLogAndFail(@"%@ could not create temp file.", self.logTag); + return nil; + } + return filePath; +} + #pragma mark - Encrypt - (nullable OWSBackupEncryptedItem *)encryptFileAsTempFile:(NSString *)srcFilePath @@ -92,7 +102,10 @@ static const compression_algorithm SignalCompressionAlgorithm = COMPRESSION_LZMA // TODO: Encrypt the data using key; NSData *encryptedData = unencryptedData; - NSString *dstFilePath = [self.jobTempDirPath stringByAppendingPathComponent:[NSUUID UUID].UUIDString]; + NSString *_Nullable dstFilePath = [self createTempFile]; + if (!dstFilePath) { + return nil; + } NSError *error; BOOL success = [encryptedData writeToFile:dstFilePath options:NSDataWritingAtomic error:&error]; if (!success || error) { diff --git a/Signal/src/util/OWSBackupJob.m b/Signal/src/util/OWSBackupJob.m index 35a0c91a0..51a4fa8e1 100644 --- a/Signal/src/util/OWSBackupJob.m +++ b/Signal/src/util/OWSBackupJob.m @@ -86,16 +86,12 @@ NSString *const kOWSBackup_KeychainService = @"kOWSBackup_KeychainService"; // might want to use a predictable directory so that repeated // import attempts can reuse downloads from previous attempts. NSString *temporaryDirectory = NSTemporaryDirectory(); - self.jobTempDirPath = [temporaryDirectory stringByAppendingString:[NSUUID UUID].UUIDString]; + self.jobTempDirPath = [temporaryDirectory stringByAppendingPathComponent:[NSUUID UUID].UUIDString]; if (![OWSFileSystem ensureDirectoryExists:self.jobTempDirPath]) { OWSProdLogAndFail(@"%@ Could not create jobTempDirPath.", self.logTag); return NO; } - if (![OWSFileSystem protectFileOrFolderAtPath:self.jobTempDirPath]) { - OWSProdLogAndFail(@"%@ Could not protect jobTempDirPath.", self.logTag); - return NO; - } return YES; } diff --git a/Signal/src/util/OWSBackupLazyRestoreJob.swift b/Signal/src/util/OWSBackupLazyRestoreJob.swift new file mode 100644 index 000000000..7c28f03da --- /dev/null +++ b/Signal/src/util/OWSBackupLazyRestoreJob.swift @@ -0,0 +1,90 @@ +// +// Copyright (c) 2018 Open Whisper Systems. All rights reserved. +// + +import Foundation +import PromiseKit +import SignalServiceKit + +@objc +public class OWSBackupLazyRestoreJob: NSObject { + + let TAG = "[OWSBackupLazyRestoreJob]" + + let primaryStorage: OWSPrimaryStorage + + private var jobTempDirPath: String? + + deinit { + if let jobTempDirPath = self.jobTempDirPath { + DispatchQueue.global().async { + OWSFileSystem.deleteFile(jobTempDirPath) + } + } + } + + @objc + public class func run() { + OWSBackupLazyRestoreJob().run() + } + + public override init() { + self.primaryStorage = OWSPrimaryStorage.shared() + } + + private func run() { + AssertIsOnMainThread() + + DispatchQueue.global().async { + self.restoreAttachments() + } + } + + private func restoreAttachments() { + let temporaryDirectory = NSTemporaryDirectory() + let jobTempDirPath = (temporaryDirectory as NSString).appendingPathComponent(NSUUID().uuidString) + // let jobTempDirPath = temporaryDirectory.appendingPathComponent(UUID().uuidString) + + guard OWSFileSystem.ensureDirectoryExists(jobTempDirPath) else { + Logger.error("\(TAG) could not create temp directory.") + return + } + + self.jobTempDirPath = jobTempDirPath + + let backupIO = OWSBackupIO(jobTempDirPath: jobTempDirPath) + + let attachmentIds = OWSBackup.shared().attachmentIdsForLazyRestore() + self.tryToRestoreNextAttachment(attachmentIds: attachmentIds, backupIO: backupIO) + } + + private func tryToRestoreNextAttachment(attachmentIds: [String], backupIO: OWSBackupIO) { + var attachmentIdsCopy = attachmentIds + guard let attachmentId = attachmentIdsCopy.last else { + // This job is done. + Logger.verbose("\(TAG) job is done.") + return + } + attachmentIdsCopy.removeLast() + guard let attachment = TSAttachmentStream.fetch(uniqueId: attachmentId) else { + Logger.warn("\(TAG) could not load attachment.") + // Not necessarily an error. + // The attachment might have been deleted since the job began. + // Continue trying to restore the other attachments. + tryToRestoreNextAttachment(attachmentIds: attachmentIds, backupIO: backupIO) + return + } + OWSBackup.shared().lazyRestoreAttachment(attachment, + backupIO: backupIO, + completion: { (success) in + if success { + Logger.info("\(self.TAG) restored attachment.") + } else { + Logger.warn("\(self.TAG) could not restore attachment.") + } + // Continue trying to restore the other attachments. + self.tryToRestoreNextAttachment(attachmentIds: attachmentIdsCopy, backupIO: backupIO) + }) + + } +} diff --git a/SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.h b/SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.h index 1349a34e0..8acb08552 100644 --- a/SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.h +++ b/SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.h @@ -34,6 +34,11 @@ NS_ASSUME_NONNULL_BEGIN @property (nonatomic, readonly) NSDate *creationTimestamp; +// Optional properties. Set only for attachments which +// need "lazy backup restore." +@property (nonatomic, readonly, nullable) NSString *backupRestoreRecordName; +@property (nonatomic, readonly, nullable) NSData *backupRestoreEncryptionKey; + #if TARGET_OS_IPHONE - (nullable UIImage *)image; - (nullable UIImage *)thumbnailImage; @@ -62,6 +67,13 @@ NS_ASSUME_NONNULL_BEGIN + (nullable NSError *)migrateToSharedData; +#pragma mark - Update With... Methods + +// Marks attachment as needing "lazy backup restore." +- (void)updateWithBackupRestoreRecordName:(NSString *)recordName encryptionKey:(NSData *)encryptionKey; +// Marks attachment as having completed "lazy backup restore." +- (void)updateWithBackupRestoreComplete; + @end NS_ASSUME_NONNULL_END diff --git a/SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m b/SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m index 1bef2f1da..05c464e13 100644 --- a/SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m +++ b/SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m @@ -26,6 +26,9 @@ NS_ASSUME_NONNULL_BEGIN // This property should only be accessed on the main thread. @property (nullable, nonatomic) NSNumber *cachedAudioDurationSeconds; +@property (nonatomic, nullable) NSString *backupRestoreRecordName; +@property (nonatomic, nullable) NSData *backupRestoreEncryptionKey; + @end #pragma mark - @@ -610,6 +613,33 @@ NS_ASSUME_NONNULL_BEGIN return audioDurationSeconds; } +#pragma mark - Update With... Methods + +- (void)updateWithBackupRestoreRecordName:(NSString *)recordName encryptionKey:(NSData *)encryptionKey +{ + OWSAssert(recordName.length > 0); + OWSAssert(encryptionKey.length > 0); + + [self.dbReadWriteConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) { + [self applyChangeToSelfAndLatestCopy:transaction + changeBlock:^(TSAttachmentStream *attachment) { + [attachment setBackupRestoreRecordName:recordName]; + [attachment setBackupRestoreEncryptionKey:encryptionKey]; + }]; + }]; +} + +- (void)updateWithBackupRestoreComplete +{ + [self.dbReadWriteConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) { + [self applyChangeToSelfAndLatestCopy:transaction + changeBlock:^(TSAttachmentStream *attachment) { + [attachment setBackupRestoreRecordName:nil]; + [attachment setBackupRestoreEncryptionKey:nil]; + }]; + }]; +} + @end NS_ASSUME_NONNULL_END diff --git a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m index e05947567..ddff2a59a 100644 --- a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m +++ b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m @@ -57,6 +57,7 @@ void runAsyncRegistrationsForStorage(OWSStorage *storage) [OWSFailedMessagesJob asyncRegisterDatabaseExtensionsWithPrimaryStorage:storage]; [OWSFailedAttachmentDownloadsJob asyncRegisterDatabaseExtensionsWithPrimaryStorage:storage]; [OWSMediaGalleryFinder asyncRegisterDatabaseExtensionsWithPrimaryStorage:storage]; + [TSDatabaseView asyncRegisterLazyRestoreAttachmentsDatabaseView:storage]; } #pragma mark - diff --git a/SignalServiceKit/src/Storage/TSDatabaseView.h b/SignalServiceKit/src/Storage/TSDatabaseView.h index 59357856f..3e83d174b 100644 --- a/SignalServiceKit/src/Storage/TSDatabaseView.h +++ b/SignalServiceKit/src/Storage/TSDatabaseView.h @@ -17,6 +17,9 @@ extern NSString *const TSUnreadDatabaseViewExtensionName; extern NSString *const TSSecondaryDevicesDatabaseViewExtensionName; +extern NSString *const TSLazyRestoreAttachmentsGroup; +extern NSString *const TSLazyRestoreAttachmentsDatabaseViewExtensionName; + @interface TSDatabaseView : NSObject - (instancetype)init NS_UNAVAILABLE; @@ -55,4 +58,6 @@ extern NSString *const TSSecondaryDevicesDatabaseViewExtensionName; + (void)asyncRegisterSecondaryDevicesDatabaseView:(OWSStorage *)storage; ++ (void)asyncRegisterLazyRestoreAttachmentsDatabaseView:(OWSStorage *)storage; + @end diff --git a/SignalServiceKit/src/Storage/TSDatabaseView.m b/SignalServiceKit/src/Storage/TSDatabaseView.m index 77d4a651f..2602c08b0 100644 --- a/SignalServiceKit/src/Storage/TSDatabaseView.m +++ b/SignalServiceKit/src/Storage/TSDatabaseView.m @@ -5,6 +5,8 @@ #import "TSDatabaseView.h" #import "OWSDevice.h" #import "OWSReadTracking.h" +#import "TSAttachment.h" +#import "TSAttachmentStream.h" #import "TSIncomingMessage.h" #import "TSInvalidIdentityKeyErrorMessage.h" #import "TSOutgoingMessage.h" @@ -28,6 +30,9 @@ NSString *const TSUnreadDatabaseViewExtensionName = @"TSUnreadDatabaseViewExtens NSString *const TSUnseenDatabaseViewExtensionName = @"TSUnseenDatabaseViewExtensionName"; NSString *const TSThreadSpecialMessagesDatabaseViewExtensionName = @"TSThreadSpecialMessagesDatabaseViewExtensionName"; NSString *const TSSecondaryDevicesDatabaseViewExtensionName = @"TSSecondaryDevicesDatabaseViewExtensionName"; +NSString *const TSLazyRestoreAttachmentsDatabaseViewExtensionName + = @"TSLazyRestoreAttachmentsDatabaseViewExtensionName"; +NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"; @interface OWSStorage (TSDatabaseView) @@ -295,30 +300,26 @@ NSString *const TSSecondaryDevicesDatabaseViewExtensionName = @"TSSecondaryDevic + (void)asyncRegisterSecondaryDevicesDatabaseView:(OWSStorage *)storage { - YapDatabaseViewGrouping *viewGrouping = - [YapDatabaseViewGrouping withObjectBlock:^NSString *_Nullable(YapDatabaseReadTransaction *_Nonnull transaction, - NSString *_Nonnull collection, - NSString *_Nonnull key, - id _Nonnull object) { - if ([object isKindOfClass:[OWSDevice class]]) { - OWSDevice *device = (OWSDevice *)object; - if (![device isPrimaryDevice]) { - return TSSecondaryDevicesGroup; - } + YapDatabaseViewGrouping *viewGrouping = [YapDatabaseViewGrouping withObjectBlock:^NSString *_Nullable( + YapDatabaseReadTransaction *transaction, NSString *collection, NSString *key, id object) { + if ([object isKindOfClass:[OWSDevice class]]) { + OWSDevice *device = (OWSDevice *)object; + if (![device isPrimaryDevice]) { + return TSSecondaryDevicesGroup; } - return nil; - }]; + } + return nil; + }]; YapDatabaseViewSorting *viewSorting = - [YapDatabaseViewSorting withObjectBlock:^NSComparisonResult(YapDatabaseReadTransaction *_Nonnull transaction, - NSString *_Nonnull group, - NSString *_Nonnull collection1, - NSString *_Nonnull key1, - id _Nonnull object1, - NSString *_Nonnull collection2, - NSString *_Nonnull key2, - id _Nonnull object2) { - + [YapDatabaseViewSorting withObjectBlock:^NSComparisonResult(YapDatabaseReadTransaction *transaction, + NSString *group, + NSString *collection1, + NSString *key1, + id object1, + NSString *collection2, + NSString *key2, + id object2) { if ([object1 isKindOfClass:[OWSDevice class]] && [object2 isKindOfClass:[OWSDevice class]]) { OWSDevice *device1 = (OWSDevice *)object1; OWSDevice *device2 = (OWSDevice *)object2; @@ -341,6 +342,58 @@ NSString *const TSSecondaryDevicesDatabaseViewExtensionName = @"TSSecondaryDevic [storage asyncRegisterExtension:view withName:TSSecondaryDevicesDatabaseViewExtensionName]; } ++ (void)asyncRegisterLazyRestoreAttachmentsDatabaseView:(OWSStorage *)storage +{ + YapDatabaseViewGrouping *viewGrouping = [YapDatabaseViewGrouping withObjectBlock:^NSString *_Nullable( + YapDatabaseReadTransaction *transaction, NSString *collection, NSString *key, id object) { + if (![object isKindOfClass:[TSAttachment class]]) { + OWSProdLogAndFail(@"%@ Unexpected entity %@ in collection: %@", self.logTag, [object class], collection); + return nil; + } + if (![object isKindOfClass:[TSAttachmentStream class]]) { + return nil; + } + TSAttachmentStream *attachmentStream = (TSAttachmentStream *)object; + if (attachmentStream.backupRestoreRecordName.length > 0 + && attachmentStream.backupRestoreEncryptionKey.length > 0) { + return TSLazyRestoreAttachmentsGroup; + } else { + return nil; + } + }]; + + YapDatabaseViewSorting *viewSorting = [YapDatabaseViewSorting withObjectBlock:^NSComparisonResult( + YapDatabaseReadTransaction *transaction, + NSString *group, + NSString *collection1, + NSString *key1, + id object1, + NSString *collection2, + NSString *key2, + id object2) { + if (![object1 isKindOfClass:[TSAttachment class]]) { + OWSProdLogAndFail(@"%@ Unexpected entity %@ in collection: %@", self.logTag, [object1 class], collection1); + return NSOrderedSame; + } + if (![object2 isKindOfClass:[TSAttachment class]]) { + OWSProdLogAndFail(@"%@ Unexpected entity %@ in collection: %@", self.logTag, [object2 class], collection2); + return NSOrderedSame; + } + + TSAttachmentStream *attachmentStream1 = (TSAttachmentStream *)object1; + TSAttachmentStream *attachmentStream2 = (TSAttachmentStream *)object2; + return [attachmentStream2.creationTimestamp compare:attachmentStream1.creationTimestamp]; + }]; + + YapDatabaseViewOptions *options = [YapDatabaseViewOptions new]; + options.isPersistent = YES; + options.allowedCollections = + [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:[TSAttachment collection]]]; + YapDatabaseView *view = + [[YapDatabaseAutoView alloc] initWithGrouping:viewGrouping sorting:viewSorting versionTag:@"1" options:options]; + [storage asyncRegisterExtension:view withName:TSLazyRestoreAttachmentsDatabaseViewExtensionName]; +} + + (id)unseenDatabaseViewExtension:(YapDatabaseReadTransaction *)transaction { OWSAssert(transaction); diff --git a/SignalServiceKit/src/Util/OWSFileSystem.h b/SignalServiceKit/src/Util/OWSFileSystem.h index e9a180289..b56e9f49d 100644 --- a/SignalServiceKit/src/Util/OWSFileSystem.h +++ b/SignalServiceKit/src/Util/OWSFileSystem.h @@ -28,6 +28,8 @@ NS_ASSUME_NONNULL_BEGIN // Returns NO IFF the directory does not exist and could not be created. + (BOOL)ensureDirectoryExists:(NSString *)dirPath; ++ (BOOL)ensureFileExists:(NSString *)filePath; + + (BOOL)deleteFile:(NSString *)filePath; + (BOOL)deleteFileIfExists:(NSString *)filePath; diff --git a/SignalServiceKit/src/Util/OWSFileSystem.m b/SignalServiceKit/src/Util/OWSFileSystem.m index fa83b41a3..88c53f091 100644 --- a/SignalServiceKit/src/Util/OWSFileSystem.m +++ b/SignalServiceKit/src/Util/OWSFileSystem.m @@ -227,6 +227,21 @@ NS_ASSUME_NONNULL_BEGIN } } ++ (BOOL)ensureFileExists:(NSString *)filePath +{ + BOOL exists = [[NSFileManager defaultManager] fileExistsAtPath:filePath]; + if (exists) { + return [self protectFileOrFolderAtPath:filePath]; + } else { + BOOL success = [[NSFileManager defaultManager] createFileAtPath:filePath contents:nil attributes:nil]; + if (!success) { + OWSFail(@"%@ Failed to create file.", self.logTag); + return NO; + } + return [self protectFileOrFolderAtPath:filePath]; + } +} + + (BOOL)deleteFile:(NSString *)filePath { NSError *error;