Use attachment pointers to restore attachments from backup.

Signal/src/util/Backup/OWSBackup.h

@@ -79,7 +79,7 @@ typedef NS_ENUM(NSUInteger, OWSBackupState) {
 
 - (NSArray<NSString *> *)attachmentIdsForLazyRestore;
 
-- (void)lazyRestoreAttachment:(TSAttachmentStream *)attachment
+- (void)lazyRestoreAttachment:(TSAttachmentPointer *)attachment
                      backupIO:(OWSBackupIO *)backupIO
                    completion:(OWSBackupBoolBlock)completion;
 

Signal/src/util/Backup/OWSBackup.m

@@ -500,18 +500,18 @@ NS_ASSUME_NONNULL_BEGIN
         [ext enumerateKeysAndObjectsInGroup:TSLazyRestoreAttachmentsGroup
                                  usingBlock:^(
                                      NSString *collection, NSString *key, id object, NSUInteger index, BOOL *stop) {
-                                     if (![object isKindOfClass:[TSAttachmentStream class]]) {
+                                     if (![object isKindOfClass:[TSAttachmentPointer class]]) {
                                          OWSFailDebug(
                                              @"Unexpected object: %@ in collection:%@", [object class], collection);
                                          return;
                                      }
-                                     TSAttachmentStream *attachmentStream = object;
+                                     TSAttachmentPointer *attachmentPointer = object;
-                                     if (!attachmentStream.lazyRestoreFragment) {
+                                     if (!attachmentPointer.lazyRestoreFragment) {
                                          OWSFailDebug(
                                              @"Invalid object: %@ in collection:%@", [object class], collection);
                                          return;
                                      }
-                                     [recordNames addObject:attachmentStream.lazyRestoreFragment.recordName];
+                                     [recordNames addObject:attachmentPointer.lazyRestoreFragment.recordName];
                                  }];
     }];
     return recordNames;
@@ -535,7 +535,7 @@ NS_ASSUME_NONNULL_BEGIN
     return attachmentIds;
 }
 
-- (void)lazyRestoreAttachment:(TSAttachmentStream *)attachment
+- (void)lazyRestoreAttachment:(TSAttachmentPointer *)attachment
                      backupIO:(OWSBackupIO *)backupIO
                    completion:(OWSBackupBoolBlock)completion
 {
@@ -543,16 +543,6 @@ NS_ASSUME_NONNULL_BEGIN
     OWSAssertDebug(backupIO);
     OWSAssertDebug(completion);
 
-    NSString *_Nullable attachmentFilePath = [attachment originalFilePath];
-    if (attachmentFilePath.length < 1) {
-        OWSLogError(@"Attachment has invalid file path.");
-        return completion(NO);
-    }
-    if ([NSFileManager.defaultManager fileExistsAtPath:attachmentFilePath]) {
-        OWSLogError(@"Attachment already has file.");
-        return completion(NO);
-    }
-
     OWSBackupFragment *_Nullable lazyRestoreFragment = attachment.lazyRestoreFragment;
     if (!lazyRestoreFragment) {
         OWSLogWarn(@"Attachment missing lazy restore metadata.");
@@ -587,13 +577,13 @@ NS_ASSUME_NONNULL_BEGIN
         }];
 }
 
-- (void)lazyRestoreAttachment:(TSAttachmentStream *)attachment
+- (void)lazyRestoreAttachment:(TSAttachmentPointer *)attachmentPointer
                      backupIO:(OWSBackupIO *)backupIO
             encryptedFilePath:(NSString *)encryptedFilePath
                 encryptionKey:(NSData *)encryptionKey
                    completion:(OWSBackupBoolBlock)completion
 {
-    OWSAssertDebug(attachment);
+    OWSAssertDebug(attachmentPointer);
     OWSAssertDebug(backupIO);
     OWSAssertDebug(encryptedFilePath.length > 0);
     OWSAssertDebug(encryptionKey.length > 0);
@@ -614,7 +604,9 @@ NS_ASSUME_NONNULL_BEGIN
         }
     }
 
