another.im-ios/AnotherXMPP/modules/roster/RosterModule.swift

// RFC 6121
// XEP-0237
import Foundation

final class RosterModule: XmppModule {
    let id = "Roseter module"

    private weak var storage: (any XMPPStorage)?
    private var isVerSupported = false
    private let registry = StanzaRegistry()

    init(_ storage: any XMPPStorage) {
        self.storage = storage
    }

    func reduce(oldState: ClientState, with _: Event) -> ClientState {
        oldState
    }

    func process(state: ClientState, with event: Event) async -> Event? {
        switch event {
        case .streamReady:
            return .requestRoster

        case .xmlInbound(let xml):
            if xml.name == "stream:features" {
                if let ver = xml.nodes.first(where: { $0.name == "ver" }), ver.xmlns == "urn:xmpp:features:rosterver" {
                    isVerSupported = true
                }
            }
            return nil

        case .requestRoster:
            var attributes: [String: String] = [:]
            if isVerSupported {
                let ver = await storage?.getRosterVer(jid: state.jid)
                attributes["ver"] = ver ?? ""
            }
            let req = Stanza.iqGet(
                from: state.jid.full,
                payload: XMLElement(
                    name: "query",
                    xmlns: "jabber:iq:roster",
                    attributes: attributes,
                    content: nil,
                    nodes: []
                )
            )
            if let req {
                await registry.enqueue(req)
                return .stanzaOutbound(req)
            } else {
                return nil
            }

        case .addRosterItem(let jidStr, let args), .updateRosterItem(let jidStr, let args):
            var attr = ["jid": jidStr]
            if let name = args["name"] {
                attr["name"] = name
            }
            let req = Stanza.iqSet(
                from: state.jid.full,
                payload: XMLElement(
                    name: "query",
                    xmlns: "jabber:iq:roster",
                    attributes: [:],
                    content: nil,
                    nodes: [
                        XMLElement(
                            name: "item",
                            xmlns: nil,
                            attributes: attr,
                            content: nil,
                            nodes: []
                        )
                    ]
                )
            )
            if let req {
                await registry.enqueue(req)
                return .stanzaOutbound(req)
            } else {
                return nil
            }

        case .deleteRosterItem(let jidStr):
            let req = Stanza.iqSet(
                from: state.jid.full,
                payload: XMLElement(
                    name: "query",
                    xmlns: "jabber:iq:roster",
                    attributes: [:],
                    content: nil,
                    nodes: [
                        XMLElement(
                            name: "item",
                            xmlns: nil,
                            attributes: ["jid": jidStr, "subscription": "remove"],
                            content: nil,
                            nodes: []
                        )
                    ]
                )
            )
            if let req {
                await registry.enqueue(req)
                return .stanzaOutbound(req)
            } else {
                return nil
            }

        case .stanzaInbound(let stanza):
            if let query = stanza.wrapped.nodes.first(where: { $0.name == "query" }), query.xmlns == "jabber:iq:roster" {
                // update version if needed
                if let ver = stanza.wrapped.nodes.first(where: { $0.name == "query" })?.attributes["ver"] {
                    await storage?.setRosterVer(jid: state.jid, version: ver)
                }

                // process stanza
                return await processInbound(state: state, stanza: stanza)
            } else {
                return nil
            }

        // if we signaling "result" to server it means roster was updated on this resource
        case .stanzaOutbound(let stanza):
            if let query = stanza.wrapped.nodes.first(where: { $0.name == "query" }), query.xmlns == "jabber:iq:roster", stanza.type == .iq(.result) {
                return .rosterUpdated
            } else {
                return nil
            }

        default:
            return nil
        }
    }
}

