From 4ce5fbd7d94198c4fc3d263dc337ffaa8e90fbef Mon Sep 17 00:00:00 2001
From: Harry <109690906+harry-signal@users.noreply.github.com>
Date: Fri, 29 Mar 2024 12:49:27 -0700
Subject: [PATCH] Remove unsafe message getter from UnpreparedOutgoingMessage

---
 .../DebugUI/DebugUIMessages.swift             |  8 +++-----
 .../UnpreparedOutgoingMessage.swift           | 12 +++++------
 .../SharingThreadPickerViewController.swift   |  2 +-
 ...SOutgoingMessage+AttachmentMultisend.swift |  2 +-
 SignalUI/Sending/ThreadUtil+SignalUI.swift    | 20 +++++++++----------
 5 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/Signal/src/ViewControllers/DebugUI/DebugUIMessages.swift b/Signal/src/ViewControllers/DebugUI/DebugUIMessages.swift
index 22af2db9f0..3bf8f05ebc 100644
--- a/Signal/src/ViewControllers/DebugUI/DebugUIMessages.swift
+++ b/Signal/src/ViewControllers/DebugUI/DebugUIMessages.swift
@@ -2923,13 +2923,12 @@ class DebugUIMessages: DebugUIPage, Dependencies {
         }
 
         databaseStorage.read { transaction in
-            let message = ThreadUtil.enqueueMessage(
+            ThreadUtil.enqueueMessage(
                 body: messageBody,
                 mediaAttachments: attachments,
                 thread: thread,
                 transaction: transaction
             )
-            Logger.debug("timestamp: \(message.timestamp)")
         }
     }
 
