Detect and handle corrupt database views.

Signal/src/ViewControllers/DebugUI/DebugUIMisc.m

@@ -120,9 +120,11 @@ NS_ASSUME_NONNULL_BEGIN
                                            }]];
 #endif
 
-    [items addObject:[OWSTableItem itemWithTitle:@"Increment database extension version suffix"
+    [items addObject:[OWSTableItem itemWithTitle:@"Increment Database Extension Versions"
                                      actionBlock:^() {
-                                         [OWSStorage incrementDatabaseExtensionVersionSuffix];
+                                         for (NSString *extensionName in ExtensionNamesForPrimaryStorage()) {
+                                             [OWSStorage incrementVersionOfDatabaseExtension:extensionName];
+                                         }
                                      }]];
 
     return [OWSTableSection sectionWithTitle:self.name items:items];

SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m

@@ -203,11 +203,7 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
     options.allowedCollections =
         [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:[OWSMessageContentJob collection]]];
 
-    return [[YapDatabaseAutoView alloc]
-        initWithGrouping:grouping
-                 sorting:sorting
-              versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:@"1"]
-                 options:options];
+    return [[YapDatabaseAutoView alloc] initWithGrouping:grouping sorting:sorting versionTag:@"1" options:options];
 }
 
 

SignalServiceKit/src/Messages/OWSDisappearingMessagesFinder.m

@@ -185,10 +185,7 @@ static NSString *const OWSDisappearingMessageFinderExpiresAtIndex = @"index_mess
             dict[OWSDisappearingMessageFinderThreadIdColumn] = message.uniqueThreadId;
         }];
 
-    return [[YapDatabaseSecondaryIndex alloc]
-        initWithSetup:setup
-              handler:handler
-           versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:nil]];
+    return [[YapDatabaseSecondaryIndex alloc] initWithSetup:setup handler:handler versionTag:nil];
 }
 
 #ifdef DEBUG

SignalServiceKit/src/Messages/OWSFailedAttachmentDownloadsJob.m

@@ -114,10 +114,7 @@ static NSString *const OWSFailedAttachmentDownloadsJobAttachmentStateIndex = @"i
             dict[OWSFailedAttachmentDownloadsJobAttachmentStateColumn] = @(attachment.state);
         }];
 
-    return [[YapDatabaseSecondaryIndex alloc]
-        initWithSetup:setup
-              handler:handler
-           versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:nil]];
+    return [[YapDatabaseSecondaryIndex alloc] initWithSetup:setup handler:handler versionTag:nil];
 }
 
 #ifdef DEBUG

SignalServiceKit/src/Messages/OWSFailedMessagesJob.m

@@ -124,10 +124,7 @@ static NSString *const OWSFailedMessagesJobMessageStateIndex = @"index_outoing_m
             dict[OWSFailedMessagesJobMessageStateColumn] = @(message.messageState);
         }];
 
-    return [[YapDatabaseSecondaryIndex alloc]
-        initWithSetup:setup
-              handler:handler
-           versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:nil]];
+    return [[YapDatabaseSecondaryIndex alloc] initWithSetup:setup handler:handler versionTag:nil];
 }
 
 #ifdef DEBUG

SignalServiceKit/src/Messages/OWSMessageReceiver.m

@@ -183,11 +183,7 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
     options.allowedCollections =
         [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:[OWSMessageDecryptJob collection]]];
 
-    return [[YapDatabaseAutoView alloc]
-        initWithGrouping:grouping
-                 sorting:sorting
-              versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:@"1"]
-                 options:options];
+    return [[YapDatabaseAutoView alloc] initWithGrouping:grouping sorting:sorting versionTag:@"1" options:options];
 }
 
 + (void)registerLegacyClasses

SignalServiceKit/src/Storage/OWSIncomingMessageFinder.m

@@ -89,10 +89,7 @@ NSString *const OWSIncomingMessageFinderColumnSourceDeviceId = @"OWSIncomingMess
 
     YapDatabaseSecondaryIndexHandler *handler = [YapDatabaseSecondaryIndexHandler withObjectBlock:block];
 
