Merge branch 'charlesmchen/attachmentCrash'

SignalServiceKit/src/Messages/Attachments/TSAttachmentPointer.m

@@ -8,6 +8,7 @@
 #import <SignalServiceKit/MimeTypeUtil.h>
 #import <SignalServiceKit/SignalServiceKit-Swift.h>
 #import <YapDatabase/YapDatabase.h>
+#import <YapDatabase/YapDatabaseTransaction.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -209,6 +210,22 @@ NS_ASSUME_NONNULL_BEGIN
                              }];
 }
 
+- (void)saveWithTransaction:(YapDatabaseReadWriteTransaction *)transaction
+{
+#ifdef DEBUG
+    if (self.uniqueId.length > 0) {
+        id _Nullable oldObject = [transaction objectForKey:self.uniqueId inCollection:TSAttachment.collection];
+        if ([oldObject isKindOfClass:[TSAttachmentStream class]]) {
+            OWSFailDebug(@"We should never overwrite a TSAttachmentStream with a TSAttachmentPointer.");
+        }
+    } else {
+        OWSFailDebug(@"Missing uniqueId.");
+    }
+#endif
+
+    [super saveWithTransaction:transaction];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m

@@ -368,11 +368,8 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
     }
 
     if (didUpdateCache) {
-        [self.dbReadWriteConnection asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-            [self applyChangeToSelfAndLatestCopy:transaction
-                                     changeBlock:^(TSAttachmentStream *attachmentStream) {
-                                         attachmentStream.isValidImageCached = @(result);
-                                     }];
+        [self applyChangeAsyncToLatestCopyWithChangeBlock:^(TSAttachmentStream *latestInstance) {
+            latestInstance.isValidImageCached = @(result);
         }];
     }
 
@@ -395,11 +392,8 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
     }
 
     if (didUpdateCache) {
-        [self.dbReadWriteConnection asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-            [self applyChangeToSelfAndLatestCopy:transaction
-                                     changeBlock:^(TSAttachmentStream *attachmentStream) {
-                                         attachmentStream.isValidVideoCached = @(result);
-                                     }];
+        [self applyChangeAsyncToLatestCopyWithChangeBlock:^(TSAttachmentStream *latestInstance) {
+            latestInstance.isValidVideoCached = @(result);
         }];
     }
 
@@ -523,28 +517,43 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
         self.cachedImageWidth = @(imageSize.width);
         self.cachedImageHeight = @(imageSize.height);
 
-        [self.dbReadWriteConnection asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-
-            NSString *collection = [[self class] collection];
-            TSAttachmentStream *latestInstance = [transaction objectForKey:self.uniqueId inCollection:collection];
-            if (latestInstance) {
-                latestInstance.cachedImageWidth = @(imageSize.width);
-                latestInstance.cachedImageHeight = @(imageSize.height);
-                [latestInstance saveWithTransaction:transaction];
-            } else {
-                // This message has not yet been saved or has been deleted; do nothing.
-                // This isn't an error per se, but these race conditions should be
-                // _very_ rare.
-                //
-                // An exception is incoming group avatar updates which we don't ever save.
-                OWSLogWarn(@"Attachment not yet saved.");
-            }
+        [self applyChangeAsyncToLatestCopyWithChangeBlock:^(TSAttachmentStream *latestInstance) {
+            latestInstance.cachedImageWidth = @(imageSize.width);
+            latestInstance.cachedImageHeight = @(imageSize.height);
         }];
 
         return imageSize;
     }
 }
 
+#pragma mark - Update With...
+
+- (void)applyChangeAsyncToLatestCopyWithChangeBlock:(void (^)(TSAttachmentStream *))changeBlock
+{
+    OWSAssertDebug(changeBlock);
+
+    [self.dbReadWriteConnection asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        NSString *collection = [TSAttachmentStream collection];
+        TSAttachmentStream *latestInstance = [transaction objectForKey:self.uniqueId inCollection:collection];
+        if (!latestInstance) {
+            // This attachment has either not yet been saved or has been deleted; do nothing.
+            // This isn't an error per se, but these race conditions should be
+            // _very_ rare.
+            //
+            // An exception is incoming group avatar updates which we don't ever save.
+            OWSLogWarn(@"Attachment not yet saved.");
+        } else if (![latestInstance isKindOfClass:[TSAttachmentStream class]]) {
+            OWSFailDebug(@"Attachment has unexpected type: %@", latestInstance.class);
+        } else {
+            changeBlock(latestInstance);
+
+            [latestInstance saveWithTransaction:transaction];
+        }
+    }];
+}
+
+#pragma mark -
+
 - (CGFloat)calculateAudioDurationSeconds
 {
     OWSAssertIsOnMainThread();
@@ -576,18 +585,8 @@ typedef void (^OWSLoadedThumbnailSuccess)(OWSLoadedThumbnail *loadedThumbnail);
     CGFloat audioDurationSeconds = [self calculateAudioDurationSeconds];
     self.cachedAudioDurationSeconds = @(audioDurationSeconds);
 
-    [self.dbReadWriteConnection asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-        NSString *collection = [[self class] collection];
-        TSAttachmentStream *latestInstance = [transaction objectForKey:self.uniqueId inCollection:collection];
-        if (latestInstance) {
-            latestInstance.cachedAudioDurationSeconds = @(audioDurationSeconds);
-            [latestInstance saveWithTransaction:transaction];
-        } else {
-            // This message has not yet been saved or has been deleted; do nothing.
-            // This isn't an error per se, but these race conditions should be
-            // _very_ rare.
-            OWSFailDebug(@"Attachment not yet saved.");
-        }
+    [self applyChangeAsyncToLatestCopyWithChangeBlock:^(TSAttachmentStream *latestInstance) {
+        latestInstance.cachedAudioDurationSeconds = @(audioDurationSeconds);
     }];
 
     return audioDurationSeconds;

SignalServiceKit/src/Storage/TSYapDatabaseObject.h

@@ -132,7 +132,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)removeWithTransaction:(YapDatabaseReadWriteTransaction *)transaction;
 - (void)remove;
 
-#pragma mark Update With...
+#pragma mark - Update With...
 
 // This method is used by "updateWith..." methods.
 //

SignalServiceKit/src/Storage/TSYapDatabaseObject.m

@@ -201,7 +201,7 @@ NS_ASSUME_NONNULL_BEGIN
     return object;
 }
 
-#pragma mark Update With...
+#pragma mark - Update With...
 
 - (void)applyChangeToSelfAndLatestCopy:(YapDatabaseReadWriteTransaction *)transaction
                            changeBlock:(void (^)(id))changeBlock