From 061ce8cb13205dfe355f84da0f217112d7fe22c7 Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Tue, 13 Mar 2018 12:30:38 -0300
Subject: [PATCH] Add database validity check.

---
 Signal/src/util/OWSBackupExportJob.m | 7 +++++++
 Signal/src/util/OWSBackupImportJob.m | 8 ++++++++
 Signal/src/util/OWSBackupJob.h       | 3 +++
 Signal/src/util/OWSBackupJob.m       | 3 +++
 4 files changed, 21 insertions(+)

diff --git a/Signal/src/util/OWSBackupExportJob.m b/Signal/src/util/OWSBackupExportJob.m
index f49039f32..4dd855c27 100644
--- a/Signal/src/util/OWSBackupExportJob.m
+++ b/Signal/src/util/OWSBackupExportJob.m
@@ -248,6 +248,10 @@ NSString *const kOWSBackup_ExportDatabaseKeySpec = @"kOWSBackup_ExportDatabaseKe
 
     [primaryDBConnection readWithBlock:^(YapDatabaseReadTransaction *srcTransaction) {
         [tempDBConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *dstTransaction) {
+            [dstTransaction setObject:@(YES)
+                               forKey:kOWSBackup_Snapshot_ValidKey
+                         inCollection:kOWSBackup_Snapshot_Collection];
+
             // Copy threads.
             [srcTransaction
                 enumerateKeysAndObjectsInCollection:[TSThread collection]
@@ -347,6 +351,9 @@ NSString *const kOWSBackup_ExportDatabaseKeySpec = @"kOWSBackup_ExportDatabaseKe
         }];
     }];
 
+    if (self.isComplete) {
+        return NO;
+    }
     // TODO: Should we do a database checkpoint?
 
     DDLogInfo(@"%@ copiedThreads: %llu", self.logTag, copiedThreads);
diff --git a/Signal/src/util/OWSBackupImportJob.m b/Signal/src/util/OWSBackupImportJob.m
index d187d318e..8d8c8c63a 100644
--- a/Signal/src/util/OWSBackupImportJob.m
+++ b/Signal/src/util/OWSBackupImportJob.m
@@ -408,6 +408,14 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
     __block BOOL aborted = NO;
     [tempDBConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *srcTransaction) {
         [primaryDBConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *dstTransaction) {
+            if (![srcTransaction boolForKey:kOWSBackup_Snapshot_ValidKey
+                               inCollection:kOWSBackup_Snapshot_Collection
+                               defaultValue:NO]) {
+                DDLogError(@"%@ invalid database.", self.logTag);
+                aborted = YES;
+                return completion(NO);
+            }
+
             for (NSString *collection in collectionsToRestore) {
                 if ([collection isEqualToString:[OWSDatabaseMigration collection]]) {
                     // It's okay if there are existing migrations; we'll clear those
diff --git a/Signal/src/util/OWSBackupJob.h b/Signal/src/util/OWSBackupJob.h
index bb723a7f7..ee5b4aa0b 100644
--- a/Signal/src/util/OWSBackupJob.h
+++ b/Signal/src/util/OWSBackupJob.h
@@ -8,6 +8,9 @@ extern NSString *const kOWSBackup_ManifestKey_DatabaseFiles;
 extern NSString *const kOWSBackup_ManifestKey_AttachmentFiles;
 extern NSString *const kOWSBackup_ManifestKey_DatabaseKeySpec;
 
+extern NSString *const kOWSBackup_Snapshot_Collection;
+extern NSString *const kOWSBackup_Snapshot_ValidKey;
+
 typedef void (^OWSBackupJobBoolCompletion)(BOOL success);
 typedef void (^OWSBackupJobCompletion)(NSError *_Nullable error);
 
diff --git a/Signal/src/util/OWSBackupJob.m b/Signal/src/util/OWSBackupJob.m
index 6c1d7697f..895845e50 100644
--- a/Signal/src/util/OWSBackupJob.m
+++ b/Signal/src/util/OWSBackupJob.m
@@ -16,6 +16,9 @@ NSString *const kOWSBackup_ManifestKey_DatabaseKeySpec = @"database_key_spec";
 
 NSString *const kOWSBackup_KeychainService = @"kOWSBackup_KeychainService";
 
+NSString *const kOWSBackup_Snapshot_Collection = @"kOWSBackup_Snapshot_Collection";
+NSString *const kOWSBackup_Snapshot_ValidKey = @"kOWSBackup_Snapshot_ValidKey";
+
 @interface OWSBackupJob ()
 
 @property (nonatomic, weak) id<OWSBackupJobDelegate> delegate;