-    return [[YapDatabaseSecondaryIndex alloc]
-        initWithSetup:setup
-              handler:handler
-           versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:nil]];
+    return [[YapDatabaseSecondaryIndex alloc] initWithSetup:setup handler:handler versionTag:nil];
 }
 
 + (NSString *)databaseExtensionName

SignalServiceKit/src/Storage/OWSMediaGalleryFinder.m

@@ -164,11 +164,7 @@ static NSString *const OWSMediaGalleryFinderExtensionName = @"OWSMediaGalleryFin
     YapDatabaseViewOptions *options = [YapDatabaseViewOptions new];
     options.allowedCollections = [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:TSMessage.collection]];
 
-    return [[YapDatabaseAutoView alloc]
-        initWithGrouping:grouping
-                 sorting:sorting
-              versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:@"1"]
-                 options:options];
+    return [[YapDatabaseAutoView alloc] initWithGrouping:grouping sorting:sorting versionTag:@"1" options:options];
 }
 
 + (BOOL)attachmentIdShouldAppearInMediaGallery:(NSString *)attachmentId transaction:(YapDatabaseReadTransaction *)transaction

SignalServiceKit/src/Storage/OWSPrimaryStorage.h

@@ -6,6 +6,8 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+NSArray<NSString *> *ExtensionNamesForPrimaryStorage();
+
 @interface OWSPrimaryStorage : OWSStorage
 
 - (instancetype)init NS_UNAVAILABLE;

SignalServiceKit/src/Storage/OWSPrimaryStorage.m

@@ -67,14 +67,12 @@ extern NSString *const TSThreadOutgoingMessageDatabaseViewExtensionName;
 extern NSString *const TSUnseenDatabaseViewExtensionName;
 extern NSString *const TSThreadSpecialMessagesDatabaseViewExtensionName;
 
