Add protocol context to protocol kit.

7 years ago · 39e353503f
parent 71782e0369
commit 39e353503f
9 changed files with 337 additions and 134 deletions
--- a/SignalMessaging/environment/AppSetup.m
+++ b/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();

--- a/SignalServiceKit/src/Account/TSAccountManager.h
+++ b/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

--- a/SignalServiceKit/src/Account/TSAccountManager.m
+++ b/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 =
-            [[self.dbConnection objectForKey:TSAccountManager_LocalRegistrationIdKey
+        uint32_t registrationID = [[transaction 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];
        }
--- a/SignalServiceKit/src/Messages/OWSIdentityManager.h
+++ b/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;
--- a/SignalServiceKit/src/Messages/OWSIdentityManager.m
+++ b/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
-//{
-//    return (int)[TSAccountManager getOrGenerateRegistrationId];
-//}
+- (int)localRegistrationId:(nullable id)protocolContext
+{
+    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,15 +280,17 @@ 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 enqueueSyncMessageForVerificationStateForRecipientId:recipientId];
+        [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];
-        }
+        [self clearSyncMessageForRecipientId:recipientId transaction:transaction];
    }

    [self fireIdentityStateChangeNotification];
@ -285,11 +298,22 @@ 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);
+
+    OWSRecipientIdentity *_Nullable currentIdentity =
+        [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];

    if (!currentIdentity) {
        // We might not know the identity for this recipient yet.
@ -298,16 +322,17 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa

    return currentIdentity.verificationState;
 }
-}

 - (nullable OWSRecipientIdentity *)recipientIdentityForRecipientId:(NSString *)recipientId
 {
    OWSAssert(recipientId.length > 0);

-    @synchronized(self)
-    {
-        return [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId];
-    }
+    __block OWSRecipientIdentity *_Nullable result;
+    // Use a read/write transaction to block on latest.
+    [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];
+    OWSRecipientIdentity *_Nullable recipientIdentity =
+        [OWSRecipientIdentity fetchObjectWithUniqueID:recipientId transaction:transaction];

    if (recipientIdentity == nil) {
        // trust on first use
@ -337,7 +361,6 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
        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;
    }
@ -430,7 +454,6 @@ NSString *const kNSNotificationName_IdentityStateDidChange = @"kNSNotificationNa
            return NO;
    }
 }
-}

 - (void)createIdentityChangeInfoMessageForRecipientId:(NSString *)recipientId
                                          transaction:(YapDatabaseReadWriteTransaction *)transaction
@ -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), ^{
-        @synchronized(self)
-        {
-            [self.dbConnection setObject:recipientId
+    [transaction 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
--- a/SignalServiceKit/src/Security/OWSRecipientIdentity.h
+++ b/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

--- a/SignalServiceKit/src/Security/OWSRecipientIdentity.m
+++ b/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,
--- a/SignalServiceKit/src/Storage/YapDatabaseReadTransaction+OWS.h
+++ b/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
--- a/SignalServiceKit/src/Storage/YapDatabaseReadTransaction+OWS.m
+++ b/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