Refine envelope processing.

Signal/src/ViewControllers/DebugUI/DebugUIMessages.m

@@ -3916,6 +3916,7 @@ typedef OWSContact * (^OWSContactBlock)(YapDatabaseReadWriteTransaction *transac
     [OWSPrimaryStorage.dbReadWriteConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
         [SSKEnvironment.shared.batchMessageProcessor enqueueEnvelopeData:envelopeData
                                                            plaintextData:plaintextData
+                                                         wasReceivedByUD:NO
                                                              transaction:transaction];
     }];
 }

SignalServiceKit/src/Messages/OWSBatchMessageProcessor.h

@@ -22,6 +22,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)enqueueEnvelopeData:(NSData *)envelopeData
               plaintextData:(NSData *_Nullable)plaintextData
+            wasReceivedByUD:(BOOL)wasReceivedByUD
                 transaction:(YapDatabaseReadWriteTransaction *)transaction;
 
 - (void)handleAnyUnprocessedEnvelopesAsync;

SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m

@@ -33,9 +33,11 @@ NS_ASSUME_NONNULL_BEGIN
 @property (nonatomic, readonly) NSDate *createdAt;
 @property (nonatomic, readonly) NSData *envelopeData;
 @property (nonatomic, readonly, nullable) NSData *plaintextData;
+@property (nonatomic, readonly) BOOL wasReceivedByUD;
 
 - (instancetype)initWithEnvelopeData:(NSData *)envelopeData
-                       plaintextData:(NSData *_Nullable)plaintextData NS_DESIGNATED_INITIALIZER;
+                       plaintextData:(NSData *_Nullable)plaintextData
+                     wasReceivedByUD:(BOOL)wasReceivedByUD NS_DESIGNATED_INITIALIZER;
 - (nullable instancetype)initWithCoder:(NSCoder *)coder NS_DESIGNATED_INITIALIZER;
 - (instancetype)initWithUniqueId:(NSString *_Nullable)uniqueId NS_UNAVAILABLE;
 
@@ -52,7 +54,9 @@ NS_ASSUME_NONNULL_BEGIN
     return @"OWSBatchMessageProcessingJob";
 }
 
-- (instancetype)initWithEnvelopeData:(NSData *)envelopeData plaintextData:(NSData *_Nullable)plaintextData
+- (instancetype)initWithEnvelopeData:(NSData *)envelopeData
+                       plaintextData:(NSData *_Nullable)plaintextData
+                     wasReceivedByUD:(BOOL)wasReceivedByUD
 {
     OWSAssertDebug(envelopeData);
 
@@ -63,6 +67,7 @@ NS_ASSUME_NONNULL_BEGIN
 
     _envelopeData = envelopeData;
     _plaintextData = plaintextData;
+    _wasReceivedByUD = wasReceivedByUD;
     _createdAt = [NSDate new];
 
     return self;
@@ -148,13 +153,15 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
 
 - (void)addJobWithEnvelopeData:(NSData *)envelopeData
                  plaintextData:(NSData *_Nullable)plaintextData
+               wasReceivedByUD:(BOOL)wasReceivedByUD
                    transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     OWSAssertDebug(envelopeData);
     OWSAssertDebug(transaction);
 
-    OWSMessageContentJob *job =
-        [[OWSMessageContentJob alloc] initWithEnvelopeData:envelopeData plaintextData:plaintextData];
+    OWSMessageContentJob *job = [[OWSMessageContentJob alloc] initWithEnvelopeData:envelopeData
+                                                                     plaintextData:plaintextData
+                                                                   wasReceivedByUD:wasReceivedByUD];
     [job saveWithTransaction:transaction];
 }
 
@@ -316,13 +323,17 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
 
 - (void)enqueueEnvelopeData:(NSData *)envelopeData
               plaintextData:(NSData *_Nullable)plaintextData
+            wasReceivedByUD:(BOOL)wasReceivedByUD
                 transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     OWSAssertDebug(envelopeData);
     OWSAssertDebug(transaction);
 
     // We need to persist the decrypted envelope data ASAP to prevent data loss.
-    [self.finder addJobWithEnvelopeData:envelopeData plaintextData:plaintextData transaction:transaction];
+    [self.finder addJobWithEnvelopeData:envelopeData
+                          plaintextData:plaintextData
+                        wasReceivedByUD:wasReceivedByUD
+                            transaction:transaction];
 }
 
 - (void)drainQueue
