Sources/SophaxChatCore/Crypto/X3DH.swift

1// X3DH.swift
2// SophaxChatCore
3//
4// Extended Triple Diffie-Hellman (X3DH) key agreement protocol.
5// Establishes a shared secret between two parties asynchronously —
6// the initiating party (Alice) can compute the secret even before
7// the responding party (Bob) is online.
8//
9// The shared secret seeds the Double Ratchet.
10//
11// Reference: https://signal.org/docs/specifications/x3dh/
12//
13// DH operations:
14//   DH1 = DH(IK_A, SPK_B)  — identity auth
15//   DH2 = DH(EK_A, IK_B)   — ephemeral x identity
16//   DH3 = DH(EK_A, SPK_B)  — ephemeral x signed prekey
17//   DH4 = DH(EK_A, OPK_B)  — ephemeral x one-time prekey (if available)
18//   SK  = KDF(DH1 || DH2 || DH3 [|| DH4])
19
20import Foundation
21import CryptoKit
22
23private extension SharedSecret {
24    var rawData: Data { withUnsafeBytes { Data($0) } }
25}
26
27public enum X3DH {
28
29    // MARK: - Sender (Alice)
30
31    public struct SenderResult {
32        /// The derived shared secret to seed the Double Ratchet.
33        public let sharedSecret: SymmetricKey
34        /// Alice's ephemeral public key — sent to Bob in the InitialMessage.
35        public let ephemeralPublicKey: Data
36        /// Which of Bob's one-time prekeys was used (if any) — sent to Bob.
37        public let usedOneTimePreKeyId: UInt32?
38        /// ML-KEM-768 encapsulated shared secret (iOS 18+ only; nil otherwise).
39        /// Must be sent to Bob alongside the X3DH ephemeral key so he can reproduce
40        /// the hybrid shared secret. nil = Bob has no PQ key or sender is on iOS < 18.
41        public let pqEncapsulatedKey: Data?
42    }
43
44    /// Perform X3DH as the initiating party (Alice).
45    ///
46    /// - Parameters:
47    ///   - senderIdentity: Alice's DH identity key pair (IK_A)
48    ///   - recipientBundle: Bob's published prekey bundle
49    /// - Returns: Shared secret + ephemeral key info to send Bob
50    public static func initiateSender(
51        senderIdentity: DHKeyPair,
52        recipientBundle: PreKeyBundle
53    ) throws -> SenderResult {
54
55        // 1. Validate bundle timestamp — only accept bundles from the past (not future-dated).
56        // Using abs() here would wrongly accept bundles up to maxPreKeyBundleAge in the future.
57        let bundleAge = Date().timeIntervalSince(recipientBundle.timestamp)
58        guard bundleAge >= 0, bundleAge < CryptoConstants.maxPreKeyBundleAge else {
59            throw SophaxError.stalePreKeyBundle
60        }
61
62        // 2. Verify signed prekey signature
63        guard try recipientBundle.verifySignedPreKey() else {
64            throw SophaxError.invalidSignature
65        }
66
67        // 3. Decode recipient keys
68        let recipientIK  = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: recipientBundle.dhIdentityKeyPublic)
69        let recipientSPK = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: recipientBundle.signedPreKeyPublic)
70
71        // 4. Generate Alice's ephemeral key pair
72        let ephemeralPair = DHKeyPair()
73
74        // 5. DH computations
75        // DH1 = DH(IK_A, SPK_B)
76        let dh1 = try senderIdentity.privateKey.sharedSecretFromKeyAgreement(with: recipientSPK)
77        // DH2 = DH(EK_A, IK_B)
78        let dh2 = try ephemeralPair.privateKey.sharedSecretFromKeyAgreement(with: recipientIK)
79        // DH3 = DH(EK_A, SPK_B)
80        let dh3 = try ephemeralPair.privateKey.sharedSecretFromKeyAgreement(with: recipientSPK)
81
82        var dhConcat = dh1.rawData + dh2.rawData + dh3.rawData
83
84        var usedOTPKId: UInt32? = nil
85        if let otpkData = recipientBundle.oneTimePreKeyPublic,
86           let otpkId   = recipientBundle.oneTimePreKeyId {
87            // DH4 = DH(EK_A, OPK_B)
88            let recipientOPK = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: otpkData)
89            dhConcat += try ephemeralPair.privateKey.sharedSecretFromKeyAgreement(with: recipientOPK).rawData
90            usedOTPKId = otpkId
91        }
92
93        // 6. Post-quantum KEM contribution (iOS 18+ only).
94        // If Bob advertised an ML-KEM-768 public key, encapsulate a fresh shared secret.
95        // The resulting (pq_shared_secret, encapsulated_key) are mixed into the HKDF IKM
96        // so an attacker needs to break BOTH the classical Curve25519 DH AND the PQ KEM.
97        var pqSharedSecret: Data? = nil
98        var pqEncapKey:     Data? = nil
99        #if swift(>=6.2)
100        if #available(iOS 19.0, macOS 26.0, *),
101           let pqPubData = recipientBundle.pqPreKeyPublic,
102           let recipientPQKey = try? MLKEM768.PublicKey(rawRepresentation: pqPubData),
103           let encResult = try? recipientPQKey.encapsulate() {
104            pqSharedSecret = encResult.sharedSecret.withUnsafeBytes { Data($0) }
105            pqEncapKey     = encResult.encapsulated
106        }
107        #endif
108
109        // 7. Derive shared secret via HKDF, then zero the DH material immediately.
110        let sharedSecret = try deriveSharedSecret(from: dhConcat, pqSharedSecret: pqSharedSecret)
111        dhConcat.withUnsafeMutableBytes { $0.initializeMemory(as: UInt8.self, repeating: 0) }
112        pqSharedSecret?.withUnsafeMutableBytes { $0.initializeMemory(as: UInt8.self, repeating: 0) }
113
114        return SenderResult(
115            sharedSecret:        sharedSecret,
116            ephemeralPublicKey:  ephemeralPair.publicKeyData,
117            usedOneTimePreKeyId: usedOTPKId,
118            pqEncapsulatedKey:   pqEncapKey
119        )
120    }
121
122    // MARK: - Receiver (Bob)
123
124    /// Perform X3DH as the receiving party (Bob).
125    ///
126    /// - Parameters:
127    ///   - recipientIdentityDH: Bob's DH identity key pair (IK_B)
128    ///   - recipientSignedPreKey: Bob's signed prekey pair (SPK_B) — used in X3DH
129    ///   - recipientOneTimePreKey: Bob's one-time prekey pair (OPK_B) — if Alice used one
130    ///   - senderIdentityDHKeyData: Alice's DH identity public key (IK_A)
131    ///   - senderEphemeralKeyData: Alice's ephemeral public key (EK_A) from the message
132    ///   - senderPQEncapsulatedKey: Alice's ML-KEM-768 encapsulated key (nil if Alice is on iOS < 18)
133    ///   - identityManager: Bob's IdentityManager — used to decapsulate the PQ key (iOS 18+ only)
134    /// - Returns: Shared secret (must match Alice's)
135    public static func initiateReceiver(
136        recipientIdentityDH: DHKeyPair,
137        recipientSignedPreKey: DHKeyPair,
138        recipientOneTimePreKey: DHKeyPair?,
139        senderIdentityDHKeyData: Data,
140        senderEphemeralKeyData: Data,
141        senderPQEncapsulatedKey: Data?    = nil,
142        identityManager: IdentityManager? = nil
143    ) throws -> SymmetricKey {
144
145        let senderIK  = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: senderIdentityDHKeyData)
146        let senderEK  = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: senderEphemeralKeyData)
147
148        // DH1 = DH(SPK_B, IK_A)   [symmetric to Alice's DH1]
149        let dh1 = try recipientSignedPreKey.privateKey.sharedSecretFromKeyAgreement(with: senderIK)
150        // DH2 = DH(IK_B, EK_A)
151        let dh2 = try recipientIdentityDH.privateKey.sharedSecretFromKeyAgreement(with: senderEK)
152        // DH3 = DH(SPK_B, EK_A)
153        let dh3 = try recipientSignedPreKey.privateKey.sharedSecretFromKeyAgreement(with: senderEK)
154
155        var dhConcat = dh1.rawData + dh2.rawData + dh3.rawData
156
157        if let otp = recipientOneTimePreKey {
158            // DH4 = DH(OPK_B, EK_A)
159            dhConcat += try otp.privateKey.sharedSecretFromKeyAgreement(with: senderEK).rawData
160        }
161
162        // Post-quantum contribution: decapsulate Alice's ML-KEM-768 ciphertext (iOS 19+ only).
163        var pqSharedSecret: Data? = nil
164        #if swift(>=6.2)
165        if #available(iOS 19.0, macOS 26.0, *),
166           let encapData = senderPQEncapsulatedKey,
167           let idMgr = identityManager {
168            let pqSS = try idMgr.pqDecapsulate(encapData)
169            pqSharedSecret = pqSS.withUnsafeBytes { Data($0) }
170        }
171        #endif
172
173        let sharedSecret = try deriveSharedSecret(from: dhConcat, pqSharedSecret: pqSharedSecret)
174        dhConcat.withUnsafeMutableBytes { $0.initializeMemory(as: UInt8.self, repeating: 0) }
175        pqSharedSecret?.withUnsafeMutableBytes { $0.initializeMemory(as: UInt8.self, repeating: 0) }
176        return sharedSecret
177    }
178
179    // MARK: - KDF
180
181    /// HKDF-SHA256 key derivation.
182    /// Follows the Signal X3DH spec: F || DH_concat is the IKM,
183    /// where F = 32 bytes of 0xFF (domain separator for non-empty use).
184    ///
185    /// When `pqSharedSecret` is non-nil (hybrid mode), it is appended to the IKM
186    /// after the DH material. HKDF's extraction step mixes all of it securely —
187    /// the session key is secure as long as either the classical DH OR the PQ KEM is secure.
188    private static func deriveSharedSecret(from dhConcat: Data, pqSharedSecret: Data? = nil) throws -> SymmetricKey {
189        guard dhConcat.count == 96 || dhConcat.count == 128 else {
190            throw SophaxError.keyAgreementFailed
191        }
192        // Per Signal X3DH spec: prepend 32 0xFF bytes as domain separator
193        let f    = Data(repeating: 0xFF, count: 32)
194        var ikm  = f + dhConcat
195        // Append PQ contribution when available — length differs from classical-only
196        // IKM, which inherently domain-separates hybrid from non-hybrid derives.
197        if let pq = pqSharedSecret {
198            ikm += pq
199        }
200        let salt = Data(repeating: 0x00, count: 32)   // 32 zero bytes as salt
201
202        return HKDF<SHA256>.deriveKey(
203            inputKeyMaterial: SymmetricKey(data: ikm),
204            salt: salt,
205            info: CryptoConstants.x3dhInfo,
206            outputByteCount: 32
207        )
208    }
209}