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

import Combine
import Foundation
import GRDB
import Martin
import MartinOMEMO

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 {
    func addRoster(_ jid: String, name: String?, groups: [String]) async throws {
        _ = try await connection.module(.roster).addItem(
            jid: JID(jid),
            name: name,
            groups: groups
        )
    }

    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()
    }
}

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
        }

        let msg = chat.createMessage(text: message.body ?? "??", id: message.id)
        msg.oob = message.oobUrl
        try await chat.send(message: msg)
    }

    func uploadFile(_ localURL: URL) async throws -> String {
        guard let data = try? Data(contentsOf: localURL) else {
            throw AppError.noData
        }
        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:
            return slot.getUri.absoluteString

        default:
            throw URLError(.badServerResponse)
        }
    }

    func fetchArchiveMessages(for roster: Roster, query: RSM.Query) async throws -> Martin.MessageArchiveManagementModule.QueryResult {
        if !discoManager.features.map({ $0.xep }).contains("XEP-0313") {
            throw AppError.featureNotSupported
        }
        let module = connection.module(MessageArchiveManagementModule.self)
        return try await module.queryItems(componentJid: JID(roster.bareJid), with: JID(roster.contactBareJid), queryId: UUID().uuidString, rsm: query)
    }
}

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()

        // register modules
        // core modules RFC 6120
        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)))

        // messaging modules RFC 6121
        client.modulesManager.register(RosterModule(rosterManager: roster))
        client.modulesManager.register(PresenceModule())

        client.modulesManager.register(PubSubModule())
        client.modulesManager.register(MessageModule(chatManager: chat))
        client.modulesManager.register(MessageArchiveManagementModule())

        client.modulesManager.register(MessageCarbonsModule())

        // file transfer modules
        client.modulesManager.register(HttpFileUploadModule())

        // extensions
        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
    }
}