Fix group member removal race condition

SignalServiceKit/src/Loki/Protocol/Closed Groups/ClosedGroupsProtocol.swift

@@ -132,6 +132,11 @@ public final class ClosedGroupsProtocol : NSObject {
             when(resolved: promises).done2 { _ in seal.fulfill(()) }.catch2 { seal.reject($0) }
             promise.done {
                 try! Storage.writeSync { transaction in
+                    let allOldRatchets = Storage.getAllClosedGroupRatchets(for: groupPublicKey)
+                    for (senderPublicKey, oldRatchet) in allOldRatchets {
+                        let collection = Storage.ClosedGroupRatchetCollectionType.old
+                        Storage.setClosedGroupRatchet(for: groupPublicKey, senderPublicKey: senderPublicKey, ratchet: oldRatchet, in: collection, using: transaction)
+                    }
                     // Delete all ratchets (it's important that this happens * after * sending out the update)
                     Storage.removeAllClosedGroupRatchets(for: groupPublicKey, using: transaction)
                     // Remove the group from the user's set of public keys to poll for if the user is leaving. Otherwise generate a new ratchet and
@@ -363,6 +368,11 @@ public final class ClosedGroupsProtocol : NSObject {
         let userPublicKey = UserDefaults.standard[.masterHexEncodedPublicKey] ?? getUserHexEncodedPublicKey()
         let wasUserRemoved = !members.contains(userPublicKey)
         if Set(members).intersection(oldMembers) != Set(oldMembers) {
+            let allOldRatchets = Storage.getAllClosedGroupRatchets(for: groupPublicKey)
+            for (senderPublicKey, oldRatchet) in allOldRatchets {
+                let collection = Storage.ClosedGroupRatchetCollectionType.old
+                Storage.setClosedGroupRatchet(for: groupPublicKey, senderPublicKey: senderPublicKey, ratchet: oldRatchet, in: collection, using: transaction)
+            }
             Storage.removeAllClosedGroupRatchets(for: groupPublicKey, using: transaction)
             if wasUserRemoved {
                 Storage.removeClosedGroupPrivateKey(for: groupPublicKey, using: transaction)

SignalServiceKit/src/Loki/Protocol/Closed Groups/SharedSenderKeysImplementation.swift

@@ -90,11 +90,12 @@ public final class SharedSenderKeysImplementation : NSObject {
     }
 
     /// - Note: Sync. Don't call from the main thread.
-    private func stepRatchet(for groupPublicKey: String, senderPublicKey: String, until targetKeyIndex: UInt, using transaction: YapDatabaseReadWriteTransaction) throws -> ClosedGroupRatchet {
+    private func stepRatchet(for groupPublicKey: String, senderPublicKey: String, until targetKeyIndex: UInt, using transaction: YapDatabaseReadWriteTransaction, isRetry: Bool = false) throws -> ClosedGroupRatchet {
         #if DEBUG
         assert(!Thread.isMainThread)
         #endif
-        guard let ratchet = Storage.getClosedGroupRatchet(for: groupPublicKey, senderPublicKey: senderPublicKey) else {
+        let collection: Storage.ClosedGroupRatchetCollectionType = (isRetry) ? .old : .current
+        guard let ratchet = Storage.getClosedGroupRatchet(for: groupPublicKey, senderPublicKey: senderPublicKey, from: collection) else {
             let error = RatchetingError.loadingFailed(groupPublicKey: groupPublicKey, senderPublicKey: senderPublicKey)
             print("[Loki] \(error.errorDescription!)")
             throw error
@@ -122,7 +123,8 @@ public final class SharedSenderKeysImplementation : NSObject {
                 }
             }
             let result = ClosedGroupRatchet(chainKey: current.chainKey, keyIndex: current.keyIndex, messageKeys: messageKeys) // Includes any skipped message keys
-            Storage.setClosedGroupRatchet(for: groupPublicKey, senderPublicKey: senderPublicKey, ratchet: result, using: transaction)
+            let collection: Storage.ClosedGroupRatchetCollectionType = (isRetry) ? .old : .current
+            Storage.setClosedGroupRatchet(for: groupPublicKey, senderPublicKey: senderPublicKey, ratchet: result, in: collection, using: transaction)
             return result
         }
     }
@@ -161,17 +163,21 @@ public final class SharedSenderKeysImplementation : NSObject {
         return try decrypt(ivAndCiphertext, for: groupPublicKey, senderPublicKey: senderPublicKey, keyIndex: keyIndex, using: transaction)
     }
 
-    public func decrypt(_ ivAndCiphertext: Data, for groupPublicKey: String, senderPublicKey: String, keyIndex: UInt, using transaction: YapDatabaseReadWriteTransaction) throws -> Data {
+    public func decrypt(_ ivAndCiphertext: Data, for groupPublicKey: String, senderPublicKey: String, keyIndex: UInt, using transaction: YapDatabaseReadWriteTransaction, isRetry: Bool = false) throws -> Data {
         let ratchet: ClosedGroupRatchet
         do {
-            ratchet = try stepRatchet(for: groupPublicKey, senderPublicKey: senderPublicKey, until: keyIndex, using: transaction)
+            ratchet = try stepRatchet(for: groupPublicKey, senderPublicKey: senderPublicKey, until: keyIndex, using: transaction, isRetry: isRetry)
         } catch {
-            // FIXME: It'd be cleaner to handle this in OWSMessageDecrypter (where all the other decryption errors are handled), but this was a lot more
-            // convenient because there's an easy way to get the sender public key from here.
-            if case RatchetingError.loadingFailed(_, _) = error {
-                ClosedGroupsProtocol.requestSenderKey(for: groupPublicKey, senderPublicKey: senderPublicKey, using: transaction)
+            if !isRetry {
+                return try decrypt(ivAndCiphertext, for: groupPublicKey, senderPublicKey: senderPublicKey, keyIndex: keyIndex, using: transaction, isRetry: true)
+            } else {
+                // FIXME: It'd be cleaner to handle this in OWSMessageDecrypter (where all the other decryption errors are handled), but this was a lot more
+                // convenient because there's an easy way to get the sender public key from here.
+                if case RatchetingError.loadingFailed(_, _) = error {
+                    ClosedGroupsProtocol.requestSenderKey(for: groupPublicKey, senderPublicKey: senderPublicKey, using: transaction)
+                }
+                throw error
             }
-            throw error
         }
         let iv = ivAndCiphertext[0..<Int(SharedSenderKeysImplementation.ivSize)]
         let ciphertext = ivAndCiphertext[Int(SharedSenderKeysImplementation.ivSize)...]
@@ -183,8 +189,12 @@ public final class SharedSenderKeysImplementation : NSObject {
         do {
             return Data(try aes.decrypt(ciphertext.bytes))
         } catch {
-            ClosedGroupsProtocol.requestSenderKey(for: groupPublicKey, senderPublicKey: senderPublicKey, using: transaction)
-            throw error
+            if !isRetry {
+                return try decrypt(ivAndCiphertext, for: groupPublicKey, senderPublicKey: senderPublicKey, keyIndex: keyIndex, using: transaction, isRetry: true)
+            } else {
+                ClosedGroupsProtocol.requestSenderKey(for: groupPublicKey, senderPublicKey: senderPublicKey, using: transaction)
+                throw error
+            }
         }
     }
 

SignalServiceKit/src/Loki/Protocol/Closed Groups/Storage+ClosedGroups.swift

@@ -1,13 +1,20 @@
 
 public extension Storage {
 
+    internal enum ClosedGroupRatchetCollectionType {
+        case old, current
+    }
+
     // MARK: Ratchets
-    internal static func getClosedGroupRatchetCollection(for groupPublicKey: String) -> String {
-        return "LokiClosedGroupRatchetCollection.\(groupPublicKey)"
+    internal static func getClosedGroupRatchetCollection(_ collection: ClosedGroupRatchetCollectionType, for groupPublicKey: String) -> String {
+        switch collection {
+        case .old: return "LokiOldClosedGroupRatchetCollection.\(groupPublicKey)"
+        case .current: return "LokiClosedGroupRatchetCollection.\(groupPublicKey)"
+        }
     }
 
-    internal static func getClosedGroupRatchet(for groupPublicKey: String, senderPublicKey: String) -> ClosedGroupRatchet? {
-        let collection = getClosedGroupRatchetCollection(for: groupPublicKey)
+    internal static func getClosedGroupRatchet(for groupPublicKey: String, senderPublicKey: String, from collection: ClosedGroupRatchetCollectionType = .current) -> ClosedGroupRatchet? {
+        let collection = getClosedGroupRatchetCollection(collection, for: groupPublicKey)
         var result: ClosedGroupRatchet?
         read { transaction in
             result = transaction.object(forKey: senderPublicKey, inCollection: collection) as? ClosedGroupRatchet
@@ -15,26 +22,31 @@ public extension Storage {
         return result
     }
 
-    internal static func setClosedGroupRatchet(for groupPublicKey: String, senderPublicKey: String, ratchet: ClosedGroupRatchet, using transaction: YapDatabaseReadWriteTransaction) {
-        let collection = getClosedGroupRatchetCollection(for: groupPublicKey)
+    internal static func setClosedGroupRatchet(for groupPublicKey: String, senderPublicKey: String, ratchet: ClosedGroupRatchet, in collection: ClosedGroupRatchetCollectionType = .current, using transaction: YapDatabaseReadWriteTransaction) {
+        let collection = getClosedGroupRatchetCollection(collection, for: groupPublicKey)
         transaction.setObject(ratchet, forKey: senderPublicKey, inCollection: collection)
     }
 
-    internal static func getAllClosedGroupSenderKeys(for groupPublicKey: String) -> Set<ClosedGroupSenderKey> {
-        let collection = getClosedGroupRatchetCollection(for: groupPublicKey)
-        var result: Set<ClosedGroupSenderKey> = []
+    internal static func getAllClosedGroupRatchets(for groupPublicKey: String, from collection: ClosedGroupRatchetCollectionType = .current) -> [(senderPublicKey: String, ratchet: ClosedGroupRatchet)] {
+        let collection = getClosedGroupRatchetCollection(collection, for: groupPublicKey)
+        var result: [(senderPublicKey: String, ratchet: ClosedGroupRatchet)] = []
         read { transaction in
             transaction.enumerateRows(inCollection: collection) { key, object, _, _ in
-                guard let publicKey = key as? String, let ratchet = object as? ClosedGroupRatchet else { return }
-                let senderKey = ClosedGroupSenderKey(chainKey: Data(hex: ratchet.chainKey), keyIndex: ratchet.keyIndex, publicKey: Data(hex: publicKey))
-                result.insert(senderKey)
+                guard let senderPublicKey = key as? String, let ratchet = object as? ClosedGroupRatchet else { return }
+                result.append((senderPublicKey: senderPublicKey, ratchet: ratchet))
             }
         }
         return result
     }
 
-    internal static func removeAllClosedGroupRatchets(for groupPublicKey: String, using transaction: YapDatabaseReadWriteTransaction) {
-        let collection = getClosedGroupRatchetCollection(for: groupPublicKey)
+    internal static func getAllClosedGroupSenderKeys(for groupPublicKey: String, from collection: ClosedGroupRatchetCollectionType = .current) -> Set<ClosedGroupSenderKey> {
+        return Set(getAllClosedGroupRatchets(for: groupPublicKey, from: collection).map { senderPublicKey, ratchet in
+            ClosedGroupSenderKey(chainKey: Data(hex: ratchet.chainKey), keyIndex: ratchet.keyIndex, publicKey: Data(hex: senderPublicKey))
+        })
+    }
+
+    internal static func removeAllClosedGroupRatchets(for groupPublicKey: String, from collection: ClosedGroupRatchetCollectionType = .current, using transaction: YapDatabaseReadWriteTransaction) {
+        let collection = getClosedGroupRatchetCollection(collection, for: groupPublicKey)
         transaction.removeAllObjects(inCollection: collection)
     }
 }