Add protocol context to protocol kit.

SignalMessaging/environment/AppSetup.m

@@ -27,9 +27,6 @@ NS_ASSUME_NONNULL_BEGIN
     dispatch_once(&onceToken, ^{
         [Environment setCurrent:[Release releaseEnvironment]];
 
-        // Encryption/Decryption mutates session state and must be synchronized on a serial queue.
-        [SessionCipher setSessionCipherDispatchQueue:[OWSDispatch sessionStoreQueue]];
-
         id<OWSCallMessageHandler> callMessageHandler = callMessageHandlerBlock();
         id<NotificationsProtocol> notificationsManager = notificationsManagerBlock();
 

SignalServiceKit/src/Account/TSAccountManager.h

@@ -1,5 +1,5 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
 #import "TSConstants.h"
@@ -13,6 +13,7 @@ extern NSString *const kNSNotificationName_LocalNumberDidChange;
 
 @class TSNetworkManager;
 @class TSStorageManager;
+@class YapDatabaseReadWriteTransaction;
 
 @interface TSAccountManager : NSObject
 
@@ -65,6 +66,7 @@ extern NSString *const kNSNotificationName_LocalNumberDidChange;
  */
 
 + (uint32_t)getOrGenerateRegistrationId;
++ (uint32_t)getOrGenerateRegistrationId:(YapDatabaseReadWriteTransaction *)transaction;
 
 #pragma mark - Register with phone number
 

SignalServiceKit/src/Account/TSAccountManager.m

@@ -194,21 +194,32 @@ NSString *const TSAccountManager_ServerSignalingKey = @"TSStorageServerSignaling
     return [[self sharedInstance] getOrGenerateRegistrationId];
 }
 
++ (uint32_t)getOrGenerateRegistrationId:(YapDatabaseReadWriteTransaction *)transaction
+{
+    return [[self sharedInstance] getOrGenerateRegistrationId:transaction];
+}
+
 - (uint32_t)getOrGenerateRegistrationId
+{
+    [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
+        [self getOrGenerateRegistrationId:transaction];
+    }];
+}
+
+- (uint32_t)getOrGenerateRegistrationId:(YapDatabaseReadWriteTransaction *)transaction
 {
     @synchronized(self)
     {
-        uint32_t registrationID =
+        uint32_t registrationID = [[transaction objectForKey:TSAccountManager_LocalRegistrationIdKey
-            [[self.dbConnection objectForKey:TSAccountManager_LocalRegistrationIdKey
+                                                inCollection:TSAccountManager_UserAccountCollection] unsignedIntValue];
-                                inCollection:TSAccountManager_UserAccountCollection] unsignedIntValue];
 
         if (registrationID == 0) {
             registrationID = (uint32_t)arc4random_uniform(16380) + 1;
             DDLogWarn(@"%@ Generated a new registrationID: %u", self.logTag, registrationID);
 
-            [self.dbConnection setObject:[NSNumber numberWithUnsignedInteger:registrationID]
+            [transaction setObject:[NSNumber numberWithUnsignedInteger:registrationID]
-                                  forKey:TSAccountManager_LocalRegistrationIdKey
+                            forKey:TSAccountManager_LocalRegistrationIdKey
-                            inCollection:TSAccountManager_UserAccountCollection];
+                      inCollection:TSAccountManager_UserAccountCollection];
         }
         return registrationID;
     }

SignalServiceKit/src/Messages/OWSIdentityManager.h

@@ -35,7 +35,8 @@ extern const NSUInteger kIdentityKeyLength;
 - (void)setVerificationState:(OWSVerificationState)verificationState
                  identityKey:(NSData *)identityKey
                  recipientId:(NSString *)recipientId
-       isUserInitiatedChange:(BOOL)isUserInitiatedChange;
+       isUserInitiatedChange:(BOOL)isUserInitiatedChange
+             protocolContext:(nullable id)protocolContext;
 
 - (OWSVerificationState)verificationStateForRecipientId:(NSString *)recipientId;
 
@@ -46,7 +47,8 @@ extern const NSUInteger kIdentityKeyLength;
  * @returns nil if the recipient does not exist, or is trusted for sending
  *          else returns the untrusted recipient.
  */
-- (nullable OWSRecipientIdentity *)untrustedIdentityForSendingToRecipientId:(NSString *)recipientId;
+- (nullable OWSRecipientIdentity *)untrustedIdentityForSendingToRecipientId:(NSString *)recipientId
+                                                            protocolContext:(nullable id)protocolContext;
 
 // This method can be called from any thread.
 - (void)processIncomingSyncMessage:(OWSSignalServiceProtosVerified *)verified;

SignalServiceKit/src/Messages/OWSIdentityManager.m

@@ -23,6 +23,7 @@
 #import "TSStorageManager.h"
 #import "TextSecureKitEnv.h"
 #import "YapDatabaseConnection+OWS.h"
+#import "YapDatabaseReadTransaction+OWS.h"
 #import <AxolotlKit/NSData+keyVersionByte.h>
 #import <Curve25519Kit/Curve25519.h>
 #import <YapDatabase/YapDatabase.h>
@@ -55,9 +56,11 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 @interface OWSIdentityManager ()
 
 @property (nonatomic, readonly) TSStorageManager *storageManager;
-@property (nonatomic, readonly) YapDatabaseConnection *dbConnection;
 @property (nonatomic, readonly) OWSMessageSender *messageSender;
 
+// TODO: Should we get rid of this property?
+@property (nonatomic, readonly) YapDatabaseConnection *dbConnection;
+
 @end
 
 #pragma mark -
@@ -136,12 +139,12 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 
     YapDatabaseReadWriteTransaction *transaction = protocolContext;
 
-    return [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId].identityKey;
+    return [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction].identityKey;
 }
 
 - (nullable ECKeyPair *)identityKeyPairWithoutProtocolContext
 {
-    ECKeyPair *_Nullable identityKeyPair = nil;
+    __block ECKeyPair *_Nullable identityKeyPair = nil;
     [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
         identityKeyPair = [self identityKeyPair:transaction];
     }];
