From 258cdab2df30a8acf210aeb531d46ec678b35b33 Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Thu, 22 Mar 2018 09:45:48 -0400
Subject: [PATCH] Don't cull CloudKit records for lazy restoring attachments.

---
 Signal/src/util/OWSBackup.h          |  2 +-
 Signal/src/util/OWSBackup.m          | 16 ++++++++++++----
 Signal/src/util/OWSBackupExportJob.m |  5 +++--
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/Signal/src/util/OWSBackup.h b/Signal/src/util/OWSBackup.h
index 924ab6672..ddcdb2552 100644
--- a/Signal/src/util/OWSBackup.h
+++ b/Signal/src/util/OWSBackup.h
@@ -75,7 +75,7 @@ typedef NS_ENUM(NSUInteger, OWSBackupState) {
 
 #pragma mark - Lazy Restore
 
-- (NSArray<TSAttachmentStream *> *)attachmentsForLazyRestore;
+- (NSArray<NSString *> *)attachmentRecordNamesForLazyRestore;
 
 - (NSArray<NSString *> *)attachmentIdsForLazyRestore;
 
diff --git a/Signal/src/util/OWSBackup.m b/Signal/src/util/OWSBackup.m
index 197602cf5..f89ddb661 100644
--- a/Signal/src/util/OWSBackup.m
+++ b/Signal/src/util/OWSBackup.m
@@ -486,9 +486,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - Lazy Restore
 
-- (NSArray<TSAttachmentStream *> *)attachmentsForLazyRestore
+- (NSArray<NSString *> *)attachmentRecordNamesForLazyRestore
 {
-    NSMutableArray<TSAttachmentStream *> *attachments = [NSMutableArray new];
+    NSMutableArray<NSString *> *recordNames = [NSMutableArray new];
     [self.dbConnection readWithBlock:^(YapDatabaseReadTransaction *transaction) {
         id ext = [transaction ext:TSMessageDatabaseViewExtensionName];
         if (!ext) {
@@ -506,10 +506,18 @@ NS_ASSUME_NONNULL_BEGIN
                                              collection);
                                          return;
                                      }
-                                     [attachments addObject:object];
+                                     TSAttachmentStream *attachmentStream = object;
+                                     if (attachmentStream.backupRestoreRecordName.length < 1) {
+                                         OWSProdLogAndFail(@"%@ Invalid object: %@ in collection:%@",
+                                             self.logTag,
+                                             [object class],
+                                             collection);
+                                         return;
+                                     }
+                                     [recordNames addObject:attachmentStream.backupRestoreRecordName];
                                  }];
     }];
-    return attachments;
+    return recordNames;
 }
 
 - (NSArray<NSString *> *)attachmentIdsForLazyRestore
diff --git a/Signal/src/util/OWSBackupExportJob.m b/Signal/src/util/OWSBackupExportJob.m
index 28d82b1f0..bb7d83c36 100644
--- a/Signal/src/util/OWSBackupExportJob.m
+++ b/Signal/src/util/OWSBackupExportJob.m
@@ -1029,9 +1029,10 @@ NS_ASSUME_NONNULL_BEGIN
     OWSAssert(![activeRecordNames containsObject:self.manifestItem.recordName]);
     [activeRecordNames addObject:self.manifestItem.recordName];
 
-    // TODO: If we implement "lazy restores" where attachments (etc.) are
-    // restored lazily, we need to include the record names for all
+    // Because we do "lazy attachment restores", we need to include the record names for all
     // records that haven't been restored yet.
+    NSArray<NSString *> *restoringRecordNames = [OWSBackup.sharedManager attachmentRecordNamesForLazyRestore];
+    [activeRecordNames addObjectsFromArray:restoringRecordNames];
 
     __weak OWSBackupExportJob *weakSelf = self;
     [OWSBackupAPI fetchAllRecordNamesWithSuccess:^(NSArray<NSString *> *recordNames) {