DRY up decryption logic.

// FREEBIE

SignalServiceKit/src/Messages/TSMessagesManager.h

@@ -34,10 +34,14 @@ typedef void (^MessageManagerCompletionBlock)();
 
 // decryptEnvelope: can be called from any thread.
 // successBlock & failureBlock may be called on any thread.
+//
+// Exactly one of successBlock & failureBlock will be called,
+// once.
 - (void)decryptEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
            successBlock:(DecryptSuccessBlock)successBlock
            failureBlock:(DecryptFailureBlock)failureBlock;
 
+// processEnvelope: can be called from any thread.
 - (void)processEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
           plaintextData:(NSData *_Nullable)plaintextData
             transaction:(YapDatabaseReadWriteTransaction *)transaction;

SignalServiceKit/src/Messages/TSMessagesManager.m

@@ -346,6 +346,7 @@ const NSUInteger kIncomingMessageBatchSize = 10;
                 // Return to avoid double-acknowledging.
                 return;
             }
+            // These message types don't have a payload to decrypt.
             case OWSSignalServiceProtosEnvelopeTypeReceipt:
             case OWSSignalServiceProtosEnvelopeTypeKeyExchange:
             case OWSSignalServiceProtosEnvelopeTypeUnknown:
@@ -364,75 +365,78 @@ const NSUInteger kIncomingMessageBatchSize = 10;
     failureBlock();
 }
 
