Merge branch 'rotate-key-completion'

9 files changed, 90 insertions, 91 deletions

diff --git a/ios/MullvadVPN/AddressCache/AddressCacheTracker.swift b/ios/MullvadVPN/AddressCache/AddressCacheTracker.swift
index de84265b1e..390fc31b4c 100644
--- a/ios/MullvadVPN/AddressCache/AddressCacheTracker.swift
+++ b/ios/MullvadVPN/AddressCache/AddressCacheTracker.swift
@@ -149,19 +149,19 @@ extension AddressCache {
             case .success(let updateResult):
                 switch updateResult {
                 case .finished:
-                    logger.debug("Finished updating address cache")
+                    logger.debug("Finished updating address cache.")
                 case .throttled:
-                    logger.debug("Address cache update was throttled")
+                    logger.debug("Address cache update was throttled.")
                 }
 
                 lastFailureAttemptDate = nil
 
             case .failure(let error):
-                logger.error(chainedError: AnyChainedError(error), message: "Failed to update address cache")
+                logger.error(chainedError: AnyChainedError(error), message: "Failed to update address cache.")
                 lastFailureAttemptDate = Date()
 
             case .cancelled:
-                logger.debug("Address cache update was cancelled")
+                logger.debug("Address cache update was cancelled.")
                 lastFailureAttemptDate = Date()
             }
         }
