session-ios/SessionMessagingKit/Jobs/MessageSendJob.swift

import SessionUtilitiesKit
import SessionSnodeKit

@objc(SNMessageSendJob)
public final class MessageSendJob : NSObject, Job, NSCoding { // NSObject/NSCoding conformance is needed for YapDatabase compatibility
    public let message: Message
    public let destination: Message.Destination
    public var delegate: JobDelegate?
    public var id: String?
    public var failureCount: UInt = 0

    // MARK: - Settings
    
    public class var collection: String { return "MessageSendJobCollection" }
    public static let maxFailureCount: UInt = 10

    // MARK: - Initialization
    
    @objc public convenience init(message: Message, publicKey: String) {
        self.init(message: message, destination: .contact(publicKey: publicKey))
    }
    
    @objc public convenience init(message: Message, groupPublicKey: String) {
        self.init(message: message, destination: .closedGroup(groupPublicKey: groupPublicKey))
    }

    public init(message: Message, destination: Message.Destination) {
        self.message = message
        self.destination = destination
    }

    // MARK: - Coding
    
    public init?(coder: NSCoder) {
        guard let message = coder.decodeObject(forKey: "message") as! Message?, var rawDestination = coder.decodeObject(forKey: "destination") as! String?, let id = coder.decodeObject(forKey: "id") as! String? else {
            return nil
        }
        
        self.message = message
        
        if rawDestination.removePrefix("contact(") {
            guard rawDestination.removeSuffix(")") else { return nil }
            
            let publicKey = rawDestination
            destination = .contact(publicKey: publicKey)
        }
        else if rawDestination.removePrefix("closedGroup(") {
            guard rawDestination.removeSuffix(")") else { return nil }
            
            let groupPublicKey = rawDestination
            destination = .closedGroup(groupPublicKey: groupPublicKey)
        }
        else if rawDestination.removePrefix("openGroup(") {
            guard rawDestination.removeSuffix(")") else { return nil }
            
            let components = rawDestination
                .split(separator: ",")
                .map { String($0).trimmingCharacters(in: .whitespacesAndNewlines) }
            
            guard components.count == 2, let channel = UInt64(components[0]) else { return nil }
            
            let server = components[1]
            destination = .legacyOpenGroup(channel: channel, server: server)
        }
        else if rawDestination.removePrefix("openGroupV2(") {
            guard rawDestination.removeSuffix(")") else { return nil }
            
            let components = rawDestination
                .split(separator: ",")
                .map { String($0).trimmingCharacters(in: .whitespacesAndNewlines) }
            
            guard components.count == 2 else { return nil }
            
            let roomToken: String = components[0]
            let server: String = components[1]
            
            destination = .openGroup(
                roomToken: roomToken,
                server: server
            )
        }
        else if rawDestination.removePrefix("openGroupV4(") {
            guard rawDestination.removeSuffix(")") else { return nil }
            
            let components = rawDestination
                .split(separator: ",")
                .map { String($0).trimmingCharacters(in: .whitespacesAndNewlines) }
            
            guard components.count == 5 else { return nil }
            
            let roomToken: String = components[0]
            let server: String = components[1]
            let whisperTo: String? = (!components[2].isEmpty ?
                components[2] :
                nil
            )
            let whisperMods: Bool = (components[3] == "true")
            let fileIdStrings: [String] = components[4]
                .replacingOccurrences(of: "[", with: "")
                .replacingOccurrences(of: "]", with: "")
                .split(separator: "|")
                .map { String($0) }
            let fileIds: [Int64]? = (fileIdStrings.isEmpty ? nil : fileIdStrings.compactMap { Int64($0) })
            
            destination = .openGroup(
                roomToken: roomToken,
                server: server,
                whisperTo: whisperTo,
                whisperMods: whisperMods,
                fileIds: fileIds
            )
        }
        else {
            return nil
        }
        
        self.id = id
        self.failureCount = coder.decodeObject(forKey: "failureCount") as! UInt? ?? 0
    }

