feat: changes needed to implement cypher session renewal on client side (#8)

* feat: changes needed to implement cypher session renewal on client side

* test: add test to cover a case with loosing session's changes

Author: AndriyVasyk

Sources/DXProtocol/DXError.swift

@@ -77,8 +77,8 @@ public enum DXError: Error, Equatable {
     case invalidSenderKeySession(distributionId: UUID, message: String)
     /// The message has already been sent.
     case duplicatedMessage(String)
-    /// The signature verification failed.
-    case verificationFailed(String)
+    /// The MAC verification failed.
+    case messageVerificationFailed(String)
     /// An error occurred in a callback.
     case callbackError(String)
     /// An unknown error occurred.

Sources/DXProtocol/Session/Session.swift

@@ -334,7 +334,6 @@ extension Session {
                 identityStore: identityStore,
                 preKeyStore: preKeyStore,
                 signedPreKeyStore: signedPreKeyStore)
-        // This is 'decrypt_message_with_record'
         let result = try session.decryptMessage(preKeyMessage.secureMessage)
 
         try sessionStore.storeSession(session, for: address)
@@ -380,7 +379,6 @@ extension Session {
             throw DXError.untrustedIdentity("Abort decrypting PreKeyMessage for untrusted identity")
         }
 
-        // This is 'decrypt_message_with_record'
         let result = try session.decryptMessage(secureMessage)
 
         try sessionStore.storeSession(session, for: address)
@@ -399,13 +397,16 @@ extension Session {
     /// - Returns: The decrypted data.
     private mutating func decryptMessage(_ message: SecureMessage) throws -> Data {
         var result = Data()
+        var errors = [Error]()
 
         do {
             // We MUST discard changes of 'state' if decryption of message fails
             var state = self.state
             result = try state.decryptMessage(message)
             self.state = state
         } catch let error {
+            errors.append(error)
+
             if case DXError.duplicatedMessage = error {
                 // This message has been already processed.
                 // Do not try previous states to decrypt it
@@ -423,6 +424,8 @@ extension Session {
                     self.promotePreviousState(state: olderState, index: index)
                     break
                 } catch let error {
+                    errors.append(error)
+
                     // This code is not covered by tests
                     if case DXError.duplicatedMessage = error {
                         // This message has been already processed.
@@ -433,6 +436,12 @@ extension Session {
             }
         }
 
+        if result.isEmpty && !errors.isEmpty {
+            let protocolErrors = errors.compactMap { $0 as? DXError }
+            let error = protocolErrors.first ?? errors[0]
+            throw error
+        }
+
         if result.isEmpty {
             throw DXError.invalidMessage("Failed to decrypt message")
         }
@@ -459,13 +468,13 @@ extension Session {
 
     /// Checks if the session contains active current state.
     /// - Returns: True if session contains active current state. False otherwise
-    func hasCurrentState() -> Bool {
+    public func hasCurrentState() -> Bool {
         return !self.state.archived
     }
 
     // TODO: - Add tests on error
     /// Archives the current active state. For example, this functionality might be used to manage stale devices
-    mutating func archiveCurrentState() {
+    public mutating func archiveCurrentState() {
         guard !self.state.archived else {
             return
         }

Sources/DXProtocol/Session/SessionState.swift

@@ -235,7 +235,7 @@ struct SessionState: Codable {
                 receiverIdentityKey: localIdentityKey,
                 macKey: messageKeys.macKey)
         guard isMacValid else {
-            throw DXError.invalidMessage("MAC verification failed")
+            throw DXError.messageVerificationFailed("MAC verification failed")
         }
 
         // Source - https://developer.apple.com/forums/thread/687212

Tests/DXProtocolTests/SessionTests/SessionTests.swift

@@ -1455,6 +1455,140 @@ final class SessionTests: XCTestCase {
             XCTFail("Invalid error type: \(error)")
         }
     }
+    
+    func testInteractionWhenLostSessionChanges() throws {
+        let senderClient = try TestClient(userId: UUID())  // Alice
+        let recipientClient = try TestClient(userId: UUID())  // Bob
+        try initializeSession(senderClient: senderClient, recipientClient: recipientClient)
+
+        let aliceClient = senderClient
+        let bobClient = recipientClient
+        
+        // Archive session state so that we have a non-empty list of previous states
+        var session = try bobClient.sessionStore.loadSession(for: aliceClient.protocolAddress)
+        session?.archiveCurrentState()
+        try bobClient.sessionStore.storeSession(
+            try XCTUnwrap(session),
+            for: aliceClient.protocolAddress)
+
+        // The code below will do the next ratchet step
+        // {
+        // Alice sends one more message
+        let data = try XCTUnwrap("From Alice".data(using: .utf8))
+        let message = try Session.encrypt(
+            data: data,
+            for: bobClient.protocolAddress,
+            sessionStore: aliceClient.sessionStore,
+            identityStore: aliceClient.identityKeyStore
+        )
+        let result = try Session.decrypt(
+            message: message,
+            from: aliceClient.protocolAddress,
+            sessionStore: bobClient.sessionStore,
+            identityStore: bobClient.identityKeyStore,
+            preKeyStore: bobClient.preKeyStore,
+            signedPreKeyStore: bobClient.signedPreKeyStore
+        )
+        XCTAssertEqual(result, data)
+        
+        // Bob sends reply to Alice
+        let replyData = try XCTUnwrap("From Bob".data(using: .utf8))
+        let replyMessage = try Session.encrypt(
+            data: replyData,
+            for: aliceClient.protocolAddress,
+            sessionStore: bobClient.sessionStore,
+            identityStore: bobClient.identityKeyStore)
+        
+        let decryptReplyMessage = try Session.decrypt(
+            message: replyMessage,
+            from: bobClient.protocolAddress,
+            sessionStore: aliceClient.sessionStore,
+            identityStore: aliceClient.identityKeyStore,
+            preKeyStore: aliceClient.preKeyStore,
+            signedPreKeyStore: aliceClient.signedPreKeyStore)
+        XCTAssertEqual(decryptReplyMessage, replyData)
+        // }
+        
+        // We will use this session to imitate loosing of session's changes
+        let outdatedSession = try bobClient.sessionStore.loadSession(for: aliceClient.protocolAddress)
+
+        // The code below will do further ratchet step (session's changes)
+        // {
+        let secondData = try XCTUnwrap("Alice second message".data(using: .utf8))
+        let secondMessage = try Session.encrypt(
+            data: secondData,
+            for: bobClient.protocolAddress,
+            sessionStore: aliceClient.sessionStore,
+            identityStore: aliceClient.identityKeyStore)
+        let decryptSecondResult = try Session.decrypt(
+            message: secondMessage,
+            from: aliceClient.protocolAddress,
+            sessionStore: bobClient.sessionStore,
+            identityStore: bobClient.identityKeyStore,
+            preKeyStore: bobClient.preKeyStore,
+            signedPreKeyStore: bobClient.signedPreKeyStore)
+        XCTAssertEqual(decryptSecondResult, secondData)
+        
+        let replySecondData = try XCTUnwrap("Bob reply on second message".data(using: .utf8))
+        let replySecondMessage = try Session.encrypt(
+            data: replySecondData,
+            for: aliceClient.protocolAddress,
+            sessionStore: bobClient.sessionStore,
+            identityStore: bobClient.identityKeyStore)
+        let decryptReplySecondResult = try Session.decrypt(
+            message: replySecondMessage,
+            from: bobClient.protocolAddress,
+            sessionStore: aliceClient.sessionStore,
+            identityStore: aliceClient.identityKeyStore,
+            preKeyStore: aliceClient.preKeyStore,
+            signedPreKeyStore: aliceClient.signedPreKeyStore)
+        XCTAssertEqual(decryptReplySecondResult, replySecondData)
+        
+        // Actually test
+        
+        // Imitate like we lost last changes made to session
+        try bobClient.sessionStore.storeSession(
+            try XCTUnwrap(outdatedSession),
+            for: aliceClient.protocolAddress)
+        
+        // Test Bob can send message to Alice after lost session changes
+        let bobDataAfterLost = try XCTUnwrap("Message From Bob after lost".data(using: .utf8))
+        let bobMessageAfterLost = try Session.encrypt(
+                        data: bobDataAfterLost,
+                        for: aliceClient.protocolAddress,
+                        sessionStore: bobClient.sessionStore,
+                        identityStore: bobClient.identityKeyStore)
+        let aliceDecryptAfterLost = try Session.decrypt(
+                message: bobMessageAfterLost,
+                from: bobClient.protocolAddress,
+                sessionStore: aliceClient.sessionStore,
+                identityStore: aliceClient.identityKeyStore,
+                preKeyStore: aliceClient.preKeyStore,
+                signedPreKeyStore: aliceClient.signedPreKeyStore)
+        XCTAssertEqual(aliceDecryptAfterLost, bobDataAfterLost)
+        
+        // Test Bob can't receive message from Alice
+        let aliceThirdMessageData = try XCTUnwrap("Alice third message".data(using: .utf8))
+        let aliceThirdMessage = try Session.encrypt(
+            data: aliceThirdMessageData,
+            for: bobClient.protocolAddress,
+            sessionStore: aliceClient.sessionStore,
+            identityStore: aliceClient.identityKeyStore)
+        do {
+            let decryptThirdMessageResult = try Session.decrypt(
+                message: aliceThirdMessage,
+                from: aliceClient.protocolAddress,
+                sessionStore: bobClient.sessionStore,
+                identityStore: bobClient.identityKeyStore,
+                preKeyStore: bobClient.preKeyStore,
+                signedPreKeyStore: bobClient.signedPreKeyStore)
+            XCTFail("Third message should not be decrypted due loosing session changes")
+            XCTAssertEqual(decryptThirdMessageResult, aliceThirdMessageData)
+        } catch DXError.messageVerificationFailed {
+        } catch {
+            XCTFail("Incorrect error thrown")
+        }
+    }
 }
 
 extension SessionTests {