@@ -208,7 +208,7 @@ extension AddressCache.Tracker {
     private func handleBackgroundTask(_ task: BGProcessingTask) {
         logger.debug("Start address cache update task")
 
-        let cancellable = updateEndpoints { result in
+        let cancellable = updateEndpoints { completion in
             do {
                 // Schedule next background task
                 try self.scheduleBackgroundTask()
@@ -216,7 +216,7 @@ extension AddressCache.Tracker {
                 self.logger.error(chainedError: AnyChainedError(error), message: "Failed to schedule next address cache update task")
             }
 
-            task.setTaskCompleted(success: result.isTaskCompleted)
+            task.setTaskCompleted(success: completion.isSuccess)
         }
 
         task.expirationHandler = {
diff --git a/ios/MullvadVPN/AddressCache/UpdateAddressCacheOperation.swift b/ios/MullvadVPN/AddressCache/UpdateAddressCacheOperation.swift
index 81940acd21..2b2d42ac33 100644
--- a/ios/MullvadVPN/AddressCache/UpdateAddressCacheOperation.swift
+++ b/ios/MullvadVPN/AddressCache/UpdateAddressCacheOperation.swift
@@ -95,14 +95,3 @@ extension AddressCache {
         }
     }
 }
-
-extension OperationCompletion where Success == AddressCache.CacheUpdateResult {
-    var isTaskCompleted: Bool {
-        switch self {
-        case .success:
-            return true
-        case .cancelled, .failure:
-            return false
-        }
-    }
-}
diff --git a/ios/MullvadVPN/AppDelegate.swift b/ios/MullvadVPN/AppDelegate.swift
index bbac93fd1a..e2ea21f452 100644
--- a/ios/MullvadVPN/AppDelegate.swift
+++ b/ios/MullvadVPN/AppDelegate.swift
@@ -181,8 +181,8 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
         let operationQueue = OperationQueue()
 
         let updateAddressCacheOperation = AsyncBlockOperation { operation in
-            let handle = self.addressCacheTracker.updateEndpoints { result in
-                addressCacheFetchResult = result.backgroundFetchResult
+            let handle = self.addressCacheTracker.updateEndpoints { completion in
+                addressCacheFetchResult = completion.backgroundFetchResult
                 operation.finish()
             }
 
@@ -195,15 +195,14 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
             let handle = RelayCache.Tracker.shared.updateRelays { completion in
                 switch completion {
                 case .success(let result):
-                    self.logger?.debug("Finished updating relays: \(result)")
+                    self.logger?.debug("Finished updating relays: \(result).")
                 case .failure(let error):
-                    self.logger?.error(chainedError: error, message: "Failed to update relays")
+                    self.logger?.error(chainedError: error, message: "Failed to update relays.")
                 case .cancelled:
                     break
                 }
 
-                relaysFetchResult = completion.result?.backgroundFetchResult
-
+                relaysFetchResult = completion.backgroundFetchResult
                 operation.finish()
             }
 
@@ -213,23 +212,17 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
         }
 
         let rotatePrivateKeyOperation = AsyncBlockOperation { operation in
-            let handle = TunnelManager.shared.rotatePrivateKey { rotationResult, error in
-                if let error = error {
-                    self.logger?.error(chainedError: error, message: "Failed to rotate the key")
-
-                    rotatePrivateKeyFetchResult = .failed
-                } else if let rotationResult = rotationResult {
-                    self.logger?.debug("Finished rotating the key: \(rotationResult)")
-
-                    switch rotationResult {
-                    case .throttled:
-                        rotatePrivateKeyFetchResult = .noData
-
-                    case .finished:
-                        rotatePrivateKeyFetchResult = .newData
-                    }
+            let handle = TunnelManager.shared.rotatePrivateKey { completion in
+                switch completion {
+                case .success(let rotationResult):
+                    self.logger?.debug("Finished rotating the key: \(rotationResult).")
+                case .failure(let error):
+                    self.logger?.error(chainedError: error, message: "Failed to rotate the key.")
+                case .cancelled:
+                    break
                 }
 
+                rotatePrivateKeyFetchResult = completion.backgroundFetchResult
                 operation.finish()
             }
 
diff --git a/ios/MullvadVPN/AppStorePaymentManager/AppStorePaymentManager.swift b/ios/MullvadVPN/AppStorePaymentManager/AppStorePaymentManager.swift
index 7bdc1ce106..88d5a7c922 100644
--- a/ios/MullvadVPN/AppStorePaymentManager/AppStorePaymentManager.swift
+++ b/ios/MullvadVPN/AppStorePaymentManager/AppStorePaymentManager.swift
@@ -175,9 +175,13 @@ class AppStorePaymentManager: NSObject, SKPaymentTransactionObserver {
     }
 
     private func sendAppStoreReceipt(accountToken: String, forceRefresh: Bool, completionHandler: @escaping (OperationCompletion<REST.CreateApplePaymentResponse, Error>) -> Void) -> Cancellable {
-        let operation = SendAppStoreReceiptOperation(restClient: REST.Client.shared, accountToken: accountToken, forceRefresh: forceRefresh, receiptProperties: nil) { completion in
-            completionHandler(completion)
-        }
+        let operation = SendAppStoreReceiptOperation(
+            restClient: REST.Client.shared,
+            accountToken: accountToken,
+            forceRefresh: forceRefresh,
+            receiptProperties: nil,
+            completionHandler: completionHandler
+        )
 
         let backgroundTaskIdentifier = UIApplication.shared.beginBackgroundTask(withName: "Send AppStore receipt") {
             operation.cancel()
diff --git a/ios/MullvadVPN/Operations/OperationCompletion.swift b/ios/MullvadVPN/Operations/OperationCompletion.swift
index 0157075319..8c5aef4b5a 100644
--- a/ios/MullvadVPN/Operations/OperationCompletion.swift
+++ b/ios/MullvadVPN/Operations/OperationCompletion.swift
@@ -13,6 +13,14 @@ enum OperationCompletion<Success, Failure: Error> {
     case success(Success)
     case failure(Failure)
 
+    var isSuccess: Bool {
+        if case .success = self {
+            return true
+        } else {
+            return false
+        }
+    }
+
     var error: Failure? {
         if case .failure(let error) = self {
             return error
diff --git a/ios/MullvadVPN/Operations/ResultOperation.swift b/ios/MullvadVPN/Operations/ResultOperation.swift
index ba2cd59c4e..3d9233ce2e 100644
--- a/ios/MullvadVPN/Operations/ResultOperation.swift
+++ b/ios/MullvadVPN/Operations/ResultOperation.swift
@@ -31,6 +31,7 @@ class ResultOperation<Success, Failure: Error>: AsyncOperation {
         super.init()
     }
 
+    @available(*, unavailable)
     override func finish() {
         // Propagate cancellation if finish() is called directly from start().
         if isCancelled {
diff --git a/ios/MullvadVPN/RelayCache/RelayCacheTracker.swift b/ios/MullvadVPN/RelayCache/RelayCacheTracker.swift
index 89ebe2bda2..9f7b9c0fc4 100644
--- a/ios/MullvadVPN/RelayCache/RelayCacheTracker.swift
+++ b/ios/MullvadVPN/RelayCache/RelayCacheTracker.swift
@@ -267,23 +267,18 @@ extension RelayCache.Tracker {
         logger.debug("Start app refresh task.")
 
         let cancellable = self.updateRelays { completion in
-            let isTaskCompleted: Bool
-
             switch completion {
             case .success(let fetchResult):
                 self.logger.debug("Finished updating relays in app refresh task: \(fetchResult).")
-                isTaskCompleted = true
 
             case .failure(let error):
                 self.logger.error(chainedError: error, message: "Failed to update relays in app refresh task.")
-                isTaskCompleted = false
 
             case .cancelled:
                 self.logger.debug("App refresh task was cancelled.")
-                isTaskCompleted = false
             }
 
-            task.setTaskCompleted(success: isTaskCompleted)
+            task.setTaskCompleted(success: completion.isSuccess)
         }
 
         task.expirationHandler = {
diff --git a/ios/MullvadVPN/Result+UIBackgroundFetchResult.swift b/ios/MullvadVPN/Result+UIBackgroundFetchResult.swift
index 40aad31c70..d1a5679bf7 100644
--- a/ios/MullvadVPN/Result+UIBackgroundFetchResult.swift
+++ b/ios/MullvadVPN/Result+UIBackgroundFetchResult.swift
@@ -21,15 +21,26 @@ extension OperationCompletion where Success == AddressCache.CacheUpdateResult {
     }
 }
 
-extension Result where Success == RelayCache.FetchResult {
+extension OperationCompletion where Success == TunnelManager.KeyRotationResult {
     var backgroundFetchResult: UIBackgroundFetchResult {
         switch self {
-        case .success(.newContent):
+        case .success(.finished):
             return .newData
-
-        case .success(.throttled), .success(.sameContent):
+        case .success(.throttled), .cancelled:
             return .noData
+        case .failure:
+            return .failed
+        }
+    }
+}
 
+extension OperationCompletion where Success == RelayCache.FetchResult {
+    var backgroundFetchResult: UIBackgroundFetchResult {
+        switch self {
+        case .success(.newContent):
+            return .newData
+        case .success(.throttled), .success(.sameContent), .cancelled:
+            return .noData
         case .failure:
             return .failed
         }
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index 74e36770c6..6f32182042 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -145,9 +145,9 @@ final class TunnelManager: TunnelManagerStateDelegate {
         timer.setEventHandler { [weak self] in
             guard let self = self else { return }
 
-            _ = self.rotatePrivateKey { rotationResult, error in
+            _ = self.rotatePrivateKey { completion in
                 self.stateQueue.async {
-                    if let scheduleDate = self.handlePrivateKeyRotationCompletion(result: rotationResult, error: error) {
+                    if let scheduleDate = self.handlePrivateKeyRotationCompletion(completion) {
                         guard self.isRunningPeriodicPrivateKeyRotation else { return }
 
                         self.schedulePrivateKeyRotationTimer(scheduleDate)
@@ -385,41 +385,36 @@ final class TunnelManager: TunnelManagerStateDelegate {
         operationQueue.addOperation(operation)
     }
 
-    func rotatePrivateKey(completionHandler: @escaping (KeyRotationResult?, TunnelManager.Error?) -> Void) -> Cancellable {
+    func rotatePrivateKey(completionHandler: @escaping (OperationCompletion<KeyRotationResult, TunnelManager.Error>) -> Void) -> Cancellable {
         let operation = ReplaceKeyOperation.operationForKeyRotation(
             queue: stateQueue,
             state: state,
             restClient: restClient,
-            rotationInterval: TunnelManagerConfiguration.privateKeyRotationInterval) { [weak self] completion in
-                guard let self = self else { return }
-
-                dispatchPrecondition(condition: .onQueue(self.stateQueue))
-
-                var rotationResult: KeyRotationResult?
-                var rotationError: TunnelManager.Error?
+            rotationInterval: TunnelManagerConfiguration.privateKeyRotationInterval
+        ) { [weak self] completion in
+            guard let self = self else { return }
 
-                switch completion {
-                case .success(let result):
-                    rotationResult = result
+            dispatchPrecondition(condition: .onQueue(self.stateQueue))
 
-                    self.reconnectTunnel {
-                        completionHandler(rotationResult, rotationError)
-                    }
+            switch completion {
+            case .success:
+                self.reconnectTunnel {
+                    completionHandler(completion)
+                }
 
-                case .failure(let error):
-                    rotationError = error
-                    self.logger.error(chainedError: error, message: "Failed to rotate private key.")
+            case .failure(let error):
+                self.logger.error(chainedError: error, message: "Failed to rotate private key.")
 
-                    DispatchQueue.main.async {
-                        completionHandler(rotationResult, rotationError)
-                    }
+                DispatchQueue.main.async {
+                    completionHandler(completion)
+                }
 
-                case .cancelled:
-                    DispatchQueue.main.async {
-                        completionHandler(rotationResult, rotationError)
-                    }
+            case .cancelled:
+                DispatchQueue.main.async {
+                    completionHandler(completion)
                 }
             }
+        }
 
         let backgroundTaskIdentifier = UIApplication.shared.beginBackgroundTask(withName: "Rotate private key") {
             operation.cancel()
@@ -787,8 +782,8 @@ extension TunnelManager {
     private func handleBackgroundTask(_ task: BGProcessingTask) {
         logger.debug("Start private key rotation task")
 
-        let request = rotatePrivateKey { rotationResult, error in
-            if let scheduleDate = self.handlePrivateKeyRotationCompletion(result: rotationResult, error: error) {
+        let cancellableTask = rotatePrivateKey { completion in
+            if let scheduleDate = self.handlePrivateKeyRotationCompletion(completion) {
                 // Schedule next background task
                 switch self.submitBackgroundTask(at: scheduleDate) {
                 case .success:
@@ -800,32 +795,35 @@ extension TunnelManager {
             }
 
             // Complete current task
-            task.setTaskCompleted(success: error == nil)
+            task.setTaskCompleted(success: completion.isSuccess)
         }
 
         task.expirationHandler = {
-            request.cancel()
+            cancellableTask.cancel()
         }
     }
 }
 
 extension TunnelManager {
-    fileprivate func handlePrivateKeyRotationCompletion(result: KeyRotationResult?, error: TunnelManager.Error?) -> Date? {
-        if let error = error {
-            logger.error(chainedError: error, message: "Failed to rotate private key")
-
-            return nextRetryScheduleDate(error)
-        } else if let result = result {
+    fileprivate func handlePrivateKeyRotationCompletion(_ completion: OperationCompletion<KeyRotationResult, TunnelManager.Error>) -> Date? {
+        switch completion {
+        case .success(let result):
             switch result {
             case .finished:
-                logger.debug("Finished private key rotation")
+                logger.debug("Finished private key rotation.")
             case .throttled:
-                logger.debug("Private key was already rotated earlier")
+                logger.debug("Private key was already rotated earlier.")
             }
 
             return nextScheduleDate(result)
-        } else {
-            logger.debug("Private key rotation was cancelled")
+
+        case .failure(let error):
+            logger.error(chainedError: error, message: "Failed to rotate private key.")
+
+            return nextRetryScheduleDate(error)
+
+        case .cancelled:
+            logger.debug("Private key rotation was cancelled.")
 
             return Date(timeIntervalSinceNow: TunnelManagerConfiguration.privateKeyRotationFailureRetryInterval)
         }