feat: implement thread-safe access to sessions (#4)

Author: AndriyVasyk

Sources/DXProtocol/Session/Session.swift

@@ -109,6 +109,10 @@ public struct Session: Codable {
 
         try identityStore.saveIdentity(theirIdentityKey, for: address)
 
+        let lock = SessionLock(address: address)
+        lock.lock()
+        defer { lock.unlock() }
+
         var session = try sessionStore.loadSession(for: address)
         if nil == session {
             session = Self(state: state)
@@ -154,6 +158,10 @@ public struct Session: Codable {
                 direction: .receiving) else {
             throw DXError.untrustedIdentity("Abort processing PreKey Message for untrusted identity")
         }
+        
+        let lock = SessionLock(address: address)
+        lock.lock()
+        defer { lock.unlock() }
 
         let theirBaseKey = message.senderBaseKey
         let messageVersion = Int(message.messageVersion)
@@ -227,6 +235,10 @@ public struct Session: Codable {
                                  for address: ProtocolAddress,
                                  sessionStore: SessionStorable,
                                  identityStore: IdentityKeyStorable) throws -> MessageContainer {
+        let lock = SessionLock(address: address)
+        lock.lock()
+        defer { lock.unlock() }
+
         let result = try self.state.encrypt(
                 data: data,
                 sessionStore: sessionStore,
@@ -311,6 +323,10 @@ extension Session {
                                 identityStore: IdentityKeyStorable,
                                 preKeyStore: PreKeyStorable,
                                 signedPreKeyStore: SignedPreKeyStorable) throws -> Data {
+        let lock = SessionLock(address: address)
+        lock.lock()
+        defer { lock.unlock() }
+
         var session = try self.processPreKeyMessage(
                 preKeyMessage,
                 from: address,
@@ -347,6 +363,10 @@ extension Session {
                                 from address: ProtocolAddress,
                                 sessionStore: SessionStorable,
                                 identityStore: IdentityKeyStorable) throws -> Data {
+        let lock = SessionLock(address: address)
+        lock.lock()
+        defer { lock.unlock() }
+
         // This code is not covered by tests
         guard var session = try sessionStore.loadSession(for: address) else {
             throw DXError.sessionNotFound("Failed to find session while decrypting message")

Sources/DXProtocol/Session/SessionLock.swift

@@ -0,0 +1,31 @@
+//
+//  SessionLock.swift
+//  
+//
+//  Created by Andriy Vasyk on 19.02.2024.
+//
+
+import Foundation
+
+struct SessionLock {
+    /// The underlying implementation of this session lock
+    private var impl: NSRecursiveLock
+    
+    /// Initialises a new lock for session that corresponds to specified protocol address
+    /// - Parameter address: The protocol address that uniquely identifies session
+    init(address: ProtocolAddress) {
+        self.impl = SessionLockStorage.shared.lock(for: address)
+    }
+    
+    // MARK: - Interface
+    
+    /// Blocks a thread’s execution until the lock can be acquired.
+    func lock() {
+        self.impl.lock()
+    }
+    
+    /// Relinquishes a previously acquired lock.
+    func unlock() {
+        self.impl.unlock()
+    }
+}

Sources/DXProtocol/Session/SessionLockStorage.swift

@@ -0,0 +1,33 @@
+//
+//  SessionLockStorage.swift
+//  
+//
+//  Created by Andriy Vasyk on 19.02.2024.
+//
+
+import Foundation
+
+final class SessionLockStorage {
+    static let shared = SessionLockStorage()
+    
+    /// Internal queue allowing to sync access to storage with locks
+    private let syncAccessQueue = DispatchQueue(label: "StorageAccessQueue")
+    
+    /// Container to store locks and associate them with corresponding protocol addresses
+    private var storage: [ProtocolAddress: NSRecursiveLock] = [:]
+    
+    /// Returns a lock for corresponding protocol addresses
+    /// - Parameter address: The protocol address that uniquely identifies session and lock for it
+    /// - Returns: A lock allowing to sync access to the session
+    func lock(for address: ProtocolAddress) -> NSRecursiveLock {
+        var result: NSRecursiveLock
+        if let lock = self.storage[address] {
+            result = lock
+        } else {
+            result = NSRecursiveLock()
+            self.storage[address] = result
+        }
+        
+        return result
+    }
+}

Tests/DXProtocolTests/SessionTests/SessionTests.swift

@@ -265,6 +265,168 @@ final class SessionTests: XCTestCase {
         let decryptedText2 = try XCTUnwrap(String(data: decrypted2, encoding: .utf8))
         XCTAssertEqual(plaintext2, decryptedText2)
     }
+    
+    func testThreadSafeSimultaneousDecrypt() async throws {
+        let senderClient = try TestClient(userId: UUID())  // Alice
+        let recipientClient = try TestClient(userId: UUID())  // Bob
+        let recipientAddress = recipientClient.protocolAddress
+        
+        try initializeSession(senderClient: senderClient, recipientClient: recipientClient)
+        
+        // Alice creates a set of messages
+        let aliceMessageCount = 50
+        var aliceMessages: [(Data, MessageContainer)] = []
+        for index in 0..<aliceMessageCount {
+            var aliceSession = try senderClient.sessionStore.loadSession(for: recipientAddress)
+            let data = try XCTUnwrap("From Alice \(index)".data(using: .utf8))
+            let message = try XCTUnwrap(
+                    try aliceSession?.encrypt(
+                            data: data,
+                            for: recipientAddress,
+                            sessionStore: senderClient.sessionStore,
+                            identityStore: senderClient.identityKeyStore)
+            )
+            let item = (data, message)
+            aliceMessages.append(item)
+        }
+        aliceMessages.shuffle()
+
+        var tasks = [Task<(Data, Data), Error>]()
+        for aliceMessage in aliceMessages {
+            let task = Task {
+                let decryptedMessage = try Session.decrypt(
+                    message: aliceMessage.1,
+                    from: senderClient.protocolAddress,
+                    sessionStore: recipientClient.sessionStore,
+                    identityStore: recipientClient.identityKeyStore,
+                    preKeyStore: recipientClient.preKeyStore,
+                    signedPreKeyStore: recipientClient.signedPreKeyStore)
+                let expectedMessage = aliceMessage.0
+                return (decryptedMessage, expectedMessage)
+            }
+            tasks.append(task)
+        }
+        
+        var results = [(decrypted: Data, expected: Data)]()
+        for task in tasks {
+            let result = try await task.value
+            results.append(result)
+        }
+
+        for result in results {
+            XCTAssertEqual(result.decrypted, result.expected)
+        }
+    }
+    
+    func testThreadSafeSimultaneousEncrypt() async throws {
+        let senderClient = try TestClient(userId: UUID())  // Alice
+        let recipientClient = try TestClient(userId: UUID())  // Bob
+        let recipientAddress = recipientClient.protocolAddress
+        
+        try initializeSession(senderClient: senderClient, recipientClient: recipientClient)
+        
+        // Actually test for Thread Safe begins here
+        
+        let plaintext = "Do not despair when your enemy attacks you."
+        let data = try XCTUnwrap(plaintext.data(using: .utf8))
+        let task1 = Task {
+            var session = try XCTUnwrap(try senderClient.sessionStore.loadSession(for: recipientAddress))
+            return try session.encrypt(
+                data: data,
+                for: recipientClient.protocolAddress,
+                sessionStore: senderClient.sessionStore,
+                identityStore: senderClient.identityKeyStore)
+        }
+        
+        let task2 = Task {
+            var session = try XCTUnwrap(try senderClient.sessionStore.loadSession(for: recipientAddress))
+            return try session.encrypt(
+                data: data,
+                for: recipientClient.protocolAddress,
+                sessionStore: senderClient.sessionStore,
+                identityStore: senderClient.identityKeyStore)
+        }
+        
+        let task3 = Task {
+            var session = try XCTUnwrap(try senderClient.sessionStore.loadSession(for: recipientAddress))
+            return try session.encrypt(
+                data: data,
+                for: recipientClient.protocolAddress,
+                sessionStore: senderClient.sessionStore,
+                identityStore: senderClient.identityKeyStore)
+        }
+        
+        let task4 = Task {
+            var session = try XCTUnwrap(try senderClient.sessionStore.loadSession(for: recipientAddress))
+            return try session.encrypt(
+                data: data,
+                for: recipientClient.protocolAddress,
+                sessionStore: senderClient.sessionStore,
+                identityStore: senderClient.identityKeyStore)
+        }
+        
+        let task5 = Task {
+            var session = try XCTUnwrap(try senderClient.sessionStore.loadSession(for: recipientAddress))
+            return try session.encrypt(
+                data: data,
+                for: recipientClient.protocolAddress,
+                sessionStore: senderClient.sessionStore,
+                identityStore: senderClient.identityKeyStore)
+        }
+        
+        let message1 = try await task1.value
+        let message2 = try await task2.value
+        let message3 = try await task3.value
+        let message4 = try await task4.value
+        let message5 = try await task5.value
+        
+        let decryptedMessage1 = try Session.decrypt(
+            message: message1,
+            from: senderClient.protocolAddress,
+            sessionStore: recipientClient.sessionStore,
+            identityStore: recipientClient.identityKeyStore,
+            preKeyStore: recipientClient.preKeyStore,
+            signedPreKeyStore: recipientClient.signedPreKeyStore)
+        
+        let decryptedMessage2 = try Session.decrypt(
+            message: message2,
+            from: senderClient.protocolAddress,
+            sessionStore: recipientClient.sessionStore,
+            identityStore: recipientClient.identityKeyStore,
+            preKeyStore: recipientClient.preKeyStore,
+            signedPreKeyStore: recipientClient.signedPreKeyStore)
+        
+        let decryptedMessage3 = try Session.decrypt(
+            message: message3,
+            from: senderClient.protocolAddress,
+            sessionStore: recipientClient.sessionStore,
+            identityStore: recipientClient.identityKeyStore,
+            preKeyStore: recipientClient.preKeyStore,
+            signedPreKeyStore: recipientClient.signedPreKeyStore)
+        
+        let decryptedMessage4 = try Session.decrypt(
+            message: message4,
+            from: senderClient.protocolAddress,
+            sessionStore: recipientClient.sessionStore,
+            identityStore: recipientClient.identityKeyStore,
+            preKeyStore: recipientClient.preKeyStore,
+            signedPreKeyStore: recipientClient.signedPreKeyStore)
+        
+        let decryptedMessage5 = try Session.decrypt(
+            message: message5,
+            from: senderClient.protocolAddress,
+            sessionStore: recipientClient.sessionStore,
+            identityStore: recipientClient.identityKeyStore,
+            preKeyStore: recipientClient.preKeyStore,
+            signedPreKeyStore: recipientClient.signedPreKeyStore)
+        
+        XCTAssertEqual(decryptedMessage1, data)
+        XCTAssertEqual(decryptedMessage2, data)
+        XCTAssertEqual(decryptedMessage3, data)
+        XCTAssertEqual(decryptedMessage4, data)
+        XCTAssertEqual(decryptedMessage5, data)
+    }
+
     // FIXME: - Need fix
     func testOutOfOrder() throws {
         let senderClient = try TestClient(userId: UUID())  // Alice