GroupMessageService.swift

GroupMessageService.swift

import Foundation
import CryptoKit
import DoubleRatchet  // https://github.com/TICESoftware/DoubleRatchet

// ────────────────────────────────────────────────
// MARK: - Errors
// ────────────────────────────────────────────────

public enum GroupMessagingError: Error, Sendable {
    case wrongGroup
    case unknownSender
    case noPairwiseSession
    case ratchetFailure(String)
    case invalidMessage
    case serializationFailed
}

// ────────────────────────────────────────────────
// MARK: - Types
// ────────────────────────────────────────────────

public struct GroupMessage: Codable, Sendable {
    public let groupID: Data
    public let senderID: String
    public let epoch: UInt64           // per-sender message counter / chain epoch
    public let ratchetHeader: Data     // serialized DoubleRatchet.Message header
    public let ciphertext: Data
}

public struct SenderKeyState: Codable, Sendable {
    public let chainKey: SymmetricKey
    public let messageCount: UInt64

    public init(chainKey: SymmetricKey, messageCount: UInt64 = 0) {
        self.chainKey = chainKey
        self.messageCount = messageCount
    }

    public func ratchetForward() throws -> (messageKey: SymmetricKey, newState: SenderKeyState) {
        let info = "message-key".data(using: .utf8)!
        let mk = try HKDF<SHA256>.deriveKey(
            inputKeyMaterial: chainKey,
            salt: nil,
            info: info,
            outputByteCount: 32
        )

        let nextInfo = "chain-key".data(using: .utf8)!
        let nextChain = try HKDF<SHA256>.deriveKey(
            inputKeyMaterial: chainKey,
            salt: nil,
            info: nextInfo,
            outputByteCount: 32
        )

        return (mk, SenderKeyState(chainKey: nextChain, messageCount: messageCount + 1))
    }

    // For persistence / export
    public func serialized() throws -> Data {
        let container = [
            "chain": chainKey.withUnsafeBytes { Data($0) },
            "count": messageCount
        ] as [String: AnyCodable]
        return try JSONEncoder().encode(container)
    }

    public static func fromSerialized(_ data: Data) throws -> SenderKeyState {
        let container = try JSONDecoder().decode([String: AnyCodable].self, from: data)
        guard let chainData = container["chain"]?.data,
              let count = container["count"]?.uint64 else {
            throw GroupMessagingError.serializationFailed
        }
        return SenderKeyState(chainKey: SymmetricKey(data: chainData), messageCount: count)
    }
}

// Helper to make Codable work with SymmetricKey
private struct AnyCodable: Codable {
    let data: Data?
    let uint64: UInt64?
    // very minimal — extend as needed
}

// ────────────────────────────────────────────────
// MARK: - EncryptedGroupSession (actor — thread-safe)
// ────────────────────────────────────────────────

