another.im-ios/AnotherIM/AppData/Client/Client.swift

import Combine
import Foundation
import GRDB
import Martin
import MartinOMEMO
import UIKit

enum ClientState: Equatable {
    enum ClientConnectionState {
        case connected
        case disconnected
    }

    case disabled
    case enabled(ClientConnectionState)
}

final class Client: ObservableObject {
    @Published private(set) var state: ClientState = .enabled(.disconnected)
    @Published private(set) var credentials: Credentials
    @Published private(set) var rosters: [Roster] = []

    private var connection: XMPPClient
    private var connectionCancellable: AnyCancellable?
    private var rostersCancellable: AnyCancellable?

    private var rosterManager = ClientMartinRosterManager()
    private var chatsManager = ClientMartinChatsManager()
    private var discoManager: ClientMartinDiscoManager
    private var messageManager: ClientMartinMessagesManager
    private var carbonsManager: ClientMartinCarbonsManager
    private var mamManager: ClientMartinMAM

    init(credentials: Credentials) {
        self.credentials = credentials
        state = credentials.isActive ? .enabled(.disconnected) : .disabled
        connection = Self.prepareConnection(credentials, rosterManager, chatsManager)
        discoManager = ClientMartinDiscoManager(connection)
        messageManager = ClientMartinMessagesManager(connection)
        carbonsManager = ClientMartinCarbonsManager(connection)
        mamManager = ClientMartinMAM(connection)
        connectionCancellable = connection.$state
            .sink { [weak self] state in
                guard let self = self else { return }
                guard self.credentials.isActive else {
                    self.state = .disabled
                    return
                }
                rostersCancellable = ValueObservation
                    .tracking { db in
                        try Roster
                            .filter(Column("bareJid") == self.credentials.bareJid)
                            .filter(Column("locallyDeleted") == false)
                            .fetchAll(db)
                    }
                    .publisher(in: Database.shared.dbQueue)
                    .catch { _ in Just([]) }
                    .sink { rosters in
                        self.rosters = rosters
                    }
                switch state {
                case .connected:
                    self.state = .enabled(.connected)

                default:
                    self.state = .enabled(.disconnected)
                }
            }
    }
}

extension Client: Identifiable {
    var id: String {
        credentials.bareJid
    }
}

extension Client {
    func updActivity(_ isActive: Bool) async {
        credentials.isActive = isActive
        Task {
            try? await credentials.setActive(flag: isActive)
            if isActive {
                self.state = .enabled(.disconnected)
            } else {
                self.state = .disabled
                self.disconnect()
            }
        }
    }

    func addRoster(_ jid: String, name: String?, groups: [String]) async throws {
        _ = try await connection.module(.roster).addItem(
            jid: JID(jid),
            name: name,
            groups: groups
        )
    }

    func addRosterLocally(_ jid: String, name: String?, groups: [String]) async throws {
        try await Roster.addRosterLocally(.init(
            bareJid: credentials.bareJid,
            contactBareJid: jid,
            name: name,
            subscription: "to",
            ask: true,
            data: .init(groups: groups, annotations: []),
            locallyDeleted: false
        ))
    }

    func deleteRoster(_ roster: Roster) async throws {
        _ = try await connection.module(.roster).removeItem(jid: JID(roster.contactBareJid))
    }

    func connect() async {
        guard credentials.isActive, state == .enabled(.disconnected) else {
            return
        }
        try? await connection.loginAndWait()
    }

    func disconnect() {
        _ = connection.disconnect(true)
    }
}

extension Client {
    func sendMessage(_ message: Message) async throws {
        guard let to = message.to else {
            return
        }
        guard let chat = connection.module(MessageModule.self).chatManager.createChat(for: connection.context, with: BareJID(to)) else {
            return
        }

        var msg = chat.createMessage(text: message.body ?? "??", id: message.id)
        msg.oob = message.oobUrl
        if message.secure {
            msg = try await encryptMessage(msg)
        }
        try await chat.send(message: msg)
    }

