From 2a609791bf7aca06f85293b3e1994417439facdb Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Wed, 24 Jul 2019 10:06:34 -0300
Subject: [PATCH] Rework message sender to assume thread is non-nil.

---
 .../OWSOutgoingSentMessageTranscript.h        |  7 ++-
 .../OWSOutgoingSentMessageTranscript.m        | 14 +++--
 .../src/Messages/OWSMessageSender.m           | 58 ++++++++++++-------
 3 files changed, 49 insertions(+), 30 deletions(-)

diff --git a/SignalServiceKit/src/Messages/DeviceSyncing/OWSOutgoingSentMessageTranscript.h b/SignalServiceKit/src/Messages/DeviceSyncing/OWSOutgoingSentMessageTranscript.h
index 0042c9b9ba..2fe554305a 100644
--- a/SignalServiceKit/src/Messages/DeviceSyncing/OWSOutgoingSentMessageTranscript.h
+++ b/SignalServiceKit/src/Messages/DeviceSyncing/OWSOutgoingSentMessageTranscript.h
@@ -17,9 +17,10 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)init NS_UNAVAILABLE;
 - (instancetype)initWithTimestamp:(uint64_t)timestamp NS_UNAVAILABLE;
 
-- (instancetype)initWithThread:(TSThread *)thread
-               outgoingMessage:(TSOutgoingMessage *)message
-             isRecipientUpdate:(BOOL)isRecipientUpdate NS_DESIGNATED_INITIALIZER;
+- (instancetype)initWithLocalThread:(TSThread *)localThread
+                      messageThread:(TSThread *)messageThread
+                    outgoingMessage:(TSOutgoingMessage *)message
+                  isRecipientUpdate:(BOOL)isRecipientUpdate NS_DESIGNATED_INITIALIZER;
 - (nullable instancetype)initWithCoder:(NSCoder *)coder NS_DESIGNATED_INITIALIZER;
 
 @end
diff --git a/SignalServiceKit/src/Messages/DeviceSyncing/OWSOutgoingSentMessageTranscript.m b/SignalServiceKit/src/Messages/DeviceSyncing/OWSOutgoingSentMessageTranscript.m
index 1b171af9eb..dd45bc4bbe 100644
--- a/SignalServiceKit/src/Messages/DeviceSyncing/OWSOutgoingSentMessageTranscript.m
+++ b/SignalServiceKit/src/Messages/DeviceSyncing/OWSOutgoingSentMessageTranscript.m
@@ -40,14 +40,17 @@ NS_ASSUME_NONNULL_BEGIN
 
 @implementation OWSOutgoingSentMessageTranscript
 
