From 7e9c3b2dac5b0188a9692f19a1a972cdc791f06c Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Tue, 22 Jan 2019 11:45:00 -0500
Subject: [PATCH] Clean up all message attachments.

---
 Signal/src/util/OWSOrphanDataCleaner.m        | 32 ++++---------------
 .../src/Messages/Interactions/TSMessage.h     |  4 +++
 .../src/Messages/Interactions/TSMessage.m     | 32 +++++++++++++------
 .../src/Messages/OWSMessageManager.m          | 13 +++-----
 4 files changed, 37 insertions(+), 44 deletions(-)

diff --git a/Signal/src/util/OWSOrphanDataCleaner.m b/Signal/src/util/OWSOrphanDataCleaner.m
index 32e685f0e..9f02179fb 100644
--- a/Signal/src/util/OWSOrphanDataCleaner.m
+++ b/Signal/src/util/OWSOrphanDataCleaner.m
@@ -1,5 +1,5 @@
 //
-//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2019 Open Whisper Systems. All rights reserved.
 //
 
 #import "OWSOrphanDataCleaner.h"
@@ -297,9 +297,7 @@ typedef void (^OrphanDataBlock)(OWSOrphanData *);
     __block NSSet *threadIds;
     // Messages
     NSMutableSet<NSString *> *orphanInteractionIds = [NSMutableSet new];
-    NSMutableSet<NSString *> *messageAttachmentIds = [NSMutableSet new];
-    NSMutableSet<NSString *> *quotedReplyThumbnailAttachmentIds = [NSMutableSet new];
-    NSMutableSet<NSString *> *contactShareAvatarAttachmentIds = [NSMutableSet new];
+    NSMutableSet<NSString *> *allMessageAttachmentIds = [NSMutableSet new];
     [databaseConnection readWithBlock:^(YapDatabaseReadTransaction *transaction) {
         [transaction
             enumerateKeysAndObjectsInCollection:TSAttachmentStream.collection
@@ -351,21 +349,7 @@ typedef void (^OrphanDataBlock)(OWSOrphanData *);
                                          }
 
                                          TSMessage *message = (TSMessage *)interaction;
-                                         if (message.attachmentIds.count > 0) {
-                                             [messageAttachmentIds addObjectsFromArray:message.attachmentIds];
-                                         }
-
-                                         TSQuotedMessage *_Nullable quotedMessage = message.quotedMessage;
-                                         if (quotedMessage) {
-                                             [quotedReplyThumbnailAttachmentIds
-                                                 addObjectsFromArray:quotedMessage.thumbnailAttachmentStreamIds];
-                                         }
-
-                                         OWSContact *_Nullable contactShare = message.contactShare;
-                                         if (contactShare && contactShare.avatarAttachmentId) {
-                                             [contactShareAvatarAttachmentIds
-                                                 addObject:contactShare.avatarAttachmentId];
-                                         }
+                                         [allMessageAttachmentIds addObjectsFromArray:message.allAttachmentIds];
                                      }];
     }];
     if (shouldAbort) {
@@ -390,15 +374,11 @@ typedef void (^OrphanDataBlock)(OWSOrphanData *);
     [self printPaths:missingAttachmentFilePaths.allObjects label:@"missing attachment file paths"];
 
     OWSLogDebug(@"attachmentIds: %zu", allAttachmentIds.count);
-    OWSLogDebug(@"messageAttachmentIds: %zu", messageAttachmentIds.count);
-    OWSLogDebug(@"quotedReplyThumbnailAttachmentIds: %zu", quotedReplyThumbnailAttachmentIds.count);
-    OWSLogDebug(@"contactShareAvatarAttachmentIds: %zu", contactShareAvatarAttachmentIds.count);
+    OWSLogDebug(@"allMessageAttachmentIds: %zu", allMessageAttachmentIds.count);
 
     NSMutableSet<NSString *> *orphanAttachmentIds = [allAttachmentIds mutableCopy];
-    [orphanAttachmentIds minusSet:messageAttachmentIds];
-    [orphanAttachmentIds minusSet:quotedReplyThumbnailAttachmentIds];
-    [orphanAttachmentIds minusSet:contactShareAvatarAttachmentIds];
-    NSMutableSet<NSString *> *missingAttachmentIds = [messageAttachmentIds mutableCopy];
+    [orphanAttachmentIds minusSet:allMessageAttachmentIds];
+    NSMutableSet<NSString *> *missingAttachmentIds = [allMessageAttachmentIds mutableCopy];
     [missingAttachmentIds minusSet:allAttachmentIds];
 
     OWSLogDebug(@"orphan attachmentIds: %zu", orphanAttachmentIds.count);
diff --git a/SignalServiceKit/src/Messages/Interactions/TSMessage.h b/SignalServiceKit/src/Messages/Interactions/TSMessage.h
index 55297e154..2050aa7c7 100644
--- a/SignalServiceKit/src/Messages/Interactions/TSMessage.h
+++ b/SignalServiceKit/src/Messages/Interactions/TSMessage.h
@@ -48,6 +48,10 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)removeAttachment:(TSAttachment *)attachment
              transaction:(YapDatabaseReadWriteTransaction *)transaction NS_SWIFT_NAME(removeAttachment(_:transaction:));
 
