From a88ce33ee0f1c3f1612459c1907168003b3bda23 Mon Sep 17 00:00:00 2001
From: Niels Andriesse <andriesseniels@gmail.com>
Date: Thu, 26 Nov 2020 16:33:42 +1100
Subject: [PATCH] WIP

---
 SessionMessagingKit/Jobs/MessageSendJob.swift             | 2 ++
 .../Messages/Signal/TSOutgoingMessage+Conversion.swift    | 2 +-
 .../Messages/Visible Messages/VisibleMessage.swift        | 8 ++++++--
 .../Sending & Receiving/MessageSender.swift               | 8 +++++++-
 .../Sending & Receiving/MessageSender+Handling.swift      | 1 +
 5 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/SessionMessagingKit/Jobs/MessageSendJob.swift b/SessionMessagingKit/Jobs/MessageSendJob.swift
index 5bb1b1f5d..61b3718ab 100644
--- a/SessionMessagingKit/Jobs/MessageSendJob.swift
+++ b/SessionMessagingKit/Jobs/MessageSendJob.swift
@@ -1,5 +1,7 @@
 import SessionUtilitiesKit
 
+// TODO: Cancel when a message is deleted
+
 @objc(SNMessageSendJob)
 public final class MessageSendJob : NSObject, Job, NSCoding { // NSObject/NSCoding conformance is needed for YapDatabase compatibility
     public let message: Message
diff --git a/SessionMessagingKit/Messages/Signal/TSOutgoingMessage+Conversion.swift b/SessionMessagingKit/Messages/Signal/TSOutgoingMessage+Conversion.swift
index 0ab7a70b4..5cf891342 100644
--- a/SessionMessagingKit/Messages/Signal/TSOutgoingMessage+Conversion.swift
+++ b/SessionMessagingKit/Messages/Signal/TSOutgoingMessage+Conversion.swift
@@ -11,7 +11,7 @@
             outgoingMessageWithTimestamp: visibleMessage.sentTimestamp!,
             in: thread,
             messageBody: visibleMessage.text,
-            attachmentIds: NSMutableArray(array: visibleMessage.attachmentIDs),
+            attachmentIds: NSMutableArray(),
             expiresInSeconds: expiration,
             expireStartedAt: 0,
             isVoiceMessage: false,
diff --git a/SessionMessagingKit/Messages/Visible Messages/VisibleMessage.swift b/SessionMessagingKit/Messages/Visible Messages/VisibleMessage.swift
index 3cd5fab53..1bc671635 100644
--- a/SessionMessagingKit/Messages/Visible Messages/VisibleMessage.swift	
+++ b/SessionMessagingKit/Messages/Visible Messages/VisibleMessage.swift	
@@ -47,6 +47,10 @@ public final class VisibleMessage : Message {
     }
 
     public override func toProto() -> SNProtoContent? {
+        preconditionFailure("Use toProto(using:) instead.")
+    }
+    
+    public func toProto(using transaction: YapDatabaseReadWriteTransaction) -> SNProtoContent? {
         let proto = SNProtoContent.builder()
         let dataMessage: SNProtoDataMessage.SNProtoDataMessageBuilder
         if let profile = profile, let profileProto = profile.toProto() {
@@ -55,13 +59,13 @@ public final class VisibleMessage : Message {
             dataMessage = SNProtoDataMessage.builder()
         }
         if let text = text { dataMessage.setBody(text) }
-        let attachments = attachmentIDs.compactMap { TSAttachmentStream.fetch(uniqueId: $0) }
+        let attachments = attachmentIDs.compactMap { TSAttachmentStream.fetch(uniqueId: $0, transaction: transaction) }
         if !attachments.allSatisfy({ $0.isUploaded }) {
             #if DEBUG
             preconditionFailure("Sending a message before all associated attachments have been uploaded.")
             #endif
         }
-        let attachmentProtos = attachments.compactMap { TSAttachmentStream.buildProto(forAttachmentId: $0.uniqueId!) }
+        let attachmentProtos = attachments.compactMap { $0.buildProto() }
         dataMessage.setAttachments(attachmentProtos)
         if let quote = quote, let quoteProto = quote.toProto() { dataMessage.setQuote(quoteProto) }
         if let linkPreview = linkPreview, let linkPreviewProto = linkPreview.toProto() { dataMessage.setPreview([ linkPreviewProto ]) }
diff --git a/SessionMessagingKit/Sending & Receiving/MessageSender.swift b/SessionMessagingKit/Sending & Receiving/MessageSender.swift
index 0154aed22..03f1f83ec 100644
--- a/SessionMessagingKit/Sending & Receiving/MessageSender.swift	
+++ b/SessionMessagingKit/Sending & Receiving/MessageSender.swift	
@@ -84,7 +84,13 @@ public final class MessageSender : NSObject {
             }
         }
         // Convert it to protobuf
-        guard let proto = message.toProto() else { seal.reject(Error.protoConversionFailed); return promise }
+        let protoOrNil: SNProtoContent?
+        if let message = message as? VisibleMessage {
+            protoOrNil = message.toProto(using: transaction as! YapDatabaseReadWriteTransaction) // Needed because of how TSAttachmentStream works
+        } else {
+            protoOrNil = message.toProto()
+        }
+        guard let proto = protoOrNil else { seal.reject(Error.protoConversionFailed); return promise }
         // Serialize the protobuf
         let plaintext: Data
         do {
diff --git a/SignalUtilitiesKit/Messaging/Sending & Receiving/MessageSender+Handling.swift b/SignalUtilitiesKit/Messaging/Sending & Receiving/MessageSender+Handling.swift
index bc2af8cc8..af5009f69 100644
--- a/SignalUtilitiesKit/Messaging/Sending & Receiving/MessageSender+Handling.swift	
+++ b/SignalUtilitiesKit/Messaging/Sending & Receiving/MessageSender+Handling.swift	
@@ -22,6 +22,7 @@ extension MessageSender : SharedSenderKeysDelegate {
             stream.save(with: transaction)
         }
         message.attachmentIDs = streams.map { $0.uniqueId! }
+        tsMessage.attachmentIds.addObjects(from: message.attachmentIDs)
     }
     
     @objc(send:withAttachments:inThread:usingTransaction:)