Add protocol context to protocol kit.

AxolotlKit/Classes/SPKAsserts.h

@@ -1,31 +1,33 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
+NS_ASSUME_NONNULL_BEGIN
+
 #ifndef SPKAssert
 
 #ifdef DEBUG
 
-#define USE_ASSERTS
+#define USE_SPK_ASSERTS
 
-#define CONVERT_TO_STRING(X) #X
-#define CONVERT_EXPR_TO_STRING(X) CONVERT_TO_STRING(X)
+#define SPK_CONVERT_TO_STRING(X) #X
+#define SPK_CONVERT_EXPR_TO_STRING(X) SPK_CONVERT_TO_STRING(X)
 
 // SPKAssert() and SPKFail() should be used in Obj-C methods.
 // SPKCAssert() and SPKCFail() should be used in free functions.
 
 #define SPKAssert(X)                                                                                                   \
     if (!(X)) {                                                                                                        \
-        DDLogError(@"%s Assertion failed: %s", __PRETTY_FUNCTION__, CONVERT_EXPR_TO_STRING(X));                        \
+        DDLogError(@"%s Assertion failed: %s", __PRETTY_FUNCTION__, SPK_CONVERT_EXPR_TO_STRING(X));                    \
         [DDLog flushLog];                                                                                              \
-        NSAssert(0, @"Assertion failed: %s", CONVERT_EXPR_TO_STRING(X));                                               \
+        NSAssert(0, @"Assertion failed: %s", SPK_CONVERT_EXPR_TO_STRING(X));                                           \
     }
 
 #define SPKCAssert(X)                                                                                                  \
     if (!(X)) {                                                                                                        \
-        DDLogError(@"%s Assertion failed: %s", __PRETTY_FUNCTION__, CONVERT_EXPR_TO_STRING(X));                        \
+        DDLogError(@"%s Assertion failed: %s", __PRETTY_FUNCTION__, SPK_CONVERT_EXPR_TO_STRING(X));                    \
         [DDLog flushLog];                                                                                              \
-        NSCAssert(0, @"Assertion failed: %s", CONVERT_EXPR_TO_STRING(X));                                              \
+        NSCAssert(0, @"Assertion failed: %s", SPK_CONVERT_EXPR_TO_STRING(X));                                          \
     }
 
 #define SPKFail(message, ...)                                                                                          \
@@ -70,3 +72,5 @@
 #endif
 
 #endif
+
+NS_ASSUME_NONNULL_END

AxolotlKit/Classes/SessionCipher.h

@@ -1,43 +1,35 @@
 //
-//  TSAxolotlRatchet.h
-//  AxolotlKit
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
-//  Created by Frederic Jacobs on 07/22/14.
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
-//
-
-#import <Foundation/Foundation.h>
 
 #import "AxolotlStore.h"
-#import "SignedPreKeyStore.h"
-#import "PreKeyStore.h"
 #import "IdentityKeyStore.h"
-#import "SessionStore.h"
+#import "PreKeyStore.h"
 #import "PreKeyWhisperMessage.h"
 #import "SessionState.h"
+#import "SessionStore.h"
+#import "SignedPreKeyStore.h"
 #import "WhisperMessage.h"
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
 
 @interface SessionCipher : NSObject
 
-/**
- * To keep Session state synchronized, encryption and decryption must happen on the same (serial) dispatch queue. If no
- * queue is specified, the main queue will be used by default. We only assert that this invariant is held. Dispatching 
- * to this thread is the responsibility of the caller.
- *
- * @param dispatchQueue    serial dispatch queue on which all encryption/decryption must be dispatched.
- */
-+ (void)setSessionCipherDispatchQueue:(dispatch_queue_t)dispatchQueue;
-+ (dispatch_queue_t)getSessionCipherDispatchQueue;
-
 - (instancetype)initWithAxolotlStore:(id<AxolotlStore>)sessionStore recipientId:(NSString*)recipientId deviceId:(int)deviceId;
 
 - (instancetype)initWithSessionStore:(id<SessionStore>)sessionStore preKeyStore:(id<PreKeyStore>)preKeyStore signedPreKeyStore:(id<SignedPreKeyStore>)signedPreKeyStore identityKeyStore:(id<IdentityKeyStore>)identityKeyStore recipientId:(NSString*)recipientId deviceId:(int)deviceId;
 
-- (id<CipherMessage>)encryptMessage:(NSData*)paddedMessage;
+// protocolContext is an optional parameter that can be used to ensure that all
+// identity and session store writes are coordinated and/or occur within a single
+// transaction.
+- (id<CipherMessage>)encryptMessage:(NSData *)paddedMessage protocolContext:(nullable id)protocolContext;
 
-- (NSData*)decrypt:(id<CipherMessage>)whisperMessage;
+- (NSData *)decrypt:(id<CipherMessage>)whisperMessage protocolContext:(nullable id)protocolContext;
 
-- (int)remoteRegistrationId;
-- (int)sessionVersion;
+- (int)remoteRegistrationId:(nullable id)protocolContext;
+- (int)sessionVersion:(nullable id)protocolContext;
 
 @end
+
+NS_ASSUME_NONNULL_END