@@ -2972,14 +2971,13 @@ class DebugUIMessages: DebugUIPage, Dependencies {
         Logger.flush()
 
         let text = "\(counter) " + randomText()
-        let message = databaseStorage.write { transaction in
-            return ThreadUtil.enqueueMessage(
+        databaseStorage.write { transaction in
+            ThreadUtil.enqueueMessage(
                 body: MessageBody(text: text, ranges: .empty),
                 thread: thread,
                 transaction: transaction
             )
         }
-        Logger.info("sendTextMessageInThread timestamp: \(message.timestamp).")
     }
 
     private static func sendOversizeTextMessageInThread(_ thread: TSThread) {
diff --git a/SignalServiceKit/Messages/OutgoingMessagePreparer/UnpreparedOutgoingMessage.swift b/SignalServiceKit/Messages/OutgoingMessagePreparer/UnpreparedOutgoingMessage.swift
index 0ab86a41d0..097bdc546a 100644
--- a/SignalServiceKit/Messages/OutgoingMessagePreparer/UnpreparedOutgoingMessage.swift
+++ b/SignalServiceKit/Messages/OutgoingMessagePreparer/UnpreparedOutgoingMessage.swift
@@ -83,18 +83,18 @@ public class UnpreparedOutgoingMessage {
         return try self._prepare(tx: tx)
     }
 
-    public var message: TSOutgoingMessage {
+    public var messageTimestampForLogging: UInt64 {
         switch messageType {
         case .persistable(let message):
-            return message.message
+            return message.message.timestamp
         case .editMessage(let message):
-            return message.messageForSending
+            return message.messageForSending.timestamp
         case .contactSync(let contactSync):
-            return contactSync.message
+            return contactSync.message.timestamp
         case .story(let story):
-            return story.message
+            return story.message.timestamp
         case .transient(let message):
-            return message
+            return message.timestamp
         }
     }
 
diff --git a/SignalShareExtension/SharingThreadPickerViewController.swift b/SignalShareExtension/SharingThreadPickerViewController.swift
index 097ecc9c28..a3f41e7a76 100644
--- a/SignalShareExtension/SharingThreadPickerViewController.swift
+++ b/SignalShareExtension/SharingThreadPickerViewController.swift
@@ -222,7 +222,7 @@ extension SharingThreadPickerViewController {
             let nonStorySendPromise = sendToOutgoingMessageThreads { thread in
                 return firstly(on: DispatchQueue.global()) { () -> Promise<Void> in
                     return self.databaseStorage.write { transaction in
-                        let unpreparedMessage = UnpreparedOutgoingMessage.build(
+                        let (_, unpreparedMessage) = UnpreparedOutgoingMessage.build(
                             thread: thread,
                             messageBody: body,
                             quotedReplyDraft: nil,
diff --git a/SignalUI/AttachmentMultisend/TSOutgoingMessage+AttachmentMultisend.swift b/SignalUI/AttachmentMultisend/TSOutgoingMessage+AttachmentMultisend.swift
index 70130de7e1..9237990676 100644
--- a/SignalUI/AttachmentMultisend/TSOutgoingMessage+AttachmentMultisend.swift
+++ b/SignalUI/AttachmentMultisend/TSOutgoingMessage+AttachmentMultisend.swift
@@ -64,7 +64,7 @@ extension TSOutgoingMessage {
         thread: TSThread,
         transaction: SDSAnyWriteTransaction
     ) throws -> PreparedOutgoingMessage {
-        let unpreparedMessage = UnpreparedOutgoingMessage.build(
+        let (_, unpreparedMessage) = UnpreparedOutgoingMessage.build(
             thread: thread,
             messageBody: messageBody,
             mediaAttachments: mediaAttachments,
diff --git a/SignalUI/Sending/ThreadUtil+SignalUI.swift b/SignalUI/Sending/ThreadUtil+SignalUI.swift
index a931f5beed..bfb62a062f 100644
--- a/SignalUI/Sending/ThreadUtil+SignalUI.swift
+++ b/SignalUI/Sending/ThreadUtil+SignalUI.swift
@@ -22,7 +22,7 @@ public extension ThreadUtil {
     ) -> TSOutgoingMessage {
         AssertIsOnMainThread()
 
-        let unpreparedMessage = UnpreparedOutgoingMessage.build(
+        let (message, unpreparedMessage) = UnpreparedOutgoingMessage.build(
             thread: thread,
             messageBody: messageBody,
             mediaAttachments: mediaAttachments,
@@ -32,27 +32,27 @@ public extension ThreadUtil {
             transaction: readTransaction
         )
 
-        return enqueueMessage(
+        enqueueMessage(
             unpreparedMessage,
             thread: thread,
             persistenceCompletionHandler: persistenceCompletion,
             transaction: readTransaction
         )
+        return message
     }
 
     // MARK: - Durable Message Enqueue
 
-    @discardableResult
     class func enqueueMessage(
         _ unpreparedMessage: UnpreparedOutgoingMessage,
         thread: TSThread,
         persistenceCompletionHandler persistenceCompletion: PersistenceCompletion? = nil,
         transaction readTransaction: SDSAnyReadTransaction
-    ) -> TSOutgoingMessage {
-        let message = unpreparedMessage.message
-        let eventId = "sendMessageMarkedAsSent-\(message.timestamp)"
+    ) {
+        let messageTimestampForLogging = unpreparedMessage.messageTimestampForLogging
+        let eventId = "sendMessageMarkedAsSent-\(messageTimestampForLogging)"
         BenchEventStart(
-            title: "Send Message Milestone: Marked as Sent (\(message.timestamp))",
+            title: "Send Message Milestone: Marked as Sent (\(messageTimestampForLogging))",
             eventId: eventId,
             logInProduction: true
         )
@@ -82,7 +82,6 @@ public extension ThreadUtil {
                 thread.donateSendMessageIntent(for: messageForIntent, transaction: writeTransaction)
             }
         }
-        return message
     }
 }
 
@@ -98,7 +97,7 @@ extension UnpreparedOutgoingMessage {
         linkPreviewDraft: OWSLinkPreviewDraft?,
         editTarget: TSOutgoingMessage?,
         transaction: SDSAnyReadTransaction
-    ) -> UnpreparedOutgoingMessage {
+    ) -> (TSOutgoingMessage, UnpreparedOutgoingMessage) {
 
         var attachments = mediaAttachments
         let truncatedText: String?
@@ -171,11 +170,12 @@ extension UnpreparedOutgoingMessage {
 
         let attachmentInfos = attachments.map { $0.buildAttachmentDataSource(message: message) }
 
-        return UnpreparedOutgoingMessage.forMessage(
+        let unpreparedMessage = UnpreparedOutgoingMessage.forMessage(
             message,
             unsavedBodyAttachments: attachmentInfos,
             linkPreviewDraft: linkPreviewDraft,
             quotedReplyDraft: quotedReplyDraft
         )
+        return (message, unpreparedMessage)
     }
 }