+// Returns ids for all attachments, including message ("body") attachments,
+// quoted reply thumbnails, contact share avatars, link preview images, etc.
+- (NSArray<NSString *> *)allAttachmentIds;
+
 - (BOOL)isMediaAlbumWithTransaction:(YapDatabaseReadTransaction *)transaction;
 
 - (void)setQuotedMessageThumbnailAttachmentStream:(TSAttachmentStream *)attachmentStream;
diff --git a/SignalServiceKit/src/Messages/Interactions/TSMessage.m b/SignalServiceKit/src/Messages/Interactions/TSMessage.m
index 9865e83b2..296348678 100644
--- a/SignalServiceKit/src/Messages/Interactions/TSMessage.m
+++ b/SignalServiceKit/src/Messages/Interactions/TSMessage.m
@@ -184,6 +184,28 @@ static const NSUInteger OWSMessageSchemaVersion = 4;
     return self.attachmentIds ? (self.attachmentIds.count > 0) : NO;
 }
 
+- (NSArray<NSString *> *)allAttachmentIds
+{
+    NSMutableArray<NSString *> *result = [NSMutableArray new];
+    if (self.attachmentIds.count > 0) {
+        [result addObjectsFromArray:self.attachmentIds];
+    }
+
+    if (self.quotedMessage) {
+        [result addObjectsFromArray:self.quotedMessage.thumbnailAttachmentStreamIds];
+    }
+
+    if (self.contactShare.avatarAttachmentId) {
+        [result addObject:self.contactShare.avatarAttachmentId];
+    }
+
+    if (self.linkPreview.imageAttachmentId) {
+        [result addObject:self.linkPreview.imageAttachmentId];
+    }
+
+    return [result copy];
+}
+
 - (NSArray<TSAttachment *> *)attachmentsWithTransaction:(YapDatabaseReadTransaction *)transaction
 {
     NSMutableArray<TSAttachment *> *attachments = [NSMutableArray new];
@@ -338,7 +360,7 @@ static const NSUInteger OWSMessageSchemaVersion = 4;
 {
     [super removeWithTransaction:transaction];
 
-    for (NSString *attachmentId in self.attachmentIds) {
+    for (NSString *attachmentId in self.allAttachmentIds) {
         // We need to fetch each attachment, since [TSAttachment removeWithTransaction:] does important work.
         TSAttachment *_Nullable attachment =
             [TSAttachment fetchObjectWithUniqueID:attachmentId transaction:transaction];
@@ -349,14 +371,6 @@ static const NSUInteger OWSMessageSchemaVersion = 4;
         [attachment removeWithTransaction:transaction];
     };
 
-    if (self.contactShare.avatarAttachmentId) {
-        [self.contactShare removeAvatarAttachmentWithTransaction:transaction];
-    }
-
-    if (self.linkPreview.imageAttachmentId) {
-        [self.linkPreview removeAttachmentWithTransaction:transaction];
-    }
-
     // Updates inbox thread preview
     [self touchThreadWithTransaction:transaction];
 }
diff --git a/SignalServiceKit/src/Messages/OWSMessageManager.m b/SignalServiceKit/src/Messages/OWSMessageManager.m
index 5635bff66..b80f1b689 100644
--- a/SignalServiceKit/src/Messages/OWSMessageManager.m
+++ b/SignalServiceKit/src/Messages/OWSMessageManager.m
@@ -1437,15 +1437,10 @@ NS_ASSUME_NONNULL_BEGIN
         [incomingMessage markAsReadAtTimestamp:envelope.timestamp sendReadReceipt:NO transaction:transaction];
     }
 
-    NSMutableArray<NSString *> *otherAttachmentIds = [NSMutableArray new];
-    if (incomingMessage.quotedMessage.thumbnailAttachmentPointerId.length > 0) {
-        [otherAttachmentIds addObject:incomingMessage.quotedMessage.thumbnailAttachmentPointerId];
-    }
-    if (incomingMessage.contactShare.avatarAttachmentId.length > 0) {
-        [otherAttachmentIds addObject:incomingMessage.contactShare.avatarAttachmentId];
-    }
-    if (incomingMessage.linkPreview.imageAttachmentId.length > 0) {
-        [otherAttachmentIds addObject:incomingMessage.linkPreview.imageAttachmentId];
+    // Download the "non-message body" attachments.
+    NSMutableArray<NSString *> *otherAttachmentIds = [incomingMessage.allAttachmentIds mutableCopy];
+    if (incomingMessage.attachmentIds) {
+        [otherAttachmentIds removeObjectsInArray:incomingMessage.attachmentIds];
     }
     for (NSString *attachmentId in otherAttachmentIds) {
         TSAttachmentPointer *_Nullable attachmentPointer =