Handle cancellation in background fetch

2 files changed, 18 insertions, 8 deletions

diff --git a/ios/MullvadVPN/AppDelegate.swift b/ios/MullvadVPN/AppDelegate.swift
index 2fd1a5999f..0ad7a204d8 100644
--- a/ios/MullvadVPN/AppDelegate.swift
+++ b/ios/MullvadVPN/AppDelegate.swift
@@ -184,14 +184,21 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
         let operationQueue = OperationQueue()
 
         let updateAddressCacheOperation = AsyncBlockOperation { operation in
-            _ = self.addressCacheTracker.updateEndpoints { result in
+            let handle = self.addressCacheTracker.updateEndpoints { result in
                 addressCacheFetchResult = result.backgroundFetchResult
                 operation.finish()
             }
+
+            operation.addCancellationBlock {
+                handle.cancel()
+            }
         }
 
         let updateRelaysOperation = AsyncBlockOperation { operation in
+            var cancellationToken: PromiseCancellationToken?
+
             RelayCache.Tracker.shared.updateRelays()
+                .storeCancellationToken(in: &cancellationToken)
                 .observe { completion in
                     switch completion.unwrappedValue {
                     case .success(let result):
@@ -206,15 +213,14 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
 
                     operation.finish()
                 }
-        }
 
-        let rotatePrivateKeyOperation = AsyncBlockOperation { operation in
-            guard !operation.isCancelled else {
-                operation.finish()
-                return
+            operation.addCancellationBlock {
+                cancellationToken?.cancel()
             }
+        }
 
-            _ = TunnelManager.shared.rotatePrivateKey { rotationResult, error in
+        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")
 
@@ -233,6 +239,10 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
 
                 operation.finish()
             }
+
+            operation.addCancellationBlock {
+                handle.cancel()
+            }
         }
 
         rotatePrivateKeyOperation.addDependency(updateRelaysOperation)
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index 22df01850c..a885bf56d4 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -124,7 +124,7 @@ class TunnelManager: TunnelManagerStateDelegate
         timer.setEventHandler { [weak self] in
             guard let self = self else { return }
 
-            self.rotatePrivateKey { rotationResult, error in
+            _ = self.rotatePrivateKey { rotationResult, error in
                 self.stateQueue.async {
                     if let scheduleDate = self.handlePrivateKeyRotationCompletion(result: rotationResult, error: error) {
                         guard self.isRunningPeriodicPrivateKeyRotation else { return }