Merge branch 'charlesmchen/backupMoreCollections'

Signal/src/util/Backup/OWSBackup.h

@@ -24,6 +24,8 @@ typedef NS_ENUM(NSUInteger, OWSBackupState) {
 NSString *NSStringForBackupExportState(OWSBackupState state);
 NSString *NSStringForBackupImportState(OWSBackupState state);
 
+NSArray<NSString *> *MiscCollectionsToBackup(void);
+
 @class OWSBackupIO;
 @class TSAttachmentPointer;
 @class TSThread;

Signal/src/util/Backup/OWSBackup.m

@@ -8,6 +8,7 @@
 #import "OWSBackupImportJob.h"
 #import "Signal-Swift.h"
 #import <SignalCoreKit/Randomness.h>
+#import <SignalServiceKit/OWSIdentityManager.h>
 #import <SignalServiceKit/YapDatabaseConnection+OWS.h>
 
 NS_ASSUME_NONNULL_BEGIN
@@ -51,6 +52,16 @@ NSString *NSStringForBackupImportState(OWSBackupState state)
     }
 }
 
+NSArray<NSString *> *MiscCollectionsToBackup(void)
+{
+    return @[
+             kOWSBlockingManager_BlockListCollection,
+             OWSUserProfile.collection,
+             SSKIncrementingIdFinder.collectionName,
+             OWSPreferencesSignalDatabaseCollection,
+    ];
+}
+
 // TODO: Observe Reachability.
 @interface OWSBackup () <OWSBackupJobDelegate>
 

Signal/src/util/Backup/OWSBackupExportJob.m

@@ -111,9 +111,14 @@ NS_ASSUME_NONNULL_BEGIN
 // It isn't strictly necessary to capture the entity type (the importer doesn't
 // use this state), but I think it'll be helpful to have around to future-proof
 // this work, help with debugging issue, etc.
-- (BOOL)writeObject:(TSYapDatabaseObject *)object entityType:(SignalIOSProtoBackupSnapshotBackupEntityType)entityType
+- (BOOL)writeObject:(NSObject *)object
+         collection:(NSString *)collection
+                key:(NSString *)key
+         entityType:(SignalIOSProtoBackupSnapshotBackupEntityType)entityType
 {
     OWSAssertDebug(object);
+    OWSAssertDebug(collection.length > 0);
+    OWSAssertDebug(key.length > 0);
 
     NSData *_Nullable data = [NSKeyedArchiver archivedDataWithRootObject:object];
     if (!data) {
@@ -126,7 +131,10 @@ NS_ASSUME_NONNULL_BEGIN
     }
 
     SignalIOSProtoBackupSnapshotBackupEntityBuilder *entityBuilder =
-        [SignalIOSProtoBackupSnapshotBackupEntity builderWithType:entityType entityData:data];
+        [SignalIOSProtoBackupSnapshotBackupEntity builderWithType:entityType
+                                                       entityData:data
+                                                       collection:collection
+                                                              key:key];
 
     NSError *error;
     SignalIOSProtoBackupSnapshotBackupEntity *_Nullable entity = [entityBuilder buildAndReturnError:&error];
@@ -496,11 +504,14 @@ NS_ASSUME_NONNULL_BEGIN
         Class,
         EntityFilter _Nullable,
         SignalIOSProtoBackupSnapshotBackupEntityType);