-- (void)decryptSecureMessage:(OWSSignalServiceProtosEnvelope *)messageEnvelope
+- (void)decryptSecureMessage:(OWSSignalServiceProtosEnvelope *)envelope
                 successBlock:(DecryptSuccessBlock)successBlock
                 failureBlock:(void (^)(NSError *_Nullable error))failureBlock
 {
-    OWSAssert(messageEnvelope);
+    OWSAssert(envelope);
     OWSAssert(successBlock);
     OWSAssert(failureBlock);
 
-    TSStorageManager *storageManager = [TSStorageManager sharedManager];
+    [self decryptEnvelope:envelope
-    NSString *recipientId = messageEnvelope.source;
+        messageTypeName:@"Secure Message"
-    int deviceId = messageEnvelope.sourceDevice;
+        missingPayloadBlock:^{
-    dispatch_async([OWSDispatch sessionStoreQueue], ^{
-        // DEPRECATED - Remove after all clients have been upgraded.
-        NSData *encryptedData = messageEnvelope.hasContent ? messageEnvelope.content : messageEnvelope.legacyMessage;
-        if (!encryptedData) {
             OWSProdFail([OWSAnalyticsEvents messageManagerErrorMessageEnvelopeHasNoContent]);
-            failureBlock(nil);
-            return;
         }
+        cipherMessageBlock:^(NSData *encryptedData) {
+            return [[WhisperMessage alloc] initWithData:encryptedData];
+        }
+        successBlock:successBlock
+        failureBlock:failureBlock];
+}
 
-        @try {
+- (void)decryptPreKeyBundle:(OWSSignalServiceProtosEnvelope *)envelope
-            WhisperMessage *message = [[WhisperMessage alloc] initWithData:encryptedData];
+               successBlock:(DecryptSuccessBlock)successBlock
-            SessionCipher *cipher = [[SessionCipher alloc] initWithSessionStore:storageManager
+               failureBlock:(void (^)(NSError *_Nullable error))failureBlock
-                                                                    preKeyStore:storageManager
+{
-                                                              signedPreKeyStore:storageManager
+    OWSAssert(envelope);
-                                                               identityKeyStore:self.identityManager
+    OWSAssert(successBlock);
-                                                                    recipientId:recipientId
+    OWSAssert(failureBlock);
-                                                                       deviceId:deviceId];
 
-            NSData *plaintextData = [[cipher decrypt:message] removePadding];
+    // Check whether we need to refresh our PreKeys every time we receive a PreKeyWhisperMessage.
-            successBlock(plaintextData);
+    [TSPreKeyManager checkPreKeys];
-        } @catch (NSException *exception) {
+
-            dispatch_async(dispatch_get_main_queue(), ^{
+    [self decryptEnvelope:envelope
-                [self processException:exception envelope:messageEnvelope];
+        messageTypeName:@"PreKey Bundle"
-                NSString *errorDescription =
+        missingPayloadBlock:^{
-                    [NSString stringWithFormat:@"Exception while decrypting: %@", exception.description];
+            OWSProdFail([OWSAnalyticsEvents messageManagerErrorPrekeyBundleEnvelopeHasNoContent]);
-                NSError *error = OWSErrorWithCodeDescription(OWSErrorCodeFailedToDecryptMessage, errorDescription);
-                failureBlock(error);
-            });
         }
-    });
+        cipherMessageBlock:^(NSData *encryptedData) {
+            return [[PreKeyWhisperMessage alloc] initWithData:encryptedData];
+        }
+        successBlock:successBlock
+        failureBlock:failureBlock];
 }
 
-- (void)decryptPreKeyBundle:(OWSSignalServiceProtosEnvelope *)preKeyEnvelope
+- (void)decryptEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
+        messageTypeName:(NSString *)messageTypeName
+    missingPayloadBlock:(void (^_Nonnull)())missingPayloadBlock
+     cipherMessageBlock:(id<CipherMessage> (^_Nonnull)(NSData *))cipherMessageBlock
            successBlock:(DecryptSuccessBlock)successBlock
            failureBlock:(void (^)(NSError *_Nullable error))failureBlock
 {
-    OWSAssert(preKeyEnvelope);
+    OWSAssert(envelope);
+    OWSAssert(messageTypeName.length > 0);
+    OWSAssert(missingPayloadBlock);
+    OWSAssert(cipherMessageBlock);
     OWSAssert(successBlock);
     OWSAssert(failureBlock);
 
     TSStorageManager *storageManager = [TSStorageManager sharedManager];
-    NSString *recipientId = preKeyEnvelope.source;
+    NSString *recipientId = envelope.source;
-    int deviceId = preKeyEnvelope.sourceDevice;
+    int deviceId = envelope.sourceDevice;
 
     // DEPRECATED - Remove after all clients have been upgraded.
-    NSData *encryptedData = preKeyEnvelope.hasContent ? preKeyEnvelope.content : preKeyEnvelope.legacyMessage;
+    NSData *encryptedData = envelope.hasContent ? envelope.content : envelope.legacyMessage;
     if (!encryptedData) {
-        OWSProdFail([OWSAnalyticsEvents messageManagerErrorPrekeyBundleEnvelopeHasNoContent]);
+        missingPayloadBlock();
         failureBlock(nil);
         return;
     }
 
     dispatch_async([OWSDispatch sessionStoreQueue], ^{
         @try {
-            // Check whether we need to refresh our PreKeys every time we receive a PreKeyWhisperMessage.
+            id<CipherMessage> cipherMessage = cipherMessageBlock(encryptedData);
-            [TSPreKeyManager checkPreKeys];
-
-            PreKeyWhisperMessage *message = [[PreKeyWhisperMessage alloc] initWithData:encryptedData];
             SessionCipher *cipher = [[SessionCipher alloc] initWithSessionStore:storageManager
                                                                     preKeyStore:storageManager
                                                               signedPreKeyStore:storageManager
@@ -440,13 +444,13 @@ const NSUInteger kIncomingMessageBatchSize = 10;
                                                                     recipientId:recipientId
                                                                        deviceId:deviceId];
 
-            NSData *plaintextData = [[cipher decrypt:message] removePadding];
+            NSData *plaintextData = [[cipher decrypt:cipherMessage] removePadding];
             successBlock(plaintextData);
         } @catch (NSException *exception) {
             dispatch_async(dispatch_get_main_queue(), ^{
-                [self processException:exception envelope:preKeyEnvelope];
+                [self processException:exception envelope:envelope];
-                NSString *errorDescription =
+                NSString *errorDescription = [NSString
-                    [NSString stringWithFormat:@"Exception while decrypting PreKey Bundle: %@", exception.description];
+                    stringWithFormat:@"Exception while decrypting %@: %@", messageTypeName, exception.description];
                 NSError *error = OWSErrorWithCodeDescription(OWSErrorCodeFailedToDecryptMessage, errorDescription);
                 failureBlock(error);
             });