Return cancellation handle for key rotation

2 files changed, 8 insertions, 4 deletions

diff --git a/ios/MullvadVPN/AppDelegate.swift b/ios/MullvadVPN/AppDelegate.swift
index 3a673b17aa..2fd1a5999f 100644
--- a/ios/MullvadVPN/AppDelegate.swift
+++ b/ios/MullvadVPN/AppDelegate.swift
@@ -214,7 +214,7 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
                 return
             }
 
-            TunnelManager.shared.rotatePrivateKey { rotationResult, error in
+            _ = TunnelManager.shared.rotatePrivateKey { rotationResult, error in
                 if let error = error {
                     self.logger?.error(chainedError: error, message: "Failed to rotate the key")
 
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index c2fe96adcd..22df01850c 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -352,7 +352,7 @@ class TunnelManager: TunnelManagerStateDelegate
         operationQueue.addOperation(operation)
     }
 
-    func rotatePrivateKey(completionHandler: @escaping (KeyRotationResult?, TunnelManager.Error?) -> Void) {
+    func rotatePrivateKey(completionHandler: @escaping (KeyRotationResult?, TunnelManager.Error?) -> Void) -> Cancellable {
         let operation = ReplaceKeyOperation.operationForKeyRotation(
             queue: stateQueue,
             state: state,
@@ -399,6 +399,10 @@ class TunnelManager: TunnelManagerStateDelegate
         exclusivityController.addOperation(operation, categories: [OperationCategory.changeTunnelSettings])
 
         operationQueue.addOperation(operation)
+
+        return AnyCancellable {
+            operation.cancel()
+        }
     }
 
     func setRelayConstraints(_ newConstraints: RelayConstraints, completionHandler: @escaping (TunnelManager.Error?) -> Void) {
@@ -666,7 +670,7 @@ extension TunnelManager {
     private func handleBackgroundTask(_ task: BGProcessingTask) {
         logger.debug("Start private key rotation task")
 
-        rotatePrivateKey { rotationResult, error in
+        let request = rotatePrivateKey { rotationResult, error in
             if let scheduleDate = self.handlePrivateKeyRotationCompletion(result: rotationResult, error: error) {
                 // Schedule next background task
                 switch self.submitBackgroundTask(at: scheduleDate) {
@@ -683,7 +687,7 @@ extension TunnelManager {
         }
 
         task.expirationHandler = {
-            // TODO: handle cancellation?
+            request.cancel()
         }
     }
 }