Lazy restore attachments.

Signal/src/util/OWSBackup.m

@@ -616,9 +616,11 @@ NS_ASSUME_NONNULL_BEGIN
         return completion(NO);
     }
 
-    if (![backupIO decryptFileAsFile:encryptedFilePath dstFilePath:decryptedFilePath encryptionKey:encryptionKey]) {
+    @autoreleasepool {
-        DDLogError(@"%@ Could not load decrypt file.", self.logTag);
+        if (![backupIO decryptFileAsFile:encryptedFilePath dstFilePath:decryptedFilePath encryptionKey:encryptionKey]) {
-        return completion(NO);
+            DDLogError(@"%@ Could not load decrypt file.", self.logTag);
+            return completion(NO);
+        }
     }
 
     NSString *_Nullable attachmentFilePath = [attachment filePath];

Signal/src/util/OWSBackupImportJob.m

@@ -255,53 +255,43 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
 {
     DDLogVerbose(@"%@ %s: %zd", self.logTag, __PRETTY_FUNCTION__, self.attachmentsItems.count);
 
-    NSString *attachmentsDirPath = [TSAttachmentStream attachmentsFolder];
+    __block NSUInteger count = 0;
-
+    [self.primaryStorage.newDatabaseConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-    NSUInteger count = 0;
+        for (OWSBackupFragment *item in self.attachmentsItems) {
-    for (OWSBackupFragment *item in self.attachmentsItems) {
+            if (self.isComplete) {
-        if (self.isComplete) {
+                return;
-            return;
+            }
-        }
+            if (item.recordName.length < 1) {
-        if (item.recordName.length < 1) {
+                DDLogError(@"%@ attachment was not downloaded.", self.logTag);
-            DDLogError(@"%@ attachment was not downloaded.", self.logTag);
+                // Attachment-related errors are recoverable and can be ignored.
-            // Attachment-related errors are recoverable and can be ignored.
-            continue;
-        }
-        if (item.relativeFilePath.length < 1) {
-            DDLogError(@"%@ attachment missing relative file path.", self.logTag);
-            // Attachment-related errors are recoverable and can be ignored.
-            continue;
-        }
-
-        count++;
-        [self updateProgressWithDescription:NSLocalizedString(@"BACKUP_IMPORT_PHASE_RESTORING_FILES",
-                                                @"Indicates that the backup import data is being restored.")
-                                   progress:@(count / (CGFloat)self.attachmentsItems.count)];
-
-        NSString *dstFilePath = [attachmentsDirPath stringByAppendingPathComponent:item.relativeFilePath];
-        if ([NSFileManager.defaultManager fileExistsAtPath:dstFilePath]) {
-            DDLogError(@"%@ skipping redundant file restore: %@.", self.logTag, dstFilePath);
-            continue;
-        }
-        NSString *dstDirPath = [dstFilePath stringByDeletingLastPathComponent];
-        if (![NSFileManager.defaultManager fileExistsAtPath:dstDirPath]) {
-            if (![OWSFileSystem ensureDirectoryExists:dstDirPath]) {
-                DDLogError(@"%@ couldn't create directory for file restore: %@.", self.logTag, dstFilePath);
                 continue;
             }
-        }
+            if (item.attachmentId.length < 1) {
-        @autoreleasepool {
+                DDLogError(@"%@ attachment missing attachment id.", self.logTag);
-            if (![self.backupIO decryptFileAsFile:item.downloadFilePath
+                // Attachment-related errors are recoverable and can be ignored.
-                                      dstFilePath:dstFilePath
+                continue;
-                                    encryptionKey:item.encryptionKey]) {
+            }
-                DDLogError(@"%@ attachment could not be restored.", self.logTag);
+            if (item.relativeFilePath.length < 1) {
+                DDLogError(@"%@ attachment missing relative file path.", self.logTag);
                 // Attachment-related errors are recoverable and can be ignored.
                 continue;
             }
+            TSAttachmentStream *_Nullable attachment =
+                [TSAttachmentStream fetchObjectWithUniqueID:item.attachmentId transaction:transaction];
+            if (!attachment) {
+                DDLogError(@"%@ attachment to restore could not be found.", self.logTag);
+                // Attachment-related errors are recoverable and can be ignored.
+                continue;
+            }
+            [attachment updateWithLazyRestoreFragment:item transaction:transaction];
+            count++;
+            [self updateProgressWithDescription:NSLocalizedString(@"BACKUP_IMPORT_PHASE_RESTORING_FILES",
+                                                    @"Indicates that the backup import data is being restored.")
+                                       progress:@(count / (CGFloat)self.attachmentsItems.count)];
         }
+    }];
 
-        DDLogError(@"%@ restored file: %@.", self.logTag, item.relativeFilePath);
+    DDLogError(@"%@ enqueued lazy restore of %zd files.", self.logTag, count);
-    }
 }
 
 - (void)restoreDatabaseWithCompletion:(OWSBackupJobBoolCompletion)completion

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.h

@@ -69,7 +69,8 @@ NS_ASSUME_NONNULL_BEGIN
 #pragma mark - Update With... Methods
 
 // Marks attachment as needing "lazy backup restore."
-- (void)updateWithLazyRestoreFragment:(OWSBackupFragment *)lazyRestoreFragment;
+- (void)updateWithLazyRestoreFragment:(OWSBackupFragment *)lazyRestoreFragment
+                          transaction:(YapDatabaseReadWriteTransaction *)transaction;
 // Marks attachment as having completed "lazy backup restore."
 - (void)updateWithLazyRestoreComplete;
 

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m

@@ -624,21 +624,21 @@ NS_ASSUME_NONNULL_BEGIN
 #pragma mark - Update With... Methods
 
 - (void)updateWithLazyRestoreFragment:(OWSBackupFragment *)lazyRestoreFragment
+                          transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     OWSAssert(lazyRestoreFragment);
+    OWSAssert(transaction);
 
-    [self.dbReadWriteConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+    if (!lazyRestoreFragment.uniqueId) {
-        if (!lazyRestoreFragment.uniqueId) {
+        // If metadata hasn't been saved yet, save now.
-            // If metadata hasn't been saved yet, save now.
+        [lazyRestoreFragment saveWithTransaction:transaction];
-            [lazyRestoreFragment saveWithTransaction:transaction];
 
-            OWSAssert(lazyRestoreFragment.uniqueId);
+        OWSAssert(lazyRestoreFragment.uniqueId);
-        }
+    }
-        [self applyChangeToSelfAndLatestCopy:transaction
+    [self applyChangeToSelfAndLatestCopy:transaction
-                                 changeBlock:^(TSAttachmentStream *attachment) {
+                             changeBlock:^(TSAttachmentStream *attachment) {
-                                     [attachment setLazyRestoreFragmentId:lazyRestoreFragment.uniqueId];
+                                 [attachment setLazyRestoreFragmentId:lazyRestoreFragment.uniqueId];
-                                 }];
+                             }];
-    }];
 }
 
 - (void)updateWithLazyRestoreComplete