-    NSString *_Nullable attachmentFilePath = [attachment originalFilePath];
+    TSAttachmentStream *stream = [[TSAttachmentStream alloc] initWithPointer:attachmentPointer];
+
+    NSString *attachmentFilePath = stream.originalFilePath;
     if (attachmentFilePath.length < 1) {
         OWSLogError(@"Attachment has invalid file path.");
         return completion(NO);
@@ -634,7 +626,10 @@ NS_ASSUME_NONNULL_BEGIN
         return completion(NO);
     }
 
-    [attachment updateWithLazyRestoreComplete];
+    [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        // This should overwrite the attachment pointer with an attachment stream.
+        [stream saveWithTransaction:transaction];
+    }];
 
     completion(YES);
 }

Signal/src/util/Backup/OWSBackupImportJob.m

@@ -281,13 +281,18 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
                 // Attachment-related errors are recoverable and can be ignored.
                 continue;
             }
-            TSAttachmentStream *_Nullable attachment =
+            TSAttachmentPointer *_Nullable attachment =
-                [TSAttachmentStream fetchObjectWithUniqueID:item.attachmentId transaction:transaction];
+                [TSAttachmentPointer fetchObjectWithUniqueID:item.attachmentId transaction:transaction];
             if (!attachment) {
                 OWSLogError(@"attachment to restore could not be found.");
                 // Attachment-related errors are recoverable and can be ignored.
                 continue;
             }