@@ -154,15 +157,19 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 
     YapDatabaseReadWriteTransaction *transaction = protocolContext;
 
-    ECKeyPair *_Nullable identityKeyPair = [self.dbConnection keyPairForKey:TSStorageManagerIdentityKeyStoreIdentityKey
+    ECKeyPair *_Nullable identityKeyPair = [transaction keyPairForKey:TSStorageManagerIdentityKeyStoreIdentityKey
-                                                               inCollection:TSStorageManagerIdentityKeyStoreCollection];
+                                                         inCollection:TSStorageManagerIdentityKeyStoreCollection];
     return identityKeyPair;
 }
 
-//- (int)localRegistrationId
+- (int)localRegistrationId:(nullable id)protocolContext
-//{
+{
-//    return (int)[TSAccountManager getOrGenerateRegistrationId];
+    OWSAssert([protocolContext isKindOfClass:[YapDatabaseReadWriteTransaction class]]);
-//}
+
+    YapDatabaseReadWriteTransaction *transaction = protocolContext;
+
+    return (int)[TSAccountManager getOrGenerateRegistrationId:transaction];
+}
 
 - (BOOL)saveRemoteIdentity:(NSData *)identityKey
                recipientId:(NSString *)recipientId
@@ -228,7 +235,7 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
         [self.storageManager archiveAllSessionsForContact:recipientId protocolContext:protocolContext];
 
         // Cancel any pending verification state sync messages for this recipient.
-        [self clearSyncMessageForRecipientId:recipientId];
+        [self clearSyncMessageForRecipientId:recipientId transaction:transaction];
 
         [self fireIdentityStateChangeNotification];
 
@@ -242,42 +249,48 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
                  identityKey:(NSData *)identityKey
                  recipientId:(NSString *)recipientId
        isUserInitiatedChange:(BOOL)isUserInitiatedChange
