Core Apps » ChatSecure:Android

//

//  SecondaryIndexes.swift

//  ChatSecureCore

//

//  Created by Chris Ballinger on 11/7/17.

//  Copyright © 2017 Chris Ballinger. All rights reserved.

//

import Foundation

import YapDatabase

extension YapDatabaseSecondaryIndexOptions {

    convenience init(whitelist: [String]) {

        let set = Set(whitelist)

        let whitelist = YapWhitelistBlacklist(whitelist: set)

        self.init()

        self.allowedCollections = whitelist

    }

}

public extension YapDatabaseSecondaryIndex {

    @objc public static var buddyIndex: YapDatabaseSecondaryIndex {

        let columns: [String:YapDatabaseSecondaryIndexType] = [

            BuddyIndexColumnName.accountKey: .text,

            BuddyIndexColumnName.username: .text

        ]

        let setup = YapDatabaseSecondaryIndexSetup(capacity: UInt(columns.count))

        columns.forEach { (key, value) in

            setup.addColumn(key, with: value)

        }

        let handler = YapDatabaseSecondaryIndexHandler.withObjectBlock { (transaction, dict, collection, key, object) in

            guard let buddy = object as? OTRXMPPBuddy else {

                return

            }

            dict[BuddyIndexColumnName.accountKey] = buddy.accountUniqueId

            dict[BuddyIndexColumnName.username] = buddy.username

        }

        let options = YapDatabaseSecondaryIndexOptions(whitelist: [OTRXMPPBuddy.collection])

        let secondaryIndex = YapDatabaseSecondaryIndex(setup: setup, handler: handler, versionTag: "2", options: options)

        return secondaryIndex

    }

    @objc public static var messageIndex: YapDatabaseSecondaryIndex {

        let columns: [String:YapDatabaseSecondaryIndexType] = [

            MessageIndexColumnName.messageKey: .text,

            MessageIndexColumnName.remoteMessageId: .text,

            MessageIndexColumnName.threadId: .text,

            MessageIndexColumnName.isMessageRead: .integer,

            MessageIndexColumnName.originId: .text,

            MessageIndexColumnName.stanzaId: .text

        ]

        let setup = YapDatabaseSecondaryIndexSetup(capacity: UInt(columns.count))

        columns.forEach { (key, value) in

            setup.addColumn(key, with: value)

        }

        let handler = YapDatabaseSecondaryIndexHandler.withObjectBlock { (transaction, dict, collection, key, object) in

            guard let message = object as? OTRMessageProtocol else {

                return

            }

            if let remoteMessageId = message.remoteMessageId,

                remoteMessageId.count > 0 {

                dict[MessageIndexColumnName.remoteMessageId] = remoteMessageId

            }

            if message.messageKey.count > 0 {

                dict[MessageIndexColumnName.messageKey] = message.messageKey

            }

            dict[MessageIndexColumnName.isMessageRead] = message.isMessageRead

            if message.threadId.count > 0 {

                dict[MessageIndexColumnName.threadId] = message.threadId

            }

            if let originId = message.originId, originId.count > 0 {

                dict[MessageIndexColumnName.originId] = originId

            }

            if let stanzaId = message.stanzaId, stanzaId.count > 0 {

                dict[MessageIndexColumnName.stanzaId] = stanzaId

            }

        }

        // These are actually the same collection

        let options = YapDatabaseSecondaryIndexOptions(whitelist: [OTRBaseMessage.collection, OTRXMPPRoomMessage.collection])

        let secondaryIndex = YapDatabaseSecondaryIndex(setup: setup, handler: handler, versionTag: "6", options: options)

        return secondaryIndex

    }

