Tests/SophaxChatCoreTests/CryptoTests.swift

1// CryptoTests.swift
2// SophaxChatCoreTests
3//
4// Tests for X3DH key agreement, Double Ratchet algorithm, and RelayRouter.
5// Uses Swift Testing framework (Swift 5.10+, no Xcode required).
6
7import Testing
8import Foundation
9import CryptoKit
10@testable import SophaxChatCore
11
12// MARK: - RelayRouter Tests
13
14@Suite("RelayRouter")
15struct RelayRouterTests {
16
17    /// Build a minimal RelayEnvelope for testing.
18    func makeEnvelope(id: String = UUID().uuidString, ttl: UInt8 = 6) -> RelayEnvelope {
19        // The inner WireMessage content doesn't matter for routing tests.
20        let fakePayload = Data("test".utf8)
21        let fakeWire = WireMessage(
22            type: .message, payload: fakePayload,
23            senderID: "alice", timestamp: Date(), signature: Data(repeating: 0, count: 64)
24        )
25        return RelayEnvelope(
26            id:           id,
27            targetPeerID: "bob",
28            originPeerID: "alice",
29            ttl:          ttl,
30            hopCount:     0,
31            message:      fakeWire
32        )
33    }
34
35    @Test("First-seen envelope is accepted")
36    func firstEnvelopeAccepted() {
37        let router   = RelayRouter()
38        let envelope = makeEnvelope()
39        #expect(router.shouldProcess(envelope) == true)
40    }
41
42    @Test("Duplicate envelope is rejected")
43    func duplicateEnvelopeRejected() {
44        let router   = RelayRouter()
45        let envelope = makeEnvelope(id: "fixed-id")
46        _ = router.shouldProcess(envelope)
47        #expect(router.shouldProcess(envelope) == false, "Relay dedup must drop already-seen IDs")
48    }
49
50    @Test("Zero-TTL envelope is rejected")
51    func zeroTTLRejected() {
52        let router   = RelayRouter()
53        let envelope = makeEnvelope(ttl: 0)
54        #expect(router.shouldProcess(envelope) == false)
55    }
56
57    @Test("Inflated TTL above maxTTL is rejected")
58    func inflatedTTLRejected() {
59        let router   = RelayRouter()
60        let envelope = makeEnvelope(ttl: RelayEnvelope.maxTTL + 1)
61        #expect(router.shouldProcess(envelope) == false)
62    }
63
64    @Test("LRU eviction allows re-entry after cap exceeded")
65    func lruEviction() {
66        let cap    = 10
67        let router = RelayRouter(maxSeen: cap)
68
69        // Fill the cache to cap
70        for _ in 0..<cap {
71            _ = router.shouldProcess(makeEnvelope())
72        }
73        #expect(router.seenCount == cap)
74
75        // Add one more — oldest should be evicted
76        _ = router.shouldProcess(makeEnvelope())
77        #expect(router.seenCount == cap, "Cache size must not exceed maxSeen")
78    }
79
80    @Test("reset() clears the seen-set")
81    func resetClearsSeen() {
82        let router = RelayRouter()
83        _ = router.shouldProcess(makeEnvelope())
84        #expect(router.seenCount == 1)
85        router.reset()
86        #expect(router.seenCount == 0)
87    }
88
89    @Test("Rate limit blocks after maxRelaysPerWindow")
90    func rateLimitBlocks() {
91        let max    = 5
92        let router = RelayRouter(maxRelaysPerWindow: max, windowSeconds: 60)
93
94        // First `max` calls should pass
95        for _ in 0..<max {
96            #expect(router.isRateLimited(senderID: "eve") == false)
97        }
98
99        // Next call should be blocked
100        #expect(router.isRateLimited(senderID: "eve") == true, "Should be rate-limited after \(max) relays")
101    }
102
103    @Test("Rate limit resets after window expires")
104    func rateLimitWindowReset() {
105        let router = RelayRouter(maxRelaysPerWindow: 2, windowSeconds: 0.01)
106
107        _ = router.isRateLimited(senderID: "eve")
108        _ = router.isRateLimited(senderID: "eve")
109        #expect(router.isRateLimited(senderID: "eve") == true)
110
111        // Wait for window to expire
112        Thread.sleep(forTimeInterval: 0.05)
113
114        // Window reset — should pass again
115        #expect(router.isRateLimited(senderID: "eve") == false, "Rate limit must reset after window expires")
116    }
117
118    @Test("Rate limit is per-sender (independent windows)")
119    func rateLimitPerSender() {
120        let max    = 3
121        let router = RelayRouter(maxRelaysPerWindow: max, windowSeconds: 60)
122
123        // Exhaust eve's quota
124        for _ in 0..<max { _ = router.isRateLimited(senderID: "eve") }
125        #expect(router.isRateLimited(senderID: "eve") == true)
126
127        // alice's quota is independent
128        #expect(router.isRateLimited(senderID: "alice") == false)
129    }
130
131    @Test("Different envelope IDs are all accepted (no false positives)")
132    func noFalsePositives() {
133        let router = RelayRouter()
134        for _ in 0..<50 {
135            let envelope = makeEnvelope() // new UUID each time
136            #expect(router.shouldProcess(envelope) == true)
137        }
138    }
139}
140
141// MARK: - X3DH Tests
142
143@Suite("X3DH Key Agreement")
144struct X3DHTests {
145
146    /// Verify that Alice and Bob derive the same shared secret via X3DH.
147    @Test("Shared secret symmetry — with one-time prekey")
148    func sharedSecretSymmetry() throws {
149        let aliceIdentityDH = DHKeyPair()
150        let bobIdentityDH   = DHKeyPair()
151        let bobSigningKey   = SigningKeyPair()
152        let bobSPK          = DHKeyPair()
153        let bobOTPK         = DHKeyPair()
154
155        let spkSignature = try bobSigningKey.privateKey.signature(for: bobSPK.publicKeyData)
156
157        let bundle = PreKeyBundle(
158            signingKeyPublic:      bobSigningKey.publicKeyData,
159            dhIdentityKeyPublic:   bobIdentityDH.publicKeyData,
160            signedPreKeyPublic:    bobSPK.publicKeyData,
161            signedPreKeySignature: spkSignature,
162            signedPreKeyId:        1,
163            oneTimePreKeyPublic:   bobOTPK.publicKeyData,
164            oneTimePreKeyId:       42,
165            username:              "bob",
166            timestamp:             Date()
167        )
168
169        let aliceResult = try X3DH.initiateSender(
170            senderIdentity: aliceIdentityDH,
171            recipientBundle: bundle
172        )
173
174        let bobSecret = try X3DH.initiateReceiver(
175            recipientIdentityDH:     bobIdentityDH,
176            recipientSignedPreKey:   bobSPK,
177            recipientOneTimePreKey:  bobOTPK,
178            senderIdentityDHKeyData: aliceIdentityDH.publicKeyData,
179            senderEphemeralKeyData:  aliceResult.ephemeralPublicKey
180        )
181
182        let aliceBytes = aliceResult.sharedSecret.withUnsafeBytes { Data($0) }
183        let bobBytes   = bobSecret.withUnsafeBytes { Data($0) }
184        #expect(aliceBytes == bobBytes, "Alice and Bob must derive the same shared secret")
185        #expect(aliceResult.usedOneTimePreKeyId == 42)
186    }
187
188    @Test("Shared secret symmetry — without one-time prekey")
189    func sharedSecretWithoutOTPK() throws {
190        let aliceIdentityDH = DHKeyPair()
191        let bobIdentityDH   = DHKeyPair()
192        let bobSigningKey   = SigningKeyPair()
193        let bobSPK          = DHKeyPair()
194
195        let spkSignature = try bobSigningKey.privateKey.signature(for: bobSPK.publicKeyData)
196
197        let bundle = PreKeyBundle(
198            signingKeyPublic:      bobSigningKey.publicKeyData,
199            dhIdentityKeyPublic:   bobIdentityDH.publicKeyData,
200            signedPreKeyPublic:    bobSPK.publicKeyData,
201            signedPreKeySignature: spkSignature,
202            signedPreKeyId:        1,
203            oneTimePreKeyPublic:   nil,
204            oneTimePreKeyId:       nil,
205            username:              "bob",
206            timestamp:             Date()
207        )
208
209        let aliceResult = try X3DH.initiateSender(senderIdentity: aliceIdentityDH, recipientBundle: bundle)
210        let bobSecret   = try X3DH.initiateReceiver(
211            recipientIdentityDH:     bobIdentityDH,
212            recipientSignedPreKey:   bobSPK,
213            recipientOneTimePreKey:  nil,
214            senderIdentityDHKeyData: aliceIdentityDH.publicKeyData,
215            senderEphemeralKeyData:  aliceResult.ephemeralPublicKey
216        )
217
218        let a = aliceResult.sharedSecret.withUnsafeBytes { Data($0) }
219        let b = bobSecret.withUnsafeBytes { Data($0) }
220        #expect(a == b)
221        #expect(aliceResult.usedOneTimePreKeyId == nil)
222    }
223
224    @Test("Stale bundle is rejected")
225    func staleBundle() throws {
226        let aliceIdentityDH = DHKeyPair()
227        let bobSigningKey   = SigningKeyPair()
228        let bobSPK          = DHKeyPair()
229        let spkSignature    = try bobSigningKey.privateKey.signature(for: bobSPK.publicKeyData)
230
231        let bundle = PreKeyBundle(
232            signingKeyPublic:      bobSigningKey.publicKeyData,
233            dhIdentityKeyPublic:   DHKeyPair().publicKeyData,
234            signedPreKeyPublic:    bobSPK.publicKeyData,
235            signedPreKeySignature: spkSignature,
236            signedPreKeyId:        1,
237            oneTimePreKeyPublic:   nil, oneTimePreKeyId: nil,
238            username:              "bob",
239            timestamp:             Date(timeIntervalSinceNow: -90_000)  // 25h ago
240        )
241
242        #expect(throws: SophaxError.stalePreKeyBundle) {
243            try X3DH.initiateSender(senderIdentity: aliceIdentityDH, recipientBundle: bundle)
244        }
245    }
246
247    @Test("Invalid SPK signature is rejected")
248    func invalidSignature() throws {
249        let aliceIdentityDH = DHKeyPair()
250        let bobSigningKey   = SigningKeyPair()
251        let bobSPK          = DHKeyPair()
252        let badSignature    = Data(repeating: 0xAA, count: 64)
253
254        let bundle = PreKeyBundle(
255            signingKeyPublic:      bobSigningKey.publicKeyData,
256            dhIdentityKeyPublic:   DHKeyPair().publicKeyData,
257            signedPreKeyPublic:    bobSPK.publicKeyData,
258            signedPreKeySignature: badSignature,
259            signedPreKeyId:        1,
260            oneTimePreKeyPublic:   nil, oneTimePreKeyId: nil,
261            username:              "bob",
262            timestamp:             Date()
263        )
264
265        #expect(throws: (any Error).self) {
266            try X3DH.initiateSender(senderIdentity: aliceIdentityDH, recipientBundle: bundle)
267        }
268    }
269
270    @Test("Different initiations produce different secrets")
271    func differentSecretsEachTime() throws {
272        let bobIdentityDH = DHKeyPair()
273        let bobSK         = SigningKeyPair()
274        let bobSPK        = DHKeyPair()
275        let spkSig        = try bobSK.privateKey.signature(for: bobSPK.publicKeyData)
276
277        let bundle = PreKeyBundle(
278            signingKeyPublic: bobSK.publicKeyData, dhIdentityKeyPublic: bobIdentityDH.publicKeyData,
279            signedPreKeyPublic: bobSPK.publicKeyData, signedPreKeySignature: spkSig,
280            signedPreKeyId: 1, oneTimePreKeyPublic: nil, oneTimePreKeyId: nil,
281            username: "bob", timestamp: Date()
282        )
283
284        let result1 = try X3DH.initiateSender(senderIdentity: DHKeyPair(), recipientBundle: bundle)
285        let result2 = try X3DH.initiateSender(senderIdentity: DHKeyPair(), recipientBundle: bundle)
286
287        let s1 = result1.sharedSecret.withUnsafeBytes { Data($0) }
288        let s2 = result2.sharedSecret.withUnsafeBytes { Data($0) }
289        #expect(s1 != s2, "Different ephemeral keys must produce different secrets")
290    }
291}
292
293// MARK: - Double Ratchet Tests
294
295@Suite("Double Ratchet")
296struct DoubleRatchetTests {
297
298    /// Helper: build matched Alice/Bob session pair from a shared secret.
299    func makeSessionPair() throws -> (alice: DoubleRatchet, bob: DoubleRatchet) {
300        let sharedSecret = SymmetricKey(size: .bits256)
301        let bobSPK       = DHKeyPair()
302        let alice = try DoubleRatchet.initAsInitiator(
303            sharedSecret: sharedSecret,
304            remoteRatchetPublicKey: bobSPK.publicKeyData
305        )
306        let bob = try DoubleRatchet.initAsResponder(
307            sharedSecret: sharedSecret,
308            ownRatchetKeyPair: bobSPK
309        )
310        return (alice, bob)
311    }
312
313    @Test("Alice → Bob basic encryption")
314    func basicEncryptDecrypt() throws {
315        let (alice, bob) = try makeSessionPair()
316        let plaintext    = Data("Hello, Bob!".utf8)
317        let encrypted    = try alice.encrypt(plaintext: plaintext)
318        let decrypted    = try bob.decrypt(message: encrypted)
319        #expect(decrypted == plaintext)
320    }
321
322    @Test("Bob replies to Alice")
323    func bobReplies() throws {
324        let (alice, bob) = try makeSessionPair()
325        let enc1 = try alice.encrypt(plaintext: Data("Hi".utf8))
326        let _    = try bob.decrypt(message: enc1)
327        let enc2 = try bob.encrypt(plaintext: Data("Hey!".utf8))
328        let dec2 = try alice.decrypt(message: enc2)
329        #expect(dec2 == Data("Hey!".utf8))
330    }
331
332    @Test("20 consecutive messages Alice → Bob")
333    func manyMessages() throws {
334        let (alice, bob) = try makeSessionPair()
335        for i in 0..<20 {
336            let text = "Message \(i)"
337            let enc  = try alice.encrypt(plaintext: Data(text.utf8))
338            let dec  = try bob.decrypt(message: enc)
339            #expect(String(data: dec, encoding: .utf8) == text)
340        }
341    }
342
343    @Test("10 round-trip bidirectional messages")
344    func bidirectional() throws {
345        let (alice, bob) = try makeSessionPair()
346        for i in 0..<10 {
347            let a2b = "A→B \(i)"
348            let b2a = "B→A \(i)"
349            let aEnc = try alice.encrypt(plaintext: Data(a2b.utf8))
350            let aDec = try bob.decrypt(message: aEnc)
351            #expect(String(data: aDec, encoding: .utf8) == a2b)
352            let bEnc = try bob.encrypt(plaintext: Data(b2a.utf8))
353            let bDec = try alice.decrypt(message: bEnc)
354            #expect(String(data: bDec, encoding: .utf8) == b2a)
355        }
356    }
357
358    @Test("Different messages produce different ciphertexts (unique keys/nonces)")
359    func uniqueKeys() throws {
360        let (alice, _) = try makeSessionPair()
361        let enc1 = try alice.encrypt(plaintext: Data("same".utf8))
362        let enc2 = try alice.encrypt(plaintext: Data("same".utf8))
363        #expect(enc1.ciphertext != enc2.ciphertext, "Each message must use a unique key+nonce")
364    }
365
366    @Test("Out-of-order message delivery")
367    func outOfOrder() throws {
368        let (alice, bob) = try makeSessionPair()
369        let pt0 = Data("msg0".utf8)
370        let pt1 = Data("msg1".utf8)
371        let pt2 = Data("msg2".utf8)
372
373        let enc0 = try alice.encrypt(plaintext: pt0)
374        let enc1 = try alice.encrypt(plaintext: pt1)
375        let enc2 = try alice.encrypt(plaintext: pt2)
376
377        // Deliver 2 → 0 → 1 (out of order)
378        let dec2 = try bob.decrypt(message: enc2)
379        let dec0 = try bob.decrypt(message: enc0)
380        let dec1 = try bob.decrypt(message: enc1)
381
382        #expect(dec0 == pt0)
383        #expect(dec1 == pt1)
384        #expect(dec2 == pt2)
385    }
386
387    @Test("Session state export/import (persistence)")
388    func stateRoundtrip() throws {
389        let (alice, bob) = try makeSessionPair()
390        let pt  = Data("hello".utf8)
391        let enc = try alice.encrypt(plaintext: pt)
392
393        // Persist and restore Bob's state
394        let stateData   = try bob.exportState()
395        let restoredBob = try DoubleRatchet.importState(stateData)
396        let dec         = try restoredBob.decrypt(message: enc)
397        #expect(dec == pt)
398    }
399
400    @Test("Wrong associated data fails decryption")
401    func wrongAD() throws {
402        let (alice, bob) = try makeSessionPair()
403        let enc = try alice.encrypt(plaintext: Data("hi".utf8), associatedData: Data("correct".utf8))
404        #expect(throws: (any Error).self) {
405            try bob.decrypt(message: enc, associatedData: Data("wrong".utf8))
406        }
407    }
408
409    @Test("KDF_CK produces distinct keys at each step")
410    func kdfCKDistinctKeys() {
411        let ck          = SymmetricKey(size: .bits256)
412        let (ck2, mk1)  = DoubleRatchet.kdfCK(ck)
413        let (ck3, mk2)  = DoubleRatchet.kdfCK(ck2)
414
415        func bytes(_ k: SymmetricKey) -> Data { k.withUnsafeBytes { Data($0) } }
416
417        #expect(bytes(ck)  != bytes(ck2),  "Chain key must change each step")
418        #expect(bytes(ck2) != bytes(ck3),  "Chain key must change each step")
419        #expect(bytes(ck2) != bytes(mk1),  "Chain key ≠ message key")
420        #expect(bytes(mk1) != bytes(mk2),  "Consecutive message keys must differ")
421    }
422}