Receiving data from Bitwarden fixed (#1841)

Task/Issue URL: https://app.asana.com/0/1177771139624306/1205069580048345/f Tech Design URL: CC: **Description**: This PR fixes receiving of long messages from Bitwarden app
duckduckgo · Nov 20, 2023 · f9ed118 · f9ed118
1 parent f5fd7b1
commit f9ed118
Showing 1 changed file with 46 additions and 28 deletions.
diff --git a/DuckDuckGo/PasswordManager/Bitwarden/Services/BWCommunicator.swift b/DuckDuckGo/PasswordManager/Bitwarden/Services/BWCommunicator.swift
@@ -122,45 +122,63 @@ final class BWCommunicator: BWCommunication {
 
     // MARK: - Receiving Messages
 
+    private let realisticMessageLength = 200000
+    private var accumulatedData = Data()
+    private let dataQueue = DispatchQueue(label: "BWCommunicator.queue")
+
     func receiveData(_ fileHandle: FileHandle) {
+        let newData = fileHandle.availableData
+        dataQueue.async {
+            self.accumulatedData.append(newData)
+            self.processAccumulatedData()
+        }
+    }
 
-        func readMessage(availableData: Data) -> (messageData: Data?, availableData: Data) {
-            guard availableData.count > 0 else { return (nil, availableData: availableData) }
+    private func processAccumulatedData() {
+        dataQueue.async {
+            repeat {
+                let (messageData, remainingData) = self.readMessage(availableData: self.accumulatedData)
+                self.accumulatedData = remainingData
 
-            // First 4 bytes of the message contain the message length
-            let dataPrefix = availableData.prefix(4)
-            guard dataPrefix.count == 4 else {
-                assertionFailure("Wrong format of the message")
-                return (nil, availableData)
-            }
+                guard let messageData = messageData else {
+                    return
+                }
 
-            let dataPrefixArray = [UInt8](dataPrefix)
-            let messageLength = fromByteArray(dataPrefixArray, UInt32.self)
+                DispatchQueue.main.async { [weak self] in
+                    guard let self = self else { return }
 
-            let dataPostfix = availableData.dropFirst(4)
-            let messageData = dataPostfix.prefix(Int(messageLength))
-            let availableData = dataPostfix.dropFirst(Int(messageLength))
-            return (messageData: messageData, availableData: availableData)
+                    self.delegate?.bitwardenCommunicator(self, didReceiveMessageData: messageData)
+                }
+            } while self.accumulatedData.count >= 2 /*EOF*/
         }
+    }
 
-        var availableData = fileHandle.availableData
-        repeat {
-            let (messageData, nextAvailableData) = readMessage(availableData: availableData)
-            availableData = nextAvailableData
+    func readMessage(availableData: Data) -> (messageData: Data?, availableData: Data) {
+        guard availableData.count > 0 else { return (nil, availableData: availableData) }
 
-            guard let messageData = messageData else {
-                if availableData.count >= 2 {
-                    assertionFailure("Wrong format of the message")
-                }
-                return
-            }
+        // First 4 bytes of the message contain the message length
+        let dataPrefix = availableData.prefix(4)
+        guard dataPrefix.count == 4 else {
+            assertionFailure("Wrong format of the message")
+            return (nil, availableData)
+        }
 
-            DispatchQueue.main.async { [weak self] in
-                guard let self = self else { return }
+        let dataPrefixArray = [UInt8](dataPrefix)
+        let messageLength = fromByteArray(dataPrefixArray, UInt32.self)
 
-                self.delegate?.bitwardenCommunicator(self, didReceiveMessageData: messageData)
+        let dataPostfix = availableData.dropFirst(4)
+
+        if messageLength > dataPostfix.count {
+            if messageLength > realisticMessageLength {
+                self.accumulatedData = Data()
+                return (nil, Data())
             }
-        } while availableData.count >= 2 /*EOF*/
+            return (nil, availableData)
+        }
+
+        let messageData = dataPostfix.prefix(Int(messageLength))
+        let availableData = dataPostfix.dropFirst(Int(messageLength))
+        return (messageData: messageData, availableData: availableData)
     }
 
     private func fromByteArray<T>(_ value: [UInt8], _: T.Type) -> T {