TunnelManager: fetch device on pubkey-in-use error

1 files changed, 44 insertions, 8 deletions

diff --git a/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift b/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
index 4c6ddace88..349487014a 100644
--- a/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
+++ b/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
@@ -12,6 +12,7 @@ import MullvadREST
 import MullvadTypes
 import Operations
 import class WireGuardKitTypes.PrivateKey
+import class WireGuardKitTypes.PublicKey
 
 enum SetAccountAction {
     /// Set new account.
@@ -348,21 +349,56 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
         let task = devicesProxy
             .createDevice(accountNumber: accountNumber, request: request, retryStrategy: .default) { [self] result in
                 dispatchQueue.async { [self] in
-                    let result = result
-                        .map { device in
-                            return NewDevice(privateKey: privateKey, device: device)
+                    // Due to retry strategy, it's possible for server to register the new key without being
+                    // able to return the acknowledgment back to client.
+                    // In that case the subsequent retry attempt will error with `.publicKeyInUse`. Fetch the device
+                    // from API when that happens.
+                    if let error = result.error as? REST.Error, error.compareErrorCode(.publicKeyInUse) {
+                        self.findDevice(accountNumber: accountNumber, publicKey: privateKey.publicKey) { result in
+                            let result = result.flatMap { device in
+                                if let device {
+                                    return .success(NewDevice(privateKey: privateKey, device: device))
+                                } else {
+                                    return .failure(error)
+                                }
+                            }
+                            completion(result)
                         }
-                        .inspectError { error in
-                            logger.error(error: error, message: "Failed to create device.")
-                        }
-
-                    completion(result)
+                    } else {
+                        completion(result.map { NewDevice(privateKey: privateKey, device: $0) })
+                    }
                 }
             }
 
         tasks.append(task)
     }
 
+    /// Find device by public key in the list of devices registered on server. The result passed to `completion` handler
+    /// may contain `nil` if such device is not found for some reason.
+    private func findDevice(
+        accountNumber: String,
+        publicKey: PublicKey,
+        completion: @escaping (Result<Device?, Error>) -> Void
+    ) {
+        let task = devicesProxy.getDevices(accountNumber: accountNumber, retryStrategy: .default) { [self] result in
+            dispatchQueue.async { [self] in
+                let result = result
+                    .flatMap { devices in
+                        return .success(devices.first { device in
+                            return device.pubkey == publicKey
+                        })
+                    }
+                    .inspectError { error in
+                        logger.error(error: error, message: "Failed to get devices.")
+                    }
+
+                completion(result)
+            }
+        }
+
+        tasks.append(task)
+    }
+
     /// Struct that holds a private key that was used for creating a new device on the API along with the successful
     /// response from the API.
     private struct NewDevice {