@@ -405,6 +416,7 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
                 } else {
                     [self.messageManager throws_processEnvelope:envelope
                                                   plaintextData:job.plaintextData
+                                                wasReceivedByUD:job.wasReceivedByUD
                                                     transaction:transaction];
                 }
             } @catch (NSException *exception) {
@@ -482,6 +494,7 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
 
 - (void)enqueueEnvelopeData:(NSData *)envelopeData
               plaintextData:(NSData *_Nullable)plaintextData
+            wasReceivedByUD:(BOOL)wasReceivedByUD
                 transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     if (envelopeData.length < 1) {
@@ -491,7 +504,10 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
     OWSAssert(transaction);
 
     // We need to persist the decrypted envelope data ASAP to prevent data loss.
-    [self.processingQueue enqueueEnvelopeData:envelopeData plaintextData:plaintextData transaction:transaction];
+    [self.processingQueue enqueueEnvelopeData:envelopeData
+                                plaintextData:plaintextData
+                              wasReceivedByUD:wasReceivedByUD
+                                  transaction:transaction];
 
     // The new envelope won't be visible to the finder until this transaction commits,
     // so drainQueue in the transaction completion.

SignalServiceKit/src/Messages/OWSMessageManager.h

@@ -21,6 +21,7 @@ NS_ASSUME_NONNULL_BEGIN
 // processEnvelope: can be called from any thread.
 - (void)throws_processEnvelope:(SSKProtoEnvelope *)envelope
                  plaintextData:(NSData *_Nullable)plaintextData
+               wasReceivedByUD:(BOOL)wasReceivedByUD
                    transaction:(YapDatabaseReadWriteTransaction *)transaction;
 
 // This should be invoked by the main app when the app is ready.

SignalServiceKit/src/Messages/OWSMessageManager.m

@@ -219,6 +219,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)throws_processEnvelope:(SSKProtoEnvelope *)envelope
                  plaintextData:(NSData *_Nullable)plaintextData
+               wasReceivedByUD:(BOOL)wasReceivedByUD
                    transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     if (!envelope) {
@@ -261,7 +262,10 @@ NS_ASSUME_NONNULL_BEGIN
                 OWSFailDebug(@"missing decrypted data for envelope: %@", [self descriptionForEnvelope:envelope]);
                 return;
             }
-            [self throws_handleEnvelope:envelope plaintextData:plaintextData transaction:transaction];
+            [self throws_handleEnvelope:envelope
+                          plaintextData:plaintextData
+                        wasReceivedByUD:wasReceivedByUD
+                            transaction:transaction];
             break;
         case SSKProtoEnvelopeTypeReceipt:
             OWSAssertDebug(!plaintextData);
@@ -352,6 +356,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)throws_handleEnvelope:(SSKProtoEnvelope *)envelope
                 plaintextData:(NSData *)plaintextData
+              wasReceivedByUD:(BOOL)wasReceivedByUD
                   transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     if (!envelope) {
@@ -406,7 +411,10 @@ NS_ASSUME_NONNULL_BEGIN
 
             [[OWSDeviceManager sharedManager] setHasReceivedSyncMessage];
         } else if (contentProto.dataMessage) {
-            [self handleIncomingEnvelope:envelope withDataMessage:contentProto.dataMessage transaction:transaction];
+            [self handleIncomingEnvelope:envelope
+                         withDataMessage:contentProto.dataMessage
+                         wasReceivedByUD:wasReceivedByUD
+                             transaction:transaction];
         } else if (contentProto.callMessage) {
             [self handleIncomingEnvelope:envelope withCallMessage:contentProto.callMessage];
         } else if (contentProto.nullMessage) {
@@ -427,7 +435,10 @@ NS_ASSUME_NONNULL_BEGIN
         }
         OWSLogInfo(@"handling message: <DataMessage: %@ />", [self descriptionForDataMessage:dataMessageProto]);
 
-        [self handleIncomingEnvelope:envelope withDataMessage:dataMessageProto transaction:transaction];
+        [self handleIncomingEnvelope:envelope
+                     withDataMessage:dataMessageProto
+                     wasReceivedByUD:wasReceivedByUD
+                         transaction:transaction];
     } else {
         OWSProdInfoWEnvelope([OWSAnalyticsEvents messageManagerErrorEnvelopeNoActionablePayload], envelope);
     }
@@ -435,6 +446,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)handleIncomingEnvelope:(SSKProtoEnvelope *)envelope
                withDataMessage:(SSKProtoDataMessage *)dataMessage
+               wasReceivedByUD:(BOOL)wasReceivedByUD
                    transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     if (!envelope) {
@@ -517,9 +529,15 @@ NS_ASSUME_NONNULL_BEGIN
     } else if ((dataMessage.flags & SSKProtoDataMessageFlagsProfileKeyUpdate) != 0) {
         [self handleProfileKeyMessageWithEnvelope:envelope dataMessage:dataMessage];
     } else if (dataMessage.attachments.count > 0) {
-        [self handleReceivedMediaWithEnvelope:envelope dataMessage:dataMessage transaction:transaction];
+        [self handleReceivedMediaWithEnvelope:envelope
+                                  dataMessage:dataMessage
+                              wasReceivedByUD:wasReceivedByUD
+                                  transaction:transaction];
     } else {
-        [self handleReceivedTextMessageWithEnvelope:envelope dataMessage:dataMessage transaction:transaction];
+        [self handleReceivedTextMessageWithEnvelope:envelope
+                                        dataMessage:dataMessage
+                                    wasReceivedByUD:wasReceivedByUD
+                                        transaction:transaction];
 
         if ([self isDataMessageGroupAvatarUpdate:dataMessage]) {
             OWSLogVerbose(@"Data message had group avatar attachment");
@@ -528,18 +546,11 @@ NS_ASSUME_NONNULL_BEGIN
     }
 
     // Send delivery receipts for "valid data" messages received via UD.
-    BOOL wasReceivedByUD = [self wasReceivedByUD:envelope];
     if (wasReceivedByUD) {
         [self.outgoingReceiptManager enqueueDeliveryReceiptForEnvelope:envelope];
     }
 }
 