    @objc public static var signalIndex: YapDatabaseSecondaryIndex {

        let columns: [String:YapDatabaseSecondaryIndexType] = [

            SignalIndexColumnName.session: .text,

            SignalIndexColumnName.preKeyId: .integer,

            SignalIndexColumnName.preKeyAccountKey: .text

        ]

        let setup = YapDatabaseSecondaryIndexSetup(capacity: UInt(columns.count))

        columns.forEach { (key, value) in

            setup.addColumn(key, with: value)

        }

        let handler = YapDatabaseSecondaryIndexHandler.withObjectBlock { (transaction, dict, collection, key, object) in

            if let session = object as? OTRSignalSession {

                if session.name.count > 0 {

                    dict[SignalIndexColumnName.session] = session.sessionKey

                }

            } else if let preKey = object as? OTRSignalPreKey {

                dict[SignalIndexColumnName.preKeyId] = preKey.keyId

                if preKey.accountKey.count > 0 {

                    dict[SignalIndexColumnName.preKeyAccountKey] = preKey.accountKey

                }

            }

        }

        let options = YapDatabaseSecondaryIndexOptions(whitelist: [OTRSignalPreKey.collection,OTRSignalSession.collection])

        let secondaryIndex = YapDatabaseSecondaryIndex(setup: setup, handler: handler, versionTag: "6", options: options)

        return secondaryIndex

    }

    @objc public static var roomOccupantIndex: YapDatabaseSecondaryIndex {

        let columns: [String:YapDatabaseSecondaryIndexType] = [

            RoomOccupantIndexColumnName.jid: .text,

            RoomOccupantIndexColumnName.realJID: .text,

            RoomOccupantIndexColumnName.roomUniqueId: .text,

            RoomOccupantIndexColumnName.buddyUniqueId: .text,

        ]

        let setup = YapDatabaseSecondaryIndexSetup(capacity: UInt(columns.count))

        columns.forEach { (key, value) in

            setup.addColumn(key, with: value)

        }

        let handler = YapDatabaseSecondaryIndexHandler.withObjectBlock { (transaction, dict, collection, key, object) in

            guard let occupant = object as? OTRXMPPRoomOccupant else {

                return

            }

            if let jid = occupant.jid, jid.count > 0 {

                dict[RoomOccupantIndexColumnName.jid] = jid

            }

            if let realJID = occupant.realJID, realJID.count > 0 {

                dict[RoomOccupantIndexColumnName.realJID] = realJID

            }

            if let roomUniqueId = occupant.roomUniqueId, roomUniqueId.count > 0 {

                dict[RoomOccupantIndexColumnName.roomUniqueId] = roomUniqueId

            }

            if let buddyUniqueId = occupant.buddyUniqueId, buddyUniqueId.count > 0 {

                dict[RoomOccupantIndexColumnName.buddyUniqueId] = buddyUniqueId

            }

        }

        let options = YapDatabaseSecondaryIndexOptions(whitelist: [OTRXMPPRoomOccupant.collection])

        let secondaryIndex = YapDatabaseSecondaryIndex(setup: setup, handler: handler, versionTag: "4", options: options)

        return secondaryIndex

    }

}

// MARK: - Extensions

public extension OTRSignalSession {

    /// "\(accountKey)-\(name)", only used for SecondaryIndex lookups

    public var sessionKey: String {

        return OTRSignalSession.sessionKey(accountKey: accountKey, name: name)

    }

    /// "\(accountKey)-\(name)", only used for SecondaryIndex lookups

    public static func sessionKey(accountKey: String, name: String) -> String {

        return "\(accountKey)-\(name)"

    }

}

public extension OTRXMPPBuddy {

    /// This function should only be used when the secondary index is not ready

    private static func slowLookup(jid: XMPPJID,

                                   accountUniqueId: String,

                                   transaction: YapDatabaseReadTransaction) -> OTRXMPPBuddy? {

        DDLogWarn("WARN: Using slow O(n) lookup for OTRXMPPBuddy: \(jid)")

        var buddy: OTRXMPPBuddy? = nil

        transaction.enumerateKeysAndObjects(inCollection: OTRXMPPBuddy.collection) { (key, object, stop) in

            if let potentialMatch = object as? OTRXMPPBuddy,

                potentialMatch.username == jid.bare {

                buddy = potentialMatch

                stop.pointee = true

            }

        }

        return buddy

    }

    /// Fetch buddy matching JID using secondary index