    public func encode(with coder: NSCoder) {
        coder.encode(message, forKey: "message")
        switch destination {
            case .contact(let publicKey):
                coder.encode("contact(\(publicKey))", forKey: "destination")
                
            case .closedGroup(let groupPublicKey):
                coder.encode("closedGroup(\(groupPublicKey))", forKey: "destination")
                
            case .legacyOpenGroup(let channel, let server):
                coder.encode("openGroup(\(channel), \(server))", forKey: "destination")
                
            case .openGroup(let room, let server, let whisperTo, let whisperMods, let fileIds):
                let whisperToString: String = (whisperTo ?? "")
                let whisperModsString: String = (whisperMods ? "true" : "false")
                let fileIdString: String = (fileIds ?? [])
                    .map { String($0) }
                    .joined(separator: "|")
                
                coder.encode(
                    "openGroupV4(\(room), \(server), \(whisperToString), \(whisperModsString), [\(fileIdString)])",
                    forKey: "destination"
                )
        }
        
        coder.encode(id, forKey: "id")
        coder.encode(failureCount, forKey: "failureCount")
    }

    // MARK: Running
    public func execute() {
        if let id = id {
            JobQueue.currentlyExecutingJobs.insert(id)
        }
        let storage = SNMessagingKitConfiguration.shared.storage
        if let message = message as? VisibleMessage {
            guard TSOutgoingMessage.find(withTimestamp: message.sentTimestamp!) != nil else { return } // The message has been deleted
            let attachments = message.attachmentIDs.compactMap { TSAttachment.fetch(uniqueId: $0) as? TSAttachmentStream }
            let attachmentsToUpload = attachments.filter { !$0.isUploaded }
            attachmentsToUpload.forEach { attachment in
                if storage.getAttachmentUploadJob(for: attachment.uniqueId!) != nil {
                    // Wait for it to finish
                } else {
                    let job = AttachmentUploadJob(attachmentID: attachment.uniqueId!, threadID: message.threadID!, message: message, messageSendJobID: id!)
                    storage.write(with: { transaction in
                        JobQueue.shared.add(job, using: transaction)
                    }, completion: { })
                }
            }
            if !attachmentsToUpload.isEmpty { return } // Wait for all attachments to upload before continuing
        }
        storage.write(with: { transaction in // Intentionally capture self
            MessageSender.send(self.message, to: self.destination, using: transaction).done(on: DispatchQueue.global(qos: .userInitiated)) {
                self.handleSuccess()
            }.catch(on: DispatchQueue.global(qos: .userInitiated)) { error in
                SNLog("Couldn't send message due to error: \(error).")
                if let error = error as? MessageSender.Error, !error.isRetryable {
                    self.handlePermanentFailure(error: error)
                } else if let error = error as? OnionRequestAPI.Error, case .httpRequestFailedAtDestination(let statusCode, _, _) = error,
                    statusCode == 429 { // Rate limited
                    self.handlePermanentFailure(error: error)
               } else {
                    self.handleFailure(error: error)
                }
            }
        }, completion: { })
    }

    private func handleSuccess() {
        delegate?.handleJobSucceeded(self)
    }
    
    private func handlePermanentFailure(error: Error) {
        delegate?.handleJobFailedPermanently(self, with: error)
    }

    private func handleFailure(error: Error) {
        SNLog("Failed to send \(type(of: message)).")
        if let message = message as? VisibleMessage {
            guard TSOutgoingMessage.find(withTimestamp: message.sentTimestamp!) != nil else { return } // The message has been deleted
        }
        delegate?.handleJobFailed(self, with: error)
    }
}

// MARK: Convenience
private extension String {

    @discardableResult
    mutating func removePrefix<T : StringProtocol>(_ prefix: T) -> Bool {
        guard hasPrefix(prefix) else { return false }
        removeFirst(prefix.count)
        return true
    }

    @discardableResult
    mutating func removeSuffix<T : StringProtocol>(_ suffix: T) -> Bool {
        guard hasSuffix(suffix) else { return false }
        removeLast(suffix.count)
        return true
    }
}