Tests/SophaxChatCoreTests/AdvancedCryptoTests.swift

1// AdvancedCryptoTests.swift
2// SophaxChatCoreTests
3//
4// Tests for:
5//   • Sealed sender (ECDH + HKDF + ChaCha20-Poly1305 envelope)
6//   • Double Ratchet header encryption opacity (encrypted headers are opaque)
7//   • Group Sender Key ratchet (HMAC-SHA256 chain advancement)
8
9import Testing
10import Foundation
11import CryptoKit
12@testable import SophaxChatCore
13
14// MARK: - Sealed Sender Tests
15
16@Suite("Sealed Sender")
17struct SealedSenderTests {
18
19    /// Replicates the sealWireMessage / unsealMessage logic from ChatManager
20    /// so we can test the crypto primitive independently of the private functions.
21    private func seal(_ payload: Data, recipientPublicKeyData: Data) throws -> (ephemeralPublicKey: Data, encryptedPayload: Data) {
22        let ephPair  = DHKeyPair()
23        let recipKey = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: recipientPublicKeyData)
24        let shared   = try ephPair.privateKey.sharedSecretFromKeyAgreement(with: recipKey)
25        var ikmData  = Data(); shared.withUnsafeBytes { ikmData.append(contentsOf: $0) }
26        let key      = HKDF<SHA256>.deriveKey(inputKeyMaterial: SymmetricKey(data: ikmData),
27                                               info: CryptoConstants.sealedSenderInfo,
28                                               outputByteCount: 32)
29        let nonce  = ChaChaPoly.Nonce()
30        let sealed = try ChaChaPoly.seal(payload, using: key, nonce: nonce, authenticating: Data())
31        return (ephPair.publicKeyData, sealed.combined)
32    }
33
34    private func unseal(ephemeralPublicKey: Data, encryptedPayload: Data,
35                        recipientPrivateKey: Curve25519.KeyAgreement.PrivateKey) throws -> Data {
36        let ephKey = try Curve25519.KeyAgreement.PublicKey(rawRepresentation: ephemeralPublicKey)
37        let shared = try recipientPrivateKey.sharedSecretFromKeyAgreement(with: ephKey)
38        var ikmData = Data(); shared.withUnsafeBytes { ikmData.append(contentsOf: $0) }
39        let key = HKDF<SHA256>.deriveKey(inputKeyMaterial: SymmetricKey(data: ikmData),
40                                          info: CryptoConstants.sealedSenderInfo,
41                                          outputByteCount: 32)
42        let box  = try ChaChaPoly.SealedBox(combined: encryptedPayload)
43        return try ChaChaPoly.open(box, using: key, authenticating: Data())
44    }
45
46    @Test("Seal → unseal round-trip restores plaintext")
47    func sealUnsealRoundtrip() throws {
48        let recipient = DHKeyPair()
49        let plaintext = Data("SophaxChat sealed sender test".utf8)
50
51        let (ephPub, ciphertext) = try seal(plaintext, recipientPublicKeyData: recipient.publicKeyData)
52        let recovered = try unseal(ephemeralPublicKey: ephPub, encryptedPayload: ciphertext,
53                                   recipientPrivateKey: recipient.privateKey)
54
55        #expect(recovered == plaintext, "Unsealed bytes must equal original plaintext")
56    }
57
58    @Test("Wrong recipient key fails to unseal")
59    func wrongKeyFails() throws {
60        let recipient = DHKeyPair()
61        let attacker  = DHKeyPair()
62        let plaintext = Data("secret".utf8)
63
64        let (ephPub, ciphertext) = try seal(plaintext, recipientPublicKeyData: recipient.publicKeyData)
65
66        #expect(throws: (any Error).self) {
67            _ = try unseal(ephemeralPublicKey: ephPub, encryptedPayload: ciphertext,
68                           recipientPrivateKey: attacker.privateKey)
69        }
70    }
71
72    @Test("Different seals of same plaintext produce different ciphertexts (fresh ephemeral key each time)")
73    func freshEphemeralEachSeal() throws {
74        let recipient = DHKeyPair()
75        let plaintext = Data("same message".utf8)
76
77        let (_, ct1) = try seal(plaintext, recipientPublicKeyData: recipient.publicKeyData)
78        let (_, ct2) = try seal(plaintext, recipientPublicKeyData: recipient.publicKeyData)
79
80        #expect(ct1 != ct2, "Each seal must use a fresh ephemeral key, producing unique ciphertext")
81    }
82
83    @Test("Ciphertext does not contain the plaintext in clear")
84    func ciphertextIsOpaque() throws {
85        let recipient = DHKeyPair()
86        let plaintext = Data("SophaxChat_visible_marker".utf8)
87
88        let (_, ciphertext) = try seal(plaintext, recipientPublicKeyData: recipient.publicKeyData)
89
90        // The raw marker string must not appear in the ciphertext bytes
91        let range = ciphertext.range(of: plaintext)
92        #expect(range == nil, "Plaintext must not appear verbatim inside ciphertext")
93    }
94
95    @Test("Tampered ciphertext fails authentication")
96    func tamperedCiphertextFails() throws {
97        let recipient = DHKeyPair()
98        let plaintext = Data("hello".utf8)
99
100        let (ephPub, var ciphertext) = try seal(plaintext, recipientPublicKeyData: recipient.publicKeyData)
101        // Flip a byte in the middle of the ciphertext (skip 12-byte nonce at start)
102        ciphertext[16] ^= 0xFF
103
104        #expect(throws: (any Error).self) {
105            _ = try unseal(ephemeralPublicKey: ephPub, encryptedPayload: ciphertext,
106                           recipientPrivateKey: recipient.privateKey)
107        }
108    }
109}
110
111// MARK: - Double Ratchet Header Encryption Tests
112
113@Suite("DR Header Encryption")
114struct DRHeaderEncryptionTests {
115
116    func makeSessionPair() throws -> (alice: DoubleRatchet, bob: DoubleRatchet) {
117        let sharedSecret = SymmetricKey(size: .bits256)
118        let bobSPK       = DHKeyPair()
119        let alice = try DoubleRatchet.initAsInitiator(
120            sharedSecret: sharedSecret,
121            remoteRatchetPublicKey: bobSPK.publicKeyData
122        )
123        let bob = try DoubleRatchet.initAsResponder(
124            sharedSecret: sharedSecret,
125            ownRatchetKeyPair: bobSPK
126        )
127        return (alice, bob)
128    }
129
130    @Test("Encrypted message does not expose message number in clear")
131    func messageNumberIsHidden() throws {
132        let (alice, _) = try makeSessionPair()
133        // Send 5 messages to advance the counter
134        var lastMessage: RatchetMessage?
135        for i in 0..<5 {
136            lastMessage = try alice.encrypt(plaintext: Data("msg\(i)".utf8))
137        }
138        guard let msg = lastMessage else { return }
139
140        // The byte sequence [0, 0, 0, 4] (message number = 4 in big-endian) must not
141        // appear unprotected in the encryptedHeader or ciphertext
142        let msgNum4 = Data([0, 0, 0, 4])
143        #expect(msg.encryptedHeader.range(of: msgNum4) == nil,
144                "Message number must not appear in plain text inside encryptedHeader")
145        #expect(msg.ciphertext.range(of: msgNum4) == nil,
146                "Message number must not appear in plain text inside ciphertext")
147    }
148
149    @Test("Decryption fails when encryptedHeader is tampered")
150    func tamperedHeaderFails() throws {
151        let (alice, bob) = try makeSessionPair()
152        var enc = try alice.encrypt(plaintext: Data("hello".utf8))
153        // Corrupt the first byte of the encrypted header
154        var hdr = enc.encryptedHeader
155        hdr[0] ^= 0xFF
156        enc = RatchetMessage(encryptedHeader: hdr, ciphertext: enc.ciphertext)
157
158        #expect(throws: (any Error).self) {
159            _ = try bob.decrypt(message: enc)
160        }
161    }
162
163    @Test("Decryption fails when ciphertext body is tampered but header is intact")
164    func tamperedBodyFails() throws {
165        let (alice, bob) = try makeSessionPair()
166        var enc = try alice.encrypt(plaintext: Data("hello".utf8))
167        // Corrupt the last byte of the body (Poly1305 tag is at the end)
168        var ct = enc.ciphertext
169        ct[ct.count - 1] ^= 0xFF
170        enc = RatchetMessage(encryptedHeader: enc.encryptedHeader, ciphertext: ct)
171
172        #expect(throws: (any Error).self) {
173            _ = try bob.decrypt(message: enc)
174        }
175    }
176
177    @Test("Header and body are bound: swapping header between two messages fails")
178    func headerBodyBinding() throws {
179        let (alice, bob) = try makeSessionPair()
180        let enc0 = try alice.encrypt(plaintext: Data("first".utf8))
181        let enc1 = try alice.encrypt(plaintext: Data("second".utf8))
182
183        // Mix header from enc1 with body from enc0 — should fail (AAD mismatch)
184        let mixed = RatchetMessage(encryptedHeader: enc1.encryptedHeader, ciphertext: enc0.ciphertext)
185
186        #expect(throws: (any Error).self) {
187            _ = try bob.decrypt(message: mixed)
188        }
189    }
190}
191
192// MARK: - Group Sender Key Ratchet Tests
193
194@Suite("Group Sender Key Ratchet")
195struct GroupSenderKeyRatchetTests {
196
197    /// Standalone HMAC-SHA256 sender key ratchet step (mirrors ChatManager.senderKeyRatchetStep).
198    private func ratchetStep(_ chainKey: Data) -> (messageKey: Data, nextChainKey: Data) {
199        let ck = SymmetricKey(data: chainKey)
200        let mk = Data(HMAC<SHA256>.authenticationCode(for: Data([0x01]), using: ck))
201        let nk = Data(HMAC<SHA256>.authenticationCode(for: Data([0x02]), using: ck))
202        return (mk, nk)
203    }
204
205    @Test("Message key and next chain key differ from each other and from input")
206    func distinctOutputs() {
207        let ck0 = SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }
208        let (mk, nk) = ratchetStep(ck0)
209        #expect(mk != nk, "Message key must differ from next chain key")
210        #expect(mk != ck0, "Message key must differ from input chain key")
211        #expect(nk != ck0, "Next chain key must differ from input chain key")
212    }
213
214    @Test("Consecutive iterations produce different message keys (no key reuse)")
215    func noKeyReuse() {
216        let ck0 = SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }
217        let (mk0, ck1) = ratchetStep(ck0)
218        let (mk1, _)   = ratchetStep(ck1)
219        #expect(mk0 != mk1, "Each ratchet step must produce a unique message key")
220    }
221
222    @Test("Chain is deterministic: same input always produces same output")
223    func deterministic() {
224        let ck = SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }
225        let (mk1, nk1) = ratchetStep(ck)
226        let (mk2, nk2) = ratchetStep(ck)
227        #expect(mk1 == mk2, "Same chain key must always produce the same message key")
228        #expect(nk1 == nk2, "Same chain key must always produce the same next chain key")
229    }
230
231    @Test("100-step chain produces 100 distinct message keys")
232    func longChain() {
233        var ck = SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }
234        var keys: [Data] = []
235        for _ in 0..<100 {
236            let (mk, nk) = ratchetStep(ck)
237            keys.append(mk)
238            ck = nk
239        }
240        let unique = Set(keys.map { $0.base64EncodedString() })
241        #expect(unique.count == 100, "100 ratchet steps must produce 100 unique message keys")
242    }
243
244    @Test("SenderKeyState encodes and decodes correctly")
245    func senderKeyStateCodable() throws {
246        let chainKey  = SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }
247        let state     = SenderKeyState(chainKey: chainKey, iteration: 42)
248        let encoded   = try JSONEncoder().encode(state)
249        let decoded   = try JSONDecoder().decode(SenderKeyState.self, from: encoded)
250        #expect(decoded.chainKey  == state.chainKey)
251        #expect(decoded.iteration == state.iteration)
252    }
253
254    @Test("Out-of-order recovery: skipping N steps then receiving in-order is consistent")
255    func skipAndRecover() {
256        // Simulate sender advancing 5 steps; receiver processes step 4 first, then 0-3
257        var ck = SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }
258        var allKeys: [Data] = []
259        var ckCopy = ck
260        for _ in 0..<5 {
261            let (mk, nk) = ratchetStep(ckCopy)
262            allKeys.append(mk)
263            ckCopy = nk
264        }
265        // "Receiver" independently derives step-4 key by advancing from the same initial ck
266        var rxCK = ck
267        for _ in 0..<4 { let (_, nk) = ratchetStep(rxCK); rxCK = nk }
268        let (mk4, _) = ratchetStep(rxCK)
269        #expect(mk4 == allKeys[4], "Receiver must derive the same message key for step 4 as the sender")
270
271        // Verify all 5 keys are unique
272        let unique = Set(allKeys.map { $0.base64EncodedString() })
273        #expect(unique.count == 5, "All 5 message keys must be distinct")
274    }
275}