private extension RosterModule {
    func processInbound(state: ClientState, stanza: Stanza) async -> Event? {
        switch stanza.type {
        // "set" type stanza from server is just "push", so change roster accordingly
        case .iq(.set):
            // sanity check (according RFC6121 skip this push if its not for current resource)
            if stanza.wrapped.attributes["to"] != state.jid.bare {
                return nil
            }

            // get exists roster items
            var existItems: [RosterItem] = []
            if let data = await storage?.getRoster(jid: state.jid), let decoded = try? JSONDecoder().decode([XMLElement].self, from: data) {
                existItems = decoded.compactMap { RosterItem(wrap: $0, owner: state.jid) }
            }

            // get item from stanza
            let xmlItem = stanza.wrapped
                .nodes
                .first(where: { $0.name == "query" })?
                .nodes
                .first(where: { $0.name == "item" })
            guard let xmlItem, let item = RosterItem(wrap: xmlItem, owner: state.jid) else { return nil }

            // filter out exists items
            existItems = existItems.filter { $0.id != item.id }

            // append updated/new item (if it not "removed")
            if item.subsription != .remove {
                existItems.append(item)
            }

            // save roster
            guard let data = try? JSONEncoder().encode(existItems.map { $0.wrapped }) else { return nil }
            await storage?.setRoster(jid: state.jid, roster: data)

            // according RFC6120 it is not necessary to answer server with "result" on push, but why not?
            let result = Stanza.iqResult(
                from: state.jid.full,
                payload: XMLElement(
                    name: "query",
                    xmlns: "jabber:iq:roster",
                    attributes: [:],
                    content: nil,
                    nodes: []
                )
            )
            if let result {
                return .stanzaOutbound(result)
            } else {
                return nil
            }

        // the "result" from server is an answer from one of our "set", or initial roster request
        case .iq(.result):
            // get request stanza from registry
            guard let respId = stanza.id, let req = await registry.deuque(for: respId) else { return nil }

            // get exists roster items
            var existItems: [RosterItem] = []
            if let data = await storage?.getRoster(jid: state.jid), let decoded = try? JSONDecoder().decode([XMLElement].self, from: data) {
                existItems = decoded.compactMap { RosterItem(wrap: $0, owner: state.jid) }
            }

            // perform changes based on stanza request type
            switch req.type {
            // for simple roster request
            case .iq(.get):
                // get items from response
                let stanzaItems = stanza.wrapped
                    .nodes
                    .first(where: { $0.name == "query" })?
                    .nodes
                    .filter { $0.name == "item" }
                    .compactMap { RosterItem(wrap: $0, owner: state.jid) } ?? []
                let stanzaItemsIds = stanzaItems.map { $0.id }

                // filter out exists items and append updated
                existItems = existItems.filter { !stanzaItemsIds.contains($0.id) }
                existItems.append(contentsOf: stanzaItems)

                // save roster
                guard let data = try? JSONEncoder().encode(existItems.map { $0.wrapped }) else { return nil }
                await storage?.setRoster(jid: state.jid, roster: data)
                return .rosterUpdated

            // for result of one of our request
            case .iq(.set):
                // get item from request stanza
                let xmlItem = req.wrapped
                    .nodes
                    .first(where: { $0.name == "query" })?
                    .nodes
                    .first(where: { $0.name == "item" })
                guard let xmlItem, let item = RosterItem(wrap: xmlItem, owner: state.jid) else { return nil }

                // filter out exists items
                existItems = existItems.filter { $0.id != item.id }

                // append updated/new item (if it not "removed")
                if item.subsription != .remove {
                    existItems.append(item)
                }

                // save roster
                guard let data = try? JSONEncoder().encode(existItems.map { $0.wrapped }) else { return nil }
                await storage?.setRoster(jid: state.jid, roster: data)
                return .rosterUpdated

            default:
                return nil
            }

        // TODO: add error handling here
        case .iq(.error):
            guard let respId = stanza.id, let req = await registry.deuque(for: respId) else { return nil }
            print("Error on request: \(req.wrapped)\n with: \(stanza.wrapped)")
            return nil

        default:
            return nil
        }
    }
}