added omemo2

Tests/LumaOMEMO2Tests.swift

@@ -0,0 +1,317 @@
+import XCTest
+import Foundation
+import CryptoKit
+@testable import Luma
+import Martin
+import MartinOMEMO
+
+final class LumaOMEMO2Tests: XCTestCase {
+
+    // MARK: - Payload crypto (XEP-0384 0.8.3, section 4.4)
+
+    func testDerivePayloadKeysHasStableStructure() throws {
+        let key = Data(repeating: 0x42, count: 32)
+        let derived = try XCTUnwrap(LumaOMEMO2Module.derivePayloadKeys(from: key))
+        XCTAssertEqual(derived.encryptionKey.count, 32)
+        XCTAssertEqual(derived.authKey.count, 32)
+        XCTAssertEqual(derived.iv.count, 16)
+
+        // Deterministic and key-dependent.
+        XCTAssertEqual(derived, LumaOMEMO2Module.derivePayloadKeys(from: key))
+        let other = try XCTUnwrap(LumaOMEMO2Module.derivePayloadKeys(from: Data(repeating: 0x43, count: 32)))
+        XCTAssertNotEqual(derived, other)
+
+        // Wrong input size is rejected.
+        XCTAssertNil(LumaOMEMO2Module.derivePayloadKeys(from: Data([1, 2, 3])))
+    }
+
+    func testCBCRoundTrip() throws {
+        let key = Data(repeating: 0x11, count: 32)
+        let iv = Data(repeating: 0x22, count: 16)
+        let plaintext = Data("OMEMO 2 payload".utf8)
+
+        let ciphertext = try XCTUnwrap(LumaOMEMO2Module.aes256CBCEncrypt(plaintext, key: key, iv: iv))
+        XCTAssertNotEqual(ciphertext, plaintext)
+        let recovered = try XCTUnwrap(LumaOMEMO2Module.aes256CBCDecrypt(ciphertext, key: key, iv: iv))
+        XCTAssertEqual(recovered, plaintext)
+
+        // A wrong key never recovers the plaintext: the padding check may
+        // either reject the result or yield garbage, but never the content.
+        // Integrity itself is guaranteed by the HMAC step in finishDecode.
+        let wrongKey = Data(repeating: 0x12, count: 32)
+        let wrong = LumaOMEMO2Module.aes256CBCDecrypt(ciphertext, key: wrongKey, iv: iv)
+        if let wrong {
+            XCTAssertNotEqual(wrong, plaintext)
+        }
+
+        // PKCS#7 padding rounds up to a whole block.
+        XCTAssertEqual(ciphertext.count % 16, 0)
+    }
+
+    func testHMACTruncationMatchesManualComputation() throws {
+        let authKey = Data(repeating: 0x33, count: 32)
+        let message = Data("authenticate me".utf8)
+        let full = Data(HMAC<SHA256>.authenticationCode(for: message, using: SymmetricKey(data: authKey)))
+        XCTAssertEqual(full.count, 32)
+        // The wire format keeps the first 16 bytes ("cutting off excess
+        // bytes from the end" in the spec).
+        let truncated = Data(full.prefix(16))
+        XCTAssertEqual(truncated.count, 16)
+        XCTAssertEqual(LumaOMEMO2Module.constantTimeEquals(truncated, full.subdata(in: 0..<16)), true)
+        XCTAssertFalse(LumaOMEMO2Module.constantTimeEquals(truncated, Data(repeating: 0, count: 16)))
+        XCTAssertFalse(LumaOMEMO2Module.constantTimeEquals(truncated, Data([1])))
+    }
+
+    // MARK: - SCE envelope (XEP-0420)
+
+    func testEnvelopeRoundTrip() throws {
+        let xml = LumaOMEMO2Module.envelopeXML(
+            body: "Привет & <тест>",
+            from: "user@example.org",
+            to: nil
+        )
+        let envelope = try XCTUnwrap(Element.from(string: xml))
+        XCTAssertEqual(envelope.name, "envelope")
+        XCTAssertEqual(envelope.xmlns, "urn:xmpp:sce:1")
+
+        let content = try XCTUnwrap(envelope.findChild(name: "content"))
+        let body = try XCTUnwrap(content.findChild(name: "body"))
+        XCTAssertEqual(body.value, "Привет & <тест>")
+        XCTAssertEqual(body.xmlns, "jabber:client")
+
+        // rpad must be present and non-empty.
+        let rpad = try XCTUnwrap(envelope.findChild(name: "rpad")?.value)
+        XCTAssertFalse(rpad.isEmpty)
+
+        // from affix is mandatory.
+        XCTAssertEqual(envelope.findChild(name: "from")?.getAttribute("jid"), "user@example.org")
+        // No <to> without a MUC.
+        XCTAssertNil(envelope.findChild(name: "to"))
+    }
+
+    func testEnvelopeIncludesToAffixForMUC() throws {
+        let xml = LumaOMEMO2Module.envelopeXML(
+            body: "hi",
+            from: "user@example.org",
+            to: "room@conference.example.org"
+        )
+        let envelope = try XCTUnwrap(Element.from(string: xml))
+        XCTAssertEqual(
+            envelope.findChild(name: "to")?.getAttribute("jid"),
+            "room@conference.example.org"
+        )
+    }
+
+    // MARK: - Bundle parsing
+
+    func testBundleParsing() throws {
+        let bundle = Element(name: "bundle", xmlns: "urn:xmpp:omemo:2")
+        bundle.addChild(Element(name: "spk", cdata: Data([1, 2, 3]).base64EncodedString(), attributes: ["id": "7"]))
+        bundle.addChild(Element(name: "spks", cdata: Data([4, 5, 6]).base64EncodedString()))
+        bundle.addChild(Element(name: "ik", cdata: Data([7, 8, 9]).base64EncodedString()))
+        let prekeys = Element(name: "prekeys")
+        prekeys.addChild(Element(name: "pk", cdata: Data([10, 11]).base64EncodedString(), attributes: ["id": "11"]))
+        prekeys.addChild(Element(name: "pk", cdata: Data([12, 13]).base64EncodedString(), attributes: ["id": "12"]))
+        bundle.addChild(prekeys)
+
+        let parsed = try XCTUnwrap(OMEMO2Bundle(from: bundle))
+        XCTAssertEqual(parsed.signedPreKeyId, 7)
+        XCTAssertEqual(parsed.signedPreKeyPublic, Data([1, 2, 3]))
+        XCTAssertEqual(parsed.signedPreKeySignature, Data([4, 5, 6]))
+        XCTAssertEqual(parsed.identityKey, Data([7, 8, 9]))
+        XCTAssertEqual(parsed.preKeys.count, 2)
+
+        // Wrong namespace is rejected.
+        bundle.xmlns = "urn:xmpp:omemo:0"
+        XCTAssertNil(OMEMO2Bundle(from: bundle))
+    }
+
+    // MARK: - Double Ratchet round trip (shared with the legacy module)
+
+    func testRatchetSessionRoundTrip() throws {
+        let alice = try makeStorage()
+        let bob = try makeStorage()
+
+        let aliceContext = try XCTUnwrap(SignalContext(withStorage: alice))
+        let bobContext = try XCTUnwrap(SignalContext(withStorage: bob))
+
+        let aliceDevice = Int32(bitPattern: alice.identities.localRegistrationId())
+        let aliceAddress = SignalAddress(name: "alice@example.org", deviceId: aliceDevice)
+        let bobAddress = SignalAddress(name: "bob@example.org", deviceId: Int32(bitPattern: bob.identities.localRegistrationId()))
+
+        // Bob builds a session with Alice from her published bundle.
+        let preKey = try XCTUnwrap(alice.preKeys.loadPreKey(withId: 1))
+        let preKeyRecord = try XCTUnwrap(SignalPreKey(fromSerializedData: preKey))
+        let signedPreKey = try XCTUnwrap(alice.signedPreKeys.loadSignedPreKey(withId: 1))
+        let signedPreKeyRecord = try XCTUnwrap(SignalSignedPreKey(fromSerializedData: signedPreKey))
+        let identityKey = try XCTUnwrap(alice.identities.keyPair()?.publicKey)
+        let bundle = try XCTUnwrap(SignalPreKeyBundle(
+            registrationId: 0,
+            deviceId: aliceDevice,
+            preKeyId: 1,
+            preKeyPublic: try XCTUnwrap(preKeyRecord.serializedPublicKey),
+            signedPreKeyId: 1,
+            signedPreKeyPublic: try XCTUnwrap(signedPreKeyRecord.publicKeyData),
+            signedPreKeySignature: signedPreKeyRecord.signature,
+            identityKey: identityKey
+        ))
+        let builder = try XCTUnwrap(SignalSessionBuilder(withAddress: aliceAddress, andContext: bobContext))
+        XCTAssertTrue(builder.processPreKeyBundle(bundle: bundle))
+
+        // The 48-byte combined key (payload key + truncated HMAC) travels
+        // through the ratchet exactly like in LumaOMEMO2Module.
+        var combined = Data(repeating: 0x5A, count: 32)
+        combined.append(Data(repeating: 0x3C, count: 16))
+
+        let bobCipher = try XCTUnwrap(SignalSessionCipher(withAddress: aliceAddress, andContext: bobContext))
+        let encryptedKey = try bobCipher.encrypt(data: combined).get()
+        XCTAssertTrue(encryptedKey.prekey)
+
+        let aliceCipher = try XCTUnwrap(SignalSessionCipher(withAddress: bobAddress, andContext: aliceContext))
+        let decrypted = try aliceCipher.decrypt(key: SignalSessionCipher.Key(
+            key: encryptedKey.key,
+            deviceId: Int32(bitPattern: bob.identities.localRegistrationId()),
+            prekey: true
+        )).get()
+        XCTAssertEqual(decrypted, combined)
+    }
+}
+
+// MARK: - In-memory signal stores
+
+private final class TestOMEMOStorage: SignalStorage {
+    let sessions = InMemorySessionStore()
+    let preKeys = InMemoryPreKeyStore()
+    let signedPreKeys = InMemorySignedPreKeyStore()
+    let identities = InMemoryIdentityKeyStore()
+    let senderKeys = InMemorySenderKeyStore()
+
+    init() {
+        super.init(
+            sessionStore: sessions,
+            preKeyStore: preKeys,
+            signedPreKeyStore: signedPreKeys,
+            identityKeyStore: identities,
+            senderKeyStore: senderKeys
+        )
+    }
+
+    override func setup(withContext context: SignalContext) {
+        identities.registrationID = context.generateRegistrationId()
+        if let pair = SignalIdentityKeyPair.generateKeyPair(context: context) {
+            identities.keyPairData = pair.serialized()
+        }
+        if let identityKeyPair = identities.keyPair(),
+            let signedPreKey = context.generateSignedPreKey(withIdentity: identityKeyPair, signedPreKeyId: 1),
+            let serialized = signedPreKey.serializedData {
+            _ = signedPreKeys.storeSignedPreKey(serialized, withId: 1)
+        }
+        let generated = context.generatePreKeys(withStartingPreKeyId: 1, count: 1)
+        for preKey in generated {
+            if let serialized = preKey.serializedData {
+                _ = preKeys.storePreKey(serialized, withId: preKey.preKeyId)
+            }
+        }
+        super.setup(withContext: context)
+    }
+}
+
+private final class InMemorySessionStore: SignalSessionStoreProtocol {
+    private var records: [String: Data] = [:]
+    func sessionRecord(forAddress address: SignalAddress) -> Data? {
+        records["\(address.name)|\(address.deviceId)"]
+    }
+    func allDevices(for name: String, activeAndTrusted: Bool) -> [Int32] { [] }
+    func storeSessionRecord(_ data: Data, forAddress address: SignalAddress) -> Bool {
+        records["\(address.name)|\(address.deviceId)"] = data
+        return true
+    }
+    func containsSessionRecord(forAddress address: SignalAddress) -> Bool {
+        sessionRecord(forAddress: address) != nil
+    }
+    func deleteSessionRecord(forAddress address: SignalAddress) -> Bool {
+        records.removeValue(forKey: "\(address.name)|\(address.deviceId)") != nil
+    }
+    func deleteAllSessions(for name: String) -> Bool {
+        records.removeAll()
+        return true
+    }
+}
+
+private final class InMemoryPreKeyStore: SignalPreKeyStoreProtocol {
+    private var keys: [UInt32: Data] = [:]
+    private var pendingDeletion: Set<UInt32> = []
+    func currentPreKeyId() -> UInt32 { keys.keys.max() ?? 0 }
+    func loadPreKey(withId id: UInt32) -> Data? { keys[id] }
+    func storePreKey(_ data: Data, withId id: UInt32) -> Bool {
+        keys[id] = data
+        return true
+    }
+    func containsPreKey(withId id: UInt32) -> Bool { keys[id] != nil }
+    func deletePreKey(withId id: UInt32) -> Bool {
+        pendingDeletion.insert(id)
+        return true
+    }
+    func flushDeletedPreKeys() -> Bool {
+        let ids = pendingDeletion
+        pendingDeletion.removeAll()
+        ids.forEach { keys.removeValue(forKey: $0) }
+        return !ids.isEmpty
+    }
+}
+
+private final class InMemorySignedPreKeyStore: SignalSignedPreKeyStoreProtocol {
+    private var keys: [UInt32: Data] = [:]
+    func countSignedPreKeys() -> Int { keys.count }
+    func loadSignedPreKey(withId id: UInt32) -> Data? { keys[id] }
+    func storeSignedPreKey(_ data: Data, withId id: UInt32) -> Bool {
+        keys[id] = data
+        return true
+    }
+    func containsSignedPreKey(withId id: UInt32) -> Bool { keys[id] != nil }
+    func deleteSignedPreKey(withId id: UInt32) -> Bool {
+        keys.removeValue(forKey: id) != nil
+    }
+}
+
+private final class InMemoryIdentityKeyStore: SignalIdentityKeyStoreProtocol {
+    var registrationID: UInt32 = 0
+    var keyPairData: Data?
+    private var identities: [String: Data] = [:]
+    func keyPair() -> SignalIdentityKeyPairProtocol? {
+        guard let keyPairData else { return nil }
+        return SignalIdentityKeyPair(fromKeyPairData: keyPairData)
+    }
+    func localRegistrationId() -> UInt32 { registrationID }
+    func save(identity: SignalAddress, key: SignalIdentityKeyProtocol?) -> Bool {
+        save(identity: identity, publicKeyData: key?.publicKey)
+    }
+    func isTrusted(identity: SignalAddress, key: SignalIdentityKeyProtocol?) -> Bool {
+        isTrusted(identity: identity, publicKeyData: key?.publicKey)
+    }
+    func save(identity: SignalAddress, publicKeyData: Data?) -> Bool {
+        guard let publicKeyData else { return false }
+        identities["\(identity.name)|\(identity.deviceId)"] = publicKeyData
+        return true
+    }
+    func isTrusted(identity: SignalAddress, publicKeyData: Data?) -> Bool { true }
+    func setStatus(_ status: IdentityStatus, forIdentity identity: SignalAddress) -> Bool { true }
+    func setStatus(active: Bool, forIdentity identity: SignalAddress) -> Bool { true }
+    func identities(forName name: String) -> [Identity] { [] }
+    func identityFingerprint(forAddress address: SignalAddress) -> String? { nil }
+}
+
+private final class InMemorySenderKeyStore: SignalSenderKeyStoreProtocol {
+    private var keys: [String: Data] = [:]
+    func storeSenderKey(_ key: Data, address: SignalAddress?, groupId: String?) -> Bool {
+        keys["\(address?.name ?? "-")|\(groupId ?? "-")"] = key
+        return true
+    }
+    func loadSenderKey(forAddress address: SignalAddress?, groupId: String?) -> Data? {
+        keys["\(address?.name ?? "-")|\(groupId ?? "-")"]
+    }
+}
+
+private func makeStorage() throws -> TestOMEMOStorage {
+    TestOMEMOStorage()
+}