-void VerifyRegistrationsForStorage(OWSStorage *storage)
+NSArray<NSString *> *ExtensionNamesForPrimaryStorage()
 {
-    OWSCAssert(storage);
-
     // This should 1:1 correspond to the database view registrations
     // done in RunSyncRegistrationsForStorage() and
     // RunAsyncRegistrationsForStorage().
-    NSArray<NSString *> *databaseViewNames = @[
+    return @[
         // We don't need to verify the cross process notifier.
         // [TSDatabaseView registerCrossProcessNotifier:storage];
 
@@ -128,20 +126,23 @@ void VerifyRegistrationsForStorage(OWSStorage *storage)
         // [TSDatabaseView asyncRegisterLazyRestoreAttachmentsDatabaseView:storage completion:completion];
         TSLazyRestoreAttachmentsDatabaseViewExtensionName,
     ];
+}
+
+void VerifyRegistrationsForPrimaryStorage(OWSStorage *storage)
+{
+    OWSCAssert(storage);
 
-    __block BOOL hasMissingDatabaseView = NO;
     [[storage newDatabaseConnection] asyncReadWithBlock:^(YapDatabaseReadTransaction *transaction) {
-        for (NSString *databaseViewName in databaseViewNames) {
-            YapDatabaseViewTransaction *_Nullable viewTransaction = [transaction ext:databaseViewName];
+        for (NSString *extensionName in ExtensionNamesForPrimaryStorage()) {
+            YapDatabaseViewTransaction *_Nullable viewTransaction = [transaction ext:extensionName];
             if (!viewTransaction) {
-                OWSProdLogAndCFail(@"VerifyRegistrationsForStorage missing database view: %@", databaseViewName);
-                hasMissingDatabaseView = YES;
+                OWSProdLogAndCFail(
+                    @"VerifyRegistrationsForPrimaryStorage missing database extension: %@", extensionName);
+
+                [OWSStorage incrementVersionOfDatabaseExtension:extensionName];
             }
         }
     }];
-    if (hasMissingDatabaseView) {
-        [OWSStorage incrementDatabaseExtensionVersionSuffix];
-    }
 }
 
 #pragma mark -
@@ -235,7 +236,7 @@ void VerifyRegistrationsForStorage(OWSStorage *storage)
 
 - (void)verifyDatabaseViews
 {
-    VerifyRegistrationsForStorage(self);
+    VerifyRegistrationsForPrimaryStorage(self);
 }
 
 + (void)protectFiles

SignalServiceKit/src/Storage/OWSStorage.h

@@ -66,8 +66,7 @@ typedef void (^OWSStorageMigrationBlock)(void);
 
 #pragma mark - Extension Registration
 
-+ (void)incrementDatabaseExtensionVersionSuffix;
-+ (nullable NSString *)appendSuffixToDatabaseExtensionVersionIfNecessary:(nullable NSString *)versionTag;
++ (void)incrementVersionOfDatabaseExtension:(NSString *)extensionName;
 
 - (BOOL)registerExtension:(YapDatabaseExtension *)extension withName:(NSString *)extensionName;
 

SignalServiceKit/src/Storage/OWSStorage.m

@@ -15,9 +15,11 @@
 #import <Curve25519Kit/Randomness.h>
 #import <SAMKeychain/SAMKeychain.h>
 #import <YapDatabase/YapDatabase.h>
+#import <YapDatabase/YapDatabaseAutoView.h>
+#import <YapDatabase/YapDatabaseCrossProcessNotification.h>
 #import <YapDatabase/YapDatabaseCryptoUtils.h>
 #import <YapDatabase/YapDatabaseSecondaryIndex.h>
-#import <YapDatabase/YapDatabaseView.h>
+#import <YapDatabase/YapDatabaseSecondaryIndexPrivate.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -39,7 +41,7 @@ const NSUInteger kDatabasePasswordLength = 30;
 typedef NSData *_Nullable (^LoadDatabaseMetadataBlock)(NSError **_Nullable);
 typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
 
-NSString *const kNSUserDefaults_DatabaseExtensionVersionSuffix = @"kNSUserDefaults_DatabaseExtensionVersionSuffix";
+NSString *const kNSUserDefaults_DatabaseExtensionVersionMap = @"kNSUserDefaults_DatabaseExtensionVersionMap";
 
 #pragma mark -
 
@@ -468,46 +470,87 @@ NSString *const kNSUserDefaults_DatabaseExtensionVersionSuffix = @"kNSUserDefaul
 
 #pragma mark - Extension Registration
 
-+ (void)incrementDatabaseExtensionVersionSuffix
++ (void)incrementVersionOfDatabaseExtension:(NSString *)extensionName
 {
     DDLogError(@"%@ %s", self.logTag, __PRETTY_FUNCTION__);
 
     NSUserDefaults *appUserDefaults = [NSUserDefaults appUserDefaults];
     OWSAssert(appUserDefaults);
-    NSNumber *_Nullable suffix = [appUserDefaults valueForKey:kNSUserDefaults_DatabaseExtensionVersionSuffix];
-    [appUserDefaults setValue:@(suffix.intValue + 1) forKey:kNSUserDefaults_DatabaseExtensionVersionSuffix];
+    NSMutableDictionary<NSString *, NSNumber *> *_Nullable versionMap =
+        [[appUserDefaults valueForKey:kNSUserDefaults_DatabaseExtensionVersionMap] mutableCopy];
+    if (!versionMap) {
+        versionMap = [NSMutableDictionary new];
+    }
+    NSNumber *_Nullable versionSuffix = versionMap[extensionName];
+    versionMap[extensionName] = @(versionSuffix.intValue + 1);
+    [appUserDefaults setValue:versionMap forKey:kNSUserDefaults_DatabaseExtensionVersionMap];
     [appUserDefaults synchronize];
 }
 
-+ (nullable NSString *)databaseExtensionVersionSuffix
+- (nullable NSString *)appendSuffixToDatabaseExtensionVersionIfNecessary:(nullable NSString *)versionTag
+                                                           extensionName:(NSString *)extensionName
 {
+    OWSAssertIsOnMainThread();
+
     NSUserDefaults *appUserDefaults = [NSUserDefaults appUserDefaults];
     OWSAssert(appUserDefaults);
-    NSNumber *_Nullable suffix = [appUserDefaults valueForKey:kNSUserDefaults_DatabaseExtensionVersionSuffix];
-    if (!suffix) {
-        return nil;
-    }
-    return [NSString stringWithFormat:@".%@", suffix];
-}
-
-+ (nullable NSString *)appendSuffixToDatabaseExtensionVersionIfNecessary:(nullable NSString *)versionTag
-{
-    OWSAssertIsOnMainThread();
+    NSDictionary<NSString *, NSNumber *> *_Nullable versionMap =
+        [appUserDefaults valueForKey:kNSUserDefaults_DatabaseExtensionVersionMap];
+    NSNumber *_Nullable versionSuffix = versionMap[extensionName];
 
-    NSString *_Nullable suffix = [self databaseExtensionVersionSuffix];
-    if (suffix) {
+    if (versionSuffix) {
         if (!versionTag) {
             versionTag = @"0";
         }
-        NSString *result = [versionTag stringByAppendingString:suffix];
-        DDLogWarn(@"%@ database extension version: %@ + %@ -> %@", self.logTag, versionTag, suffix, result);
+        NSString *result = [NSString stringWithFormat:@"%@.%@", versionTag, versionSuffix];
+        DDLogWarn(@"%@ database extension version: %@ + %@ -> %@", self.logTag, versionTag, versionSuffix, result);
         return result;
     }
     return versionTag;
 }
 
+- (YapDatabaseExtension *)updateExtensionVersion:(YapDatabaseExtension *)extension withName:(NSString *)extensionName
+{
+    OWSAssert(extension);
+    OWSAssert(extensionName.length > 0);
+
+    if ([extension isKindOfClass:[YapDatabaseAutoView class]]) {
+        YapDatabaseAutoView *databaseView = (YapDatabaseAutoView *)extension;
+        YapDatabaseAutoView *databaseViewCopy = [[YapDatabaseAutoView alloc]
+            initWithGrouping:databaseView.grouping
+                     sorting:databaseView.sorting
+                  versionTag:[self appendSuffixToDatabaseExtensionVersionIfNecessary:databaseView.versionTag
+                                                                       extensionName:extensionName]
+                     options:databaseView.options];
+        return databaseViewCopy;
+    } else if ([extension isKindOfClass:[YapDatabaseSecondaryIndex class]]) {
+        YapDatabaseSecondaryIndex *secondaryIndex = (YapDatabaseSecondaryIndex *)extension;
+        OWSAssert(secondaryIndex->setup);
+        OWSAssert(secondaryIndex->handler);
+        YapDatabaseSecondaryIndex *secondaryIndexCopy = [[YapDatabaseSecondaryIndex alloc]
+            initWithSetup:secondaryIndex->setup
+                  handler:secondaryIndex->handler
+               versionTag:[self appendSuffixToDatabaseExtensionVersionIfNecessary:secondaryIndex.versionTag
+                                                                    extensionName:extensionName]
+                  options:secondaryIndex->options];
+        return secondaryIndexCopy;
+    } else if ([extension isKindOfClass:[YapDatabaseCrossProcessNotification class]]) {
+        // versionTag doesn't matter for YapDatabaseCrossProcessNotification.
+        return extension;
+    } else {
+        // This method needs to be able to update the versionTag of all extensions.
+        // If we start using other extension types, we need to modify this method to
+        // handle them as well.
+        OWSProdLogAndFail(@"%@ Unknown extension type: %@", self.logTag, [extension class]);
+
+        return extension;
+    }
+}
+
 - (BOOL)registerExtension:(YapDatabaseExtension *)extension withName:(NSString *)extensionName
 {
+    extension = [self updateExtensionVersion:extension withName:extensionName];
+
     return [self.database registerExtension:extension withName:extensionName];
 }
 
@@ -521,23 +564,7 @@ NSString *const kNSUserDefaults_DatabaseExtensionVersionSuffix = @"kNSUserDefaul
                       withName:(NSString *)extensionName
                     completion:(nullable dispatch_block_t)completion
 {
-    if ([extension isKindOfClass:[YapDatabaseView class]]) {
-        YapDatabaseView *databaseView = (YapDatabaseView *)extension;
-
-        NSString *_Nullable databaseExtensionVersionSuffix = [OWSStorage databaseExtensionVersionSuffix];
-        if (databaseExtensionVersionSuffix) {
-            OWSAssert([databaseView.versionTag hasSuffix:databaseExtensionVersionSuffix]);
-        }
-    } else if ([extension isKindOfClass:[YapDatabaseSecondaryIndex class]]) {
-        YapDatabaseSecondaryIndex *secondaryIndex = (YapDatabaseSecondaryIndex *)extension;
-
-        NSString *_Nullable databaseExtensionVersionSuffix = [OWSStorage databaseExtensionVersionSuffix];
-        if (databaseExtensionVersionSuffix) {
-            OWSAssert([secondaryIndex.versionTag hasSuffix:databaseExtensionVersionSuffix]);
-        }
-    } else {
-        OWSProdLogAndFail(@"%@ Unknown extension type: %@", self.logTag, [extension class]);
-    }
+    extension = [self updateExtensionVersion:extension withName:extensionName];
 
     [self.database asyncRegisterExtension:extension
                                  withName:extensionName

SignalServiceKit/src/Storage/TSDatabaseSecondaryIndexes.m

@@ -26,10 +26,8 @@
 
     YapDatabaseSecondaryIndexHandler *handler = [YapDatabaseSecondaryIndexHandler withObjectBlock:block];
 
-    YapDatabaseSecondaryIndex *secondaryIndex = [[YapDatabaseSecondaryIndex alloc]
-        initWithSetup:setup
-              handler:handler
-           versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:nil]];
+    YapDatabaseSecondaryIndex *secondaryIndex =
+        [[YapDatabaseSecondaryIndex alloc] initWithSetup:setup handler:handler versionTag:nil];
 
     return secondaryIndex;
 }

SignalServiceKit/src/Storage/TSDatabaseView.m

@@ -78,11 +78,10 @@ NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"
     options.allowedCollections =
         [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:[TSInteraction collection]]];
 
-    YapDatabaseView *view = [[YapDatabaseAutoView alloc]
-        initWithGrouping:viewGrouping
-                 sorting:viewSorting
-              versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:version]
-                 options:options];
+    YapDatabaseView *view = [[YapDatabaseAutoView alloc] initWithGrouping:viewGrouping
+                                                                  sorting:viewSorting
+                                                               versionTag:version
+                                                                  options:options];
     [storage asyncRegisterExtension:view withName:viewName];
 }
 