+    NSMutableSet<NSString *> *exportedCollections = [NSMutableSet new];
     ExportBlock exportEntities = ^(YapDatabaseReadTransaction *transaction,
         NSString *collection,
         Class expectedClass,
         EntityFilter _Nullable filter,
         SignalIOSProtoBackupSnapshotBackupEntityType entityType) {
+        [exportedCollections addObject:collection];
+
         __block NSUInteger count = 0;
         [transaction enumerateKeysAndObjectsInCollection:collection
                                               usingBlock:^(NSString *key, id object, BOOL *stop) {
@@ -515,7 +526,7 @@ NS_ASSUME_NONNULL_BEGIN
                                                       OWSFailDebug(@"unexpected class: %@", [object class]);
                                                       return;
                                                   }
-                                                  TSYapDatabaseObject *entity = object;
+                                                  NSObject *entity = object;
                                                   count++;
 
                                                   if ([entity isKindOfClass:[TSAttachmentStream class]]) {
@@ -529,7 +540,10 @@ NS_ASSUME_NONNULL_BEGIN
                                                       entity = attachmentPointer;
                                                   }
 
-                                                  if (![exportStream writeObject:entity entityType:entityType]) {
+                                                  if (![exportStream writeObject:entity
+                                                                      collection:collection
+                                                                             key:key
+                                                                      entityType:entityType]) {
                                                       *stop = YES;
                                                       aborted = YES;
                                                       return;
@@ -542,6 +556,7 @@ NS_ASSUME_NONNULL_BEGIN
     __block NSUInteger copiedInteractions = 0;
     __block NSUInteger copiedAttachments = 0;
     __block NSUInteger copiedMigrations = 0;
+    __block NSUInteger copiedMisc = 0;
     self.unsavedAttachmentExports = [NSMutableArray new];
     [dbConnection readWithBlock:^(YapDatabaseReadTransaction *transaction) {
         copiedThreads = exportEntities(transaction,
@@ -615,6 +630,30 @@ NS_ASSUME_NONNULL_BEGIN
             [OWSDatabaseMigration class],
             nil,
             SignalIOSProtoBackupSnapshotBackupEntityTypeMigration);
+        if (aborted) {
+            return;
+        }
+
+        for (NSString *collection in MiscCollectionsToBackup()) {
+            copiedMisc += exportEntities(
+                transaction, collection, [NSObject class], nil, SignalIOSProtoBackupSnapshotBackupEntityTypeMisc);
+            if (aborted) {
+                return;
+            }
+        }
+
+        for (NSString *collection in [exportedCollections.allObjects sortedArrayUsingSelector:@selector(compare:)]) {
+            OWSLogVerbose(@"Exported collection: %@", collection);
+        }
+        OWSLogVerbose(@"Exported collections: %lu", (unsigned long)exportedCollections.count);
+
+        NSSet<NSString *> *allCollections = [NSSet setWithArray:transaction.allCollections];
+        NSMutableSet *unexportedCollections = [allCollections mutableCopy];
+        [unexportedCollections minusSet:exportedCollections];
+        for (NSString *collection in [unexportedCollections.allObjects sortedArrayUsingSelector:@selector(compare:)]) {
+            OWSLogVerbose(@"Unexported collection: %@", collection);
+        }
+        OWSLogVerbose(@"Unexported collections: %lu", (unsigned long)unexportedCollections.count);
     }];
 
     if (aborted || self.isComplete) {
@@ -636,6 +675,7 @@ NS_ASSUME_NONNULL_BEGIN
     OWSLogInfo(@"copiedMessages: %zd", copiedInteractions);
     OWSLogInfo(@"copiedAttachments: %zd", copiedAttachments);
     OWSLogInfo(@"copiedMigrations: %zd", copiedMigrations);
+    OWSLogInfo(@"copiedMisc: %zd", copiedMisc);
     OWSLogInfo(@"copiedEntities: %zd", exportStream.totalItemCount);
 
     return YES;

Signal/src/util/Backup/OWSBackupImportJob.m

@@ -44,6 +44,18 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
     return SSKEnvironment.shared.primaryStorage;
 }
 
+- (OWSProfileManager *)profileManager
+{
+    return [OWSProfileManager sharedManager];
+}
+
+- (TSAccountManager *)tsAccountManager
+{
+    OWSAssertDebug(SSKEnvironment.shared.tsAccountManager);
+
+    return SSKEnvironment.shared.tsAccountManager;
+}
+
 #pragma mark -
 
 - (void)startAsync
@@ -173,6 +185,10 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
                                     return;
                                 }
 
+                                [weakSelf.profileManager fetchLocalUsersProfile];
+
+                                [weakSelf.tsAccountManager updateAccountAttributes];
+
                                 // Kick off lazy restore.
                                 [OWSBackupLazyRestoreJob runAsync];
 
@@ -432,7 +448,23 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
                         return completion(NO);
                     }
 
-                    __block TSYapDatabaseObject *object = nil;
+                    NSString *_Nullable collection = entity.collection;
+                    if (collection.length < 1) {
+                        OWSLogError(@"missing collection.");
+                        // Database-related errors are unrecoverable.
+                        aborted = YES;
+                        return completion(NO);
+                    }
+
+                    NSString *_Nullable key = entity.key;
+                    if (key.length < 1) {
+                        OWSLogError(@"missing key.");
+                        // Database-related errors are unrecoverable.
+                        aborted = YES;
+                        return completion(NO);
+                    }
+
+                    __block NSObject *object = nil;
                     @try {
                         NSKeyedUnarchiver *unarchiver = [[NSKeyedUnarchiver alloc] initForReadingWithData:entityData];
                         object = [unarchiver decodeObjectForKey:@"root"];
@@ -449,9 +481,8 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
                         return completion(NO);
                     }
 
-                    [object saveWithTransaction:transaction];
+                    [transaction setObject:object forKey:key inCollection:collection];
                     copiedEntities++;
-                    NSString *collection = [object.class collection];
                     NSUInteger restoredEntityCount = restoredEntityCounts[collection].unsignedIntValue;
                     restoredEntityCounts[collection] = @(restoredEntityCount + 1);
                 }

SignalMessaging/profiles/OWSProfileManager.h

@@ -51,6 +51,8 @@ extern const NSUInteger kOWSProfileManager_MaxAvatarDiameter;
 
 - (BOOL)isProfileNameTooLong:(nullable NSString *)profileName;
 
+- (void)fetchLocalUsersProfile;
+
 #pragma mark - Profile Whitelist
 
 // These methods are for debugging.

SignalServiceKit/protobuf/SignalIOS.proto

@@ -14,16 +14,21 @@ package IOSProtos;
 message BackupSnapshot {
     message BackupEntity {
         enum Type {
-            UNKNOWN = 0;
+            UNKNOWN     = 0;
-            MIGRATION = 1;
+            MIGRATION   = 1;
-            THREAD = 2;
+            THREAD      = 2;
             INTERACTION = 3;
-            ATTACHMENT = 4;
+            ATTACHMENT  = 4;
+            MISC        = 5;
         }
         // @required
-        optional Type type = 1;
+        optional Type type         = 1;
         // @required
-        optional bytes entityData = 2;
+        optional bytes entityData  = 2;
+        // @required
+        optional string collection = 3;
+        // @required
+        optional string key        = 4;
     }
 
     repeated BackupEntity entity = 1;

SignalServiceKit/src/Messages/OWSBlockingManager.h

@@ -10,6 +10,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 extern NSString *const kNSNotificationName_BlockListDidChange;
 
+extern NSString *const kOWSBlockingManager_BlockListCollection;
+
 // This class can be safely accessed and used from any thread.
 @interface OWSBlockingManager : NSObject
 

SignalServiceKit/src/Messages/OWSIdentityManager.h

@@ -7,6 +7,8 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+extern NSString *const OWSPrimaryStorageIdentityKeyStoreCollection;
+
 extern NSString *const OWSPrimaryStorageTrustedKeysCollection;
 
 // This notification will be fired whenever identities are created

SignalServiceKit/src/Protos/Generated/SignalIOS.pb.swift

@@ -61,6 +61,26 @@ struct IOSProtos_BackupSnapshot {
     /// Clears the value of `entityData`. Subsequent reads from it will return its default value.
     mutating func clearEntityData() {self._entityData = nil}
 
+    /// @required
+    var collection: String {
+      get {return _collection ?? String()}
+      set {_collection = newValue}
+    }
+    /// Returns true if `collection` has been explicitly set.
+    var hasCollection: Bool {return self._collection != nil}
+    /// Clears the value of `collection`. Subsequent reads from it will return its default value.
+    mutating func clearCollection() {self._collection = nil}
+
+    /// @required
+    var key: String {
+      get {return _key ?? String()}
+      set {_key = newValue}
+    }
+    /// Returns true if `key` has been explicitly set.
+    var hasKey: Bool {return self._key != nil}
+    /// Clears the value of `key`. Subsequent reads from it will return its default value.
+    mutating func clearKey() {self._key = nil}
+
     var unknownFields = SwiftProtobuf.UnknownStorage()
 
     enum TypeEnum: SwiftProtobuf.Enum {
@@ -70,6 +90,7 @@ struct IOSProtos_BackupSnapshot {
       case thread // = 2
       case interaction // = 3
       case attachment // = 4
+      case misc // = 5
 
       init() {
         self = .unknown
@@ -82,6 +103,7 @@ struct IOSProtos_BackupSnapshot {
         case 2: self = .thread
         case 3: self = .interaction
         case 4: self = .attachment
+        case 5: self = .misc
         default: return nil
         }
       }
@@ -93,6 +115,7 @@ struct IOSProtos_BackupSnapshot {
         case .thread: return 2
         case .interaction: return 3
         case .attachment: return 4
+        case .misc: return 5
         }
       }
 
@@ -102,6 +125,8 @@ struct IOSProtos_BackupSnapshot {
 
     fileprivate var _type: IOSProtos_BackupSnapshot.BackupEntity.TypeEnum? = nil
     fileprivate var _entityData: Data? = nil
+    fileprivate var _collection: String? = nil
+    fileprivate var _key: String? = nil
   }
 
   init() {}
@@ -145,6 +170,8 @@ extension IOSProtos_BackupSnapshot.BackupEntity: SwiftProtobuf.Message, SwiftPro
   static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
     1: .same(proto: "type"),
     2: .same(proto: "entityData"),
+    3: .same(proto: "collection"),
+    4: .same(proto: "key"),
   ]
 
   mutating func decodeMessage<D: SwiftProtobuf.Decoder>(decoder: inout D) throws {
@@ -152,6 +179,8 @@ extension IOSProtos_BackupSnapshot.BackupEntity: SwiftProtobuf.Message, SwiftPro
       switch fieldNumber {
       case 1: try decoder.decodeSingularEnumField(value: &self._type)
       case 2: try decoder.decodeSingularBytesField(value: &self._entityData)
+      case 3: try decoder.decodeSingularStringField(value: &self._collection)
+      case 4: try decoder.decodeSingularStringField(value: &self._key)
       default: break
       }
     }
@@ -164,12 +193,20 @@ extension IOSProtos_BackupSnapshot.BackupEntity: SwiftProtobuf.Message, SwiftPro
     if let v = self._entityData {
       try visitor.visitSingularBytesField(value: v, fieldNumber: 2)
     }
+    if let v = self._collection {
+      try visitor.visitSingularStringField(value: v, fieldNumber: 3)
+    }
+    if let v = self._key {
+      try visitor.visitSingularStringField(value: v, fieldNumber: 4)
+    }
     try unknownFields.traverse(visitor: &visitor)
   }
 
   static func ==(lhs: IOSProtos_BackupSnapshot.BackupEntity, rhs: IOSProtos_BackupSnapshot.BackupEntity) -> Bool {
     if lhs._type != rhs._type {return false}
     if lhs._entityData != rhs._entityData {return false}
+    if lhs._collection != rhs._collection {return false}
+    if lhs._key != rhs._key {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
   }
@@ -182,5 +219,6 @@ extension IOSProtos_BackupSnapshot.BackupEntity.TypeEnum: SwiftProtobuf._ProtoNa
     2: .same(proto: "THREAD"),
     3: .same(proto: "INTERACTION"),
     4: .same(proto: "ATTACHMENT"),
+    5: .same(proto: "MISC"),
   ]
 }

SignalServiceKit/src/Protos/Generated/SignalIOSProto.swift

@@ -22,6 +22,7 @@ public enum SignalIOSProtoError: Error {
         case thread = 2
         case interaction = 3
         case attachment = 4
+        case misc = 5
     }
 
     private class func SignalIOSProtoBackupSnapshotBackupEntityTypeWrap(_ value: IOSProtos_BackupSnapshot.BackupEntity.TypeEnum) -> SignalIOSProtoBackupSnapshotBackupEntityType {
@@ -31,6 +32,7 @@ public enum SignalIOSProtoError: Error {
         case .thread: return .thread
         case .interaction: return .interaction
         case .attachment: return .attachment
+        case .misc: return .misc
         }
     }
 
@@ -41,18 +43,19 @@ public enum SignalIOSProtoError: Error {
         case .thread: return .thread
         case .interaction: return .interaction
         case .attachment: return .attachment
+        case .misc: return .misc
         }
     }
 
     // MARK: - SignalIOSProtoBackupSnapshotBackupEntityBuilder
 