AxolotlKit/Classes/SessionCipher.m

@@ -1,34 +1,30 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
 #import "SessionCipher.h"
+#import "AES-CBC.h"
+#import "AxolotlExceptions.h"
+#import "AxolotlParameters.h"
+#import "ChainKey.h"
+#import "MessageKeys.h"
+#import "NSData+keyVersionByte.h"
+#import "PreKeyStore.h"
+#import "RootKey.h"
+#import "SessionBuilder.h"
+#import "SessionState.h"
+#import "SessionStore.h"
+#import "SignedPreKeyStore.h"
+#import "WhisperMessage.h"
 #import <Curve25519Kit/Curve25519.h>
 #import <Curve25519Kit/Ed25519.h>
-
-#import "NSData+keyVersionByte.h"
-
-#import "AxolotlExceptions.h"
-#import "SessionBuilder.h"
-#import "SessionStore.h"
-#import "AES-CBC.h"
-#import "AxolotlParameters.h"
-#import "MessageKeys.h"
-#import "SessionState.h"
-#import "ChainKey.h"
-#import "RootKey.h"
-#import "WhisperMessage.h"
-
-#import "SignedPreKeyStore.h"
-#import "PreKeyStore.h"
-
 #import <HKDFKit/HKDFKit.h>
 
+NS_ASSUME_NONNULL_BEGIN
+
 #define SYSTEM_VERSION_GREATER_THAN_OR_EQUAL_TO(major, minor) \
     ([[NSProcessInfo processInfo] isOperatingSystemAtLeastVersion:(NSOperatingSystemVersion){.majorVersion = major, .minorVersion = minor, .patchVersion = 0}])
 
-static dispatch_queue_t _sessionCipherDispatchQueue;
-
 @interface SessionCipher ()
 
 @property (nonatomic, readonly) NSString *recipientId;
@@ -41,10 +37,10 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
 
 @end
 
+#pragma mark -
 
 @implementation SessionCipher
 
-
 - (instancetype)initWithAxolotlStore:(id<AxolotlStore>)sessionStore recipientId:(NSString*)recipientId deviceId:(int)deviceId{
     return [self initWithSessionStore:sessionStore
                           preKeyStore:sessionStore
@@ -78,34 +74,12 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     return self;
 }
 
-#pragma mark - dispatch queue 
-
-+ (dispatch_queue_t)getSessionCipherDispatchQueue;
+- (id<CipherMessage>)encryptMessage:(NSData *)paddedMessage protocolContext:(nullable id)protocolContext
 {
-    if (_sessionCipherDispatchQueue) {
-        return _sessionCipherDispatchQueue;
-    } else {
-        return dispatch_get_main_queue();
-    }
-}
+    SPKAssert(paddedMessage);
 
-+ (void)setSessionCipherDispatchQueue:(dispatch_queue_t)dispatchQueue
-{
-    _sessionCipherDispatchQueue = dispatchQueue;
-}
-
-- (void)assertOnSessionCipherDispatchQueue
-{
-#ifdef DEBUG
-    if (SYSTEM_VERSION_GREATER_THAN_OR_EQUAL_TO(10, 0)) {
-        dispatch_assert_queue([[self class] getSessionCipherDispatchQueue]);
-    } // else, skip assert as it's a development convenience.
-#endif
-}
-
-- (id<CipherMessage>)encryptMessage:(NSData*)paddedMessage{
-    [self assertOnSessionCipherDispatchQueue];
-    SessionRecord *sessionRecord = [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId];
+    SessionRecord *sessionRecord =
+        [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId protocolContext:protocolContext];
     SessionState *sessionState = sessionRecord.sessionState;
     ChainKey *chainKey = sessionState.senderChainKey;
     MessageKeys *messageKeys     = chainKey.messageKeys;
@@ -115,7 +89,8 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
 
     if (![self.identityKeyStore isTrustedIdentityKey:sessionState.remoteIdentityKey
                                          recipientId:self.recipientId
-                                           direction:TSMessageDirectionOutgoing]) {
+                                           direction:TSMessageDirectionOutgoing
+                                     protocolContext:protocolContext]) {
         DDLogWarn(
             @"%@ Previously known identity key for while encrypting for recipient: %@", self.tag, self.recipientId);
         @throw [NSException exceptionWithName:UntrustedIdentityKeyException
@@ -123,7 +98,9 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
                                      userInfo:@{}];
     }
 
-    [self.identityKeyStore saveRemoteIdentity:sessionState.remoteIdentityKey recipientId:self.recipientId];
+    [self.identityKeyStore saveRemoteIdentity:sessionState.remoteIdentityKey
+                                  recipientId:self.recipientId
+                              protocolContext:protocolContext];
 
     NSData *ciphertextBody = [AES_CBC encryptCBCMode:paddedMessage withKey:messageKeys.cipherKey withIV:messageKeys.iv];
 
@@ -153,27 +130,42 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     }
 
     [sessionState setSenderChainKey:[chainKey nextChainKey]];