+             protocolContext:(nullable id)protocolContext
 {
     OWSAssert(identityKey.length == kStoredIdentityKeyLength);
     OWSAssert(recipientId.length > 0);
+    OWSAssert([protocolContext isKindOfClass:[YapDatabaseReadWriteTransaction class]]);
 
-    @synchronized(self)
+    YapDatabaseReadWriteTransaction *transaction = protocolContext;
-    {
-        // Ensure a remote identity exists for this key. We may be learning about
-        // it for the first time.
-        [self saveRemoteIdentity:identityKey recipientId:recipientId protocolContext:protocolContext];
-
-        OWSRecipientIdentity *recipientIdentity = [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId];
 
-        if (recipientIdentity == nil) {
+    // TODO: Remove all @synchronized
-            OWSFail(@"Missing expected identity: %@", recipientId);
+    // Ensure a remote identity exists for this key. We may be learning about
-            return;
+    // it for the first time.
-        }
+    [self saveRemoteIdentity:identityKey recipientId:recipientId protocolContext:protocolContext];
 
-        if (recipientIdentity.verificationState == verificationState) {
+    OWSRecipientIdentity *recipientIdentity =
-            return;
+        [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];
-        }
 
-        DDLogInfo(@"%@ setVerificationState: %@ (%@ -> %@)",
+    if (recipientIdentity == nil) {
-            self.logTag,
+        OWSFail(@"Missing expected identity: %@", recipientId);
-            recipientId,
+        return;
-            OWSVerificationStateToString(recipientIdentity.verificationState),
+    }
-            OWSVerificationStateToString(verificationState));
 
-        [recipientIdentity updateWithVerificationState:verificationState];
+    if (recipientIdentity.verificationState == verificationState) {
+        return;
+    }
 
-        if (isUserInitiatedChange) {
+    DDLogInfo(@"%@ setVerificationState: %@ (%@ -> %@)",
-            [self saveChangeMessagesForRecipientId:recipientId verificationState:verificationState isLocalChange:YES];
+        self.logTag,
-            [self enqueueSyncMessageForVerificationStateForRecipientId:recipientId];
+        recipientId,
-        } else {
+        OWSVerificationStateToString(recipientIdentity.verificationState),
-            // Cancel any pending verification state sync messages for this recipient.
+        OWSVerificationStateToString(verificationState));
-            [self clearSyncMessageForRecipientId:recipientId];
+
-        }
+    [recipientIdentity updateWithVerificationState:verificationState transaction:transaction];
+
+    if (isUserInitiatedChange) {
+        [self saveChangeMessagesForRecipientId:recipientId
+                             verificationState:verificationState
+                                 isLocalChange:YES
+                                   transaction:transaction];
+        [self enqueueSyncMessageForVerificationStateForRecipientId:recipientId transaction:transaction];
+    } else {
+        // Cancel any pending verification state sync messages for this recipient.
+        [self clearSyncMessageForRecipientId:recipientId transaction:transaction];
     }
 
     [self fireIdentityStateChangeNotification];
@@ -285,29 +298,41 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 
 - (OWSVerificationState)verificationStateForRecipientId:(NSString *)recipientId
 {
-    OWSAssert(recipientId.length > 0);
+    __block OWSVerificationState result;
+    // Use a read/write transaction to block on latest.
+    [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
+        result = [self verificationStateForRecipientId:recipientId transaction:transaction];
+    }];
+    return result;
+}
 
-    @synchronized(self)
+- (OWSVerificationState)verificationStateForRecipientId:(NSString *)recipientId
-    {
+                                            transaction:(YapDatabaseReadWriteTransaction *)transaction
-        OWSRecipientIdentity *_Nullable currentIdentity = [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId];
+{
+    OWSAssert(recipientId.length > 0);
+    OWSAssert(transaction);
 
-        if (!currentIdentity) {
+    OWSRecipientIdentity *_Nullable currentIdentity =
-            // We might not know the identity for this recipient yet.
+        [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];
-            return OWSVerificationStateDefault;
-        }
 
-        return currentIdentity.verificationState;
+    if (!currentIdentity) {
+        // We might not know the identity for this recipient yet.
+        return OWSVerificationStateDefault;
     }
+
+    return currentIdentity.verificationState;
 }
 
 - (nullable OWSRecipientIdentity *)recipientIdentityForRecipientId:(NSString *)recipientId
 {
     OWSAssert(recipientId.length > 0);
 
-    @synchronized(self)
+    __block OWSRecipientIdentity *_Nullable result;
-    {
+    // Use a read/write transaction to block on latest.
-        return [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId];
+    [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
-    }
+        result = [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];
+    }];
+    return result;
 }
 
 - (nullable OWSRecipientIdentity *)untrustedIdentityForSendingToRecipientId:(NSString *)recipientId
@@ -318,24 +343,22 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 
     YapDatabaseReadWriteTransaction *transaction = protocolContext;
 
-    @synchronized(self)
+    OWSRecipientIdentity *_Nullable recipientIdentity =
-    {
+        [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];
-        OWSRecipientIdentity *_Nullable recipientIdentity = [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId];
 
-        if (recipientIdentity == nil) {
+    if (recipientIdentity == nil) {
-            // trust on first use
+        // trust on first use
-            return nil;
+        return nil;
-        }
+    }
 
-        BOOL isTrusted = [self isTrustedIdentityKey:recipientIdentity.identityKey
+    BOOL isTrusted = [self isTrustedIdentityKey:recipientIdentity.identityKey
-                                        recipientId:recipientId
+                                    recipientId:recipientId
-                                          direction:TSMessageDirectionOutgoing
+                                      direction:TSMessageDirectionOutgoing
-                                    protocolContext:protocolContext];
+                                protocolContext:protocolContext];
-        if (isTrusted) {
+    if (isTrusted) {
-            return nil;
+        return nil;
-        } else {
+    } else {
-            return recipientIdentity;
+        return recipientIdentity;
-        }
     }
 }
 
@@ -361,13 +384,15 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
     @synchronized(self)
     {
         if ([[TSAccountManager localNumber] isEqualToString:recipientId]) {
-            if ([[self identityKeyPair].publicKey isEqualToData:identityKey]) {
+            ECKeyPair *_Nullable localIdentityKeyPair = [self identityKeyPair:protocolContext];
+
+            if ([localIdentityKeyPair.publicKey isEqualToData:identityKey]) {
                 return YES;
             } else {
                 OWSFail(@"%@ Wrong identity: %@ for local key: %@, recipientId: %@",
                     self.logTag,
                     identityKey,
-                    [self identityKeyPair].publicKey,
+                    localIdentityKeyPair.publicKey,
                     recipientId);
                 return NO;
             }
@@ -378,7 +403,8 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
                 return YES;
             }
             case TSMessageDirectionOutgoing: {
-                OWSRecipientIdentity *existingIdentity = [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId];
+                OWSRecipientIdentity *existingIdentity =
+                    [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];
                 return [self isTrustedKey:identityKey forSendingToIdentity:existingIdentity];
             }
             default: {
@@ -393,42 +419,39 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 {
     OWSAssert(identityKey.length == kStoredIdentityKeyLength);
 
-    @synchronized(self)
+    if (recipientIdentity == nil) {
-    {
+        return YES;
-        if (recipientIdentity == nil) {
+    }
-            return YES;
-        }
 
-        OWSAssert(recipientIdentity.identityKey.length == kStoredIdentityKeyLength);
+    OWSAssert(recipientIdentity.identityKey.length == kStoredIdentityKeyLength);
-        if (![recipientIdentity.identityKey isEqualToData:identityKey]) {
+    if (![recipientIdentity.identityKey isEqualToData:identityKey]) {
-            DDLogWarn(@"%@ key mismatch for recipient: %@", self.logTag, recipientIdentity.recipientId);
+        DDLogWarn(@"%@ key mismatch for recipient: %@", self.logTag, recipientIdentity.recipientId);
-            return NO;
+        return NO;
-        }
+    }
 
-        if ([recipientIdentity isFirstKnownKey]) {
+    if ([recipientIdentity isFirstKnownKey]) {
-            return YES;
+        return YES;
-        }
+    }
 
-        switch (recipientIdentity.verificationState) {
+    switch (recipientIdentity.verificationState) {
-            case OWSVerificationStateDefault: {
+        case OWSVerificationStateDefault: {
-                BOOL isNew = (fabs([recipientIdentity.createdAt timeIntervalSinceNow])
+            BOOL isNew = (fabs([recipientIdentity.createdAt timeIntervalSinceNow])
-                    < kIdentityKeyStoreNonBlockingSecondsThreshold);
+                < kIdentityKeyStoreNonBlockingSecondsThreshold);
-                if (isNew) {
+            if (isNew) {
-                    DDLogWarn(
+                DDLogWarn(
-                        @"%@ not trusting new identity for recipient: %@", self.logTag, recipientIdentity.recipientId);
+                    @"%@ not trusting new identity for recipient: %@", self.logTag, recipientIdentity.recipientId);
-                    return NO;
-                } else {
-                    return YES;
-                }
-            }
-            case OWSVerificationStateVerified:
-                return YES;
-            case OWSVerificationStateNoLongerVerified:
-                DDLogWarn(@"%@ not trusting no longer verified identity for recipient: %@",
-                    self.logTag,
-                    recipientIdentity.recipientId);
                 return NO;
+            } else {
+                return YES;
+            }
         }
+        case OWSVerificationStateVerified:
+            return YES;
+        case OWSVerificationStateNoLongerVerified:
+            DDLogWarn(@"%@ not trusting no longer verified identity for recipient: %@",
+                self.logTag,
+                recipientIdentity.recipientId);
+            return NO;
     }
 }
 
@@ -460,20 +483,17 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 }
 
 - (void)enqueueSyncMessageForVerificationStateForRecipientId:(NSString *)recipientId
+                                                 transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     OWSAssert(recipientId.length > 0);
+    OWSAssert(transaction);
 
-    dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
+    [transaction setObject:recipientId
-        @synchronized(self)
+                    forKey:recipientId
-        {
+              inCollection:OWSIdentityManager_QueuedVerificationStateSyncMessages];
-            [self.dbConnection setObject:recipientId
-                                  forKey:recipientId
-                            inCollection:OWSIdentityManager_QueuedVerificationStateSyncMessages];
-        }
 
-        dispatch_async(dispatch_get_main_queue(), ^{
+    dispatch_async(dispatch_get_main_queue(), ^{
-            [self tryToSyncQueuedVerificationStates];
+        [self tryToSyncQueuedVerificationStates];
-        });
     });
 }
 
@@ -544,10 +564,12 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
     });
 }
 
+// TODO: Keep converting.
 - (void)sendSyncVerificationStateMessage:(OWSVerificationStateSyncMessage *)message
 {
     OWSAssert(message);
     OWSAssert(message.verificationForRecipientId.length > 0);
+    OWSAssertIsOnMainThread();
 
     TSContactThread *contactThread = [TSContactThread getOrCreateThreadWithContactId:message.verificationForRecipientId];
     
@@ -738,7 +760,10 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 
             [recipientIdentity updateWithVerificationState:verificationState];
 
-            [self saveChangeMessagesForRecipientId:recipientId verificationState:verificationState isLocalChange:NO];
+            [self saveChangeMessagesForRecipientId:recipientId
+                                 verificationState:verificationState
+                                     isLocalChange:NO
+                                       transaction:transaction];
         }
     }
 }