-    @objc public class func builder(type: SignalIOSProtoBackupSnapshotBackupEntityType, entityData: Data) -> SignalIOSProtoBackupSnapshotBackupEntityBuilder {
+    @objc public class func builder(type: SignalIOSProtoBackupSnapshotBackupEntityType, entityData: Data, collection: String, key: String) -> SignalIOSProtoBackupSnapshotBackupEntityBuilder {
-        return SignalIOSProtoBackupSnapshotBackupEntityBuilder(type: type, entityData: entityData)
+        return SignalIOSProtoBackupSnapshotBackupEntityBuilder(type: type, entityData: entityData, collection: collection, key: key)
     }
 
     // asBuilder() constructs a builder that reflects the proto's contents.
     @objc public func asBuilder() -> SignalIOSProtoBackupSnapshotBackupEntityBuilder {
-        let builder = SignalIOSProtoBackupSnapshotBackupEntityBuilder(type: type, entityData: entityData)
+        let builder = SignalIOSProtoBackupSnapshotBackupEntityBuilder(type: type, entityData: entityData, collection: collection, key: key)
         return builder
     }
 
@@ -62,11 +65,13 @@ public enum SignalIOSProtoError: Error {
 
         @objc fileprivate override init() {}
 
-        @objc fileprivate init(type: SignalIOSProtoBackupSnapshotBackupEntityType, entityData: Data) {
+        @objc fileprivate init(type: SignalIOSProtoBackupSnapshotBackupEntityType, entityData: Data, collection: String, key: String) {
             super.init()
 
             setType(type)
             setEntityData(entityData)
+            setCollection(collection)
+            setKey(key)
         }
 
         @objc public func setType(_ valueParam: SignalIOSProtoBackupSnapshotBackupEntityType) {
@@ -77,6 +82,14 @@ public enum SignalIOSProtoError: Error {
             proto.entityData = valueParam
         }
 
+        @objc public func setCollection(_ valueParam: String) {
+            proto.collection = valueParam
+        }
+
+        @objc public func setKey(_ valueParam: String) {
+            proto.key = valueParam
+        }
+
         @objc public func build() throws -> SignalIOSProtoBackupSnapshotBackupEntity {
             return try SignalIOSProtoBackupSnapshotBackupEntity.parseProto(proto)
         }
@@ -92,12 +105,20 @@ public enum SignalIOSProtoError: Error {
 
     @objc public let entityData: Data
 
+    @objc public let collection: String
+
+    @objc public let key: String
+
     private init(proto: IOSProtos_BackupSnapshot.BackupEntity,
                  type: SignalIOSProtoBackupSnapshotBackupEntityType,
-                 entityData: Data) {
+                 entityData: Data,
+                 collection: String,
+                 key: String) {
         self.proto = proto
         self.type = type
         self.entityData = entityData
+        self.collection = collection
+        self.key = key
     }
 
     @objc
@@ -121,13 +142,25 @@ public enum SignalIOSProtoError: Error {
         }
         let entityData = proto.entityData
 
+        guard proto.hasCollection else {
+            throw SignalIOSProtoError.invalidProtobuf(description: "\(logTag) missing required field: collection")
+        }
+        let collection = proto.collection
+
+        guard proto.hasKey else {
+            throw SignalIOSProtoError.invalidProtobuf(description: "\(logTag) missing required field: key")
+        }
+        let key = proto.key
+
         // MARK: - Begin Validation Logic for SignalIOSProtoBackupSnapshotBackupEntity -
 
         // MARK: - End Validation Logic for SignalIOSProtoBackupSnapshotBackupEntity -
 
         let result = SignalIOSProtoBackupSnapshotBackupEntity(proto: proto,
                                                               type: type,
-                                                              entityData: entityData)
+                                                              entityData: entityData,
+                                                              collection: collection,
+                                                              key: key)
         return result
     }
 

SignalServiceKit/src/Storage/SSKIncrementingIdFinder.swift

@@ -7,7 +7,8 @@ import Foundation
 @objc
 public class SSKIncrementingIdFinder: NSObject {
 
-    private static let collectionName = "IncrementingIdCollection"
+    @objc
+    public static let collectionName = "IncrementingIdCollection"
 
     @objc
     public class func previousId(key: String, transaction: YapDatabaseReadTransaction) -> UInt64 {