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,19 +194,30 @@ 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];
 
         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
                       inCollection:TSAccountManager_UserAccountCollection];
         }

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];
     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,17 +249,21 @@ 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;
-    {
+
+    // TODO: Remove all @synchronized
     // 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];
+    OWSRecipientIdentity *recipientIdentity =
+        [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];
 
     if (recipientIdentity == nil) {
         OWSFail(@"Missing expected identity: %@", recipientId);
@@ -269,27 +280,40 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
         OWSVerificationStateToString(recipientIdentity.verificationState),
         OWSVerificationStateToString(verificationState));
 
-        [recipientIdentity updateWithVerificationState:verificationState];
+    [recipientIdentity updateWithVerificationState:verificationState transaction:transaction];
 
     if (isUserInitiatedChange) {
-            [self saveChangeMessagesForRecipientId:recipientId verificationState:verificationState isLocalChange:YES];
+        [self saveChangeMessagesForRecipientId:recipientId
-            [self enqueueSyncMessageForVerificationStateForRecipientId: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];
+        [self clearSyncMessageForRecipientId:recipientId transaction:transaction];
-        }
     }
 
     [self fireIdentityStateChangeNotification];
 }
 
 - (OWSVerificationState)verificationStateForRecipientId:(NSString *)recipientId
+{
+    __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;
+}
+
+- (OWSVerificationState)verificationStateForRecipientId:(NSString *)recipientId
+                                            transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     OWSAssert(recipientId.length > 0);
+    OWSAssert(transaction);
 
-    @synchronized(self)
+    OWSRecipientIdentity *_Nullable currentIdentity =
-    {
+        [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];
-        OWSRecipientIdentity *_Nullable currentIdentity = [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId];
 
     if (!currentIdentity) {
         // We might not know the identity for this recipient yet.
@@ -297,17 +321,18 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
     }
 
     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,9 +343,8 @@ 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) {
         // trust on first use
@@ -336,7 +360,6 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
     } else {
         return recipientIdentity;
     }
-    }
 }
 
 - (void)fireIdentityStateChangeNotification
@@ -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,8 +419,6 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
 {
     OWSAssert(identityKey.length == kStoredIdentityKeyLength);
 
-    @synchronized(self)
-    {
     if (recipientIdentity == nil) {
         return YES;
     }
@@ -429,7 +453,6 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
                 recipientIdentity.recipientId);
             return NO;
     }
-    }
 }
 
 - (void)createIdentityChangeInfoMessageForRecipientId:(NSString *)recipientId
@@ -460,21 +483,18 @@ 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)
-        {
-            [self.dbConnection setObject:recipientId
                     forKey:recipientId
               inCollection:OWSIdentityManager_QueuedVerificationStateSyncMessages];
-        }
 
     dispatch_async(dispatch_get_main_queue(), ^{
         [self tryToSyncQueuedVerificationStates];
     });
-    });
 }
 
 - (void)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) {
         [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