-- (instancetype)initWithThread:(TSThread *)thread
-               outgoingMessage:(TSOutgoingMessage *)message
-             isRecipientUpdate:(BOOL)isRecipientUpdate
+- (instancetype)initWithLocalThread:(TSThread *)localThread
+                      messageThread:(TSThread *)messageThread
+                    outgoingMessage:(TSOutgoingMessage *)message
+                  isRecipientUpdate:(BOOL)isRecipientUpdate
 {
-    OWSAssertDebug(message);
+    OWSAssertDebug(message != nil);
+    OWSAssertDebug(localThread != nil);
+    OWSAssertDebug(messageThread != nil);
 
     // The sync message's timestamp must match the original outgoing message's timestamp.
-    self = [super initWithTimestamp:message.timestamp thread:thread];
+    self = [super initWithTimestamp:message.timestamp thread:localThread];
 
     if (!self) {
         return self;
@@ -56,7 +59,6 @@ NS_ASSUME_NONNULL_BEGIN
     _message = message;
     _isRecipientUpdate = isRecipientUpdate;
 
-    TSThread *_Nullable messageThread = message.threadWithSneakyTransaction;
     if ([messageThread isKindOfClass:[TSContactThread class]]) {
         TSContactThread *contactThread = (TSContactThread *)messageThread;
         _sentRecipientAddress = contactThread.contactAddress;
diff --git a/SignalServiceKit/src/Messages/OWSMessageSender.m b/SignalServiceKit/src/Messages/OWSMessageSender.m
index f01d84004b..6016727a6c 100644
--- a/SignalServiceKit/src/Messages/OWSMessageSender.m
+++ b/SignalServiceKit/src/Messages/OWSMessageSender.m
@@ -852,24 +852,30 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
 {
     __block TSThread *_Nullable thread = nil;
     [self.databaseStorage writeWithBlock:^(SDSAnyWriteTransaction *transaction) {
-        thread = [message threadWithTransaction:transaction];
-        OWSAssertDebug(thread != nil);
-
-        // For some legacy sync messages, thread may be nil.
-        // In this case, we should try to use the "local" thread.
-        BOOL isSyncMessage = [message isKindOfClass:[OWSOutgoingSyncMessage class]];
-        if (thread == nil && isSyncMessage) {
-            thread = [TSAccountManager getOrCreateLocalThreadWithTransaction:transaction];
-            if (thread == nil) {
-                OWSFailDebug(@"Could not restore thread for sync message.");
-            } else {
-                OWSLogInfo(@"Thread restored for sync message.");
-            }
-        }
+        thread = [self threadForMessage:message transaction:transaction];
     }];
     return thread;
 }
 
+- (nullable TSThread *)threadForMessage:(TSMessage *)message transaction:(SDSAnyWriteTransaction *)transaction
+{
+    TSThread *_Nullable thread = [message threadWithTransaction:transaction];
+    OWSAssertDebug(thread != nil);
+
+    // For some legacy sync messages, thread may be nil.
+    // In this case, we should try to use the "local" thread.
+    BOOL isSyncMessage = [message isKindOfClass:[OWSOutgoingSyncMessage class]];
+    if (thread == nil && isSyncMessage) {
+        thread = [TSAccountManager getOrCreateLocalThreadWithTransaction:transaction];
+        if (thread == nil) {
+            OWSFailDebug(@"Could not restore thread for sync message.");
+        } else {
+            OWSLogInfo(@"Thread restored for sync message.");
+        }
+    }
+    return thread;
+}
+
 - (void)unregisteredRecipient:(SignalRecipient *)recipient
                       message:(TSOutgoingMessage *)message
                        thread:(TSThread *)thread
@@ -1524,25 +1530,35 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
                              failure:(RetryableFailureHandler)failure
 {
     SignalServiceAddress *localAddress = self.tsAccountManager.localAddress;
-    __block TSThread *_Nullable thread;
+    // After sending a message to its "message thread",
+    // we send a sync transcript to the "local thread".
+    __block TSThread *_Nullable localThread;
+    __block TSThread *_Nullable messageThread;
     __block SignalRecipient *recipient;
     [self.databaseStorage writeWithBlock:^(SDSAnyWriteTransaction *transaction) {
-        thread = [TSAccountManager getOrCreateLocalThreadWithTransaction:transaction];
+        localThread = [TSAccountManager getOrCreateLocalThreadWithTransaction:transaction];
+
+        messageThread = [self threadForMessage:message transaction:transaction];
 
         recipient = [SignalRecipient markRecipientAsRegisteredAndGet:localAddress transaction:transaction];
     }];
-    if (thread == nil) {
+    if (localThread == nil) {
         OWSFailDebug(@"Missing local thread.");
         return;
     }
+    if (messageThread == nil) {
+        OWSFailDebug(@"Missing message thread.");
+        return;
+    }
 
     OWSOutgoingSentMessageTranscript *sentMessageTranscript =
-        [[OWSOutgoingSentMessageTranscript alloc] initWithThread:thread
-                                                 outgoingMessage:message
-                                               isRecipientUpdate:isRecipientUpdate];
+        [[OWSOutgoingSentMessageTranscript alloc] initWithLocalThread:localThread
+                                                        messageThread:messageThread
+                                                      outgoingMessage:message
+                                                    isRecipientUpdate:isRecipientUpdate];
 
     OWSMessageSend *messageSend = [[OWSMessageSend alloc] initWithMessage:sentMessageTranscript
-        thread:thread
+        thread:localThread
         recipient:recipient
         senderCertificate:nil
         udAccess:nil