public actor EncryptedGroupSession {

    public let groupID: Data
    public let myUserID: String

    // My current sender key chain
    private var mySenderState: SenderKeyState

    // Pairwise Double Ratchet sessions (used only for key distribution)
    private var pairwiseDR: [String: DoubleRatchet]

    // Known sender keys from other participants
    private var knownSenderKeys: [String: SenderKeyState]

    // Last known message count per sender (helps detect replays / gaps)
    private var lastSeen: [String: UInt64]

    public init(
        groupID: Data,
        myUserID: String,
        initialSenderKey: SymmetricKey? = nil
    ) throws {
        self.groupID = groupID
        self.myUserID = myUserID

        if let key = initialSenderKey {
            self.mySenderState = SenderKeyState(chainKey: key)
        } else {
            let fresh = SymmetricKey(size: .bits256)
            self.mySenderState = SenderKeyState(chainKey: fresh)
        }

        self.pairwiseDR = [:]
        self.knownSenderKeys = [:]
        self.lastSeen = [:]
    }

    // ─── Membership ──────────────────────────────────────

    /// Establish pairwise Double Ratchet (usually via shared secret from invite link / QR / etc.)
    public func establishPairwise(with userID: String, sharedSecret: Data) throws {
        let dr = try DoubleRatchet(
            ourKeyPair: nil,           // passive / receiving side — adjust as needed
            theirPublicKey: nil,
            sharedSecret: sharedSecret,
            maxSkip: 1000,
            info: "groupkeydist-\(groupID.hexString)"
        )
        pairwiseDR[userID] = dr
    }

    /// Distribute current sender key to one participant (usually called after establishPairwise)
    public func distributeMySenderKey(to userID: String) throws -> Data {
        guard let dr = pairwiseDR[userID] else {
            throw GroupMessagingError.noPairwiseSession
        }

        let material = try mySenderState.serialized()
        let message = try dr.encrypt(plaintext: material)

        return try JSONEncoder().encode([
            "group": groupID.base64EncodedString(),
            "from": myUserID,
            "type": "sender-key-update",
            "payload": message
        ])
    }

    /// Receive sender key update from another member
    public func processSenderKeyUpdate(from userID: String, payload: Data) throws {
        guard let dr = pairwiseDR[userID] else {
            throw GroupMessagingError.noPairwiseSession
        }

        let json = try JSONDecoder().decode([String: String].self, from: payload)
        guard let encodedMessage = json["payload"]?.data(using: .base64Encoded)! else {
            throw GroupMessagingError.invalidMessage
        }

        let message = try DoubleRatchet.Message(from: encodedMessage)
        let plaintext = try dr.decrypt(message: message)

        let newKeyState = try SenderKeyState.fromSerialized(plaintext)
        knownSenderKeys[userID] = newKeyState
        lastSeen[userID] = newKeyState.messageCount - 1  // expect next to be current +1
    }

    // ─── Sending ─────────────────────────────────────────

    public func encrypt(_ plaintext: Data) throws -> GroupMessage {
        let (mk, newState) = try mySenderState.ratchetForward()
        mySenderState = newState

        let sealed = try ChaChaPoly.seal(plaintext, using: mk)

        return GroupMessage(
            groupID: groupID,
            senderID: myUserID,
            epoch: newState.messageCount,
            ratchetHeader: Data(), // can be empty or include public header if needed
            ciphertext: sealed.combined!
        )
    }

    // ─── Receiving ───────────────────────────────────────

    public func decrypt(_ message: GroupMessage) throws -> Data {
        guard message.groupID == groupID else {
            throw GroupMessagingError.wrongGroup
        }

        guard let senderKeyState = knownSenderKeys[message.senderID] else {
            throw GroupMessagingError.unknownSender
        }

        // Very basic replay / ordering check
        if let last = lastSeen[message.senderID], message.epoch <= last {
            throw GroupMessagingError.invalidMessage // replay or old message
        }

        let (mk, newState) = try senderKeyState.ratchetForward()
        knownSenderKeys[message.senderID] = newState
        lastSeen[message.senderID] = message.epoch

        let box = try ChaChaPoly.SealedBox(combined: message.ciphertext)
        let plaintext = try ChaChaPoly.open(box, using: mk)

        return plaintext
    }

    // ─── Forward Secrecy Rotation (recommended on leave / compromise suspicion) ───

    public func rotateMySenderKeyAndRedistribute() throws -> [String: Data] {
        let newKey = SymmetricKey(size: .bits256)
        mySenderState = SenderKeyState(chainKey: newKey)

        var distributions: [String: Data] = [:]

        for userID in pairwiseDR.keys {
            let payload = try distributeMySenderKey(to: userID)
            distributions[userID] = payload
        }

        return distributions
    }

    // ─── Persistence ─────────────────────────────────────

    public func saveState() throws -> Data {
        let state: [String: Any] = [
            "mySender": try mySenderState.serialized(),
            "knownSenders": knownSenderKeys.mapValues { try $0.serialized() }
            // pairwiseDR state would need DoubleRatchet's own serialization
        ]
        return try JSONSerialization.data(withJSONObject: state)
    }

    public static func restore(from data: Data, groupID: Data, myUserID: String) throws -> EncryptedGroupSession {
        let json = try JSONSerialization.jsonObject(with: data) as? [String: Any]
        guard let myData = json?["mySender"] as? Data,
              let myState = try? SenderKeyState.fromSerialized(myData) else {
            throw GroupMessagingError.serializationFailed
        }

        let session = try EncryptedGroupSession(groupID: groupID, myUserID: myUserID, initialSenderKey: myState.chainKey)
        // restore knownSenders, pairwiseDR, etc. — omitted for brevity
        return session
    }
}

GroupMessageTest.swift

import XCTest
@testable import YourModuleName

final class EncryptedGroupSessionTests: XCTestCase {

    func testBasicSendReceive() throws {
        let groupID = Data("test-group-123".utf8)

        let alice = try EncryptedGroupSession(groupID: groupID, myUserID: "alice")
        let bob   = try EncryptedGroupSession(groupID: groupID, myUserID: "bob")

        // Simulate key distribution (in reality: via invite link / pairwise X3DH)
        let shared = SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }
        try alice.establishPairwise(with: "bob", sharedSecret: shared)
        try bob.establishPairwise(with: "alice", sharedSecret: shared)

        let keyUpdateForBob = try alice.distributeMySenderKey(to: "bob")
        try bob.processSenderKeyUpdate(from: "alice", payload: keyUpdateForBob)

        // Alice sends
        let plaintext = Data("Hello secure group!".utf8)
        let encrypted = try alice.encrypt(plaintext)

        // Bob receives
        let decrypted = try bob.decrypt(encrypted)

        XCTAssertEqual(decrypted, plaintext)
    }

    func testRotation() throws {
        let group = try EncryptedGroupSession(groupID: Data("grp".utf8), myUserID: "test")

        let oldChain = group.mySenderState.chainKey

        let _ = try group.rotateMySenderKeyAndRedistribute()

        XCTAssertNotEqual(oldChain, group.mySenderState.chainKey)
    }

    func testReplayProtection() throws {
        let groupID = Data("replay-test".utf8)
        let alice = try EncryptedGroupSession(groupID: groupID, myUserID: "alice")
        let bob   = try EncryptedGroupSession(groupID: groupID, myUserID: "bob")

        // ... establish pairwise and distribute key ...

        let msg1 = try alice.encrypt(Data("msg1".utf8))
        _ = try bob.decrypt(msg1)

        // Replay msg1
        XCTAssertThrowsError(try bob.decrypt(msg1)) { error in
            XCTAssertTrue((error as? GroupMessagingError) == .invalidMessage)
        }
    }
}