Merge branch 'change-retry-behavior-for-key-rotation-ios-79'

8 files changed, 82 insertions, 115 deletions

diff --git a/ios/MullvadVPN/AppDelegate.swift b/ios/MullvadVPN/AppDelegate.swift
index fcc5f2fe92..126109663f 100644
--- a/ios/MullvadVPN/AppDelegate.swift
+++ b/ios/MullvadVPN/AppDelegate.swift
@@ -186,7 +186,7 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
             forTaskWithIdentifier: ApplicationConfiguration.privateKeyRotationTaskIdentifier,
             using: nil
         ) { task in
-            let handle = self.tunnelManager.rotatePrivateKey(forceRotate: false) { completion in
+            let handle = self.tunnelManager.rotatePrivateKey { completion in
                 self.scheduleKeyRotationTask()
 
                 task.setTaskCompleted(success: completion.isSuccess)
diff --git a/ios/MullvadVPN/SettingsManager/Migrations/MigrationFromV1ToV2.swift b/ios/MullvadVPN/SettingsManager/Migrations/MigrationFromV1ToV2.swift
index a7f8883c0f..c5898ee516 100644
--- a/ios/MullvadVPN/SettingsManager/Migrations/MigrationFromV1ToV2.swift
+++ b/ios/MullvadVPN/SettingsManager/Migrations/MigrationFromV1ToV2.swift
@@ -140,6 +140,7 @@ final class MigrationFromV1ToV2: Migration {
                     ipv6Address: device.ipv6Address,
                     wgKeyData: StoredWgKeyData(
                         creationDate: privateKeyWithMetadata.creationDate,
+                        lastRotationAttemptDate: nil,
                         privateKey: privateKeyWithMetadata.privateKey
                     )
                 )
diff --git a/ios/MullvadVPN/SettingsManager/TunnelSettingsV2.swift b/ios/MullvadVPN/SettingsManager/TunnelSettingsV2.swift
index de9308c84a..a305839bb9 100644
--- a/ios/MullvadVPN/SettingsManager/TunnelSettingsV2.swift
+++ b/ios/MullvadVPN/SettingsManager/TunnelSettingsV2.swift
@@ -114,6 +114,24 @@ struct StoredWgKeyData: Codable, Equatable {
     /// Private key creation date.
     var creationDate: Date
 
+    /// Last date a rotation was attempted. Nil if last attempt was successful.
+    var lastRotationAttemptDate: Date?
+
     /// Private key.
     var privateKey: PrivateKey
 }
+
+extension StoredWgKeyData {
+    struct KeyRotationConfiguration {
+        let rotationInterval: TimeInterval = 60 * 60 * 24 * 14
+        let retryInterval: TimeInterval = 60 * 60 * 24
+    }
+
+    func getNextRotationDate(for configuration: KeyRotationConfiguration) -> Date {
+        return max(
+            Date(),
+            lastRotationAttemptDate?.addingTimeInterval(configuration.retryInterval) ?? creationDate
+                .addingTimeInterval(configuration.rotationInterval)
+        )
+    }
+}
diff --git a/ios/MullvadVPN/TunnelManager/RotateKeyOperation.swift b/ios/MullvadVPN/TunnelManager/RotateKeyOperation.swift
index 1793453107..e6059c5320 100644
--- a/ios/MullvadVPN/TunnelManager/RotateKeyOperation.swift
+++ b/ios/MullvadVPN/TunnelManager/RotateKeyOperation.swift
@@ -19,18 +19,18 @@ class RotateKeyOperation: ResultOperation<Bool> {
     private let devicesProxy: REST.DevicesProxy
     private var task: Cancellable?
 
-    private let rotationInterval: TimeInterval?
+    private let keyRotationConfiguration: StoredWgKeyData.KeyRotationConfiguration
     private let logger = Logger(label: "ReplaceKeyOperation")
 
     init(
         dispatchQueue: DispatchQueue,
         interactor: TunnelInteractor,
         devicesProxy: REST.DevicesProxy,
-        rotationInterval: TimeInterval?
+        keyRotationConfiguration: StoredWgKeyData.KeyRotationConfiguration
     ) {
         self.interactor = interactor
         self.devicesProxy = devicesProxy
-        self.rotationInterval = rotationInterval
+        self.keyRotationConfiguration = keyRotationConfiguration
 
         super.init(
             dispatchQueue: dispatchQueue,
@@ -40,27 +40,24 @@ class RotateKeyOperation: ResultOperation<Bool> {
     }
 
     override func main() {
-        guard case let .loggedIn(accountData, deviceData) = interactor.deviceState else {
+        guard case .loggedIn(let accountData, var deviceData) = interactor.deviceState else {
             finish(result: .failure(InvalidDeviceStateError()))
             return
         }
 
-        if let rotationInterval = rotationInterval {
-            let creationDate = deviceData.wgKeyData.creationDate
-            let nextRotationDate = creationDate.addingTimeInterval(rotationInterval)
+        let nextRotationDate = deviceData.wgKeyData.getNextRotationDate(for: keyRotationConfiguration)
+        if nextRotationDate > Date() {
+            logger.debug("Throttle private key rotation.")
 
-            if nextRotationDate > Date() {
-                logger.debug("Throttle private key rotation.")
-
-                finish(result: .success(false))
-                return
-            } else {
-                logger.debug("Private key is old enough, rotate right away.")
-            }
+            finish(result: .success(false))
+            return
         } else {
-            logger.debug("Rotate private key right away.")
+            logger.debug("Private key is old enough, rotate right away.")
         }
 
+        deviceData.wgKeyData.lastRotationAttemptDate = Date()
+        interactor.setDeviceState(.loggedIn(accountData, deviceData), persist: true)
+
         logger.debug("Replacing old key with new key on server...")
 
         let newPrivateKey = PrivateKey()
@@ -93,8 +90,10 @@ class RotateKeyOperation: ResultOperation<Bool> {
             switch interactor.deviceState {
             case .loggedIn(let accountData, var deviceData):
                 deviceData.update(from: device)
+
                 deviceData.wgKeyData = StoredWgKeyData(
                     creationDate: Date(),
+                    lastRotationAttemptDate: nil,
                     privateKey: newPrivateKey
                 )
 
@@ -118,6 +117,17 @@ class RotateKeyOperation: ResultOperation<Bool> {
                     message: "Failed to rotate device key."
                 )
             }
+
+            switch interactor.deviceState {
+            case .loggedIn(let accountData, var deviceData):
+                deviceData.wgKeyData.lastRotationAttemptDate = Date()
+                interactor.setDeviceState(.loggedIn(accountData, deviceData), persist: true)
+
+            default:
+                finish(result: .failure(InvalidDeviceStateError()))
+            }
+
+            interactor.handleRestError(error)
             finish(result: .failure(error))
         }
     }
diff --git a/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift b/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
index d07170bdc2..fd8e7008ec 100644
--- a/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
+++ b/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
@@ -412,6 +412,7 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
                     ipv6Address: device.ipv6Address,
                     wgKeyData: StoredWgKeyData(
                         creationDate: Date(),
+                        lastRotationAttemptDate: nil,
                         privateKey: input.privateKey
                     )
                 )
diff --git a/ios/MullvadVPN/TunnelManager/TunnelInteractor.swift b/ios/MullvadVPN/TunnelManager/TunnelInteractor.swift
index 2f95622fbc..e22e33cf16 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelInteractor.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelInteractor.swift
@@ -34,6 +34,7 @@ protocol TunnelInteractor {
     func setConfigurationLoaded()
     func setSettings(_ settings: TunnelSettingsV2, persist: Bool)
     func setDeviceState(_ deviceState: DeviceState, persist: Bool)
+    func handleRestError(_ error: Error)
 
     func startTunnel()
     func prepareForVPNConfigurationDeletion()
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index d72c314c3e..8b4bd53c84 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -27,12 +27,6 @@ private let establishingTunnelStatusPollInterval: TimeInterval = 3
 /// is established.
 private let establishedTunnelStatusPollInterval: TimeInterval = 5
 
-/// Private key rotation interval (in seconds).
-private let privateKeyRotationInterval: TimeInterval = 60 * 60 * 24 * 14
-
-/// Private key rotation retry interval (in seconds).
-private let privateKeyRotationFailureRetryInterval: TimeInterval = 60 * 60 * 24
-
 /// A class that provides a convenient interface for VPN tunnels configuration, manipulation and
 /// monitoring.
 final class TunnelManager: StorePaymentObserver {
@@ -65,10 +59,6 @@ final class TunnelManager: StorePaymentObserver {
     private let observerList = ObserverList<TunnelObserver>()
 
     private var privateKeyRotationTimer: DispatchSourceTimer?
-    private var lastKeyRotationData: (
-        attempt: Date,
-        completion: Result<Bool, Error>
-    )?
     private var isRunningPeriodicPrivateKeyRotation = false
 
     private var tunnelStatusPollTimer: DispatchSourceTimer?
@@ -139,37 +129,7 @@ final class TunnelManager: StorePaymentObserver {
         nslock.lock()
         defer { nslock.unlock() }
 
-        guard case let .loggedIn(_, deviceData) = deviceState else {
-            return nil
-        }
-
-        if case .some(let (lastAttemptDate, completion)) = lastKeyRotationData {
-            if completion.error is InvalidDeviceStateError {
-                return nil
-            }
-
-            // Do not rotate the key if account or device is not found.
-            if let restError = completion.error as? REST.Error,
-               restError.compareErrorCode(.invalidAccount) ||
-               restError.compareErrorCode(.deviceNotFound)
-            {
-                return nil
-            }
-
-            // Retry at equal interval if failed or cancelled.
-            if !completion.isSuccess {
-                let date = lastAttemptDate.addingTimeInterval(
-                    privateKeyRotationFailureRetryInterval
-                )
-
-                return max(date, Date())
-            }
-        }
-
-        // Rotate at long intervals otherwise.
-        let date = deviceData.wgKeyData.creationDate.addingTimeInterval(privateKeyRotationInterval)
-
-        return max(date, Date())
+        return deviceState.deviceData?.wgKeyData.getNextRotationDate(for: StoredWgKeyData.KeyRotationConfiguration())
     }
 
     private func updatePrivateKeyRotationTimer() {
@@ -185,7 +145,7 @@ final class TunnelManager: StorePaymentObserver {
         let timer = DispatchSource.makeTimerSource(queue: .main)
 
         timer.setEventHandler { [weak self] in
-            _ = self?.rotatePrivateKey(forceRotate: false) { _ in
+            _ = self?.rotatePrivateKey { _ in
                 // no-op
             }
         }
@@ -198,22 +158,6 @@ final class TunnelManager: StorePaymentObserver {
         logger.debug("Schedule next private key rotation at \(scheduleDate.logFormatDate()).")
     }
 
-    private func setFinishedKeyRotation(_ result: Result<Bool, Error>) {
-        nslock.lock()
-        defer { nslock.unlock() }
-
-        lastKeyRotationData = (Date(), result)
-        updatePrivateKeyRotationTimer()
-    }
-
-    private func resetKeyRotationData() {
-        nslock.lock()
-        defer { nslock.unlock() }
-
-        lastKeyRotationData = nil
-        updatePrivateKeyRotationTimer()
-    }
-
     // MARK: - Public methods
 
     func loadConfiguration(completionHandler: @escaping (Error?) -> Void) {
@@ -394,7 +338,7 @@ final class TunnelManager: StorePaymentObserver {
 
         operation.completionQueue = .main
         operation.completionHandler = { [weak self] result in
-            self?.resetKeyRotationData()
+            self?.updatePrivateKeyRotationTimer()
 
             completionHandler(result)
         }
@@ -465,15 +409,8 @@ final class TunnelManager: StorePaymentObserver {
         )
 
         operation.completionQueue = .main
-        operation.completionHandler = { [weak self] completion in
-            guard let self = self else { return }
-
-            let error = completion.error
-            if let error = error {
-                self.checkIfDeviceRevoked(error)
-            }
-
-            completionHandler?(error)
+        operation.completionHandler = { completion in
+            completionHandler?(completion.error)
         }
 
         operation.addObserver(
@@ -491,27 +428,19 @@ final class TunnelManager: StorePaymentObserver {
         operationQueue.addOperation(operation)
     }
 
-    func rotatePrivateKey(
-        forceRotate: Bool,
-        completionHandler: @escaping (Result<Bool, Error>) -> Void
-    ) -> Cancellable {
-        var rotationInterval: TimeInterval?
-        if !forceRotate {
-            rotationInterval = privateKeyRotationInterval
-        }
-
+    func rotatePrivateKey(completionHandler: @escaping (Result<Bool, Error>) -> Void) -> Cancellable {
         let operation = RotateKeyOperation(
             dispatchQueue: internalQueue,
             interactor: TunnelInteractorProxy(self),
             devicesProxy: devicesProxy,
-            rotationInterval: rotationInterval
+            keyRotationConfiguration: StoredWgKeyData.KeyRotationConfiguration()
         )
 
         operation.completionQueue = .main
         operation.completionHandler = { [weak self] result in
             guard let self = self else { return }
 
-            self.setFinishedKeyRotation(result)
+            self.updatePrivateKeyRotationTimer()
 
             switch result {
             case .success:
@@ -519,7 +448,7 @@ final class TunnelManager: StorePaymentObserver {
 
             case let .failure(error):
                 if !error.isOperationCancellationError {
-                    self.checkIfDeviceRevoked(error)
+                    self.handleRestError(error)
                 }
 
                 completionHandler(result)
@@ -700,8 +629,13 @@ final class TunnelManager: StorePaymentObserver {
            deviceCheck.identifier != lastDeviceCheckIdentifier
         {
             if deviceCheck.isDeviceRevoked ?? false {
-                didDetectDeviceRevoked()
-
+                scheduleDeviceStateUpdate(
+                    taskName: "Set device revoked",
+                    modificationBlock: { deviceState in
+                        deviceState = .revoked
+                    },
+                    completionHandler: nil
+                )
             } else if let accountExpiry = deviceCheck.accountExpiry {
                 scheduleDeviceStateUpdate(
                     taskName: "Update account expiry",
@@ -827,22 +761,6 @@ final class TunnelManager: StorePaymentObserver {
         lastMapConnectionStatusOperation = nil
     }
 
-    private func checkIfDeviceRevoked(_ error: Error) {
-        if let error = error as? REST.Error, error.compareErrorCode(.deviceNotFound) {
-            didDetectDeviceRevoked()
-        }
-    }
-
-    private func didDetectDeviceRevoked() {
-        scheduleDeviceStateUpdate(
-            taskName: "Set device revoked",
-            modificationBlock: { deviceState in
-                deviceState = .revoked
-            },
-            completionHandler: nil
-        )
-    }
-
     private func didReconnectTunnel(error: Error?) {
         nslock.lock()
         defer { nslock.unlock() }
@@ -1025,6 +943,16 @@ final class TunnelManager: StorePaymentObserver {
         tunnelStatusPollTimer = nil
         isPolling = false
     }
+
+    func handleRestError(_ error: Error) {
+        guard let restError = error as? REST.Error else { return }
+
+        if restError.compareErrorCode(.deviceNotFound) {
+            setDeviceState(.revoked, persist: true)
+        } else if restError.compareErrorCode(.invalidAccount) {
+            setDeviceState(.loggedOut, persist: true)
+        }
+    }
 }
 
 private struct TunnelInteractorProxy: TunnelInteractor {
@@ -1093,4 +1021,8 @@ private struct TunnelInteractorProxy: TunnelInteractor {
     func selectRelay() throws -> RelaySelectorResult {
         return try tunnelManager.selectRelay()
     }
+
+    func handleRestError(_ error: Error) {
+        tunnelManager.handleRestError(error)
+    }
 }
diff --git a/ios/MullvadVPN/TunnelManager/UpdateDeviceDataOperation.swift b/ios/MullvadVPN/TunnelManager/UpdateDeviceDataOperation.swift
index 2f0f276153..c58e32eca9 100644
--- a/ios/MullvadVPN/TunnelManager/UpdateDeviceDataOperation.swift
+++ b/ios/MullvadVPN/TunnelManager/UpdateDeviceDataOperation.swift
@@ -68,6 +68,10 @@ class UpdateDeviceDataOperation: ResultOperation<StoredDeviceData> {
             }
         }
 
+        if let error = result.error {
+            interactor.handleRestError(error)
+        }
+
         finish(result: result)
     }
 }