Avoid lost messages by acknowledges message receipt after the message is processed.

// FREEBIE

src/Messages/InvalidKeyMessages/TSInvalidIdentityKeyReceivingErrorMessage.m

@@ -89,7 +89,7 @@ NS_ASSUME_NONNULL_BEGIN
                 [self.thread receivedMessagesForInvalidKey:newKey];
 
             for (TSInvalidIdentityKeyReceivingErrorMessage *errorMessage in messagesToDecrypt) {
-                [[TSMessagesManager sharedManager] handleReceivedEnvelope:errorMessage.envelope];
+                [[TSMessagesManager sharedManager] handleReceivedEnvelope:errorMessage.envelope completion:nil];
 
                 // Here we remove the existing error message because handleReceivedEnvelope will either
                 //  1.) succeed and create a new successful message in the thread or...

src/Messages/TSMessagesManager.h

@@ -18,6 +18,8 @@ NS_ASSUME_NONNULL_BEGIN
 @protocol ContactsManagerProtocol;
 @protocol OWSCallMessageHandler;
 
+typedef void (^MessageManagerCompletionBlock)();
+
 @interface TSMessagesManager : NSObject
 
 - (instancetype)init NS_UNAVAILABLE;
@@ -27,18 +29,8 @@ NS_ASSUME_NONNULL_BEGIN
 @property (nonatomic, readonly) TSNetworkManager *networkManager;
 @property (nonatomic, readonly) ContactsUpdater *contactsUpdater;
 
-- (void)handleReceivedEnvelope:(OWSSignalServiceProtosEnvelope *)envelope;
-
-/**
- * Processes all kinds of incoming envelopes with a data message, along with any attachments.
- *
- * @returns
- *   If an incoming message is created, it will be returned. If it is, for example, a group update,
- *   no incoming message is created, so nil will be returned.
- */
-- (TSIncomingMessage *)handleReceivedEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
-                              withDataMessage:(OWSSignalServiceProtosDataMessage *)dataMessage
-                                attachmentIds:(NSArray<NSString *> *)attachmentIds;
+- (void)handleReceivedEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
+                    completion:(nullable MessageManagerCompletionBlock)completion;
 
 /**
  * @returns

src/Messages/TSMessagesManager.m

@@ -153,15 +153,27 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)handleReceivedEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
+                    completion:(nullable MessageManagerCompletionBlock)completionHandler
 {
     OWSAssert([NSThread isMainThread]);
 
+    // Ensure that completionHandler is called on the main thread,
+    // and handle the nil case.
+    MessageManagerCompletionBlock completion = ^{
+        dispatch_async(dispatch_get_main_queue(), ^{
+            if (completionHandler) {
+                completionHandler();
+            }
+        });
+    };
+
     DDLogInfo(@"%@ received envelope: %@", self.tag, [self descriptionForEnvelope:envelope]);
 
     OWSAssert(envelope.source.length > 0);
     BOOL isEnvelopeBlocked = [_blockingManager.blockedPhoneNumbers containsObject:envelope.source];
     if (isEnvelopeBlocked) {
         DDLogInfo(@"%@ ignoring blocked envelope: %@", self.tag, envelope.source);
+        completion();
         return;
     }
 
@@ -175,8 +187,10 @@ NS_ASSUME_NONNULL_BEGIN
                                             DDLogError(
                                                 @"%@ handling secure message failed with error: %@", self.tag, error);
                                         }
+                                        completion();
                                     }];
-                break;
+                // Return to avoid double-acknowledging.
+                return;
             }
             case OWSSignalServiceProtosEnvelopeTypePrekeyBundle: {
                 [self handlePreKeyBundleAsync:envelope
@@ -186,8 +200,10 @@ NS_ASSUME_NONNULL_BEGIN
                                            DDLogError(
                                                @"%@ handling pre-key bundle failed with error: %@", self.tag, error);
                                        }
+                                       completion();
                                    }];
-                break;
+                // Return to avoid double-acknowledging.
+                return;
             }
             case OWSSignalServiceProtosEnvelopeTypeReceipt:
                 DDLogInfo(@"Received a delivery receipt");
@@ -207,8 +223,10 @@ NS_ASSUME_NONNULL_BEGIN
                 break;
         }
     } @catch (NSException *exception) {
-        DDLogWarn(@"Received an incorrectly formatted protocol buffer: %@", exception.debugDescription);
+        DDLogError(@"Received an incorrectly formatted protocol buffer: %@", exception.debugDescription);
     }
+
+    completion();
 }
 
 - (void)handleDeliveryReceipt:(OWSSignalServiceProtosEnvelope *)envelope
@@ -241,6 +259,8 @@ NS_ASSUME_NONNULL_BEGIN
                         [TSErrorMessage missingSessionWithEnvelope:messageEnvelope withTransaction:transaction];
                     [errorMessage saveWithTransaction:transaction];
                 }];
+                DDLogError(@"Skipping message envelope for unknown session.");
+                completion(nil);
                 return;
             }
 
@@ -248,7 +268,8 @@ NS_ASSUME_NONNULL_BEGIN
             NSData *encryptedData
                 = messageEnvelope.hasContent ? messageEnvelope.content : messageEnvelope.legacyMessage;
             if (!encryptedData) {
-                DDLogError(@"Skipping message envelope which had no encrypted data");
+                DDLogError(@"Skipping message envelope which had no encrypted data.");
+                completion(nil);
                 return;
             }
 
@@ -297,6 +318,7 @@ NS_ASSUME_NONNULL_BEGIN
         NSData *encryptedData = preKeyEnvelope.hasContent ? preKeyEnvelope.content : preKeyEnvelope.legacyMessage;
         if (!encryptedData) {
             DDLogError(@"Skipping message envelope which had no encrypted data");
+            completion(nil);
             return;
         }
 
@@ -327,6 +349,7 @@ NS_ASSUME_NONNULL_BEGIN
 
             dispatch_async(dispatch_get_main_queue(), ^{
                 [self handleEnvelope:preKeyEnvelope plaintextData:plaintextData];
+                completion(nil);
             });
         });
     }

src/Network/WebSockets/TSSocketManager.m

@@ -368,8 +368,6 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 
     DDLogInfo(@"Got message with verb: %@ and path: %@", message.verb, message.path);
 
-    [self sendWebSocketMessageAcknowledgement:message];
-
     // If we receive a message over the socket while the app is in the background,
     // prolong how long the socket stays open.
     [self requestSocketAliveForAtLeastSeconds:kBackgroundKeepSocketAliveDurationSeconds];
@@ -381,15 +379,22 @@ NSString *const SocketConnectingNotification = @"SocketConnectingNotification";
 
         if (!decryptedPayload) {
             DDLogWarn(@"Failed to decrypt incoming payload or bad HMAC");
+            [self sendWebSocketMessageAcknowledgement:message];
             return;
         }
 
         OWSSignalServiceProtosEnvelope *envelope = [OWSSignalServiceProtosEnvelope parseFromData:decryptedPayload];
 
-        [[TSMessagesManager sharedManager] handleReceivedEnvelope:envelope];
-
+        [[TSMessagesManager sharedManager] handleReceivedEnvelope:envelope
+                                                       completion:^{
+                                                           // Don't acknowledge delivery until the envelope has been
+                                                           // processed.
+                                                           [self sendWebSocketMessageAcknowledgement:message];
+                                                       }];
     } else {
         DDLogWarn(@"Unsupported WebSocket Request");
+
+        [self sendWebSocketMessageAcknowledgement:message];
     }
 }