Narrow the scope of code run on SessionCipher queue

And run all non-cipher code on the main thread.

Note: Running encryption on the sessionCipher queue is more about
serializing access to session mutations than it is about any performance
gains.

// FREEBIE

src/Messages/TSMessagesManager.m

@@ -1,5 +1,6 @@
-//  Created by Frederic Jacobs on 11/11/14.
+//
-//  Copyright (c) 2014 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//
 
 #import "TSMessagesManager.h"
 #import "ContactsManagerProtocol.h"
@@ -11,6 +12,8 @@
 #import "OWSDisappearingConfigurationUpdateInfoMessage.h"
 #import "OWSDisappearingMessagesConfiguration.h"
 #import "OWSDisappearingMessagesJob.h"
+#import "OWSDispatch.h"
+#import "OWSError.h"
 #import "OWSIncomingSentMessageTranscript.h"
 #import "OWSMessageSender.h"
 #import "OWSReadReceiptsProcessor.h"
@@ -100,14 +103,31 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)handleReceivedEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
 {
+    OWSAssert([NSThread isMainThread]);
     @try {
         switch (envelope.type) {
-            case OWSSignalServiceProtosEnvelopeTypeCiphertext:
+            case OWSSignalServiceProtosEnvelopeTypeCiphertext: {
-                [self handleSecureMessage:envelope];
+                [self handleSecureMessageAsync:envelope
+                                    completion:^(NSError *_Nullable error) {
+                                        DDLogDebug(@"%@ handled secure message.", self.tag);
+                                        if (error) {
+                                            DDLogError(
+                                                @"%@ handling secure message failed with error: %@", self.tag, error);
+                                        }
+                                    }];
                 break;
-            case OWSSignalServiceProtosEnvelopeTypePrekeyBundle:
+            }
-                [self handlePreKeyBundle:envelope];
+            case OWSSignalServiceProtosEnvelopeTypePrekeyBundle: {
+                [self handlePreKeyBundleAsync:envelope
+                                   completion:^(NSError *_Nullable error) {
+                                       DDLogDebug(@"%@ handled pre-key bundle", self.tag);
+                                       if (error) {
+                                           DDLogError(
+                                               @"%@ handling pre-key bundle failed with error: %@", self.tag, error);
+                                       }
+                                   }];
                 break;
+            }
             case OWSSignalServiceProtosEnvelopeTypeReceipt:
                 DDLogInfo(@"Received a delivery receipt");
                 [self handleDeliveryReceipt:envelope];
@@ -132,6 +152,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)handleDeliveryReceipt:(OWSSignalServiceProtosEnvelope *)envelope
 {
+    OWSAssert([NSThread isMainThread]);
     [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
         TSInteraction *interaction =
             [TSInteraction interactionForTimestamp:envelope.timestamp withTransaction:transaction];
@@ -144,8 +165,10 @@ NS_ASSUME_NONNULL_BEGIN
     }];
 }
 
-- (void)handleSecureMessage:(OWSSignalServiceProtosEnvelope *)messageEnvelope
+- (void)handleSecureMessageAsync:(OWSSignalServiceProtosEnvelope *)messageEnvelope
+                      completion:(void (^)(NSError *_Nullable error))completion
 {
+    OWSAssert([NSThread isMainThread]);
     @synchronized(self) {
         TSStorageManager *storageManager = [TSStorageManager sharedManager];
         NSString *recipientId = messageEnvelope.source;
@@ -167,7 +190,9 @@ NS_ASSUME_NONNULL_BEGIN
             return;
         }
 
+        dispatch_async([OWSDispatch sessionCipher], ^{
             NSData *plaintextData;
+
             @try {
                 WhisperMessage *message = [[WhisperMessage alloc] initWithData:encryptedData];
                 SessionCipher *cipher = [[SessionCipher alloc] initWithSessionStore:storageManager
@@ -179,16 +204,28 @@ NS_ASSUME_NONNULL_BEGIN
 
                 plaintextData = [[cipher decrypt:message] removePadding];
             } @catch (NSException *exception) {
+                dispatch_async(dispatch_get_main_queue(), ^{
                     [self processException:exception envelope:messageEnvelope];
+                    NSString *errorDescription =
+                        [NSString stringWithFormat:@"Exception while decrypting: %@", exception.description];
+                    NSError *error = OWSErrorWithCodeDescription(OWSErrorCodeFailedToDecryptMessage, errorDescription);
+                    completion(error);
+                });
                 return;
             }
 
+            dispatch_async(dispatch_get_main_queue(), ^{
                 [self handleEnvelope:messageEnvelope plaintextData:plaintextData];
+                completion(nil);
+            });
+        });
     }
 }
 
-- (void)handlePreKeyBundle:(OWSSignalServiceProtosEnvelope *)preKeyEnvelope
+- (void)handlePreKeyBundleAsync:(OWSSignalServiceProtosEnvelope *)preKeyEnvelope
+                     completion:(void (^)(NSError *_Nullable error))completion
 {
+    OWSAssert([NSThread isMainThread]);
     @synchronized(self) {
         TSStorageManager *storageManager = [TSStorageManager sharedManager];
         NSString *recipientId = preKeyEnvelope.source;
@@ -201,6 +238,7 @@ NS_ASSUME_NONNULL_BEGIN
             return;
         }
 
+        dispatch_async([OWSDispatch sessionCipher], ^{
             NSData *plaintextData;
             @try {
                 PreKeyWhisperMessage *message = [[PreKeyWhisperMessage alloc] initWithData:encryptedData];
@@ -213,16 +251,25 @@ NS_ASSUME_NONNULL_BEGIN
 
                 plaintextData = [[cipher decrypt:message] removePadding];
             } @catch (NSException *exception) {
+                dispatch_async(dispatch_get_main_queue(), ^{
                     [self processException:exception envelope:preKeyEnvelope];
+                    NSString *errorDescription = [NSString stringWithFormat:@"Exception while decrypting PreKey Bundle: %@", exception.description];
+                    NSError *error = OWSErrorWithCodeDescription(OWSErrorCodeFailedToDecryptMessage, errorDescription);
+                    completion(error);
+                });
                 return;
             }
 
+            dispatch_async(dispatch_get_main_queue(), ^{
                 [self handleEnvelope:preKeyEnvelope plaintextData:plaintextData];
+            });
+        });
     }
 }
 
 - (void)handleEnvelope:(OWSSignalServiceProtosEnvelope *)envelope plaintextData:(NSData *)plaintextData
 {
+    OWSAssert([NSThread isMainThread]);
     if (envelope.hasContent) {
         OWSSignalServiceProtosContent *content = [OWSSignalServiceProtosContent parseFromData:plaintextData];
         if (content.hasSyncMessage) {
@@ -244,6 +291,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)handleIncomingEnvelope:(OWSSignalServiceProtosEnvelope *)incomingEnvelope
                withDataMessage:(OWSSignalServiceProtosDataMessage *)dataMessage
 {
+    OWSAssert([NSThread isMainThread]);
     if (dataMessage.hasGroup) {
         __block BOOL ignoreMessage = NO;
         [self.dbConnection readWithBlock:^(YapDatabaseReadTransaction *transaction) {
@@ -282,6 +330,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)handleReceivedGroupAvatarUpdateWithEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
                                         dataMessage:(OWSSignalServiceProtosDataMessage *)dataMessage
 {
+    OWSAssert([NSThread isMainThread]);
     TSGroupThread *groupThread = [TSGroupThread getOrCreateThreadWithGroupIdData:dataMessage.group.id];
     OWSAttachmentsProcessor *attachmentsProcessor =
         [[OWSAttachmentsProcessor alloc] initWithAttachmentProtos:@[ dataMessage.group.avatar ]
@@ -309,6 +358,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)handleReceivedMediaWithEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
                             dataMessage:(OWSSignalServiceProtosDataMessage *)dataMessage
 {
+    OWSAssert([NSThread isMainThread]);
     TSThread *thread = [self threadForEnvelope:envelope dataMessage:dataMessage];
     OWSAttachmentsProcessor *attachmentsProcessor =
         [[OWSAttachmentsProcessor alloc] initWithAttachmentProtos:dataMessage.attachments
@@ -339,6 +389,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)handleIncomingEnvelope:(OWSSignalServiceProtosEnvelope *)messageEnvelope
                withSyncMessage:(OWSSignalServiceProtosSyncMessage *)syncMessage
 {
+    OWSAssert([NSThread isMainThread]);
     if (syncMessage.hasSent) {
         DDLogInfo(@"%@ Received `sent` syncMessage, recording message transcript.", self.tag);
         OWSIncomingSentMessageTranscript *transcript =
@@ -409,6 +460,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)handleEndSessionMessageWithEnvelope:(OWSSignalServiceProtosEnvelope *)endSessionEnvelope
                                 dataMessage:(OWSSignalServiceProtosDataMessage *)dataMessage
 {
+    OWSAssert([NSThread isMainThread]);
     [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
         TSContactThread *thread =
             [TSContactThread getOrCreateThreadWithContactId:endSessionEnvelope.source transaction:transaction];
@@ -427,6 +479,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)handleExpirationTimerUpdateMessageWithEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
                                            dataMessage:(OWSSignalServiceProtosDataMessage *)dataMessage
 {
+    OWSAssert([NSThread isMainThread]);
     TSThread *thread = [self threadForEnvelope:envelope dataMessage:dataMessage];
 
     OWSDisappearingMessagesConfiguration *disappearingMessagesConfiguration;
@@ -459,6 +512,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)handleReceivedTextMessageWithEnvelope:(OWSSignalServiceProtosEnvelope *)textMessageEnvelope
                                   dataMessage:(OWSSignalServiceProtosDataMessage *)dataMessage
 {
+    OWSAssert([NSThread isMainThread]);
     [self handleReceivedEnvelope:textMessageEnvelope withDataMessage:dataMessage attachmentIds:@[]];
 }
 
@@ -466,6 +520,7 @@ NS_ASSUME_NONNULL_BEGIN
                               withDataMessage:(OWSSignalServiceProtosDataMessage *)dataMessage
                                 attachmentIds:(NSArray<NSString *> *)attachmentIds
 {
+    OWSAssert([NSThread isMainThread]);
     uint64_t timestamp = envelope.timestamp;
     NSString *body = dataMessage.body;
     NSData *groupId = dataMessage.hasGroup ? dataMessage.group.id : nil;
@@ -595,6 +650,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)processException:(NSException *)exception envelope:(OWSSignalServiceProtosEnvelope *)envelope
 {
+    OWSAssert([NSThread isMainThread]);
     DDLogError(@"%@ Got exception: %@ of type: %@", self.tag, exception.description, exception.name);
     [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
       TSErrorMessage *errorMessage;

src/Network/WebSockets/TSSocketManager.m

@@ -219,9 +219,7 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
     [self keepAliveBackground];
 
     if ([message.path isEqualToString:@"/api/v1/message"] && [message.verb isEqualToString:@"PUT"]) {
-        // SessionCipher state is mutated in a couple places, during decryption, in handleReceivedEnvelope.
+
-        // To ensure consistent state, that decryption mutation must occur on the same queue as encryption.
-        dispatch_async([OWSDispatch sessionCipher], ^{
         NSData *decryptedPayload =
             [Cryptography decryptAppleMessagePayload:message.body withSignalingKey:TSStorageManager.signalingKey];
 
@@ -233,7 +231,7 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
         OWSSignalServiceProtosEnvelope *envelope = [OWSSignalServiceProtosEnvelope parseFromData:decryptedPayload];
 
         [[TSMessagesManager sharedManager] handleReceivedEnvelope:envelope];
-        });
+
     } else {
         DDLogWarn(@"Unsupported WebSocket Request");
     }

src/Util/OWSDispatch.h

@@ -13,7 +13,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 /**
  * Signal protocol session state must be coordinated on a serial queue. This is sometimes used synchronously,
- * so it's recommended that you avoid dispatching sync *from* this queue to avoid deadlock.
+ * so never dispatching sync *from* this queue to avoid deadlock.
  */
 + (dispatch_queue_t)sessionCipher;
 

src/Util/OWSError.h

@@ -1,4 +1,6 @@
-//  Copyright © 2016 Open Whisper Systems. All rights reserved.
+//
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -14,6 +16,7 @@ typedef NS_ENUM(NSInteger, OWSErrorCode) {
     OWSErrorCodeUntrustedIdentityKey = 25,
     OWSErrorCodeFailedToSendOutgoingMessage = 30,
     OWSErrorCodeFailedToDecryptMessage = 100,
+    OWSErrorCodeFailedToEncryptMessage = 110,
     OWSErrorCodeSignalServiceFailure = 1001,
     OWSErrorCodeUserError = 2001,
 };