Core Apps » ChatSecure:Android

//

//  OTROMEMOStorageTest.swift

//  ChatSecure

//

//  Created by David Chiles on 9/16/16.

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

//

import XCTest

@testable import ChatSecureCore

extension OTROMEMOSignalCoordinator {

    convenience init(accountYapKey: String, databaseConnection: YapDatabaseConnection) throws {

        let capsStorage = XMPPCapabilitiesCoreDataStorage(inMemoryStore: ())!

        let serverCaps = OTRServerCapabilities()

        let caps = XMPPCapabilities(capabilitiesStorage: capsStorage)

        let file = FileTransferManager(connection: databaseConnection, serverCapabilities: serverCaps, sessionConfiguration: URLSessionConfiguration.ephemeral)

        let messageStorage = MessageStorage(connection: databaseConnection, capabilities: caps, fileTransfer: file)

        try self.init(accountYapKey: accountYapKey, databaseConnection: databaseConnection, messageStorage: messageStorage)

    }

}

class OTROmemoStorageTest: XCTestCase {

    var databaseManager:OTRDatabaseManager!

    var omemoStorage:OTROMEMOStorageManager!

    var signalStorage:OTRSignalStorageManager!

    var signalCoordinator:OTROMEMOSignalCoordinator!

    var accountKey:String!

    var accountCollection:String!

    let initialDevices:[NSNumber] = [1,2,3]

    let secondDeviceNumbers:[NSNumber] = [1,2,4,5]

    override func setUp() {

        super.setUp()

        FileManager.default.clearDirectory(OTRTestDatabaseManager.yapDatabaseDirectory())

    }

    override func tearDown() {

        // Put teardown code here. This method is called after the invocation of each test method in the class.

        super.tearDown()

    }

    /** 

     * This creates teh necessary database and signal managers. Also creates an account which is required for all the above objects

     *

     */

    func setupDatabase(_ name:String) {

        let account = TestXMPPAccount(username: "", accountType: .jabber)!

        self.accountKey = account.uniqueId

        self.accountCollection = OTRXMPPAccount.collection

        self.databaseManager = OTRTestDatabaseManager()

        self.databaseManager.setDatabasePassphrase("help", remember: false, error: nil)

        self.databaseManager.setupDatabase(withName: name, withMediaStorage: false)

        self.omemoStorage = OTROMEMOStorageManager(accountKey: accountKey, accountCollection:accountCollection, databaseConnection: databaseManager.readWriteDatabaseConnection!)

        self.signalStorage = OTRSignalStorageManager(accountKey: accountKey, databaseConnection: databaseManager.readWriteDatabaseConnection!, delegate: nil)

        self.signalCoordinator = try! OTROMEMOSignalCoordinator(accountYapKey: accountKey, databaseConnection: databaseManager.readWriteDatabaseConnection!)

        databaseManager.readWriteDatabaseConnection?.readWrite( { (transaction) in

            account.save(with: transaction)

        })

    }

    /**

     * Store the inital device list.

     * Ensure the correct trust level is set and that it's attached to the correct account.

     */

    func storeInitialDevices() {

        self.omemoStorage.storeOurDevices(self.initialDevices)

        let firstStoredDevices = omemoStorage.getDevicesForOurAccount(nil)

        XCTAssertEqual(firstStoredDevices.count, self.initialDevices.count)

        firstStoredDevices.forEach { (device) in

            XCTAssert(device.trustLevel == .trustedTofu)

            XCTAssertEqual(device.parentKey, self.accountKey)

            XCTAssertEqual(device.parentCollection, self.accountCollection)

        }

        let firstStoredDeviceId = firstStoredDevices.map { (device) -> NSNumber in

            return device.deviceId

        }

        XCTAssertEqual(Set(self.initialDevices), Set(firstStoredDeviceId))

    }

    /**

     * Store the second device list. This should override the inital device list and remove devie 3 and add device 4 and 5.

     * Checks to make sure the trust level is set correctly.

     */