+            if (![attachment isKindOfClass:[TSAttachmentPointer class]]) {
+                OWSFailDebug(@"attachment has unexpected type: %@.", attachment.class);
+                // Attachment-related errors are recoverable and can be ignored.
+                continue;
+            }
             [attachment markForLazyRestoreWithFragment:item transaction:transaction];
             count++;
             [self updateProgressWithDescription:NSLocalizedString(@"BACKUP_IMPORT_PHASE_RESTORING_FILES",

SignalServiceKit/src/Messages/Attachments/TSAttachment.h

@@ -27,7 +27,7 @@ typedef NS_ENUM(NSUInteger, TSAttachmentType) {
 // TSAttachmentPointer, which can be distinguished by the isDownloaded
 // property.
 @property (atomic, readwrite) UInt64 serverId;
-@property (atomic, readwrite) NSData *encryptionKey;
+@property (atomic, readwrite, nullable) NSData *encryptionKey;
 @property (nonatomic, readonly) NSString *contentType;
 @property (atomic, readwrite) BOOL isDownloaded;
 @property (nonatomic) TSAttachmentType attachmentType;
@@ -62,6 +62,13 @@ typedef NS_ENUM(NSUInteger, TSAttachmentType) {
                          caption:(nullable NSString *)caption
                   albumMessageId:(nullable NSString *)albumMessageId;
 
+// This constructor is used for new instances of TSAttachmentPointer,
+// i.e. undownloaded restoring attachments.
+- (instancetype)initWithContentType:(NSString *)contentType
+                     sourceFilename:(nullable NSString *)sourceFilename
+                            caption:(nullable NSString *)caption
+                     albumMessageId:(nullable NSString *)albumMessageId;
+
 // This constructor is used for new instances of TSAttachmentStream
 // that represent new, un-uploaded outgoing attachments.
 - (instancetype)initWithContentType:(NSString *)contentType

SignalServiceKit/src/Messages/Attachments/TSAttachment.m

@@ -65,6 +65,35 @@ NSUInteger const TSAttachmentSchemaVersion = 4;
     return self;
 }
 
+// This constructor is used for new instances of TSAttachmentPointer,
+// i.e. undownloaded restoring attachments.
+- (instancetype)initWithContentType:(NSString *)contentType
+                     sourceFilename:(nullable NSString *)sourceFilename
+                            caption:(nullable NSString *)caption
+                     albumMessageId:(nullable NSString *)albumMessageId
+{
+    if (contentType.length < 1) {
+        OWSLogWarn(@"incoming attachment has invalid content type");
+
+        contentType = OWSMimeTypeApplicationOctetStream;
+    }
+    OWSAssertDebug(contentType.length > 0);
+
+    self = [super init];
+    if (!self) {
+        return self;
+    }
+
+    _contentType = contentType;
+    _sourceFilename = sourceFilename;
+    _caption = caption;
+    _albumMessageId = albumMessageId;
+
+    _attachmentSchemaVersion = TSAttachmentSchemaVersion;
+
+    return self;
+}
+
 // This constructor is used for new instances of TSAttachmentStream
 // that represent new, un-uploaded outgoing attachments.
 - (instancetype)initWithContentType:(NSString *)contentType

SignalServiceKit/src/Messages/Attachments/TSAttachmentPointer.h

@@ -6,9 +6,15 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+@class OWSBackupFragment;
 @class SSKProtoAttachmentPointer;
 @class TSMessage;
 
+typedef NS_ENUM(NSUInteger, TSAttachmentPointerType) {
+    TSAttachmentPointerTypeIncoming = 0,
+    TSAttachmentPointerTypeRestoring = 2,
+};
+
 typedef NS_ENUM(NSUInteger, TSAttachmentPointerState) {
     TSAttachmentPointerStateEnqueued = 0,
     TSAttachmentPointerStateDownloading = 1,
@@ -20,6 +26,17 @@ typedef NS_ENUM(NSUInteger, TSAttachmentPointerState) {
  */
 @interface TSAttachmentPointer : TSAttachment
 
+@property (nonatomic) TSAttachmentPointerType pointerType;
+@property (atomic) TSAttachmentPointerState state;
+@property (nullable, atomic) NSString *mostRecentFailureLocalizedText;
+
+// Though now required, `digest` may be null for pre-existing records or from
+// messages received from other clients
+@property (nullable, nonatomic, readonly) NSData *digest;
+
+// Non-nil for attachments which need "lazy backup restore."
+- (nullable OWSBackupFragment *)lazyRestoreFragment;
+
 - (nullable instancetype)initWithCoder:(NSCoder *)coder NS_DESIGNATED_INITIALIZER;
 
 - (instancetype)initWithServerId:(UInt64)serverId
@@ -32,6 +49,12 @@ typedef NS_ENUM(NSUInteger, TSAttachmentPointerState) {
                   albumMessageId:(nullable NSString *)albumMessageId
                   attachmentType:(TSAttachmentType)attachmentType NS_DESIGNATED_INITIALIZER;
 
+- (instancetype)initForRestoreWithContentType:(NSString *)contentType
+                               sourceFilename:(nullable NSString *)sourceFilename
+                                      caption:(nullable NSString *)caption
+                               albumMessageId:(nullable NSString *)albumMessageId
+                               attachmentType:(TSAttachmentType)attachmentType NS_DESIGNATED_INITIALIZER;
+
 + (nullable TSAttachmentPointer *)attachmentPointerFromProto:(SSKProtoAttachmentPointer *)attachmentProto
                                                 albumMessage:(nullable TSMessage *)message;
 
@@ -39,12 +62,11 @@ typedef NS_ENUM(NSUInteger, TSAttachmentPointerState) {
                                         (NSArray<SSKProtoAttachmentPointer *> *)attachmentProtos
                                                     albumMessage:(TSMessage *)message;
 
-@property (atomic) TSAttachmentPointerState state;
+#pragma mark - Update With... Methods
-@property (nullable, atomic) NSString *mostRecentFailureLocalizedText;
 
-// Though now required, `digest` may be null for pre-existing records or from
+// Marks attachment as needing "lazy backup restore."
-// messages received from other clients
+- (void)markForLazyRestoreWithFragment:(OWSBackupFragment *)lazyRestoreFragment
-@property (nullable, nonatomic, readonly) NSData *digest;
+                           transaction:(YapDatabaseReadWriteTransaction *)transaction;
 
 @end
 

SignalServiceKit/src/Messages/Attachments/TSAttachmentPointer.m

@@ -3,11 +3,22 @@
 //
 
 #import "TSAttachmentPointer.h"
+#import "OWSBackupFragment.h"
 #import <SignalServiceKit/MimeTypeUtil.h>
 #import <SignalServiceKit/SignalServiceKit-Swift.h>
+#import <YapDatabase/YapDatabase.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
+@interface TSAttachmentPointer ()
+
+// Optional property.  Only set for attachments which need "lazy backup restore."
+@property (nonatomic, nullable) NSString *lazyRestoreFragmentId;
+
+@end
+
+#pragma mark -
+
 @implementation TSAttachmentPointer
 
 - (nullable instancetype)initWithCoder:(NSCoder *)coder
@@ -24,6 +35,10 @@ NS_ASSUME_NONNULL_BEGIN
         _state = TSAttachmentPointerStateFailed;
     }
 
+    if (![coder containsValueForKey:@"pointerType"]) {
+        _pointerType = TSAttachmentPointerTypeIncoming;
+    }
+
     return self;
 }
 
@@ -51,6 +66,28 @@ NS_ASSUME_NONNULL_BEGIN
     _digest = digest;
     _state = TSAttachmentPointerStateEnqueued;
     self.attachmentType = attachmentType;
+    _pointerType = TSAttachmentPointerTypeIncoming;
+
+    return self;
+}
+
+- (instancetype)initForRestoreWithContentType:(NSString *)contentType
+                               sourceFilename:(nullable NSString *)sourceFilename
+                                      caption:(nullable NSString *)caption
+                               albumMessageId:(nullable NSString *)albumMessageId
+                               attachmentType:(TSAttachmentType)attachmentType
+{
+    self = [super initWithContentType:contentType
+                       sourceFilename:sourceFilename
+                              caption:caption
+                       albumMessageId:albumMessageId];
+    if (!self) {
+        return self;
+    }
+
+    _state = TSAttachmentPointerStateEnqueued;
+    self.attachmentType = attachmentType;
+    _pointerType = TSAttachmentPointerTypeRestoring;
 
     return self;
 }
@@ -141,6 +178,34 @@ NS_ASSUME_NONNULL_BEGIN
     }
 }
 
+- (nullable OWSBackupFragment *)lazyRestoreFragment
+{
+    if (!self.lazyRestoreFragmentId) {
+        return nil;
+    }
+    return [OWSBackupFragment fetchObjectWithUniqueID:self.lazyRestoreFragmentId];
+}
+
+#pragma mark - Update With... Methods
+
+- (void)markForLazyRestoreWithFragment:(OWSBackupFragment *)lazyRestoreFragment
+                           transaction:(YapDatabaseReadWriteTransaction *)transaction
+{
+    OWSAssertDebug(lazyRestoreFragment);
+    OWSAssertDebug(transaction);
+
+    if (!lazyRestoreFragment.uniqueId) {
+        // If metadata hasn't been saved yet, save now.
+        [lazyRestoreFragment saveWithTransaction:transaction];
+
+        OWSAssertDebug(lazyRestoreFragment.uniqueId);
+    }
+    [self applyChangeToSelfAndLatestCopy:transaction
+                             changeBlock:^(TSAttachmentPointer *attachment) {
+                                 [attachment setLazyRestoreFragmentId:lazyRestoreFragment.uniqueId];
+                             }];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.h

@@ -3,7 +3,6 @@
 //
 
 #import "DataSource.h"
-#import "OWSBackupFragment.h"
 #import "TSAttachment.h"
 
 #if TARGET_OS_IPHONE
@@ -70,9 +69,6 @@ typedef void (^OWSThumbnailFailure)(void);
 
 + (nullable NSError *)migrateToSharedData;
 
-// Non-nil for attachments which need "lazy backup restore."
-- (nullable OWSBackupFragment *)lazyRestoreFragment;
-
 #pragma mark - Thumbnails
 
 // On cache hit, the thumbnail will be returned synchronously and completion will never be invoked.
@@ -99,12 +95,6 @@ typedef void (^OWSThumbnailFailure)(void);
 
 #pragma mark - Update With... Methods
 
-// Marks attachment as needing "lazy backup restore."
-- (void)markForLazyRestoreWithFragment:(OWSBackupFragment *)lazyRestoreFragment
-                           transaction:(YapDatabaseReadWriteTransaction *)transaction;
-// Marks attachment as having completed "lazy backup restore."
-- (void)updateWithLazyRestoreComplete;
-
 - (nullable TSAttachmentStream *)cloneAsThumbnail;
 
 #pragma mark - Protobuf

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m

@@ -39,9 +39,6 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
 // This property should only be accessed on the main thread.
 @property (nullable, nonatomic) NSNumber *cachedAudioDurationSeconds;
 
-// Optional property.  Only set for attachments which need "lazy backup restore."
-@property (nonatomic, nullable) NSString *lazyRestoreFragmentId;
-
 @property (atomic, nullable) NSNumber *isValidImageCached;
 @property (atomic, nullable) NSNumber *isValidVideoCached;
 
@@ -358,10 +355,6 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
 {
     OWSAssertDebug(self.isImage || self.isAnimated);
 
-    if (self.lazyRestoreFragment) {
-        return NO;
-    }
-
     @synchronized(self) {
         if (!self.isValidImageCached) {
             self.isValidImageCached =
@@ -375,10 +368,6 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
 {
     OWSAssertDebug(self.isVideo);
 
-    if (self.lazyRestoreFragment) {
-        return NO;
-    }
-
     @synchronized(self) {
         if (!self.isValidVideoCached) {
             self.isValidVideoCached = @([OWSMediaUtils isValidVideoWithPath:self.originalFilePath]);
@@ -629,14 +618,6 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
     return audioDurationSeconds;
 }
 
-- (nullable OWSBackupFragment *)lazyRestoreFragment
-{
-    if (!self.lazyRestoreFragmentId) {
-        return nil;
-    }
-    return [OWSBackupFragment fetchObjectWithUniqueID:self.lazyRestoreFragmentId];
-}
-
 #pragma mark - Thumbnails
 
 - (nullable UIImage *)thumbnailImageWithSizeHint:(CGSize)sizeHint
@@ -819,34 +800,6 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
 
 #pragma mark - Update With... Methods
 
-- (void)markForLazyRestoreWithFragment:(OWSBackupFragment *)lazyRestoreFragment
-                           transaction:(YapDatabaseReadWriteTransaction *)transaction
-{
-    OWSAssertDebug(lazyRestoreFragment);
-    OWSAssertDebug(transaction);
-
-    if (!lazyRestoreFragment.uniqueId) {
-        // If metadata hasn't been saved yet, save now.
-        [lazyRestoreFragment saveWithTransaction:transaction];
-
-        OWSAssertDebug(lazyRestoreFragment.uniqueId);
-    }
-    [self applyChangeToSelfAndLatestCopy:transaction
-                             changeBlock:^(TSAttachmentStream *attachment) {
-                                 [attachment setLazyRestoreFragmentId:lazyRestoreFragment.uniqueId];
-                             }];
-}
-
-- (void)updateWithLazyRestoreComplete
-{
-    [self.dbReadWriteConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-        [self applyChangeToSelfAndLatestCopy:transaction
-                                 changeBlock:^(TSAttachmentStream *attachment) {
-                                     [attachment setLazyRestoreFragmentId:nil];
-                                 }];
-    }];
-}
-
 - (nullable TSAttachmentStream *)cloneAsThumbnail
 {
     NSData *_Nullable thumbnailData = self.thumbnailDataSmallSync;

SignalServiceKit/src/Storage/TSDatabaseView.m

@@ -6,7 +6,7 @@
 #import "OWSDevice.h"
 #import "OWSReadTracking.h"
 #import "TSAttachment.h"
-#import "TSAttachmentStream.h"
+#import "TSAttachmentPointer.h"
 #import "TSIncomingMessage.h"
 #import "TSInvalidIdentityKeyErrorMessage.h"
 #import "TSOutgoingMessage.h"
@@ -369,47 +369,47 @@ NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"
             OWSFailDebug(@"Unexpected entity %@ in collection: %@", [object class], collection);
             return nil;
         }
-        if (![object isKindOfClass:[TSAttachmentStream class]]) {
+        if (![object isKindOfClass:[TSAttachmentPointer class]]) {
             return nil;
         }
-        TSAttachmentStream *attachmentStream = (TSAttachmentStream *)object;
+        TSAttachmentPointer *attachmentPointer = (TSAttachmentPointer *)object;
-        if (attachmentStream.lazyRestoreFragment) {
+        if (attachmentPointer.lazyRestoreFragment) {
             return TSLazyRestoreAttachmentsGroup;
         } else {
             return nil;
         }
     }];
 
-    YapDatabaseViewSorting *viewSorting = [YapDatabaseViewSorting withObjectBlock:^NSComparisonResult(
+    YapDatabaseViewSorting *viewSorting =
-        YapDatabaseReadTransaction *transaction,
+        [YapDatabaseViewSorting withObjectBlock:^NSComparisonResult(YapDatabaseReadTransaction *transaction,
-        NSString *group,
+            NSString *group,
-        NSString *collection1,
+            NSString *collection1,
-        NSString *key1,
+            NSString *key1,
-        id object1,
+            id object1,
-        NSString *collection2,
+            NSString *collection2,
-        NSString *key2,
+            NSString *key2,
-        id object2) {
+            id object2) {
-        if (![object1 isKindOfClass:[TSAttachmentStream class]]) {
+            if (![object1 isKindOfClass:[TSAttachmentPointer class]]) {
-            OWSFailDebug(@"Unexpected entity %@ in collection: %@", [object1 class], collection1);
+                OWSFailDebug(@"Unexpected entity %@ in collection: %@", [object1 class], collection1);
-            return NSOrderedSame;
+                return NSOrderedSame;
-        }
+            }
-        if (![object2 isKindOfClass:[TSAttachmentStream class]]) {
+            if (![object2 isKindOfClass:[TSAttachmentPointer class]]) {
-            OWSFailDebug(@"Unexpected entity %@ in collection: %@", [object2 class], collection2);
+                OWSFailDebug(@"Unexpected entity %@ in collection: %@", [object2 class], collection2);
-            return NSOrderedSame;
+                return NSOrderedSame;
-        }
+            }
 
-        // Specific ordering doesn't matter; we just need a stable ordering.
+            // Specific ordering doesn't matter; we just need a stable ordering.
-        TSAttachmentStream *attachmentStream1 = (TSAttachmentStream *)object1;
+            TSAttachmentPointer *attachmentPointer1 = (TSAttachmentPointer *)object1;
-        TSAttachmentStream *attachmentStream2 = (TSAttachmentStream *)object2;
+            TSAttachmentPointer *attachmentPointer2 = (TSAttachmentPointer *)object2;
-        return [attachmentStream2.creationTimestamp compare:attachmentStream1.creationTimestamp];
+            return [attachmentPointer1.uniqueId compare:attachmentPointer2.uniqueId];
-    }];
+        }];
 
     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:@"3" options:options];
+        [[YapDatabaseAutoView alloc] initWithGrouping:viewGrouping sorting:viewSorting versionTag:@"4" options:options];
     [storage asyncRegisterExtension:view withName:TSLazyRestoreAttachmentsDatabaseViewExtensionName];
 }