Bind capability state to Noise generations

bitchatTests/Noise/NoiseCoverageTests.swift

@@ -542,8 +542,12 @@ struct NoiseCoverageTests {
         let aliceManager = NoiseSessionManager(localStaticKey: aliceStaticKey, keychain: keychain)
         let bobManager = NoiseSessionManager(localStaticKey: bobStaticKey, keychain: keychain)
 
-        aliceManager.onSessionEstablished = establishedRecorder.recordEstablished(peerID:remoteKey:)
-        bobManager.onSessionEstablished = establishedRecorder.recordEstablished(peerID:remoteKey:)
+        aliceManager.onSessionEstablished = establishedRecorder.recordEstablished(
+            peerID:remoteKey:sessionGeneration:
+        )
+        bobManager.onSessionEstablished = establishedRecorder.recordEstablished(
+            peerID:remoteKey:sessionGeneration:
+        )
 
         try establishManagerSessions(aliceManager: aliceManager, bobManager: bobManager)
 
@@ -658,6 +662,95 @@ struct NoiseCoverageTests {
         #expect(rekeyedSession.getState() == .handshaking)
     }
 
+    @Test("A stale decrypt generation cannot commit across session promotion")
+    func staleDecryptGenerationCannotCommitAcrossPromotion() throws {
+        let aliceManager = NoiseSessionManager(
+            localStaticKey: aliceStaticKey,
+            keychain: keychain,
+            sessionFactory: { peerID, role in
+                BlockingDecryptNoiseSession(
+                    peerID: peerID,
+                    role: role,
+                    keychain: self.keychain,
+                    localStaticKey: self.aliceStaticKey
+                )
+            }
+        )
+        let bobManager = NoiseSessionManager(localStaticKey: bobStaticKey, keychain: keychain)
+        try establishManagerSessions(aliceManager: aliceManager, bobManager: bobManager)
+
+        let oldSession = try #require(
+            aliceManager.getSession(for: alicePeerID) as? BlockingDecryptNoiseSession
+        )
+        let oldGeneration = try #require(aliceManager.sessionGeneration(for: alicePeerID))
+
+        // Prepare a fully authenticated responder candidate without promoting
+        // it yet. Its final XX message is the exact operation that replaces
+        // the old `sessions[peerID]` entry.
+        let replacementInitiator = NoiseSession(
+            peerID: bobPeerID,
+            role: .initiator,
+            keychain: keychain,
+            localStaticKey: bobStaticKey
+        )
+        let message1 = try replacementInitiator.startHandshake()
+        let message2 = try #require(
+            try aliceManager.handleIncomingHandshake(from: alicePeerID, message: message1)
+        )
+        let message3 = try #require(try replacementInitiator.processHandshakeMessage(message2))
+
+        let ciphertext = try bobManager.encrypt(Data("old session".utf8), for: bobPeerID)
+        oldSession.pauseNextDecrypt()
+
+        let decryptResult = ConcurrentTestResult<(plaintext: Data, sessionGeneration: UUID)>()
+        DispatchQueue.global().async {
+            decryptResult.capture {
+                try aliceManager.decryptWithSessionGeneration(ciphertext, from: self.alicePeerID)
+            }
+        }
+        #expect(oldSession.waitForDecryptStart(timeout: 2))
+
+        let promotionStarted = DispatchSemaphore(value: 0)
+        let promotionResult = ConcurrentTestResult<Data?>()
+        DispatchQueue.global().async {
+            promotionStarted.signal()
+            promotionResult.capture {
+                try aliceManager.handleIncomingHandshake(from: self.alicePeerID, message: message3)
+            }
+        }
+        #expect(promotionStarted.wait(timeout: .now() + 2) == .success)
+        #expect(
+            promotionResult.wait(timeout: 0.05) == nil,
+            "Promotion must wait for the exact decrypting-session lease"
+        )
+
+        oldSession.resumeDecrypt()
+        let decrypted = try #require(decryptResult.wait(timeout: 2)).get()
+        _ = try #require(promotionResult.wait(timeout: 2)).get()
+
+        #expect(decrypted.plaintext == Data("old session".utf8))
+        #expect(decrypted.sessionGeneration == oldGeneration)
+        #expect(aliceManager.sessionGeneration(for: alicePeerID) != oldGeneration)
+        #expect(throws: NoiseEncryptionError.sessionNotEstablished) {
+            try aliceManager.encrypt(
+                Data("stale send".utf8),
+                for: alicePeerID,
+                expectedSessionGeneration: oldGeneration
+            )
+        }
+
+        var staleCommitRan = false
+        let staleCommit = aliceManager.withCurrentSessionGeneration(
+            for: alicePeerID,
+            expected: decrypted.sessionGeneration
+        ) {
+            staleCommitRan = true
+            return true
+        }
+        #expect(staleCommit == nil)
+        #expect(!staleCommitRan)
+    }
+
     @Test("Secure noise sessions enforce limits and renegotiation thresholds")
     func secureNoiseSessionsEnforceLimitsAndThresholds() throws {
         let initiator = SecureNoiseSession(
@@ -852,7 +945,11 @@ private final class SessionCallbackRecorder: @unchecked Sendable {
         return establishedEntries.map(\.0)
     }
 
-    func recordEstablished(peerID: PeerID, remoteKey: Curve25519.KeyAgreement.PublicKey) {
+    func recordEstablished(
+        peerID: PeerID,
+        remoteKey: Curve25519.KeyAgreement.PublicKey,
+        sessionGeneration _: UUID
+    ) {
         lock.lock()
         establishedEntries.append((peerID, remoteKey.rawRepresentation))
         lock.unlock()
@@ -874,3 +971,58 @@ private final class FailingNoiseSession: NoiseSession {
         throw Error.synthetic
     }
 }
+
+private final class BlockingDecryptNoiseSession: NoiseSession, @unchecked Sendable {
+    private let controlLock = NSLock()
+    private var shouldPauseNextDecrypt = false
+    private let decryptStarted = DispatchSemaphore(value: 0)
+    private let resumeDecryptSemaphore = DispatchSemaphore(value: 0)
+
+    func pauseNextDecrypt() {
+        controlLock.lock()
+        shouldPauseNextDecrypt = true
+        controlLock.unlock()
+    }
+
+    func waitForDecryptStart(timeout: TimeInterval) -> Bool {
+        decryptStarted.wait(timeout: .now() + timeout) == .success
+    }
+
+    func resumeDecrypt() {
+        resumeDecryptSemaphore.signal()
+    }
+
+    override func decrypt(_ ciphertext: Data) throws -> Data {
+        controlLock.lock()
+        let shouldPause = shouldPauseNextDecrypt
+        shouldPauseNextDecrypt = false
+        controlLock.unlock()
+
+        if shouldPause {
+            decryptStarted.signal()
+            resumeDecryptSemaphore.wait()
+        }
+        return try super.decrypt(ciphertext)
+    }
+}
+
+private final class ConcurrentTestResult<Value>: @unchecked Sendable {
+    private let lock = NSLock()
+    private let completed = DispatchSemaphore(value: 0)
+    private var storedResult: Result<Value, Error>?
+
+    func capture(_ operation: () throws -> Value) {
+        let result = Result(catching: operation)
+        lock.lock()
+        storedResult = result
+        lock.unlock()
+        completed.signal()
+    }
+
+    func wait(timeout: TimeInterval) -> Result<Value, Error>? {
+        guard completed.wait(timeout: .now() + timeout) == .success else { return nil }
+        lock.lock()
+        defer { lock.unlock() }
+        return storedResult
+    }
+}

bitchatTests/Services/BLENoisePacketHandlerTests.swift

@@ -9,6 +9,7 @@ struct BLENoisePacketHandlerTests {
     private final class Recorder {
         var handshakeResult: Result<Data?, Error> = .success(nil)
         var hasSession = false
+        let sessionGeneration = UUID()
         var decryptResult: Result<Data, Error> = .success(Data())
 
         var processedHandshakes: [(peerID: PeerID, message: Data)] = []
@@ -18,7 +19,7 @@ struct BLENoisePacketHandlerTests {
         var lastSeenUpdates: [PeerID] = []
         var decryptCalls: [(payload: Data, peerID: PeerID)] = []
         var clearedSessions: [PeerID] = []
-        var authenticatedPeerStates: [(peerID: PeerID, payload: Data)] = []
+        var authenticatedPeerStates: [(peerID: PeerID, payload: Data, generation: UUID)] = []
         var deliveries: [(peerID: PeerID, type: NoisePayloadType, payload: Data, timestamp: Date)] = []
         /// Ordered side-effect log to assert recovery sequencing.
         var events: [String] = []
@@ -57,14 +58,17 @@ struct BLENoisePacketHandlerTests {
             },
             decrypt: { payload, peerID in
                 recorder.decryptCalls.append((payload, peerID))
-                return try recorder.decryptResult.get()
+                return BLENoiseDecryptionResult(
+                    plaintext: try recorder.decryptResult.get(),
+                    sessionGeneration: recorder.sessionGeneration
+                )
             },
             clearSession: { peerID in
                 recorder.clearedSessions.append(peerID)
                 recorder.events.append("clearSession")
             },
-            handleAuthenticatedPeerState: { peerID, payload in
-                recorder.authenticatedPeerStates.append((peerID, payload))
+            handleAuthenticatedPeerState: { peerID, payload, generation in
+                recorder.authenticatedPeerStates.append((peerID, payload, generation))
             },
             deliverNoisePayload: { peerID, type, payload, timestamp in
                 recorder.deliveries.append((peerID, type, payload, timestamp))
@@ -240,6 +244,7 @@ struct BLENoisePacketHandlerTests {
         #expect(recorder.authenticatedPeerStates.count == 1)
         #expect(recorder.authenticatedPeerStates.first?.peerID == remotePeerID)
         #expect(recorder.authenticatedPeerStates.first?.payload == Data([0x01, 0x02, 0x03]))
+        #expect(recorder.authenticatedPeerStates.first?.generation == recorder.sessionGeneration)
         #expect(recorder.deliveries.isEmpty)
     }
 

bitchatTests/Services/NoiseEncryptionServiceTests.swift

@@ -96,13 +96,22 @@ struct NoiseEncryptionServiceTests {
         let recorder = AuthenticationRecorder()
 
         #expect(alice.onPeerAuthenticated == nil)
+        #expect(bob.onPeerAuthenticatedWithGeneration == nil)
         alice.addOnPeerAuthenticatedHandler(recorder.record(peerID:fingerprint:))
         bob.onPeerAuthenticated = recorder.record(peerID:fingerprint:)
+        bob.onPeerAuthenticatedWithGeneration = recorder.record(
+            peerID:fingerprint:sessionGeneration:
+        )
 
         try establishSessions(alice: alice, bob: bob)
 
         let authenticated = await TestHelpers.waitUntil({ recorder.count >= 2 }, timeout: 5.0)
         #expect(authenticated)
+        let generationAuthenticated = await TestHelpers.waitUntil(
+            { recorder.generationCount >= 1 },
+            timeout: 5.0
+        )
+        #expect(generationAuthenticated)
         #expect(alice.hasEstablishedSession(with: bobPeerID))
         #expect(bob.hasEstablishedSession(with: alicePeerID))
         #expect(alice.hasSession(with: bobPeerID))
@@ -111,6 +120,7 @@ struct NoiseEncryptionServiceTests {
         #expect(bob.getPeerPublicKeyData(alicePeerID)?.count == 32)
         #expect(alice.getPeerFingerprint(bobPeerID) != nil)
         #expect(bob.getPeerFingerprint(alicePeerID) != nil)
+        #expect(recorder.generation(for: alicePeerID) == bob.sessionGeneration(for: alicePeerID))
 
         let plaintext = Data("secret payload".utf8)
         let ciphertext = try alice.encrypt(plaintext, for: bobPeerID)
@@ -244,6 +254,17 @@ struct NoiseEncryptionServiceTests {
         let alicePeerID = PeerID(publicKey: alice.getStaticPublicKeyData())
         let bobPeerID = PeerID(publicKey: bob.getStaticPublicKeyData())
         try establishSessions(alice: alice, bob: bob)
+        let originalGeneration = try #require(alice.sessionGeneration(for: bobPeerID))
+        var leaseRan = false
+        let leased = alice.withCurrentSessionGeneration(
+            for: bobPeerID,
+            expected: originalGeneration
+        ) {
+            leaseRan = true
+            return true
+        }
+        #expect(leased == true)
+        #expect(leaseRan)
 
         var emittedPeerID: PeerID?
         var emittedMessage: Data?
@@ -254,6 +275,17 @@ struct NoiseEncryptionServiceTests {
         try alice._test_initiateAutomaticRekey(for: bobPeerID)
 
         #expect(emittedPeerID == bobPeerID)
+        #expect(alice.sessionGeneration(for: bobPeerID) == nil)
+        leaseRan = false
+        let staleLease = alice.withCurrentSessionGeneration(
+            for: bobPeerID,
+            expected: originalGeneration
+        ) {
+            leaseRan = true
+            return true
+        }
+        #expect(staleLease == nil)
+        #expect(!leaseRan)
         let message1 = try #require(emittedMessage)
         #expect(!message1.isEmpty)
         #expect(alice.hasSession(with: bobPeerID))
@@ -269,6 +301,7 @@ struct NoiseEncryptionServiceTests {
 
         #expect(alice.hasEstablishedSession(with: bobPeerID))
         #expect(bob.hasEstablishedSession(with: alicePeerID))
+        #expect(alice.sessionGeneration(for: bobPeerID) != originalGeneration)
     }
 
     @Test("Large private-file payloads use the bounded Noise extension")
@@ -411,6 +444,7 @@ struct NoiseEncryptionServiceTests {
 private final class AuthenticationRecorder: @unchecked Sendable {
     private let lock = NSLock()
     private var entries: [(PeerID, String)] = []
+    private var generationEntries: [(PeerID, UUID)] = []
 
     var count: Int {
         lock.lock()
@@ -418,9 +452,27 @@ private final class AuthenticationRecorder: @unchecked Sendable {
         return entries.count
     }
 
+    var generationCount: Int {
+        lock.lock()
+        defer { lock.unlock() }
+        return generationEntries.count
+    }
+
     func record(peerID: PeerID, fingerprint: String) {
         lock.lock()
         entries.append((peerID, fingerprint))
         lock.unlock()
     }
+
+    func record(peerID: PeerID, fingerprint _: String, sessionGeneration: UUID) {
+        lock.lock()
+        generationEntries.append((peerID, sessionGeneration))
+        lock.unlock()
+    }
+
+    func generation(for peerID: PeerID) -> UUID? {
+        lock.lock()
+        defer { lock.unlock() }
+        return generationEntries.last { $0.0 == peerID }?.1
+    }
 }