-- (BOOL)wasReceivedByUD:(SSKProtoEnvelope *)envelope
-{
-    return (
-        envelope.type == SSKProtoEnvelopeTypeUnidentifiedSender && (!envelope.hasSource || envelope.source.length < 1));
-}
-
 - (void)sendGroupInfoRequest:(NSData *)groupId
                     envelope:(SSKProtoEnvelope *)envelope
                  transaction:(YapDatabaseReadWriteTransaction *)transaction
@@ -702,6 +713,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)handleReceivedMediaWithEnvelope:(SSKProtoEnvelope *)envelope
                             dataMessage:(SSKProtoDataMessage *)dataMessage
+                        wasReceivedByUD:(BOOL)wasReceivedByUD
                             transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     if (!envelope) {
@@ -734,6 +746,7 @@ NS_ASSUME_NONNULL_BEGIN
 
     TSIncomingMessage *_Nullable createdMessage = [self handleReceivedEnvelope:envelope
                                                                withDataMessage:dataMessage
+                                                               wasReceivedByUD:wasReceivedByUD
                                                                  attachmentIds:attachmentsProcessor.attachmentIds
                                                                    transaction:transaction];
 
@@ -983,6 +996,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)handleReceivedTextMessageWithEnvelope:(SSKProtoEnvelope *)envelope
                                   dataMessage:(SSKProtoDataMessage *)dataMessage
+                              wasReceivedByUD:(BOOL)wasReceivedByUD
                                   transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     if (!envelope) {
@@ -998,7 +1012,11 @@ NS_ASSUME_NONNULL_BEGIN
         return;
     }
 
-    [self handleReceivedEnvelope:envelope withDataMessage:dataMessage attachmentIds:@[] transaction:transaction];
+    [self handleReceivedEnvelope:envelope
+                 withDataMessage:dataMessage
+                 wasReceivedByUD:wasReceivedByUD
+                   attachmentIds:@[]
+                     transaction:transaction];
 }
 
 - (void)sendGroupUpdateForThread:(TSGroupThread *)gThread message:(TSOutgoingMessage *)message
@@ -1106,6 +1124,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (TSIncomingMessage *_Nullable)handleReceivedEnvelope:(SSKProtoEnvelope *)envelope
                                        withDataMessage:(SSKProtoDataMessage *)dataMessage
+                                       wasReceivedByUD:(BOOL)wasReceivedByUD
                                          attachmentIds:(NSArray<NSString *> *)attachmentIds
                                            transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
@@ -1127,7 +1146,6 @@ NS_ASSUME_NONNULL_BEGIN
     NSData *groupId = dataMessage.group ? dataMessage.group.id : nil;
     OWSContact *_Nullable contact = [OWSContacts contactForDataMessage:dataMessage transaction:transaction];
     NSNumber *_Nullable serverTimestamp = (envelope.hasServerTimestamp ? @(envelope.serverTimestamp) : nil);
-    BOOL wasReceivedByUD = [self wasReceivedByUD:envelope];
 
     if (dataMessage.group.type == SSKProtoGroupContextTypeRequestInfo) {
         [self handleGroupInfoRequest:envelope dataMessage:dataMessage transaction:transaction];

SignalServiceKit/src/Messages/OWSMessageReceiver.m

@@ -333,6 +333,12 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
           }];
 }
 
+- (BOOL)wasReceivedByUD:(SSKProtoEnvelope *)envelope
+{
+    return (
+        envelope.type == SSKProtoEnvelopeTypeUnidentifiedSender && (!envelope.hasSource || envelope.source.length < 1));
+}
+
 - (void)processJob:(OWSMessageDecryptJob *)job completion:(void (^)(BOOL))completion
 {
     AssertOnDispatchQueue(self.serialQueue);
@@ -356,6 +362,10 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
         return;
     }
 
+    // We use the original envelope for this check;
+    // the decryption process might rewrite the envelope.
+    BOOL wasReceivedByUD = [self wasReceivedByUD:envelope];
+
     [self.messageDecrypter decryptEnvelope:envelope
         envelopeData:job.envelopeData
         successBlock:^(OWSMessageDecryptResult *result, YapDatabaseReadWriteTransaction *transaction) {
@@ -369,6 +379,7 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
             // since the envelope may be altered by the decryption process in the UD case.
             [self.batchMessageProcessor enqueueEnvelopeData:result.envelopeData
                                               plaintextData:result.plaintextData
+                                            wasReceivedByUD:wasReceivedByUD
                                                 transaction:transaction];
 
             dispatch_async(self.serialQueue, ^{