@@ -748,12 +773,15 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 - (void)saveChangeMessagesForRecipientId:(NSString *)recipientId
                        verificationState:(OWSVerificationState)verificationState
                            isLocalChange:(BOOL)isLocalChange
+                             transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     OWSAssert(recipientId.length > 0);
+    OWSAssert(transaction);
 
     NSMutableArray<TSMessage *> *messages = [NSMutableArray new];
 
-    TSContactThread *contactThread = [TSContactThread getOrCreateThreadWithContactId:recipientId];
+    TSContactThread *contactThread =
+        [TSContactThread getOrCreateThreadWithContactId:recipientId transaction:transaction];
     OWSAssert(contactThread);
     [messages addObject:[[OWSVerificationStateChangeMessage alloc] initWithTimestamp:[NSDate ows_millisecondTimeStamp]
                                                                               thread:contactThread
@@ -761,7 +789,8 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
                                                                    verificationState:verificationState
                                                                        isLocalChange:isLocalChange]];
 
-    for (TSGroupThread *groupThread in [TSGroupThread groupThreadsWithRecipientId:recipientId]) {
+    for (TSGroupThread *groupThread in
+        [TSGroupThread groupThreadsWithRecipientId:recipientId transaction:transaction]) {
         [messages
             addObject:[[OWSVerificationStateChangeMessage alloc] initWithTimestamp:[NSDate ows_millisecondTimeStamp]
                                                                             thread:groupThread
@@ -770,11 +799,9 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
                                                                      isLocalChange:isLocalChange]];
     }
 
-    [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+    for (TSMessage *message in messages) {
-        for (TSMessage *message in messages) {
+        [message saveWithTransaction:transaction];
-            [message saveWithTransaction:transaction];
+    }
-        }
-    }];
 }
 
 #pragma mark - Debug