@@ -228,11 +227,8 @@ NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"
     options.allowedCollections =
         [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:[TSThread collection]]];
 
-    YapDatabaseView *databaseView = [[YapDatabaseAutoView alloc]
-        initWithGrouping:viewGrouping
-                 sorting:viewSorting
-              versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:@"3"]
-                 options:options];
+    YapDatabaseView *databaseView =
+        [[YapDatabaseAutoView alloc] initWithGrouping:viewGrouping sorting:viewSorting versionTag:@"3" options:options];
 
     [storage asyncRegisterExtension:databaseView withName:TSThreadDatabaseViewExtensionName];
 }
@@ -340,11 +336,8 @@ NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"
     NSSet *deviceCollection = [NSSet setWithObject:[OWSDevice collection]];
     options.allowedCollections = [[YapWhitelistBlacklist alloc] initWithWhitelist:deviceCollection];
 
-    YapDatabaseView *view = [[YapDatabaseAutoView alloc]
-        initWithGrouping:viewGrouping
-                 sorting:viewSorting
-              versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:@"3"]
-                 options:options];
+    YapDatabaseView *view =
+        [[YapDatabaseAutoView alloc] initWithGrouping:viewGrouping sorting:viewSorting versionTag:@"3" options:options];
 
     [storage asyncRegisterExtension:view withName:TSSecondaryDevicesDatabaseViewExtensionName];
 }
@@ -397,11 +390,8 @@ NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"
     options.isPersistent = YES;
     options.allowedCollections =
         [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:[TSAttachment collection]]];
-    YapDatabaseView *view = [[YapDatabaseAutoView alloc]
-        initWithGrouping:viewGrouping
-                 sorting:viewSorting
-              versionTag:[OWSStorage appendSuffixToDatabaseExtensionVersionIfNecessary:@"3"]
-                 options:options];
+    YapDatabaseView *view =
+        [[YapDatabaseAutoView alloc] initWithGrouping:viewGrouping sorting:viewSorting versionTag:@"3" options:options];
     [storage asyncRegisterExtension:view
                            withName:TSLazyRestoreAttachmentsDatabaseViewExtensionName
                          completion:completion];