From eadb64b75aa908b25afbcc66253c7e7e7aa4bebd Mon Sep 17 00:00:00 2001
From: Matthew Chen <charlesmchen@gmail.com>
Date: Tue, 23 Jan 2018 10:16:36 -0500
Subject: [PATCH] Elaborate test suite around database conversion.

---
 Signal/test/util/OWSDatabaseConverterTest.m  | 51 ++++++++++++++++++--
 SignalMessaging/utils/OWSDatabaseConverter.h |  1 +
 SignalMessaging/utils/OWSDatabaseConverter.m |  2 +-
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/Signal/test/util/OWSDatabaseConverterTest.m b/Signal/test/util/OWSDatabaseConverterTest.m
index 8173ffe10..1d671191c 100644
--- a/Signal/test/util/OWSDatabaseConverterTest.m
+++ b/Signal/test/util/OWSDatabaseConverterTest.m
@@ -35,6 +35,11 @@ NS_ASSUME_NONNULL_BEGIN
     return [Randomness generateRandomBytes:30];
 }
 
+- (NSData *)randomDatabaseSalt
+{
+    return [Randomness generateRandomBytes:kSQLCipherSaltLength];
+}
+
 // * Open a YapDatabase.
 // * Do some work with a block.
 // * Close the database.
@@ -177,6 +182,15 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (nullable NSString *)createUnconvertedDatabase:(NSData *)databasePassword
+{
+    return [self createDatabase:databasePassword
+                   databaseSalt:nil];
+}
+
+// If databaseSalt is nil, creates a non-converted database.
+// Otherwise creates a pre-converted database.
+- (nullable NSString *)createDatabase:(NSData *)databasePassword
+                         databaseSalt:(NSData *_Nullable)databaseSalt
 {
     NSString *temporaryDirectory = NSTemporaryDirectory();
     NSString *filename = [[NSUUID UUID].UUIDString stringByAppendingString:@".sqlite"];
@@ -186,7 +200,7 @@ NS_ASSUME_NONNULL_BEGIN
 
     [self createTestDatabase:databaseFilePath databasePassword:databasePassword];
 
-    BOOL isValid = [self verifyTestDatabase:databaseFilePath databasePassword:databasePassword databaseSalt:nil];
+    BOOL isValid = [self verifyTestDatabase:databaseFilePath databasePassword:databasePassword databaseSalt:databaseSalt];
     OWSAssert(isValid);
 
     return databaseFilePath;
@@ -204,17 +218,18 @@ NS_ASSUME_NONNULL_BEGIN
     // TODO: When we can create converted databases.
 }
 
+// Verifies that legacy users with non-converted databases can convert.
 - (void)testDatabaseConversion
 {
     NSData *databasePassword = [self randomDatabasePassword];
     NSString *_Nullable databaseFilePath = [self createUnconvertedDatabase:databasePassword];
     XCTAssertTrue([OWSDatabaseConverter doesDatabaseNeedToBeConverted:databaseFilePath]);
-
+    
     __block NSData *_Nullable databaseSalt = nil;
     OWSDatabaseSaltBlock saltBlock = ^(NSData *saltData) {
         OWSAssert(!databaseSalt);
         OWSAssert(saltData);
-
+        
         databaseSalt = saltData;
     };
     NSError *_Nullable error = [OWSDatabaseConverter convertDatabaseIfNecessary:databaseFilePath
@@ -225,9 +240,37 @@ NS_ASSUME_NONNULL_BEGIN
     }
     XCTAssertNil(error);
     XCTAssertFalse([OWSDatabaseConverter doesDatabaseNeedToBeConverted:databaseFilePath]);
+    
+    BOOL isValid =
+    [self verifyTestDatabase:databaseFilePath databasePassword:databasePassword databaseSalt:databaseSalt];
+    XCTAssertTrue(isValid);
+}
 
+// Verifies new users who create new pre-converted databases.
+- (void)testDatabaseCreation
+{
+    NSData *databasePassword = [self randomDatabasePassword];
+    NSData *databaseSalt = [self randomDatabaseSalt];
+    NSString *_Nullable databaseFilePath = [self createDatabase:databasePassword
+                                                   databaseSalt:databaseSalt];
+    XCTAssertFalse([OWSDatabaseConverter doesDatabaseNeedToBeConverted:databaseFilePath]);
+    
+    OWSDatabaseSaltBlock saltBlock = ^(NSData *saltData) {
+        OWSAssert(saltData);
+        
+        XCTFail(@"%s No conversion should be necessary", __PRETTY_FUNCTION__);
+    };
+    NSError *_Nullable error = [OWSDatabaseConverter convertDatabaseIfNecessary:databaseFilePath
+                                                               databasePassword:databasePassword
+                                                                      saltBlock:saltBlock];
+    if (error) {
+        DDLogError(@"%s error: %@", __PRETTY_FUNCTION__, error);
+    }
+    XCTAssertNil(error);
+    XCTAssertFalse([OWSDatabaseConverter doesDatabaseNeedToBeConverted:databaseFilePath]);
+    
     BOOL isValid =
-        [self verifyTestDatabase:databaseFilePath databasePassword:databasePassword databaseSalt:databaseSalt];
+    [self verifyTestDatabase:databaseFilePath databasePassword:databasePassword databaseSalt:databaseSalt];
     XCTAssertTrue(isValid);
 }
 
diff --git a/SignalMessaging/utils/OWSDatabaseConverter.h b/SignalMessaging/utils/OWSDatabaseConverter.h
index d255876ee..c62275fbf 100644
--- a/SignalMessaging/utils/OWSDatabaseConverter.h
+++ b/SignalMessaging/utils/OWSDatabaseConverter.h
@@ -5,6 +5,7 @@
 NS_ASSUME_NONNULL_BEGIN
 
 extern const NSUInteger kSqliteHeaderLength;
+extern const NSUInteger kSQLCipherSaltLength;
 
 typedef void (^OWSDatabaseSaltBlock)(NSData *saltData);
 
diff --git a/SignalMessaging/utils/OWSDatabaseConverter.m b/SignalMessaging/utils/OWSDatabaseConverter.m
index 3e9d1dad5..7c75242e6 100644
--- a/SignalMessaging/utils/OWSDatabaseConverter.m
+++ b/SignalMessaging/utils/OWSDatabaseConverter.m
@@ -13,6 +13,7 @@
 NS_ASSUME_NONNULL_BEGIN
 
 const NSUInteger kSqliteHeaderLength = 32;
+const NSUInteger kSQLCipherSaltLength = 16;
 
 @interface OWSStorage (OWSDatabaseConverter)
 
@@ -122,7 +123,6 @@ const NSUInteger kSqliteHeaderLength = 32;
         NSData *headerData = [self readFirstNBytesOfDatabaseFile:databaseFilePath byteCount:kSqliteHeaderLength];
         OWSAssert(headerData);
 
-        const NSUInteger kSQLCipherSaltLength = 16;
         OWSAssert(headerData.length >= kSQLCipherSaltLength);
         sqlCipherSaltData = [headerData subdataWithRange:NSMakeRange(0, kSQLCipherSaltLength)];