    func uploadFile(_ localURL: URL, needEncrypt: Bool) async throws -> String {
        // get data from file
        guard var data = try? Data(contentsOf: localURL) else {
            throw AppError.noData
        }

        // encrypt data if needed
        var key = Data()
        var iv = Data()
        if needEncrypt {
            key = try AESGSMEngine.generateKey()
            iv = try AESGSMEngine.generateIV()
            var encrypted = Data()
            var tag = Data()
            guard AESGSMEngine.shared.encrypt(iv: iv, key: key, message: data, output: &encrypted, tag: &tag) else {
                throw AppError.securityError
            }

            // attach tag to end of encrypted data
            encrypted.append(tag)
            data = encrypted
        }

        // upload
        let httpModule = connection.module(HttpFileUploadModule.self)
        let components = try await httpModule.findHttpUploadComponents()
        guard let component = components.first(where: { $0.maxSize > data.count }) else {
            throw AppError.fileTooBig
        }

        let slot = try await httpModule.requestUploadSlot(
            componentJid: component.jid,
            filename: localURL.lastPathComponent,
            size: data.count,
            contentType: localURL.mimeType
        )
        var request = URLRequest(url: slot.putUri)
        for (key, value) in slot.putHeaders {
            request.addValue(value, forHTTPHeaderField: key)
        }
        request.httpMethod = "PUT"
        request.httpBody = data
        request.addValue(String(data.count), forHTTPHeaderField: "Content-Length")
        request.addValue(localURL.mimeType, forHTTPHeaderField: "Content-Type")
        let (_, response) = try await URLSession.shared.data(for: request)
        switch response {
        case let httpResponse as HTTPURLResponse where httpResponse.statusCode == 201:
            if needEncrypt {
                guard var parts = URLComponents(url: slot.getUri, resolvingAgainstBaseURL: true) else {
                    throw URLError(.badServerResponse)
                }
                parts.scheme = "aesgcm"
                parts.fragment = (iv + key).map { String(format: "%02x", $0) }.joined()
                guard let shareUrl = parts.url else {
                    throw URLError(.badServerResponse)
                }
                return shareUrl.absoluteString
            } else {
                return slot.getUri.absoluteString
            }

        default:
            throw URLError(.badServerResponse)
        }
    }

    func fetchArchiveMessages(for roster: Roster, query: RSM.Query) async throws -> Martin.MessageArchiveManagementModule.QueryResult {
        let module = connection.module(MessageArchiveManagementModule.self)
        return try await module.queryItems(componentJid: JID(roster.bareJid), with: JID(roster.contactBareJid), queryId: UUID().uuidString, rsm: query)
    }
}

private extension Client {
    func encryptMessage(_ message: Martin.Message) async throws -> Martin.Message {
        try await withCheckedThrowingContinuation { continuation in
            connection.module(.omemo).encode(message: message, completionHandler: { result in
                switch result {
                case .successMessage(let encodedMessage, _):
                    // guard connection.isConnected else {
                    //     continuation.resume(returning: message)
                    //     return
                    // }
                    continuation.resume(returning: encodedMessage)

                case .failure(let error):
                    var errorMessage = NSLocalizedString("It was not possible to send encrypted message due to encryption error", comment: "message encryption failure")
                    switch error {
                    case .noSession:
                        errorMessage = NSLocalizedString("There is no trusted device to send message to", comment: "message encryption failure")

                    default:
                        break
                    }
                    continuation.resume(throwing: XMPPError.unexpected_request(errorMessage))
                }
            })
        }
    }
}

extension Client {
    static func tryLogin(with credentials: Credentials) async throws -> Client {
        let client = Client(credentials: credentials)
        try await client.connection.loginAndWait()
        return client
    }
}

private extension Client {
    static func prepareConnection(_ credentials: Credentials, _ roster: RosterManager, _ chat: ChatManager) -> XMPPClient {
        let client = XMPPClient()
        client.connectionConfiguration.resource = UIDevice.current.name

        // register modules
        client.modulesManager.register(StreamFeaturesModule())
        client.modulesManager.register(SaslModule())
        client.modulesManager.register(AuthModule())
        client.modulesManager.register(SessionEstablishmentModule())
        client.modulesManager.register(ResourceBinderModule())
        client.modulesManager.register(DiscoveryModule(identity: .init(category: "client", type: "iOS", name: Const.appName)))

        client.modulesManager.register(RosterModule(rosterManager: roster))

        client.modulesManager.register(PubSubModule())
        client.modulesManager.register(PEPUserAvatarModule())
        client.modulesManager.register(PEPBookmarksModule())

        client.modulesManager.register(MessageModule(chatManager: chat))
        client.modulesManager.register(MessageArchiveManagementModule())
        client.modulesManager.register(MessageCarbonsModule())
        // client.modulesManager.register(MessageDeliveryReceiptsModule()).sendReceived = false

        client.modulesManager.register(HttpFileUploadModule())

        client.modulesManager.register(PresenceModule())
        client.modulesManager.register(SoftwareVersionModule())
        client.modulesManager.register(PingModule())
        client.connectionConfiguration.userJid = .init(credentials.bareJid)
        client.connectionConfiguration.credentials = .password(password: credentials.pass)

        // OMEMO
        let omemoManager = ClientMartinOMEMO(credentials)
        let (signalStorage, signalContext) = omemoManager.signal
        client.modulesManager.register(OMEMOModule(aesGCMEngine: AESGSMEngine.shared, signalContext: signalContext, signalStorage: signalStorage))

        // group chats
        // client.modulesManager.register(MucModule(roomManager: manager))

        // channels
        // client.modulesManager.register(MixModule(channelManager: manager))

        return client
    }
}