SignalServiceKit/src/Security/OWSRecipientIdentity.h

@@ -1,5 +1,5 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
 #import "OWSSignalServiceProtos.pb.h"
@@ -27,7 +27,8 @@ OWSSignalServiceProtosVerifiedState OWSVerificationStateToProtoState(OWSVerifica
 
 @property (atomic, readonly) OWSVerificationState verificationState;
 
-- (void)updateWithVerificationState:(OWSVerificationState)verificationState;
+- (void)updateWithVerificationState:(OWSVerificationState)verificationState
+                        transaction:(YapDatabaseReadWriteTransaction *)transaction;
 
 #pragma mark - Initializers
 

SignalServiceKit/src/Security/OWSRecipientIdentity.m

@@ -80,13 +80,35 @@ OWSSignalServiceProtosVerifiedState OWSVerificationStateToProtoState(OWSVerifica
 }
 
 - (void)updateWithVerificationState:(OWSVerificationState)verificationState
+                        transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
+    OWSAssert(transaction);
+
     // Ensure changes are persisted without clobbering any work done on another thread or instance.
     [self updateWithChangeBlock:^(OWSRecipientIdentity *_Nonnull obj) {
         obj.verificationState = verificationState;
-    }];
+    }
+                    transaction:transaction];
+}
+
+- (void)updateWithChangeBlock:(void (^)(OWSRecipientIdentity *obj))changeBlock
+                  transaction:(YapDatabaseReadWriteTransaction *)transaction
+{
+    OWSAssert(transaction);
+
+    changeBlock(self);
+
+    OWSRecipientIdentity *latest = [[self class] fetchObjectWithUniqueID:self.uniqueId transaction:transaction];
+    if (latest == nil) {
+        [self saveWithTransaction:transaction];
+        return;
+    }
+
+    changeBlock(latest);
+    [latest saveWithTransaction:transaction];
 }
 