    @objc public static func fetchBuddy(jid: XMPPJID,

                                        accountUniqueId: String,

                                        transaction: YapDatabaseReadTransaction) -> OTRXMPPBuddy? {

        guard let indexTransaction = transaction.ext(SecondaryIndexName.buddy) as? YapDatabaseSecondaryIndexTransaction else {

            DDLogError("Error looking up OTRXMPPBuddy via SecondaryIndex: Extension not ready.")

            return self.slowLookup(jid:jid, accountUniqueId:accountUniqueId, transaction: transaction)

        }

        let queryString = "Where \(BuddyIndexColumnName.accountKey) == ? AND \(BuddyIndexColumnName.username) == ?"

        let query = YapDatabaseQuery(string: queryString, parameters: [accountUniqueId, jid.bare])

        var matchingBuddies: [OTRXMPPBuddy] = []

        let success = indexTransaction.enumerateKeysAndObjects(matching: query) { (collection, key, object, stop) in

            if let matchingBuddy = object as? OTRXMPPBuddy {

                matchingBuddies.append(matchingBuddy)

            }

        }

        if !success {

            DDLogError("Error looking up OTRXMPPBuddy with query \(query) \(jid) \(accountUniqueId)")

            return nil

        }

        if matchingBuddies.count > 1 {

            DDLogWarn("WARN: More than one OTRXMPPBuddy matching query \(query) \(jid) \(accountUniqueId): \(matchingBuddies.count)")

        }

//        #if DEBUG

//            if matchingBuddies.count == 0 {

//                DDLogWarn("WARN: No OTRXMPPBuddy matching query \(jid) \(accountUniqueId)")

//                let buddy = slowLookup(jid: jid, accountUniqueId: accountUniqueId, transaction: transaction)

//                if buddy != nil {

//                    DDLogWarn("WARN: Found buddy using O(n) lookup that wasn't found in secondary index: \(jid) \(accountUniqueId)")

//                }

//            }

//        #endif

        return matchingBuddies.first

    }

}

// MARK: - Constants

/// YapDatabase extension names for Secondary Indexes

@objc public class SecondaryIndexName: NSObject {

    @objc public static let messages = "OTRMessagesSecondaryIndex"

    @objc public static let signal = "OTRYapDatabseMessageIdSecondaryIndexExtension"

    @objc public static let roomOccupants = "SecondaryIndexName_roomOccupantIndex"

    @objc public static let buddy = "SecondaryIndexName_buddy"

}

@objc public class BuddyIndexColumnName: NSObject {

    @objc public static let accountKey = "BuddyIndexColumnName_accountKey"

    @objc public static let username = "BuddyIndexColumnName_username"

}

@objc public class MessageIndexColumnName: NSObject {

    @objc public static let messageKey = "OTRYapDatabaseMessageIdSecondaryIndexColumnName"

    @objc public static let remoteMessageId = "OTRYapDatabaseRemoteMessageIdSecondaryIndexColumnName"

    @objc public static let threadId = "OTRYapDatabaseMessageThreadIdSecondaryIndexColumnName"

    @objc public static let isMessageRead = "OTRYapDatabaseUnreadMessageSecondaryIndexColumnName"

    /// XEP-0359 origin-id

    @objc public static let originId = "SecondaryIndexNameOriginId"

    /// XEP-0359 stanza-id

    @objc public static let stanzaId = "SecondaryIndexNameStanzaId"

}

@objc public class RoomOccupantIndexColumnName: NSObject {

    /// jid

    @objc public static let jid = "OTRYapDatabaseRoomOccupantJidSecondaryIndexColumnName"

    @objc public static let realJID = "RoomOccupantIndexColumnName_realJID"

    @objc public static let roomUniqueId = "RoomOccupantIndexColumnName_roomUniqueId"

    @objc public static let buddyUniqueId = "RoomOccupantIndexColumnName_buddyUniqueId"

}

@objc public class SignalIndexColumnName: NSObject {

    @objc public static let session = "OTRYapDatabaseSignalSessionSecondaryIndexColumnName"

    @objc public static let preKeyId = "OTRYapDatabaseSignalPreKeyIdSecondaryIndexColumnName"

    @objc public static let preKeyAccountKey = "OTRYapDatabaseSignalPreKeyAccountKeySecondaryIndexColumnName"

}