added omemo2

Sources/Shared/Services/PhoneWatchBridge.swift

@@ -19,6 +19,24 @@ struct PhoneWatchSnapshot: Codable {
     let chats: [Chat]
 }
 
+/// Plain-value snapshots of the SwiftData models, taken on the main actor
+/// before the background queue builds the watch payload. Reading @Model
+/// properties off the main actor crashes with EXC_BAD_ACCESS.
+private struct PhoneWatchChatData {
+    let jid: String
+    let name: String
+    let unread: Int
+}
+
+private struct PhoneWatchMessageData {
+    let id: String
+    let conversationID: String
+    let body: String
+    let timestamp: Date
+    let outgoing: Bool
+    let encrypted: Bool
+}
+
 struct WatchVoiceMessage: Sendable {
     let transferID: String
     let jid: String
@@ -59,18 +77,33 @@ final class PhoneWatchBridge: NSObject, WCSessionDelegate {
         session?.activate()
     }
 
+    @MainActor
     func update(conversations: [Conversation], messages: [ChatMessage]) {
         guard let session else { return }
-        // Copy-on-write snapshots are cheap here; the filtering, sorting and
-        // JSON encoding are not. Keep that work away from SwiftUI's main run
-        // loop, especially immediately after a MAM page is committed.
+        // SwiftData models must never be read off the main actor: snapshot
+        // the fields the watch needs here, then hand plain values to the
+        // background queue for the filtering, sorting and JSON encoding.
+        let chats = conversations.map { conversation in
+            PhoneWatchChatData(
+                jid: conversation.jid,
+                name: conversation.displayName,
+                unread: conversation.unreadCount
+            )
+        }
+        let entries = messages.map { message in
+            PhoneWatchMessageData(
+                id: message.clientID,
+                conversationID: message.conversationID,
+                body: message.previewText,
+                timestamp: message.timestamp,
+                outgoing: message.direction == .outgoing,
+                encrypted: message.security == .omemo
+            )
+        }
         snapshotQueue.async { [weak self, weak session] in
             guard let self,
                   let session,
-                  let data = Self.makeSnapshotData(
-                    conversations: conversations,
-                    messages: messages
-                  ),
+                  let data = Self.makeSnapshotData(chats: chats, entries: entries),
                   data != self.latestSnapshot else { return }
             self.latestSnapshot = data
             self.publishLatestSnapshot(using: session)
@@ -78,57 +111,56 @@ final class PhoneWatchBridge: NSObject, WCSessionDelegate {
     }
 
     private static func makeSnapshotData(
-        conversations: [Conversation],
-        messages: [ChatMessage]
+        chats: [PhoneWatchChatData],
+        entries: [PhoneWatchMessageData]
     ) -> Data? {
-        let visibleConversations = Array(conversations.prefix(20))
-        let visibleConversationIDs = Set(visibleConversations.map(\.jid))
-        var recentMessagesByConversation: [String: [ChatMessage]] = [:]
+        let visibleChats = Array(chats.prefix(20))
+        let visibleChatIDs = Set(visibleChats.map(\.jid))
+        var recentMessagesByChat: [String: [PhoneWatchMessageData]] = [:]
 
         // Keep a bounded, chronologically sorted window. MAM can append an old
         // overlapping stanza after a newer live message, so array order alone
         // cannot identify the latest messages reliably.
-        for message in messages where visibleConversationIDs.contains(message.conversationID) {
-            var recent = recentMessagesByConversation[message.conversationID] ?? []
-            let insertionIndex = Self.insertionIndex(for: message, in: recent)
-            recent.insert(message, at: insertionIndex)
+        for entry in entries where visibleChatIDs.contains(entry.conversationID) {
+            var recent = recentMessagesByChat[entry.conversationID] ?? []
+            let insertionIndex = Self.insertionIndex(for: entry, in: recent)
+            recent.insert(entry, at: insertionIndex)
             if recent.count > 20 {
                 recent.removeFirst(recent.count - 20)
             }
-            recentMessagesByConversation[message.conversationID] = recent
+            recentMessagesByChat[entry.conversationID] = recent
         }
 
-        let chats = visibleConversations.map { conversation in
+        let snapshotChats = visibleChats.map { chat in
             PhoneWatchSnapshot.Chat(
-                jid: conversation.jid,
-                name: conversation.displayName,
-                unread: conversation.unreadCount,
-                messages: (recentMessagesByConversation[conversation.jid] ?? [])
-                    .map {
-                        PhoneWatchSnapshot.Message(
-                            id: $0.clientID,
-                            body: $0.previewText,
-                            timestamp: $0.timestamp,
-                            outgoing: $0.direction == .outgoing,
-                            encrypted: $0.security == .omemo
-                        )
-                    }
+                jid: chat.jid,
+                name: chat.name,
+                unread: chat.unread,
+                messages: (recentMessagesByChat[chat.jid] ?? []).map { entry in
+                    PhoneWatchSnapshot.Message(
+                        id: entry.id,
+                        body: entry.body,
+                        timestamp: entry.timestamp,
+                        outgoing: entry.outgoing,
+                        encrypted: entry.encrypted
+                    )
+                }
             )
         }
-        return try? JSONEncoder.watchEncoder.encode(PhoneWatchSnapshot(chats: chats))
+        return try? JSONEncoder.watchEncoder.encode(PhoneWatchSnapshot(chats: snapshotChats))
     }
 
     private static func insertionIndex(
-        for message: ChatMessage,
-        in sortedMessages: [ChatMessage]
+        for entry: PhoneWatchMessageData,
+        in sorted: [PhoneWatchMessageData]
     ) -> Int {
         var lowerBound = 0
-        var upperBound = sortedMessages.count
+        var upperBound = sorted.count
         while lowerBound < upperBound {
             let middle = lowerBound + (upperBound - lowerBound) / 2
-            let existing = sortedMessages[middle]
-            let existingComesFirst = existing.timestamp < message.timestamp
-                || (existing.timestamp == message.timestamp && existing.clientID < message.clientID)
+            let existing = sorted[middle]
+            let existingComesFirst = existing.timestamp < entry.timestamp
+                || (existing.timestamp == entry.timestamp && existing.id < entry.id)
             if existingComesFirst {
                 lowerBound = middle + 1
             } else {

Sources/Shared/UI/ServerInfoView.swift

@@ -8,13 +8,13 @@ struct ServerInfoView: View {
     @State private var isLoading = false
 
     var body: some View {
-        List {
+        Form {
             if let info = model.serverInformation {
-                Section {
+                Section("Это статистика вашего соединения.") {
                     statisticsEntries(info)
-                } header: {
-                    Text("Это статистика вашего соединения.")
-                }
+                } //header: {
+//                    Text("Это статистика вашего соединения.")
+//                }
 
                 Section {
                     softwareEntry(info)
@@ -255,7 +255,7 @@ struct ServerInfoView: View {
                 .font(.caption)
                 .foregroundStyle(.secondary)
         }
-        .frame(maxWidth: .infinity, alignment: .leading)
+//        .frame(maxWidth: .infinity, alignment: .leading)
         .listRowBackground(status.color)
     }
 

Sources/Shared/XMPP/LumaOMEMO2Module.swift

@@ -0,0 +1,839 @@
+import Foundation
+import Combine
+import CommonCrypto
+import CryptoKit
+import os
+import Martin
+import MartinOMEMO
+
+extension XmppModuleIdentifier {
+    static var omemo2: XmppModuleIdentifier<LumaOMEMO2Module> {
+        LumaOMEMO2Module.IDENTIFIER
+    }
+}
+
+/// OMEMO 2 (XEP-0384 0.8.3, urn:xmpp:omemo:2) support built on the
+/// MartinOMEMO signal stack. Wire format implemented here:
+/// - payload: AES-256-CBC + HMAC-SHA-256 keyed via HKDF ("OMEMO Payload"),
+///   plaintext is a Stanza Content Encryption `<envelope>` (XEP-0420);
+/// - header: `<keys jid>` groups with `<key rid kex>` elements;
+/// - device list node `urn:xmpp:omemo:2:devices` (`<devices>` element);
+/// - bundle node `urn:xmpp:omemo:2:bundles` (one item per device id,
+///   `<bundle>` with `<spk>/<spks>/<ik>/<prekeys><pk>`).
+/// The Double Ratchet sessions are shared with the legacy OMEMO module
+/// through the same SignalStorage.
+final class LumaOMEMO2Module: AbstractPEPModule, XmppModule {
+    public static let ID = "omemo2"
+    public static let IDENTIFIER = XmppModuleIdentifier<LumaOMEMO2Module>()
+    public static let XMLNS = "urn:xmpp:omemo:2"
+    public static let DEVICES_LIST_NODE = "urn:xmpp:omemo:2:devices"
+    public static let BUNDLES_NODE = "urn:xmpp:omemo:2:bundles"
+    public static let SCE_XMLNS = "urn:xmpp:sce:1"
+    public static let HKDF_INFO = "OMEMO Payload"
+
+    public let id: String = ID
+    public let criteria = Criteria.empty()
+    public let features: [String] = [LumaOMEMO2Module.DEVICES_LIST_NODE + "+notify"]
+
+    public let signalContext: SignalContext
+    public let storage: SignalStorage
+    private let devicesQueue = DispatchQueue(label: "app.luma.omemo2.devices")
+    private var devices: [BareJID: [Int32]] = [:]
+    private var devicesFetchError: [BareJID: [Int32]] = [:]
+
+    /// True once our device id has been published to the OMEMO 2 device
+    /// list and our bundle is on the server. Drives the UI readiness flag.
+    @Published public private(set) var isReady: Bool = false
+
+    public override var isPepAvailable: Bool {
+        didSet {
+            if isPepAvailable {
+                publishBundleIfNeeded { [weak self] in
+                    self?.publishDeviceListIfNeeded()
+                }
+            }
+        }
+    }
+
+    public init(signalContext: SignalContext, signalStorage: SignalStorage) {
+        self.signalContext = signalContext
+        self.storage = signalStorage
+        super.init()
+    }
+
+    public func process(stanza: Stanza) throws {
+        // Decryption is driven explicitly by XMPPService decode calls;
+        // PEP notifications arrive through onItemNotification.
+    }
+
+    // MARK: - Device list
+
+    public func devices(for jid: BareJID) -> [Int32]? {
+        let known = devicesQueue.sync { devices[jid] }
+        guard let known else { return nil }
+        guard let failed = devicesFetchError[jid] else { return known }
+        return known.filter { !failed.contains($0) }
+    }
+
+    public func isAvailable(for jid: BareJID) -> Bool {
+        !(devicesQueue.sync { devices[jid] }?.isEmpty ?? true)
+            || !storage.sessionStore.allDevices(for: jid.stringValue, activeAndTrusted: true).isEmpty
+    }
+
+    private func publishDeviceListIfNeeded() {
+        guard isPepAvailable, let context else { return }
+        let pepJid = context.userBareJid
+        context.module(.pubsub).retrieveItems(
+            from: pepJid,
+            for: Self.DEVICES_LIST_NODE,
+            limit: .lastItems(1)
+        ) { [weak self] result in
+            switch result {
+            case .success(let items):
+                self?.processDeviceList(jid: pepJid, payload: items.items.first?.payload)
+            case .failure(let error):
+                guard error.error == .item_not_found || error.error == .internal_server_error() else { return }
+                self?.processDeviceList(jid: pepJid, payload: nil)
+            }
+        }
+    }
+
+    private func processDeviceList(jid: BareJID, payload input: Element?) {
+        guard let context else { return }
+        var list = input
+        if list?.name != "devices" || list?.xmlns != Self.XMLNS {
+            list = Element(name: "devices", xmlns: Self.XMLNS)
+        }
+        guard let list else { return }
+
+        let isOwn = jid == context.userBareJid
+        let ourID = String(storage.identityKeyStore.localRegistrationId())
+        var changed = false
+        if isOwn,
+            list.findChild(where: { $0.name == "device" && $0.getAttribute("id") == ourID }) == nil
+        {
+            list.addChild(Element(name: "device", attributes: ["id": ourID, "label": "Luma"]))
+            changed = true
+        }
+
+        if isOwn, changed {
+            let options = PubSubNodeConfig()
+            options.accessModel = .open
+            context.module(.pubsub).publishItem(
+                at: jid,
+                to: Self.DEVICES_LIST_NODE,
+                itemId: "current",
+                payload: list,
+                publishOptions: options
+            ) { [weak self] result in
+                switch result {
+                case .success:
+                    Logger(subsystem: "Luma", category: "omemo2")
+                        .info("device list published jid=\(jid.stringValue)")
+                    self?.isReady = true
+                case .failure(let error):
+                    Logger(subsystem: "Luma", category: "omemo2")
+                        .warning("device list publish failed: \(error.error)")
+                    guard error.error == .conflict() else { return }
+                    context.module(.pubsub).retrieveNodeConfiguration(
+                        from: jid,
+                        node: Self.DEVICES_LIST_NODE
+                    ) { result in
+                        guard case .success(let form) = result else { return }
+                        form.accessModel = .open
+                        context.module(.pubsub).configureNode(
+                            at: jid,
+                            node: Self.DEVICES_LIST_NODE,
+                            with: form
+                        ) { _ in
+                            context.module(.pubsub).publishItem(
+                                at: jid,
+                                to: Self.DEVICES_LIST_NODE,
+                                itemId: "current",
+                                payload: list,
+                                publishOptions: options
+                            ) { result in
+                                if case .success = result { self?.isReady = true }
+                            }
+                        }
+                    }
+                }
+            }
+        } else if isOwn {
+            isReady = true
+        }
+
+        // Track known devices and align identity states with the list.
+        let known = list.mapChildren(transform: { $0.getAttribute("id").flatMap(Int32.init) })
+        devicesQueue.async { [weak self] in self?.devices[jid] = known }
+        let active = storage.sessionStore.allDevices(for: jid.stringValue, activeAndTrusted: true)
+        active.filter { !known.contains($0) }.forEach {
+            _ = storage.identityKeyStore.setStatus(active: false, forIdentity: SignalAddress(name: jid.stringValue, deviceId: $0))
+        }
+        known.filter { !active.contains($0) }.forEach {
+            _ = storage.identityKeyStore.setStatus(active: true, forIdentity: SignalAddress(name: jid.stringValue, deviceId: $0))
+        }
+    }
+
+    override func onItemNotification(notification: PubSubModule.ItemNotification) {
+        guard notification.node == Self.DEVICES_LIST_NODE, let context else { return }
+        switch notification.action {
+        case .published(let item):
+            let from = notification.message.from?.bareJid ?? context.userBareJid
+            processDeviceList(jid: from, payload: item.payload)
+        default:
+            break
+        }
+    }
+
+    // MARK: - Bundle
+
+    func publishBundleIfNeeded(completionHandler: (() -> Void)?) {
+        guard isPepAvailable, let context else {
+            completionHandler?()
+            return
+        }
+        let deviceID = String(storage.identityKeyStore.localRegistrationId())
+        context.module(.pubsub).retrieveItems(
+            from: context.userBareJid,
+            for: Self.BUNDLES_NODE,
+            limit: .items(withIds: [deviceID])
+        ) { [weak self] result in
+            switch result {
+            case .success(let items):
+                self?.publishBundle(current: items.items.first?.payload, completionHandler: completionHandler)
+            case .failure(let error):
+                guard error.error == .item_not_found || error.error == .internal_server_error() else {
+                    completionHandler?()
+                    return
+                }
+                self?.publishBundle(current: nil, completionHandler: completionHandler)
+            }
+        }
+    }
+
+    private func signedPreKey(regenerate: Bool = false) -> SignalSignedPreKey? {
+        let signedPreKeyId = storage.signedPreKeyStore.countSignedPreKeys()
+        var signedPreKey: SignalSignedPreKey?
+        if !regenerate, signedPreKeyId != 0,
+            let data = signalContext.storage.signedPreKeyStore.loadSignedPreKey(withId: UInt32(signedPreKeyId))
+        {
+            signedPreKey = SignalSignedPreKey(fromSerializedData: data)
+        }
+        if signedPreKey == nil {
+            guard let identityKeyPair = storage.identityKeyStore.keyPair() else { return nil }
+            signedPreKey = signalContext.generateSignedPreKey(
+                withIdentity: identityKeyPair,
+                signedPreKeyId: UInt32(signedPreKeyId + 1)
+            )
+            guard let signedPreKey, let serialized = signedPreKey.serializedData else { return nil }
+            guard signalContext.storage.signedPreKeyStore.storeSignedPreKey(serialized, withId: signedPreKey.preKeyId) else {
+                return nil
+            }
+        }
+        return signedPreKey
+    }
+
+    private func publishBundle(current input: Element?, completionHandler: (() -> Void)?) {
+        guard let identityKeyPair = storage.identityKeyStore.keyPair(),
+            let identityPublicKey = identityKeyPair.publicKey
+        else {
+            completionHandler?()
+            return
+        }
+
+        var flush = input == nil
+        if !flush {
+            flush = identityPublicKey.base64EncodedString() != input?.findChild(name: "ik")?.value
+        }
+
+        guard let signedPreKey = signedPreKey(regenerate: flush) else {
+            completionHandler?()
+            return
+        }
+        let signedPublicBase64 = signedPreKey.publicKeyData?.base64EncodedString() ?? ""
+        let signatureBase64 = signedPreKey.signature.base64EncodedString()
+        var changed = flush
+            || signedPublicBase64 != input?.findChild(name: "spk")?.value
+            || signatureBase64 != input?.findChild(name: "spks")?.value
+
+        let currentIDs = input?.findChild(name: "prekeys")?.mapChildren(transform: {
+            $0.getAttribute("id").flatMap(UInt32.init)
+        }) ?? []
+        var validKeys = currentIDs.compactMap { id -> SignalPreKey? in
+            guard let data = storage.preKeyStore.loadPreKey(withId: id) else { return nil }
+            return SignalPreKey(fromSerializedData: data)
+        }
+        let needKeys = 100 - validKeys.count
+        if needKeys > 0 {
+            changed = true
+            let start = storage.preKeyStore.currentPreKeyId() + 1
+            let newKeys = signalContext.generatePreKeys(withStartingPreKeyId: start, count: UInt32(needKeys))
+            validKeys += newKeys.filter { key in
+                guard let serialized = key.serializedData else { return false }
+                return storage.preKeyStore.storePreKey(serialized, withId: key.preKeyId)
+            }
+        }
+
+        if changed {
+            publishBundle(
+                signedPreKey: signedPreKey,
+                identityKey: identityPublicKey,
+                preKeys: validKeys,
+                completionHandler: completionHandler
+            )
+        } else {
+            completionHandler?()
+        }
+    }
+
+    private func publishBundle(
+        signedPreKey: SignalSignedPreKey,
+        identityKey: Data,
+        preKeys: [SignalPreKey],
+        completionHandler: (() -> Void)?
+    ) {
+        guard let context, let signedPreKeyPublic = signedPreKey.publicKeyData else {
+            completionHandler?()
+            return
+        }
+
+        let bundleEl = Element(name: "bundle", xmlns: Self.XMLNS)
+        bundleEl.addChild(Element(
+            name: "spk",
+            cdata: signedPreKeyPublic.base64EncodedString(),
+            attributes: ["id": String(signedPreKey.preKeyId)]
+        ))
+        bundleEl.addChild(Element(name: "spks", cdata: signedPreKey.signature.base64EncodedString()))
+        bundleEl.addChild(Element(name: "ik", cdata: identityKey.base64EncodedString()))
+        let pkElements = preKeys.map { preKey -> Element in
+            let publicData = preKey.serializedPublicKey?.base64EncodedString() ?? ""
+            return Element(name: "pk", cdata: publicData, attributes: ["id": String(preKey.preKeyId)])
+        }
+        bundleEl.addChild(Element(name: "prekeys", children: pkElements))
+
+        let options = PubSubNodeConfig()
+        options.accessModel = .open
+        options.maxItems = .max
+
+        let deviceID = String(storage.identityKeyStore.localRegistrationId())
+        context.module(.pubsub).publishItem(
+            at: nil,
+            to: Self.BUNDLES_NODE,
+            itemId: deviceID,
+            payload: bundleEl,
+            publishOptions: options
+        ) { result in
+            switch result {
+            case .success:
+                Logger(subsystem: "Luma", category: "omemo2")
+                    .info("bundle published deviceID=\(deviceID)")
+            case .failure(let error):
+                Logger(subsystem: "Luma", category: "omemo2")
+                    .warning("bundle publish failed: \(error.error)")
+            }
+            completionHandler?()
+        }
+    }
+
+    // MARK: - Session building
+
+    func buildSession(forAddress address: SignalAddress, completionHandler: (() -> Void)? = nil) {
+        guard let context else {
+            completionHandler?()
+            return
+        }
+        let pepJid = BareJID(address.name)
+        context.module(.pubsub).retrieveItems(
+            from: pepJid,
+            for: Self.BUNDLES_NODE,
+            limit: .items(withIds: [String(address.deviceId)])
+        ) { [weak self] result in
+            defer { completionHandler?() }
+            guard let self, case .success(let items) = result,
+                let bundle = OMEMO2Bundle(from: items.items.first?.payload)
+            else {
+                self?.markDeviceAsFailed(for: pepJid, andDeviceId: address.deviceId)
+                return
+            }
+            guard let preKey = bundle.preKeys.randomElement(),
+                let preKeyBundle = SignalPreKeyBundle(
+                    registrationId: 0,
+                    deviceId: address.deviceId,
+                    preKeyId: preKey.id,
+                    preKeyPublic: preKey.data,
+                    signedPreKeyId: bundle.signedPreKeyId,
+                    signedPreKeyPublic: bundle.signedPreKeyPublic,
+                    signedPreKeySignature: bundle.signedPreKeySignature,
+                    identityKey: bundle.identityKey
+                ),
+                let builder = SignalSessionBuilder(withAddress: address, andContext: self.signalContext)
+            else {
+                self.markDeviceAsFailed(for: pepJid, andDeviceId: address.deviceId)
+                return
+            }
+            _ = builder.processPreKeyBundle(bundle: preKeyBundle)
+        }
+    }
+
+    private func markDeviceAsFailed(for jid: BareJID, andDeviceId deviceId: Int32) {
+        var failed = devicesFetchError[jid] ?? []
+        if !failed.contains(deviceId) {
+            failed.append(deviceId)
+            devicesFetchError[jid] = failed
+        }
+    }
+
+    // MARK: - Decode
+
+    public func decode(message: Message, from: BareJID, serverMsgId: String? = nil) -> DecryptionResult<Message, SignalError> {
+        guard context != nil else { return .failure(.unknown) }
+        guard let encryptedEl = message.findChild(name: "encrypted", xmlns: Self.XMLNS),
+            let headerEl = encryptedEl.findChild(name: "header"),
+            let sid = UInt32(headerEl.getAttribute("sid") ?? "")
+        else {
+            return .failure(.notEncrypted)
+        }
+
+        let localDeviceID = String(storage.identityKeyStore.localRegistrationId())
+        let ownJID = context!.userBareJid.stringValue.lowercased()
+
+        // Collect candidate keys for our device: prefer the 0.8.3
+        // `<keys jid>` grouping, fall back to direct `<key>` children.
+        var keyElements: [Element] = []
+        for keysEl in headerEl.getChildren(where: { $0.name == "keys" }) {
+            guard keysEl.getAttribute("jid")?.lowercased() == ownJID else { continue }
+            keyElements.append(contentsOf: keysEl.getChildren(where: { $0.name == "key" }))
+        }
+        if keyElements.isEmpty {
+            keyElements = headerEl.getChildren(where: { $0.name == "key" })
+        }
+        let ours = keyElements.filter { $0.getAttribute("rid") == localDeviceID }
+        guard !ours.isEmpty else {
+            Logger(subsystem: "Luma", category: "omemo2")
+                .warning("message not encrypted for us: from=\(from.stringValue) sid=\(sid) ourRid=\(localDeviceID) keyCount=\(keyElements.count) hasPayload=\(encryptedEl.findChild(name: "payload") != nil)")
+            guard context!.userBareJid != from || sid != storage.identityKeyStore.localRegistrationId() else {
+                return .failure(.duplicateMessage)
+            }
+            guard encryptedEl.findChild(name: "payload") != nil else {
+                return .failure(.duplicateMessage)
+            }
+            return .failure(.invalidMessage)
+        }
+
+        let address = SignalAddress(name: from.stringValue, deviceId: Int32(bitPattern: sid))
+        var lastError: SignalError = .unknown
+        for keyEl in ours {
+            guard let base64 = keyEl.value, let keyData = Data(base64Encoded: base64),
+                let cipher = SignalSessionCipher(withAddress: address, andContext: signalContext)
+            else {
+                Logger(subsystem: "Luma", category: "omemo2")
+                    .warning("unusable key element from=\(from.stringValue) sid=\(sid)")
+                continue
+            }
+            let kexRaw = keyEl.getAttribute("kex") ?? ""
+            let isKex = kexRaw == "true" || kexRaw == "1"
+                || keyEl.getAttribute("prekey") == "true" || keyEl.getAttribute("prekey") == "1"
+            // Trust the kex attribute first; some clients mislabel the key
+            // element, so retry with the opposite interpretation before
+            // giving up on the message.
+            var elementError: SignalError = .unknown
+            for prekey in [isKex, !isKex] {
+                let result = cipher.decrypt(key: SignalSessionCipher.Key(
+                    key: keyData,
+                    deviceId: Int32(bitPattern: sid),
+                    prekey: prekey
+                ))
+                switch result {
+                case .success(let combined):
+                    Logger(subsystem: "Luma", category: "omemo2")
+                        .info("key decrypted from=\(from.stringValue) sid=\(sid) prekey=\(prekey)")
+                    return finishDecode(
+                        message: message,
+                        encryptedEl: encryptedEl,
+                        combined: combined,
+                        address: address,
+                        isKex: prekey
+                    )
+                case .failure(let error):
+                    elementError = error
+                }
+            }
+            lastError = elementError
+            Logger(subsystem: "Luma", category: "omemo2")
+                .warning("key decrypt failed: error=\(elementError) from=\(from.stringValue) sid=\(sid) kexAttr=\(kexRaw)")
+        }
+
+        Logger(subsystem: "Luma", category: "omemo2")
+            .warning("decode failed: error=\(lastError) from=\(from.stringValue) sid=\(sid)")
+        if (lastError == .noSession || lastError == .invalidMessage), serverMsgId != nil {
+            buildSession(forAddress: address)
+        }
+        return .failure(lastError)
+    }
+
+    private func finishDecode(
+        message: Message,
+        encryptedEl: Element,
+        combined: Data,
+        address: SignalAddress,
+        isKex: Bool
+    ) -> DecryptionResult<Message, SignalError> {
+        message.removeChild(encryptedEl)
+
+        // A consumed prekey must leave our published bundle.
+        if isKex, storage.preKeyStore.flushDeletedPreKeys() {
+            publishBundleIfNeeded(completionHandler: nil)
+        }
+
+        guard let payloadValue = encryptedEl.findChild(name: "payload")?.value,
+            let ciphertext = Data(base64Encoded: payloadValue)
+        else {
+            // Empty message: session management, nothing to display.
+            return .successTransportKey(combined, iv: Data())
+        }
+
+        // combined = key (32) || truncated HMAC (16)
+        guard combined.count >= 48 else { return .failure(.invalidMac) }
+        let key = combined.subdata(in: 0..<32)
+        let expectedMAC = combined.subdata(in: 32..<48)
+        guard let derived = Self.derivePayloadKeys(from: key) else { return .failure(.invalidMac) }
+        let computedMAC = Data(HMAC<SHA256>.authenticationCode(
+            for: ciphertext,
+            using: SymmetricKey(data: derived.authKey)
+        ).prefix(16))
+        guard Self.constantTimeEquals(computedMAC, expectedMAC) else {
+            Logger(subsystem: "Luma", category: "omemo2")
+                .warning("payload HMAC mismatch from=\(address.name) device=\(address.deviceId)")
+            return .failure(.invalidMac)
+        }
+
+        guard let plaintext = Self.aes256CBCDecrypt(ciphertext, key: derived.encryptionKey, iv: derived.iv),
+            let plaintextString = String(data: plaintext, encoding: .utf8),
+            let envelope = Element.from(string: plaintextString),
+            envelope.name == "envelope"
+        else {
+            Logger(subsystem: "Luma", category: "omemo2")
+                .warning("payload envelope parse failed from=\(address.name) device=\(address.deviceId)")
+            return .failure(.invalidMessage)
+        }
+
+        if let content = envelope.findChild(name: "content"),
+            let bodyEl = content.findChild(name: "body"),
+            let body = bodyEl.value, !body.isEmpty
+        {
+            message.body = body
+        }
+
+        Logger(subsystem: "Luma", category: "omemo2")
+            .info("decoded message from=\(address.name) device=\(address.deviceId) bodyLen=\(message.body?.count ?? 0)")
+        _ = storage.identityKeyStore.setStatus(active: true, forIdentity: address)
+        return .successMessage(
+            message,
+            fingerprint: storage.identityKeyStore.identityFingerprint(forAddress: address)
+        )
+    }
+
+    // MARK: - Encode
+
+    public func encode(
+        message: Message,
+        withStoreHint: Bool = true,
+        completionHandler: @escaping (EncryptionResult<Message, SignalError>) -> Void
+    ) {
+        guard let jid = message.to?.bareJid else {
+            completionHandler(.failure(.noDestination))
+            return
+        }
+        encode(message: message, for: [jid], withStoreHint: withStoreHint, completionHandler: completionHandler)
+    }
+
+    public func addresses(
+        for jids: [BareJID],
+        completionHandler: @escaping (Result<[SignalAddress], SignalError>) -> Void
+    ) {
+        guard let pubsub = context?.module(.pubsub) else {
+            completionHandler(.failure(.unknown))
+            return
+        }
+        let group = DispatchGroup()
+        var addresses: [SignalAddress] = []
+        for jid in jids {
+            if let known = devices(for: jid) {
+                let knownAddresses = known.map { SignalAddress(name: jid.stringValue, deviceId: $0) }
+                addresses.append(contentsOf: knownAddresses)
+                for address in knownAddresses where !storage.sessionStore.containsSessionRecord(forAddress: address) {
+                    group.enter()
+                    buildSession(forAddress: address) { group.leave() }
+                }
+            } else {
+                group.enter()
+                pubsub.retrieveItems(from: jid, for: Self.DEVICES_LIST_NODE, limit: .lastItems(1)) { [weak self] result in
+                    defer { group.leave() }
+                    guard let self, case .success(let items) = result,
+                        let listEl = items.items.first?.payload,
+                        listEl.name == "devices", listEl.xmlns == Self.XMLNS
+                    else { return }
+                    let known = listEl.mapChildren(transform: { $0.getAttribute("id").flatMap(Int32.init) })
+                    self.devicesQueue.async { self.devices[jid] = known }
+                    let knownAddresses = known.map { SignalAddress(name: jid.stringValue, deviceId: $0) }
+                    addresses.append(contentsOf: knownAddresses)
+                    for address in knownAddresses where !self.storage.sessionStore.containsSessionRecord(forAddress: address) {
+                        group.enter()
+                        self.buildSession(forAddress: address) { group.leave() }
+                    }
+                }
+            }
+        }
+        group.notify(queue: .main) {
+            completionHandler(.success(addresses))
+        }
+    }
+
+    public func encode(
+        message: Message,
+        for jids: [BareJID],
+        withStoreHint: Bool = true,
+        completionHandler: @escaping (EncryptionResult<Message, SignalError>) -> Void
+    ) {
+        addresses(for: jids) { result in
+            switch result {
+            case .failure(let error):
+                completionHandler(.failure(error))
+            case .success(let addresses):
+                if addresses.isEmpty {
+                    completionHandler(.failure(.noSession))
+                } else {
+                    self.encode(message: message, forAddresses: addresses, withStoreHint: withStoreHint, completionHandler: completionHandler)
+                }
+            }
+        }
+    }
+
+    public func encode(
+        message: Message,
+        forAddresses addresses: [SignalAddress],
+        roomJID: BareJID? = nil,
+        withStoreHint: Bool = true,
+        completionHandler: @escaping (EncryptionResult<Message, SignalError>) -> Void
+    ) {
+        let result = encodeMessage(message: message, for: addresses, roomJID: roomJID)
+        switch result {
+        case .successMessage(let encoded, _):
+            if withStoreHint {
+                encoded.addChild(Element(name: "store", xmlns: "urn:xmpp:hints"))
+            }
+        default:
+            break
+        }
+        completionHandler(result)
+    }
+
+    private func encodeMessage(
+        message: Message,
+        for remoteAddresses: [SignalAddress],
+        roomJID: BareJID?
+    ) -> EncryptionResult<Message, SignalError> {
+        guard let context else { return .failure(.unknown) }
+
+        let localAddresses = storage.sessionStore
+            .allDevices(for: context.userBareJid.stringValue, activeAndTrusted: true)
+            .map { SignalAddress(name: context.userBareJid.stringValue, deviceId: $0) }
+        let destinations = Set(remoteAddresses + localAddresses)
+
+        // Stanza Content Encryption envelope (XEP-0420) as the plaintext.
+        let envelopeXML = Self.envelopeXML(
+            body: message.body,
+            from: context.userBareJid.stringValue,
+            to: roomJID?.stringValue
+        )
+        guard let plaintext = envelopeXML.data(using: .utf8) else { return .failure(.unknown) }
+
+        var key = Data(count: 32)
+        key.withUnsafeMutableBytes { bytes in
+            _ = SecRandomCopyBytes(kSecRandomDefault, 32, bytes.baseAddress!)
+        }
+        guard let derived = Self.derivePayloadKeys(from: key) else { return .failure(.unknown) }
+        guard let ciphertext = Self.aes256CBCEncrypt(plaintext, key: derived.encryptionKey, iv: derived.iv) else {
+            return .failure(.unknown)
+        }
+        let mac = Data(HMAC<SHA256>.authenticationCode(
+            for: ciphertext,
+            using: SymmetricKey(data: derived.authKey)
+        ).prefix(16))
+        var combinedKey = key
+        combinedKey.append(mac)
+
+        let encryptedEl = Element(name: "encrypted", xmlns: Self.XMLNS)
+        let header = Element(name: "header")
+        header.setAttribute("sid", value: String(storage.identityKeyStore.localRegistrationId()))
+        encryptedEl.addChild(header)
+
+        var keysByJID: [String: [Element]] = [:]
+        for address in destinations {
+            guard let cipher = SignalSessionCipher(withAddress: address, andContext: signalContext) else { continue }
+            switch cipher.encrypt(data: combinedKey) {
+            case .success(let output):
+                let keyEl = Element(name: "key", cdata: output.key.base64EncodedString())
+                keyEl.setAttribute("rid", value: String(output.deviceId))
+                if output.prekey { keyEl.setAttribute("kex", value: "true") }
+                keysByJID[address.name.lowercased(), default: []].append(keyEl)
+            case .failure:
+                break
+            }
+        }
+        for (jid, keys) in keysByJID {
+            let keysEl = Element(name: "keys")
+            keysEl.setAttribute("jid", value: jid)
+            keysEl.addChildren(keys)
+            header.addChild(keysEl)
+        }
+
+        encryptedEl.addChild(Element(name: "payload", cdata: ciphertext.base64EncodedString()))
+        message.body = nil
+        message.addChild(encryptedEl)
+
+        let fingerprint = storage.identityKeyStore.identityFingerprint(
+            forAddress: SignalAddress(
+                name: context.userBareJid.stringValue,
+                deviceId: Int32(bitPattern: storage.identityKeyStore.localRegistrationId())
+            )
+        )
+        return .successMessage(message, fingerprint: fingerprint)
+    }
+
+    // MARK: - Crypto helpers (internal for unit tests)
+
+    struct PayloadKeys: Equatable {
+        let encryptionKey: Data
+        let authKey: Data
+        let iv: Data
+    }
+
+    static func derivePayloadKeys(from key: Data) -> PayloadKeys? {
+        guard key.count == 32 else { return nil }
+        let salt = Data(count: 32) // 256 zero bits
+        let material = HKDF<SHA256>.deriveKey(
+            inputKeyMaterial: SymmetricKey(data: key),
+            salt: salt,
+            info: Data(HKDF_INFO.utf8),
+            outputByteCount: 80
+        )
+        let bytes = material.withUnsafeBytes { Data($0) }
+        return PayloadKeys(
+            encryptionKey: bytes.subdata(in: 0..<32),
+            authKey: bytes.subdata(in: 32..<64),
+            iv: bytes.subdata(in: 64..<80)
+        )
+    }
+
+    static func aes256CBCEncrypt(_ plaintext: Data, key: Data, iv: Data) -> Data? {
+        crypt(operation: CCOperation(kCCEncrypt), data: plaintext, key: key, iv: iv)
+    }
+
+    static func aes256CBCDecrypt(_ ciphertext: Data, key: Data, iv: Data) -> Data? {
+        crypt(operation: CCOperation(kCCDecrypt), data: ciphertext, key: key, iv: iv)
+    }
+
+    private static func crypt(operation: CCOperation, data: Data, key: Data, iv: Data) -> Data? {
+        guard key.count == kCCKeySizeAES256, iv.count == kCCBlockSizeAES128 else { return nil }
+        let bufferSize = data.count + kCCBlockSizeAES128
+        var buffer = Data(count: bufferSize)
+        var moved = 0
+        let status = key.withUnsafeBytes { keyBytes in
+            iv.withUnsafeBytes { ivBytes in
+                data.withUnsafeBytes { dataBytes in
+                    buffer.withUnsafeMutableBytes { outBytes in
+                        CCCrypt(
+                            operation,
+                            CCAlgorithm(kCCAlgorithmAES),
+                            CCOptions(kCCOptionPKCS7Padding),
+                            keyBytes.baseAddress, key.count,
+                            ivBytes.baseAddress,
+                            dataBytes.baseAddress, data.count,
+                            outBytes.baseAddress, bufferSize,
+                            &moved
+                        )
+                    }
+                }
+            }
+        }
+        guard status == kCCSuccess else { return nil }
+        buffer.removeSubrange(moved..<buffer.count)
+        return buffer
+    }
+
+    static func constantTimeEquals(_ lhs: Data, _ rhs: Data) -> Bool {
+        guard lhs.count == rhs.count else { return false }
+        var difference: UInt8 = 0
+        for index in 0..<lhs.count {
+            difference |= lhs[index] ^ rhs[index]
+        }
+        return difference == 0
+    }
+
+    static func envelopeXML(body: String?, from fromJID: String, to toJID: String?) -> String {
+        var content = ""
+        if let body, !body.isEmpty {
+            content += "<body xmlns='jabber:client'>" + escapeXML(body) + "</body>"
+        }
+        let alphabet = Array("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
+        let rpadLength = Int.random(in: 12...96)
+        let rpad = String((0..<rpadLength).map { _ in alphabet.randomElement()! })
+        var result = "<envelope xmlns='urn:xmpp:sce:1'><content>" + content + "</content>"
+        result += "<rpad>" + rpad + "</rpad>"
+        result += "<from jid='" + escapeXML(fromJID) + "'/>"
+        if let toJID {
+            result += "<to jid='" + escapeXML(toJID) + "'/>"
+        }
+        result += "</envelope>"
+        return result
+    }
+
+    static func escapeXML(_ value: String) -> String {
+        value
+            .replacingOccurrences(of: "&", with: "&amp;")
+            .replacingOccurrences(of: "<", with: "&lt;")
+            .replacingOccurrences(of: ">", with: "&gt;")
+    }
+}
+
+/// Parsed `urn:xmpp:omemo:2` bundle (`<spk>/<spks>/<ik>/<prekeys><pk>`).
+struct OMEMO2Bundle {
+    let signedPreKeyId: UInt32
+    let signedPreKeyPublic: Data
+    let signedPreKeySignature: Data
+    let identityKey: Data
+    let preKeys: [(id: UInt32, data: Data)]
+
+    init?(from element: Element?) {
+        guard let element,
+            element.name == "bundle", element.xmlns == LumaOMEMO2Module.XMLNS,
+            let spk = element.findChild(name: "spk"),
+            let spkID = spk.getAttribute("id").flatMap(UInt32.init),
+            let spkValue = spk.value,
+            let spkData = Data(base64Encoded: spkValue),
+            let spks = element.findChild(name: "spks")?.value,
+            let spksData = Data(base64Encoded: spks),
+            let ik = element.findChild(name: "ik")?.value,
+            let ikData = Data(base64Encoded: ik),
+            let preKeysEl = element.findChild(name: "prekeys")
+        else { return nil }
+
+        let parsedPreKeys = preKeysEl.mapChildren(transform: { pk -> (UInt32, Data)? in
+            guard pk.name == "pk",
+                let id = pk.getAttribute("id").flatMap(UInt32.init),
+                let value = pk.value,
+                let data = Data(base64Encoded: value)
+            else { return nil }
+            return (id, data)
+        })
+        guard !parsedPreKeys.isEmpty else { return nil }
+
+        signedPreKeyId = spkID
+        signedPreKeyPublic = spkData
+        signedPreKeySignature = spksData
+        identityKey = ikData
+        preKeys = parsedPreKeys
+    }
+}

Sources/Shared/XMPP/XMPPService.swift

@@ -854,14 +854,28 @@ final class XMPPService {
             addReply(replyTo, fallback: replyFallback, to: message)
             if !encrypted, let outOfBandURL { message.oob = outOfBandURL }
             if encrypted {
-                guard let omemo = client.moduleOrNil(.omemo), omemo.isReady else {
-                    throw LumaXMPPError.omemoNotReady
-                }
                 let recipients = try await groupOMEMORecipients(
                     in: room,
                     using: muc,
                     client: client
                 )
+                // OMEMO 2 only when every member publishes OMEMO 2 devices;
+                // otherwise the legacy protocol keeps legacy members readable.
+                if let omemo2 = client.moduleOrNil(.omemo2),
+                    omemo2.isReady,
+                    recipients.allSatisfy({ omemo2.devices(for: $0)?.isEmpty == false }) {
+                    let encryptedMessage = try await encrypt(
+                        message,
+                        forGroupRecipients: recipients,
+                        using: omemo2,
+                        roomJID: roomJID
+                    )
+                    try await room.send(message: encryptedMessage.message)
+                    return encryptedMessage.fingerprint
+                }
+                guard let omemo = client.moduleOrNil(.omemo), omemo.isReady else {
+                    throw LumaXMPPError.omemoNotReady
+                }
                 let encryptedMessage = try await encrypt(
                     message,
                     forGroupRecipients: recipients,
@@ -893,6 +907,15 @@ final class XMPPService {
             message.oob = outOfBandURL
         }
         if encrypted {
+            // Prefer OMEMO 2 when the peer publishes OMEMO 2 devices; legacy
+            // OMEMO stays as the fallback for legacy-only contacts.
+            if let omemo2 = client.moduleOrNil(.omemo2),
+                omemo2.isReady,
+                omemo2.devices(for: peer)?.isEmpty == false {
+                let encryptedMessage = try await encrypt(message, using: omemo2)
+                try await chat.send(message: encryptedMessage.message)
+                return encryptedMessage.fingerprint
+            }
             guard let omemo = client.moduleOrNil(.omemo), omemo.isReady else {
                 throw LumaXMPPError.omemoNotReady
             }
@@ -1699,6 +1722,12 @@ final class XMPPService {
                 signalStorage: omemoStorage
             )
         )
+        // OMEMO 2 (urn:xmpp:omemo:2) shares the signal storage with the
+        // legacy module: Double Ratchet sessions are per-device and
+        // namespace-agnostic.
+        _ = client.modulesManager.register(
+            LumaOMEMO2Module(signalContext: signalContext, signalStorage: omemoStorage)
+        )
     }
 
     private func configureConnection(
@@ -1873,14 +1902,20 @@ final class XMPPService {
             }
             .store(in: &cancellables)
 
-        client.module(.omemo).$isReady
-            .removeDuplicates()
-            .receive(on: DispatchQueue.main)
-            .sink { [weak self] ready in
-                self?.eventHandler?(
-                    .omemo(ready: ready, ownFingerprint: omemoStorage.ownFingerprint))
-            }
-            .store(in: &cancellables)
+        // The UI is ready for encrypted messaging once either protocol
+        // has published its keys.
+        Publishers.CombineLatest(
+            client.module(.omemo).$isReady,
+            client.module(.omemo2).$isReady
+        )
+        .map { legacyReady, omemo2Ready in legacyReady || omemo2Ready }
+        .removeDuplicates()
+        .receive(on: DispatchQueue.main)
+        .sink { [weak self] ready in
+            self?.eventHandler?(
+                .omemo(ready: ready, ownFingerprint: omemoStorage.ownFingerprint))
+        }
+        .store(in: &cancellables)
     }
 
     private func deliverOrDelayDirect(_ message: Message, timestamp: Date) async {
@@ -1936,6 +1971,22 @@ final class XMPPService {
         }
     }
 
+    /// OMEMO 2 twin of `decodeOmemoOffMain`.
+    private func decodeOmemo2OffMain(
+        _ message: Message,
+        from sender: BareJID,
+        serverMsgId: String?,
+        module: LumaOMEMO2Module
+    ) async -> DecryptionResult<Message, SignalError> {
+        let queue = omemoDecodeQueue
+        return await withCheckedContinuation { continuation in
+            queue.async {
+                let result = module.decode(message: message, from: sender, serverMsgId: serverMsgId)
+                continuation.resume(returning: result)
+            }
+        }
+    }
+
     private func handle(state: XMPPClient.State) {
         switch state {
         case .connecting, .disconnecting:
@@ -1988,12 +2039,22 @@ final class XMPPService {
         let security: ChatMessage.Security
         let fingerprint: String?
         let contentMessage: Message?
-        switch await decodeOmemoOffMain(
-            message,
-            from: sender,
-            serverMsgId: archiveID,
-            module: client.module(.omemo)
-        ) {
+        let decryptionResult = await (
+            Self.isOMEMO2Payload(message)
+                ? decodeOmemo2OffMain(
+                    message,
+                    from: sender,
+                    serverMsgId: archiveID,
+                    module: client.module(.omemo2)
+                )
+                : decodeOmemoOffMain(
+                    message,
+                    from: sender,
+                    serverMsgId: archiveID,
+                    module: client.module(.omemo)
+                )
+        )
+        switch decryptionResult {
         case .successMessage(let decodedMessage, let value):
             security = .omemo
             fingerprint = value
@@ -2003,18 +2064,9 @@ final class XMPPService {
         case .failure(let error):
             switch error {
             case .notEncrypted:
-                if Self.isOMEMO2Payload(message), message.body?.isEmpty != false {
-                    // OMEMO 2 (urn:xmpp:omemo:2) is not implemented by the
-                    // pinned MartinOMEMO library; show the honest
-                    // undecryptable state instead of an empty bubble.
-                    security = .decryptionFailed
-                    fingerprint = nil
-                    contentMessage = nil
-                } else {
-                    security = .plaintext
-                    fingerprint = nil
-                    contentMessage = message
-                }
+                security = .plaintext
+                fingerprint = nil
+                contentMessage = message
             default:
                 // Some clients include a plaintext <body> fallback alongside the
                 // OMEMO <encrypted> payload. When decryption fails, prefer that
@@ -2231,6 +2283,7 @@ final class XMPPService {
             : (stanzaID ?? message.originId ?? message.id ?? UUID().uuidString)
 
         let encrypted = message.firstChild(name: "encrypted", xmlns: OMEMOModule.XMLNS) != nil
+            || message.firstChild(name: "encrypted", xmlns: LumaOMEMO2Module.XMLNS) != nil
         let realSender: BareJID?
         if outgoing {
             realSender = client.userBareJid
@@ -2252,12 +2305,22 @@ final class XMPPService {
         let fingerprint: String?
         let contentMessage: Message?
         if let realSender {
-            switch await decodeOmemoOffMain(
-                message,
-                from: realSender,
-                serverMsgId: stanzaID,
-                module: client.module(.omemo)
-            ) {
+            let decryptionResult = await (
+                Self.isOMEMO2Payload(message)
+                    ? decodeOmemo2OffMain(
+                        message,
+                        from: realSender,
+                        serverMsgId: stanzaID,
+                        module: client.module(.omemo2)
+                    )
+                    : decodeOmemoOffMain(
+                        message,
+                        from: realSender,
+                        serverMsgId: stanzaID,
+                        module: client.module(.omemo)
+                    )
+            )
+            switch decryptionResult {
             case .successMessage(let decodedMessage, let value):
                 security = .omemo
                 fingerprint = value
@@ -2267,15 +2330,9 @@ final class XMPPService {
             case .failure(let error):
                 switch error {
                 case .notEncrypted:
-                    if Self.isOMEMO2Payload(message), message.body?.isEmpty != false {
-                        security = .decryptionFailed
-                        fingerprint = nil
-                        contentMessage = nil
-                    } else {
-                        security = .plaintext
-                        fingerprint = nil
-                        contentMessage = message
-                    }
+                    security = .plaintext
+                    fingerprint = nil
+                    contentMessage = message
                 case .duplicateMessage:
                     emitGroupEchoIfPossible(
                         roomJID: roomJID,
@@ -2871,6 +2928,113 @@ final class XMPPService {
         }
     }
 
+    /// OMEMO 2 twin of `encrypt(_:using:)`.
+    private func encrypt(
+        _ message: Message,
+        using omemo: LumaOMEMO2Module
+    ) async throws -> (message: Message, fingerprint: String?) {
+        try await withCheckedThrowingContinuation { continuation in
+            omemo.encode(message: message, withStoreHint: true) { result in
+                switch result {
+                case .successMessage(let message, let fingerprint):
+                    continuation.resume(returning: (message, fingerprint))
+                case .failure(let error):
+                    continuation.resume(
+                        throwing: LumaXMPPError.omemoEncryptionFailed(error.rawValue))
+                }
+            }
+        }
+    }
+
+    /// OMEMO 2 twin of `encrypt(_:forGroupRecipients:using:)`. Keys are
+    /// nested in `<keys jid>` groups, so the per-device validation walks the
+    /// header's groups.
+    private func encrypt(
+        _ message: Message,
+        forGroupRecipients recipients: [BareJID],
+        using omemo: LumaOMEMO2Module,
+        roomJID: BareJID
+    ) async throws -> (message: Message, fingerprint: String?) {
+        let fetchedAddresses: [SignalAddress] = try await withCheckedThrowingContinuation {
+            continuation in
+            omemo.addresses(for: recipients) { result in
+                switch result {
+                case .success(let addresses):
+                    continuation.resume(returning: addresses)
+                case .failure(let error):
+                    continuation.resume(
+                        throwing: LumaXMPPError.omemoEncryptionFailed(error.rawValue))
+                }
+            }
+        }
+        let addresses = Array(Set(fetchedAddresses))
+
+        let availableJIDs = Set(addresses.map { $0.name.lowercased() })
+        let missingJIDs =
+            recipients
+            .map(\.stringValue)
+            .filter { !availableJIDs.contains($0.lowercased()) }
+        guard missingJIDs.isEmpty else {
+            throw LumaXMPPError.groupOMEMODevicesUnavailable(missingJIDs.sorted())
+        }
+
+        let addressesWithoutSessions = addresses.filter {
+            omemoStorage?.hasSession(for: $0) != true
+        }
+        guard addressesWithoutSessions.isEmpty else {
+            throw LumaXMPPError.groupOMEMOSessionsUnavailable(
+                addressesWithoutSessions
+                    .map { "\($0.name)/\($0.deviceId)" }
+                    .sorted()
+            )
+        }
+
+        let encryptedMessage: (message: Message, fingerprint: String?) =
+            try await withCheckedThrowingContinuation { continuation in
+                omemo.encode(
+                    message: message,
+                    forAddresses: addresses,
+                    roomJID: roomJID,
+                    withStoreHint: true
+                ) { result in
+                    switch result {
+                    case .successMessage(let message, let fingerprint):
+                        continuation.resume(returning: (message, fingerprint))
+                    case .failure(let error):
+                        continuation.resume(
+                            throwing: LumaXMPPError.omemoEncryptionFailed(error.rawValue))
+                    }
+                }
+            }
+
+        var encryptedKeyCounts: [Int32: Int] = [:]
+        if let encrypted = encryptedMessage.message.firstChild(
+            name: "encrypted", xmlns: LumaOMEMO2Module.XMLNS
+        ), let header = encrypted.findChild(name: "header") {
+            var keyElements: [Element] = []
+            for keysEl in header.getChildren(where: { $0.name == "keys" }) {
+                keyElements.append(contentsOf: keysEl.getChildren(where: { $0.name == "key" }))
+            }
+            keyElements
+                .compactMap { $0.getAttribute("rid").flatMap(Int32.init) }
+                .forEach { encryptedKeyCounts[$0, default: 0] += 1 }
+        }
+
+        var unavailableAddresses: [String] = []
+        for address in addresses {
+            let availableCount = encryptedKeyCounts[address.deviceId, default: 0]
+            if availableCount > 0 {
+                encryptedKeyCounts[address.deviceId] = availableCount - 1
+            } else {
+                unavailableAddresses.append("\(address.name)/\(address.deviceId)")
+            }
+        }
+        guard unavailableAddresses.isEmpty else {
+            throw LumaXMPPError.groupOMEMOSessionsUnavailable(unavailableAddresses.sorted())
+        }
+        return encryptedMessage
+    }
+
     private func encrypt(
         _ message: Message,
         forGroupRecipients recipients: [BareJID],