+// TODO: Is this method obsolete?
 - (void)updateWithChangeBlock:(void (^)(OWSRecipientIdentity *obj))changeBlock
 {
     changeBlock(self);
@@ -139,6 +161,7 @@ OWSSignalServiceProtosVerifiedState OWSVerificationStateToProtoState(OWSVerifica
     return self.dbReadWriteConnection;
 }
 
+// TODO: Replace with protocol connection?
 /**
  * Override to disable the object cache to better enforce transaction semantics on the store.
  * Note that it's still technically possible to access this collection from a different collection,

SignalServiceKit/src/Storage/YapDatabaseReadTransaction+OWS.h

@@ -0,0 +1,28 @@
+//
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
+//
+
+#import <YapDatabase/YapDatabaseTransaction.h>
+
+@class ECKeyPair;
+@class PreKeyRecord;
+@class SignedPreKeyRecord;
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface YapDatabaseReadTransaction (OWS)
+
+- (BOOL)boolForKey:(NSString *)key inCollection:(NSString *)collection;
+- (BOOL)boolForKey:(NSString *)key inCollection:(NSString *)collection defaultValue:(BOOL)defaultValue;
+- (int)intForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable NSDate *)dateForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable NSDictionary *)dictionaryForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable NSString *)stringForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable NSData *)dataForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable ECKeyPair *)keyPairForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable PreKeyRecord *)preKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable SignedPreKeyRecord *)signedPreKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection;
+
+@end
+
+NS_ASSUME_NONNULL_END

SignalServiceKit/src/Storage/YapDatabaseReadTransaction+OWS.m

@@ -0,0 +1,112 @@
+//
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
+//
+
+#import "YapDatabaseReadTransaction+OWS.h"
+#import <AxolotlKit/PreKeyRecord.h>
+#import <AxolotlKit/SignedPrekeyRecord.h>
+#import <Curve25519Kit/Curve25519.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@implementation YapDatabaseReadTransaction (OWS)
+
+- (nullable id)objectForKey:(NSString *)key inCollection:(NSString *)collection ofExpectedType:(Class) class {
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    id _Nullable value = [self objectForKey:key inCollection:collection];
+    OWSAssert(!value || [value isKindOfClass:class]);
+    return value;
+}
+
+    - (nullable NSDictionary *)dictionaryForKey : (NSString *)key inCollection : (NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    return [self objectForKey:key inCollection:collection ofExpectedType:[NSDictionary class]];
+}
+
+- (nullable NSString *)stringForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    return [self objectForKey:key inCollection:collection ofExpectedType:[NSString class]];
+}
+
+- (BOOL)boolForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    return [self boolForKey:key inCollection:collection defaultValue:NO];
+}
+
+- (BOOL)boolForKey:(NSString *)key inCollection:(NSString *)collection defaultValue:(BOOL)defaultValue
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    NSNumber *_Nullable value = [self objectForKey:key inCollection:collection ofExpectedType:[NSNumber class]];
+    return value ? [value boolValue] : defaultValue;
+}
+
+- (nullable NSData *)dataForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    return [self objectForKey:key inCollection:collection ofExpectedType:[NSData class]];
+}
+
+- (nullable ECKeyPair *)keyPairForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    return [self objectForKey:key inCollection:collection ofExpectedType:[ECKeyPair class]];
+}
+
+- (nullable PreKeyRecord *)preKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    return [self objectForKey:key inCollection:collection ofExpectedType:[PreKeyRecord class]];
+}
+
+- (nullable SignedPreKeyRecord *)signedPreKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    return [self objectForKey:key inCollection:collection ofExpectedType:[SignedPreKeyRecord class]];
+}
+
+- (int)intForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    NSNumber *_Nullable number = [self objectForKey:key inCollection:collection ofExpectedType:[NSNumber class]];
+    return [number intValue];
+}
+
+- (nullable NSDate *)dateForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    OWSAssert(key.length > 0);
+    OWSAssert(collection.length > 0);
+
+    NSNumber *_Nullable value = [self objectForKey:key inCollection:collection ofExpectedType:[NSNumber class]];
+    if (value) {
+        return [NSDate dateWithTimeIntervalSince1970:value.doubleValue];
+    } else {
+        return nil;
+    }
+}
+
+@end
+
+NS_ASSUME_NONNULL_END