    func storeSecondDeviceList() {

        //Now simulate getting a new set of devices where one was removed and two were added.

        omemoStorage.storeOurDevices(self.secondDeviceNumbers)

        let secondStoredDevices = omemoStorage.getDevicesForOurAccount(nil)

        XCTAssertEqual(secondStoredDevices.count, 5)

        secondStoredDevices.forEach { (device) in

            XCTAssertEqual(device.parentKey, accountKey)

            XCTAssertEqual(device.parentCollection, accountCollection)

            switch device.deviceId {

            case 3:

                XCTAssert(device.trustLevel == .removed)

            case 4:

                fallthrough

            case 5:

                XCTAssert(device.trustLevel == .untrustedNew)

            default:

                XCTAssert(device.trustLevel == .trustedTofu,"Device \(device.deviceId) Should be Tofu")

            }

        }

        let secondStoredDeviceId = secondStoredDevices.map { (device) -> NSNumber in

            return device.deviceId

        }

        XCTAssertEqual(Set([1,2,3,4,5]), Set(secondStoredDeviceId))

    }

    func testOmemoFirstDevicesStorage() {

        self.setupDatabase(#function)

        self.storeInitialDevices()

    }

    func testOmemoSecondDeviceStorage() {

        self.setupDatabase(#function)

        self.storeInitialDevices()

        self.storeSecondDeviceList()

    }

    /**

     * Do the above test and then remove all devices and ensure there are no devices left.

     */

    func testOmemoRemoveDevices() {

        self.setupDatabase(#function)

        self.storeInitialDevices()

        self.storeSecondDeviceList()

        let thirdDeviceNumbers = [NSNumber]()

        omemoStorage.storeOurDevices(thirdDeviceNumbers)

        let thirdStoredDevices = omemoStorage.getDevicesForParentYapKey(accountKey, yapCollection: accountCollection, trusted:true)

        XCTAssertEqual(thirdStoredDevices.count, 0)

    }

    /**

     * 1. Check to make sure there are no pre-keys currently stored.

     * 2. Add 3 preKeys and check that the current max is correct.

     * 3. Delete a pre-key and check that values are correct

     */

    func testPreKeyStorage() {

        self.setupDatabase(#function)

        let none = self.signalStorage.currentMaxPreKeyId()

        XCTAssertNil(none)

        //Save a few random prekeys with id.

        //Normally these should be sequential but shouldn't matter

        let keyIds:[UInt32] =  [1,2,100]

        keyIds.forEach { (id) in

            let result = self.signalStorage.storePreKey(Data(), preKeyId: id)

            XCTAssertTrue(result)

        }

        XCTAssertEqual(self.signalStorage.currentMaxPreKeyId(),100)

        // Delete a prekey. In this case it's the max id but that shouldn't matter either.

        XCTAssertTrue(self.signalStorage.deletePreKey(withId:100))

        XCTAssertFalse(self.signalStorage.containsPreKey(withId:100))

        XCTAssertFalse(self.signalStorage.containsPreKey(withId:999))

        XCTAssertTrue(self.signalStorage.containsPreKey(withId:2))

        let result = self.signalStorage.currentMaxPreKeyId()!

        XCTAssertEqual(result, 100)

        //Fetch all the remaining prekeys should get back 1,2

        let allRemainingPrekeys = self.signalStorage.fetchAllPreKeys(false)

        let remainingPreKeyIds = allRemainingPrekeys.map { (prekey) -> UInt32 in

            return prekey.keyId

        }

        //Check that it has the two id expected and not the one that was deleted

        XCTAssertTrue(remainingPreKeyIds.contains(1))

        XCTAssertTrue(remainingPreKeyIds.contains(2))

        XCTAssertFalse(remainingPreKeyIds.contains(100))

        XCTAssertEqual(allRemainingPrekeys.count, 2)

    }

    /**

     * NOTES: This test is failing because the stored bundle cannot

     * be validated after re-fetching, so it regenerates a completely

     * new bundle. Why does store/fetch corrupt the bundle?

     *

     * 1. Test generating a bundle like on first launch.

     * 2. Remove a few pre keys. This simulates what will happen when the signal library uses up some pre keys from incoming bundle messages

     * 3. Fetch our own bundle again. This time it should all come from the database and the only new information is two new pre-keys to replace teh ones deleted

     */

    func testOurBundleStorage() {

        self.setupDatabase(#function)

        let firstFetch = self.signalCoordinator.fetchMyBundle()!

        XCTAssertNotNil(firstFetch)

        var firstPreKeyFetch = [UInt32:Data]()

        firstFetch.preKeys.forEach { (preKey) in

            firstPreKeyFetch.updateValue(preKey.publicKey, forKey: preKey.preKeyId)

        }

        //Remove some pre keys

        XCTAssertTrue(self.signalStorage.deletePreKey(withId:22))

        XCTAssertTrue(self.signalStorage.deletePreKey(withId:25))

        //Fetch again

        let secondFetch = self.signalCoordinator.fetchMyBundle()!

        XCTAssertNotNil(secondFetch)

        var secondPreKeyFetch = [UInt32:Data]()

        secondFetch.preKeys.forEach { (preKey) in

            secondPreKeyFetch.updateValue(preKey.publicKey, forKey: preKey.preKeyId)

        }

        XCTAssertEqual(firstFetch.deviceId, secondFetch.deviceId,"Should be the same device id")

        XCTAssertEqual(firstFetch.identityKey, secondFetch.identityKey,"Same Identity Key")

        XCTAssertEqual(firstFetch.signedPreKey.signature, secondFetch.signedPreKey.signature,"Same signature")

        XCTAssertEqual(firstFetch.signedPreKey.preKeyId, secondFetch.signedPreKey.preKeyId,"Same prekey Id")

        XCTAssertEqual(firstFetch.signedPreKey.publicKey, secondFetch.signedPreKey.publicKey,"Same prekey public key")

        //Checking Pre Keys

        let firstIdArray = Array(firstPreKeyFetch.keys).sorted()

        let secondIdArray = Array(secondPreKeyFetch.keys).sorted()

        //The two deleted keys should not show up in the second key fetch

        XCTAssertFalse(secondIdArray.contains(22),"Should not contain this key id")

        XCTAssertFalse(secondIdArray.contains(25),"Should not contain this key id")

        // Both times it should fetch 100 keys

        XCTAssertEqual(firstIdArray.count, 100,"Should have fetched 100 keys")

        XCTAssertEqual(secondIdArray.count, 100,"Should have fetched 100 keys")

        XCTAssertEqual(firstIdArray.first!, 1,"Should start with id 1")

        XCTAssertEqual(secondIdArray.first!, 1,"Should start with id 1")

        XCTAssertEqual(firstIdArray.last!, 100,"Should end with id 100")

        XCTAssertEqual(secondIdArray.last!, 102,"Should start with id 1")

        //Make sure all the data is the same as previous attempt

        secondPreKeyFetch.forEach { (id,secondData) in

            if let firstData = firstPreKeyFetch[id] {

                XCTAssertEqual(firstData, secondData,"Public key information should be the same")

            } else {

                //These should not be in the first set because they were created after we deleted two keys

                let excluded:Set<UInt32> = Set([101,102])

                XCTAssertTrue(excluded.contains(id))

            }

        }

        //Double check to make sure that the key information is there or not there depending on added and removed keys.

        XCTAssertNil(self.signalStorage.loadPreKey(withId:22))

        XCTAssertNil(self.signalStorage.loadPreKey(withId:25))

        XCTAssertNotNil(self.signalStorage.loadPreKey(withId:101))

        XCTAssertNotNil(self.signalStorage.loadPreKey(withId:102))

    }

}