-    [self.sessionStore storeSession:self.recipientId deviceId:self.deviceId session:sessionRecord];
-    
+    [self.sessionStore storeSession:self.recipientId
+                           deviceId:self.deviceId
+                            session:sessionRecord
+                    protocolContext:protocolContext];
+
     return cipherMessage;
 }
 
-- (NSData*)decrypt:(id<CipherMessage>)whisperMessage{
-    [self assertOnSessionCipherDispatchQueue];
+- (NSData *)decrypt:(id<CipherMessage>)whisperMessage protocolContext:(nullable id)protocolContext
+{
+    SPKAssert(whisperMessage);
+
     if ([whisperMessage isKindOfClass:[PreKeyWhisperMessage class]]) {
-        return [self decryptPreKeyWhisperMessage:(PreKeyWhisperMessage*)whisperMessage];
+        return
+            [self decryptPreKeyWhisperMessage:(PreKeyWhisperMessage *)whisperMessage protocolContext:protocolContext];
     } else{
-        return [self decryptWhisperMessage:whisperMessage];
+        return [self decryptWhisperMessage:whisperMessage protocolContext:protocolContext];
     }
 }
 
-- (NSData*)decryptPreKeyWhisperMessage:(PreKeyWhisperMessage*)preKeyWhisperMessage{
-    [self assertOnSessionCipherDispatchQueue];
-    SessionRecord *sessionRecord = [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId];
-    int unsignedPreKeyId         = [self.sessionBuilder processPrekeyWhisperMessage:preKeyWhisperMessage withSession:sessionRecord];
-    NSData *plaintext            = [self decryptWithSessionRecord:sessionRecord whisperMessage:preKeyWhisperMessage.message];
+- (NSData *)decryptPreKeyWhisperMessage:(PreKeyWhisperMessage *)preKeyWhisperMessage
+                        protocolContext:(nullable id)protocolContext
+{
+    SPKAssert(preKeyWhisperMessage);
 
-    [self.sessionStore storeSession:self.recipientId deviceId:self.deviceId session:sessionRecord];
+    SessionRecord *sessionRecord =
+        [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId protocolContext:protocolContext];
+    int unsignedPreKeyId         = [self.sessionBuilder processPrekeyWhisperMessage:preKeyWhisperMessage withSession:sessionRecord protocolContext:protocolContext];
+    NSData *plaintext = [self decryptWithSessionRecord:sessionRecord
+                                        whisperMessage:preKeyWhisperMessage.message
+                                       protocolContext:protocolContext];
+
+    [self.sessionStore storeSession:self.recipientId
+                           deviceId:self.deviceId
+                            session:sessionRecord
+                    protocolContext:protocolContext];
 
     // If there was an unsigned PreKey
     if (unsignedPreKeyId >= 0) {
@@ -183,15 +175,19 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     return plaintext;
 }
 
-- (NSData*)decryptWhisperMessage:(WhisperMessage*)message{
-    [self assertOnSessionCipherDispatchQueue];
+- (NSData *)decryptWhisperMessage:(WhisperMessage *)whisperMessage protocolContext:(nullable id)protocolContext
+{
+    SPKAssert(whisperMessage);
 
-    SessionRecord  *sessionRecord  = [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId];
-    NSData         *plaintext      = [self decryptWithSessionRecord:sessionRecord whisperMessage:message];
+    SessionRecord *sessionRecord =
+        [self.sessionStore loadSession:self.recipientId deviceId:self.deviceId protocolContext:protocolContext];
+    NSData *plaintext =
+        [self decryptWithSessionRecord:sessionRecord whisperMessage:whisperMessage protocolContext:protocolContext];
 
     if (![self.identityKeyStore isTrustedIdentityKey:sessionRecord.sessionState.remoteIdentityKey
                                          recipientId:self.recipientId
-                                           direction:TSMessageDirectionIncoming]) {
+                                           direction:TSMessageDirectionIncoming
+                                     protocolContext:protocolContext]) {
         DDLogWarn(
             @"%@ Previously known identity key for while decrypting from recipient: %@", self.tag, self.recipientId);
         @throw [NSException exceptionWithName:UntrustedIdentityKeyException
@@ -200,20 +196,30 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     }
 
     [self.identityKeyStore saveRemoteIdentity:sessionRecord.sessionState.remoteIdentityKey
-                                  recipientId:self.recipientId];
-    [self.sessionStore storeSession:self.recipientId deviceId:self.deviceId session:sessionRecord];
-    
+                                  recipientId:self.recipientId
+                              protocolContext:protocolContext];
+    [self.sessionStore storeSession:self.recipientId
+                           deviceId:self.deviceId
+                            session:sessionRecord
+                    protocolContext:protocolContext];
+
     return plaintext;
 }
 
 
--(NSData*)decryptWithSessionRecord:(SessionRecord*)sessionRecord whisperMessage:(WhisperMessage*)message{
-    [self assertOnSessionCipherDispatchQueue];
+- (NSData *)decryptWithSessionRecord:(SessionRecord *)sessionRecord
+                      whisperMessage:(WhisperMessage *)whisperMessage
+                     protocolContext:(nullable id)protocolContext
+{
+    SPKAssert(sessionRecord);
+    SPKAssert(whisperMessage);
+
     SessionState   *sessionState   = [sessionRecord sessionState];
     NSMutableArray *exceptions     = [NSMutableArray array];
     
     @try {
-        NSData *decryptedData = [self decryptWithSessionState:sessionState whisperMessage:message];
+        NSData *decryptedData =
+            [self decryptWithSessionState:sessionState whisperMessage:whisperMessage protocolContext:protocolContext];
         DDLogDebug(@"%@ successfully decrypted with current session state: %@", self.tag, sessionState);
         return decryptedData;
     }
@@ -232,7 +238,9 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     [[sessionRecord previousSessionStates]
         enumerateObjectsUsingBlock:^(SessionState *_Nonnull previousState, NSUInteger idx, BOOL *_Nonnull stop) {
             @try {
-                decryptedData = [self decryptWithSessionState:previousState whisperMessage:message];
+                decryptedData = [self decryptWithSessionState:previousState
+                                               whisperMessage:whisperMessage
+                                              protocolContext:protocolContext];
                 DDLogInfo(@"%@ successfully decrypted with PREVIOUS session state: %@", self.tag, previousState);
                 NSAssert(decryptedData != nil, @"Expected exception or non-nil data");
                 stateToPromoteIdx = idx;
@@ -252,7 +260,8 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
         return decryptedData;
     }
 
-    BOOL containsActiveSession = [self.sessionStore containsSession:self.recipientId deviceId:self.deviceId];
+    BOOL containsActiveSession =
+        [self.sessionStore containsSession:self.recipientId deviceId:self.deviceId protocolContext:protocolContext];
     DDLogError(@"%@ No valid session for recipient: %@ containsActiveSession: %@, previousStates: %lu",
         self.tag,
         self.recipientId,
@@ -273,35 +282,51 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     }
 }
 
--(NSData*)decryptWithSessionState:(SessionState*)sessionState whisperMessage:(WhisperMessage*)message{
-    [self assertOnSessionCipherDispatchQueue];
+- (NSData *)decryptWithSessionState:(SessionState *)sessionState
+                     whisperMessage:(WhisperMessage *)whisperMessage
+                    protocolContext:(nullable id)protocolContext
+{
+    SPKAssert(sessionState);
+    SPKAssert(whisperMessage);
+
     if (![sessionState hasSenderChain]) {
         @throw [NSException exceptionWithName:InvalidMessageException reason:@"Uninitialized session!" userInfo:nil];
     }
-    
-    if (message.version != sessionState.version) {
-        @throw [NSException exceptionWithName:InvalidMessageException reason:[NSString stringWithFormat:@"Got message version %d but was expecting %d", message.version, sessionState.version] userInfo:nil];
+
+    if (whisperMessage.version != sessionState.version) {
+        @throw [NSException exceptionWithName:InvalidMessageException
+                                       reason:[NSString stringWithFormat:@"Got message version %d but was expecting %d",
+                                                        whisperMessage.version,
+                                                        sessionState.version]
+                                     userInfo:nil];
     }
 
-    int messageVersion       = message.version;
-    NSData *theirEphemeral   = message.senderRatchetKey.removeKeyType;
-    int counter              = message.counter;
+    int messageVersion = whisperMessage.version;
+    NSData *theirEphemeral = whisperMessage.senderRatchetKey.removeKeyType;
+    int counter = whisperMessage.counter;
     ChainKey *chainKey       = [self getOrCreateChainKeys:sessionState theirEphemeral:theirEphemeral];
     SPKAssert(chainKey);
     MessageKeys *messageKeys = [self getOrCreateMessageKeysForSession:sessionState theirEphemeral:theirEphemeral chainKey:chainKey counter:counter];
     SPKAssert(messageKeys);
 
-    [message verifyMacWithVersion:messageVersion senderIdentityKey:sessionState.remoteIdentityKey receiverIdentityKey:sessionState.localIdentityKey macKey:messageKeys.macKey];
-    
-    NSData *plaintext = [AES_CBC decryptCBCMode:message.cipherText withKey:messageKeys.cipherKey withIV:messageKeys.iv];
-    
+    [whisperMessage verifyMacWithVersion:messageVersion
+                       senderIdentityKey:sessionState.remoteIdentityKey
+                     receiverIdentityKey:sessionState.localIdentityKey
+                                  macKey:messageKeys.macKey];
+
+    NSData *plaintext =
+        [AES_CBC decryptCBCMode:whisperMessage.cipherText withKey:messageKeys.cipherKey withIV:messageKeys.iv];
+
     [sessionState clearUnacknowledgedPreKeyMessage];
     
     return plaintext;
 }
 
-- (ChainKey*)getOrCreateChainKeys:(SessionState*)sessionState theirEphemeral:(NSData*)theirEphemeral{
-    [self assertOnSessionCipherDispatchQueue];
+- (ChainKey *)getOrCreateChainKeys:(SessionState *)sessionState
+                    theirEphemeral:(NSData *)theirEphemeral
+{
+    SPKAssert(sessionState);
+    SPKAssert(theirEphemeral);
     SPKAssert(theirEphemeral.length == ECCKeyLength);
 
     @try {
@@ -339,8 +364,15 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     }
 }
 
-- (MessageKeys*)getOrCreateMessageKeysForSession:(SessionState*)sessionState theirEphemeral:(NSData*)theirEphemeral chainKey:(ChainKey*)chainKey counter:(int)counter{
-    [self assertOnSessionCipherDispatchQueue];
+- (MessageKeys *)getOrCreateMessageKeysForSession:(SessionState *)sessionState
+                                   theirEphemeral:(NSData *)theirEphemeral
+                                         chainKey:(ChainKey *)chainKey
+                                          counter:(int)counter
+{
+    SPKAssert(sessionState);
+    SPKAssert(theirEphemeral);
+    SPKAssert(chainKey);
+
     if (chainKey.index > counter) {
         if ([sessionState hasMessageKeys:theirEphemeral counter:counter]) {
             return [sessionState removeMessageKeys:theirEphemeral counter:counter];
@@ -388,10 +420,12 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
 }
 
 
-- (int)remoteRegistrationId{
-    [self assertOnSessionCipherDispatchQueue];
-    SessionRecord *record = [self.sessionStore loadSession:self.recipientId deviceId:_deviceId];
-    
+- (int)remoteRegistrationId:(nullable id)protocolContext
+{
+
+    SessionRecord *record =
+        [self.sessionStore loadSession:self.recipientId deviceId:_deviceId protocolContext:protocolContext];
+
     if (!record) {
         @throw [NSException exceptionWithName:NoSessionException reason:@"Trying to get registration Id of a non-existing session." userInfo:nil];
     }
@@ -399,10 +433,12 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
     return record.sessionState.remoteRegistrationId;
 }
 
-- (int)sessionVersion{
-    [self assertOnSessionCipherDispatchQueue];
-    SessionRecord *record = [self.sessionStore loadSession:self.recipientId deviceId:_deviceId];
-    
+- (int)sessionVersion:(nullable id)protocolContext
+{
+
+    SessionRecord *record =
+        [self.sessionStore loadSession:self.recipientId deviceId:_deviceId protocolContext:protocolContext];
+
     if (!record) {
         @throw [NSException exceptionWithName:NoSessionException reason:@"Trying to get the version of a non-existing session." userInfo:nil];
     }
@@ -423,3 +459,5 @@ static dispatch_queue_t _sessionCipherDispatchQueue;
 }
 
 @end
+
+NS_ASSUME_NONNULL_END

AxolotlKit/Classes/Sessions/SessionBuilder.h

@@ -1,5 +1,5 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
 #import "AxolotlStore.h"
@@ -10,6 +10,8 @@
 #import "SignedPreKeyStore.h"
 #import <Foundation/Foundation.h>
 
+NS_ASSUME_NONNULL_BEGIN
+
 @class PreKeyWhisperMessage;
 
 extern const int kPreKeyOfLastResortId;
@@ -27,8 +29,12 @@ extern const int kPreKeyOfLastResortId;
                          recipientId:(NSString *)recipientId
                             deviceId:(int)deviceId;
 
-- (void)processPrekeyBundle:(PreKeyBundle *)preKeyBundle;
+- (void)processPrekeyBundle:(PreKeyBundle *)preKeyBundle protocolContext:(nullable id)protocolContext;
+
 - (int)processPrekeyWhisperMessage:(PreKeyWhisperMessage *)message
-                       withSession:(SessionRecord *)sessionRecord;
+                       withSession:(SessionRecord *)sessionRecord
+                   protocolContext:(nullable id)protocolContext;
 
 @end
+
+NS_ASSUME_NONNULL_END

AxolotlKit/Classes/Sessions/SessionBuilder.m

@@ -1,24 +1,22 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
+#import "SessionBuilder.h"
+#import "AliceAxolotlParameters.h"
 #import "AxolotlExceptions.h"
 #import "AxolotlParameters.h"
-#import "AliceAxolotlParameters.h"
-#import "BobAxolotlParameters.h"
-
-#import "NSData+keyVersionByte.h"
-
 #import "AxolotlStore.h"
-#import "SessionState.h"
-#import "SessionBuilder.h"
+#import "BobAxolotlParameters.h"
+#import "NSData+keyVersionByte.h"
 #import "PreKeyWhisperMessage.h"
+#import "PrekeyBundle.h"
 #import "RatchetingSession.h"
-
+#import "SessionState.h"
 #import <Curve25519Kit/Curve25519.h>
 #import <Curve25519Kit/Ed25519.h>
 
-#import "PrekeyBundle.h"
+NS_ASSUME_NONNULL_BEGIN
 
 #define CURRENT_VERSION 3
 #define MINUMUM_VERSION 3
@@ -69,32 +67,40 @@ const int kPreKeyOfLastResortId = 0xFFFFFF;
     return self;
 }
 
-- (void)processPrekeyBundle:(PreKeyBundle*)preKeyBundle{
+- (void)processPrekeyBundle:(PreKeyBundle *)preKeyBundle protocolContext:(nullable id)protocolContext
+{
+    SPKAssert(preKeyBundle);
+
     NSData *theirIdentityKey  = preKeyBundle.identityKey.removeKeyType;
     NSData *theirSignedPreKey = preKeyBundle.signedPreKeyPublic.removeKeyType;
-    
-    if (![self.identityStore isTrustedIdentityKey:theirIdentityKey recipientId:self.recipientId direction:TSMessageDirectionOutgoing]) {
+
+    if (![self.identityStore isTrustedIdentityKey:theirIdentityKey
+                                      recipientId:self.recipientId
+                                        direction:TSMessageDirectionOutgoing
+                                  protocolContext:protocolContext]) {
         @throw [NSException exceptionWithName:UntrustedIdentityKeyException reason:@"Identity key is not valid" userInfo:@{}];
     }
 
     if (![Ed25519 verifySignature:preKeyBundle.signedPreKeySignature publicKey:theirIdentityKey data:preKeyBundle.signedPreKeyPublic]) {
         @throw [NSException exceptionWithName:InvalidKeyException reason:@"KeyIsNotValidlySigned" userInfo:nil];
     }
-    
-    SessionRecord *sessionRecord       = [self.sessionStore loadSession:self.recipientId deviceId:preKeyBundle.deviceId];
+
+    SessionRecord *sessionRecord =
+        [self.sessionStore loadSession:self.recipientId deviceId:preKeyBundle.deviceId protocolContext:protocolContext];
     ECKeyPair     *ourBaseKey          = [Curve25519 generateKeyPair];
     NSData        *theirOneTimePreKey  = preKeyBundle.preKeyPublic.removeKeyType;
     int           theirOneTimePreKeyId = preKeyBundle.preKeyId;
     int           theirSignedPreKeyId  = preKeyBundle.signedPreKeyId;
-    
-    
-    AliceAxolotlParameters *params = [[AliceAxolotlParameters alloc] initWithIdentityKey:[self.identityStore identityKeyPair]
-                                                                        theirIdentityKey:theirIdentityKey
-                                                                              ourBaseKey:ourBaseKey
-                                                                       theirSignedPreKey:theirSignedPreKey
-                                                                      theirOneTimePreKey:theirOneTimePreKey
-                                                                         theirRatchetKey:theirSignedPreKey];
-    
+
+
+    AliceAxolotlParameters *params =
+        [[AliceAxolotlParameters alloc] initWithIdentityKey:[self.identityStore identityKeyPair:protocolContext]
+                                           theirIdentityKey:theirIdentityKey
+                                                 ourBaseKey:ourBaseKey
+                                          theirSignedPreKey:theirSignedPreKey
+                                         theirOneTimePreKey:theirOneTimePreKey
+                                            theirRatchetKey:theirSignedPreKey];
+
     if (!sessionRecord.isFresh) {
         [sessionRecord archiveCurrentState];
     }
@@ -104,13 +110,14 @@ const int kPreKeyOfLastResortId = 0xFFFFFF;
     DDLogInfo(@"setUnacknowledgedPreKeyMessage for: %@ with preKeyId: %d", self.recipientId, theirOneTimePreKeyId);
 
     [sessionRecord.sessionState setUnacknowledgedPreKeyMessage:theirOneTimePreKeyId signedPreKey:theirSignedPreKeyId baseKey:ourBaseKey.publicKey];
-    [sessionRecord.sessionState setLocalRegistrationId:self.identityStore.localRegistrationId];
+    [sessionRecord.sessionState setLocalRegistrationId:[self.identityStore localRegistrationId:protocolContext]];
     [sessionRecord.sessionState setRemoteRegistrationId:preKeyBundle.registrationId];
     [sessionRecord.sessionState setAliceBaseKey:ourBaseKey.publicKey];
 
     // Saving invalidates any existing sessions, so be sure to save *before* storing the new session.
-    BOOL previousIdentityExisted =
-        [self.identityStore saveRemoteIdentity:theirIdentityKey recipientId:self.recipientId];
+    BOOL previousIdentityExisted = [self.identityStore saveRemoteIdentity:theirIdentityKey
+                                                              recipientId:self.recipientId
+                                                          protocolContext:protocolContext];
     if (previousIdentityExisted) {
         DDLogInfo(@"%@ PKBundle removing previous session states for changed identity for recipient:%@",
             self.tag,
@@ -118,15 +125,26 @@ const int kPreKeyOfLastResortId = 0xFFFFFF;
         [sessionRecord removePreviousSessionStates];
     }
 
-    [self.sessionStore storeSession:self.recipientId deviceId:self.deviceId session:sessionRecord];
+    [self.sessionStore storeSession:self.recipientId
+                           deviceId:self.deviceId
+                            session:sessionRecord
+                    protocolContext:protocolContext];
 }
 
-- (int)processPrekeyWhisperMessage:(PreKeyWhisperMessage*)message withSession:(SessionRecord*)sessionRecord{
+- (int)processPrekeyWhisperMessage:(PreKeyWhisperMessage *)message
+                       withSession:(SessionRecord *)sessionRecord
+                   protocolContext:(nullable id)protocolContext
+{
+    SPKAssert(message);
+    SPKAssert(sessionRecord);
 
     int    messageVersion    = message.version;
     NSData *theirIdentityKey = message.identityKey.removeKeyType;
 
-    if (![self.identityStore isTrustedIdentityKey:theirIdentityKey recipientId:self.recipientId direction:TSMessageDirectionIncoming]) {
+    if (![self.identityStore isTrustedIdentityKey:theirIdentityKey
+                                      recipientId:self.recipientId
+                                        direction:TSMessageDirectionIncoming
+                                  protocolContext:protocolContext]) {
         @throw [NSException exceptionWithName:UntrustedIdentityKeyException reason:@"There is a previously known identity key." userInfo:@{}];
     }
     
@@ -134,20 +152,24 @@ const int kPreKeyOfLastResortId = 0xFFFFFF;
     
     switch (messageVersion) {
         case 3:
-            unSignedPrekeyId = [self processPrekeyV3:message withSession:sessionRecord];
+            unSignedPrekeyId = [self processPrekeyV3:message withSession:sessionRecord protocolContext:protocolContext];
             break;
         default:
             @throw [NSException exceptionWithName:InvalidVersionException reason:@"Trying to initialize with unknown version" userInfo:@{}];
             break;
     }
 
-    [self.identityStore saveRemoteIdentity:theirIdentityKey recipientId:self.recipientId];
+    [self.identityStore saveRemoteIdentity:theirIdentityKey
+                               recipientId:self.recipientId
+                           protocolContext:protocolContext];
 
     return unSignedPrekeyId;
 }
 
-- (int)processPrekeyV3:(PreKeyWhisperMessage*)message withSession:(SessionRecord*)sessionRecord{
-    
+- (int)processPrekeyV3:(PreKeyWhisperMessage *)message
+           withSession:(SessionRecord *)sessionRecord
+       protocolContext:(nullable id)protocolContext
+{
     NSData *baseKey = message.baseKey.removeKeyType;
     
     if ([sessionRecord hasSessionState:message.version baseKey:baseKey]) {
@@ -163,20 +185,21 @@ const int kPreKeyOfLastResortId = 0xFFFFFF;
         DDLogWarn(@"%@ Processing PreKey message which had no one-time prekey.", self.tag);
     }
 
-    BobAxolotlParameters *params = [[BobAxolotlParameters alloc] initWithMyIdentityKeyPair:self.identityStore.identityKeyPair
-                                                                          theirIdentityKey:message.identityKey.removeKeyType
-                                                                           ourSignedPrekey:ourSignedPrekey
-                                                                             ourRatchetKey:ourSignedPrekey
-                                                                          ourOneTimePrekey:ourOneTimePreKey
-                                                                              theirBaseKey:baseKey];
-    
+    BobAxolotlParameters *params =
+        [[BobAxolotlParameters alloc] initWithMyIdentityKeyPair:[self.identityStore identityKeyPair:protocolContext]
+                                               theirIdentityKey:message.identityKey.removeKeyType
+                                                ourSignedPrekey:ourSignedPrekey
+                                                  ourRatchetKey:ourSignedPrekey
+                                               ourOneTimePrekey:ourOneTimePreKey
+                                                   theirBaseKey:baseKey];
+
     if (!sessionRecord.isFresh) {
         [sessionRecord archiveCurrentState];
     }
     
     [RatchetingSession initializeSession:sessionRecord.sessionState sessionVersion:message.version BobParameters:params];
     
-    [sessionRecord.sessionState setLocalRegistrationId:self.identityStore.localRegistrationId];
+    [sessionRecord.sessionState setLocalRegistrationId:[self.identityStore localRegistrationId:protocolContext]];
     [sessionRecord.sessionState setRemoteRegistrationId:message.registrationId];
     [sessionRecord.sessionState setAliceBaseKey:baseKey];
 
@@ -201,3 +224,5 @@ const int kPreKeyOfLastResortId = 0xFFFFFF;
 }
 
 @end
+
+NS_ASSUME_NONNULL_END

AxolotlKit/Classes/State/AxolotlStore.h

@@ -1,16 +1,14 @@
 //
-//  SessionBuilder.h
-//  AxolotlKit
-//
-//  Created by Frederic Jacobs on 23/07/14.
-//  Copyright (c) 2014 Frederic Jacobs. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
-#import <Foundation/Foundation.h>
 #import "IdentityKeyStore.h"
 #import "PreKeyStore.h"
 #import "SessionStore.h"
 #import "SignedPreKeyStore.h"
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
 
 /**
  *  The Session Store defines the interface of the storage of sesssions.
@@ -19,3 +17,5 @@
 @protocol AxolotlStore <SessionStore, IdentityKeyStore, PreKeyStore, SessionStore, SignedPreKeyStore>
 
 @end
+
+NS_ASSUME_NONNULL_END

AxolotlKit/Classes/State/IdentityKeyStore.h

@@ -1,5 +1,5 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
 #import <Foundation/Foundation.h>
@@ -14,10 +14,12 @@ typedef NS_ENUM(NSInteger, TSMessageDirection) {
     TSMessageDirectionOutgoing
 };
 
+// See a discussion of the protocolContext in SessionCipher.h.
 @protocol IdentityKeyStore <NSObject>
 
-- (nullable ECKeyPair *)identityKeyPair;
-- (int)localRegistrationId;
+- (nullable ECKeyPair *)identityKeyPair:(nullable id)protocolContext;
+
+- (int)localRegistrationId:(nullable id)protocolContext;
 
 /**
  * Record a recipients identity key
@@ -28,7 +30,9 @@ typedef NS_ENUM(NSInteger, TSMessageDirection) {
  * @returns YES if we are replacing an existing known identity key for recipientId.
  *          NO  if there was no previously stored identity key for the recipient.
  */
-- (BOOL)saveRemoteIdentity:(NSData *)identityKey recipientId:(NSString *)recipientId;
+- (BOOL)saveRemoteIdentity:(NSData *)identityKey
+               recipientId:(NSString *)recipientId
+           protocolContext:(nullable id)protocolContext;
 
 /**
  * @param   identityKey key data used to identify the recipient
@@ -41,7 +45,8 @@ typedef NS_ENUM(NSInteger, TSMessageDirection) {
  */
 - (BOOL)isTrustedIdentityKey:(NSData *)identityKey
                  recipientId:(NSString *)recipientId
-                   direction:(TSMessageDirection)direction;
+                   direction:(TSMessageDirection)direction
+             protocolContext:(nullable id)protocolContext;
 
 @end
 

AxolotlKit/Classes/State/PreKeyStore.h

@@ -1,22 +1,22 @@
 //
-//  PreKeyStore.h
-//  AxolotlKit
-//
-//  Created by Frederic Jacobs on 12/10/14.
-//  Copyright (c) 2014 Frederic Jacobs. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
-#import <Foundation/Foundation.h>
 #import "PreKeyRecord.h"
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
 
 @protocol PreKeyStore <NSObject>
 
-- (PreKeyRecord*)loadPreKey:(int)preKeyId;
+- (PreKeyRecord *)loadPreKey:(int)preKeyId;
 
-- (void)storePreKey:(int)preKeyId preKeyRecord:(PreKeyRecord*)record;
+- (void)storePreKey:(int)preKeyId preKeyRecord:(PreKeyRecord *)record;
 
 - (BOOL)containsPreKey:(int)preKeyId;
 
 - (void)removePreKey:(int)preKeyId;
 
 @end
+
+NS_ASSUME_NONNULL_END

AxolotlKit/Classes/State/SessionStore.h

@@ -1,18 +1,15 @@
 //
-//  SessionStore.h
-//  AxolotlKit
-//
-//  Created by Frederic Jacobs on 12/10/14.
-//  Copyright (c) 2014 Frederic Jacobs. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
-#import <Foundation/Foundation.h>
 #import "SessionRecord.h"
+#import <Foundation/Foundation.h>
 
+NS_ASSUME_NONNULL_BEGIN
+
+// See a discussion of the protocolContext in SessionCipher.h.
 @protocol SessionStore <NSObject>
 
-
-
 /**
  *  Returns a copy of the SessionRecord corresponding to the recipientId + deviceId tuple or a new SessionRecord if one does not currently exist.
  *
@@ -21,17 +18,28 @@
  *
  *  @return a copy of the SessionRecord corresponding to the recipientId + deviceId tuple.
  */
+- (SessionRecord *)loadSession:(NSString *)contactIdentifier
+                      deviceId:(int)deviceId
+               protocolContext:(nullable id)protocolContext;
 
-- (SessionRecord*)loadSession:(NSString*)contactIdentifier deviceId:(int)deviceId;
+// Deprecated.
+- (NSArray *)subDevicesSessions:(NSString *)contactIdentifier protocolContext:(nullable id)protocolContext;
 
-- (NSArray*)subDevicesSessions:(NSString*)contactIdentifier;
+- (void)storeSession:(NSString *)contactIdentifier
+            deviceId:(int)deviceId
+             session:(SessionRecord *)session
+     protocolContext:(nullable id)protocolContext;
 
-- (void)storeSession:(NSString*)contactIdentifier deviceId:(int)deviceId session:(SessionRecord*)session;
+- (BOOL)containsSession:(NSString *)contactIdentifier
+               deviceId:(int)deviceId
+        protocolContext:(nullable id)protocolContext;
 
-- (BOOL)containsSession:(NSString*)contactIdentifier deviceId:(int)deviceId;
+- (void)deleteSessionForContact:(NSString *)contactIdentifier
+                       deviceId:(int)deviceId
+                protocolContext:(nullable id)protocolContext;
 
-- (void)deleteSessionForContact:(NSString*)contactIdentifier deviceId:(int)deviceId;
-
-- (void)deleteAllSessionsForContact:(NSString*)contactIdentifier;
+- (void)deleteAllSessionsForContact:(NSString *)contactIdentifier protocolContext:(nullable id)protocolContext;
 
 @end
+
+NS_ASSUME_NONNULL_END

AxolotlKit/Classes/State/SignedPreKeyStore.h

@@ -1,9 +1,9 @@
 //
-//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
-#import <Foundation/Foundation.h>
 #import "SignedPrekeyRecord.h"
+#import <Foundation/Foundation.h>
 
 NS_ASSUME_NONNULL_BEGIN