Store next private key when rotating the key

6 files changed, 240 insertions, 135 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index c4410329e7..06d9d34cc6 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -283,8 +283,7 @@
 		58F2E144276A13F300A79513 /* StartTunnelOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F2E143276A13F300A79513 /* StartTunnelOperation.swift */; };
 		58F2E146276A2C9900A79513 /* StopTunnelOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F2E145276A2C9900A79513 /* StopTunnelOperation.swift */; };
 		58F2E148276A307400A79513 /* MapConnectionStatusOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F2E147276A307400A79513 /* MapConnectionStatusOperation.swift */; };
-		58F2E14A276A43AA00A79513 /* RegenerateTunnelPrivateKeyOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F2E149276A43AA00A79513 /* RegenerateTunnelPrivateKeyOperation.swift */; };
-		58F2E14C276A61C000A79513 /* RotatePrivateKeyOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F2E14B276A61C000A79513 /* RotatePrivateKeyOperation.swift */; };
+		58F2E14C276A61C000A79513 /* ReplaceKeyOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F2E14B276A61C000A79513 /* ReplaceKeyOperation.swift */; };
 		58F3C0A4249CB069003E76BE /* HeaderBarView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F3C0A3249CB069003E76BE /* HeaderBarView.swift */; };
 		58F3C0A624A50157003E76BE /* relays.json in Resources */ = {isa = PBXBuildFile; fileRef = 58F3C0A524A50155003E76BE /* relays.json */; };
 		58F3C0A724A50C02003E76BE /* relays.json in Resources */ = {isa = PBXBuildFile; fileRef = 58F3C0A524A50155003E76BE /* relays.json */; };
@@ -572,8 +571,7 @@
 		58F2E143276A13F300A79513 /* StartTunnelOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StartTunnelOperation.swift; sourceTree = "<group>"; };
 		58F2E145276A2C9900A79513 /* StopTunnelOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StopTunnelOperation.swift; sourceTree = "<group>"; };
 		58F2E147276A307400A79513 /* MapConnectionStatusOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MapConnectionStatusOperation.swift; sourceTree = "<group>"; };
-		58F2E149276A43AA00A79513 /* RegenerateTunnelPrivateKeyOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RegenerateTunnelPrivateKeyOperation.swift; sourceTree = "<group>"; };
-		58F2E14B276A61C000A79513 /* RotatePrivateKeyOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = RotatePrivateKeyOperation.swift; sourceTree = "<group>"; };
+		58F2E14B276A61C000A79513 /* ReplaceKeyOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ReplaceKeyOperation.swift; sourceTree = "<group>"; };
 		58F3C0A3249CB069003E76BE /* HeaderBarView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = HeaderBarView.swift; sourceTree = "<group>"; };
 		58F3C0A524A50155003E76BE /* relays.json */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text.json; path = relays.json; sourceTree = "<group>"; };
 		58F558DC2695B85E00F630D0 /* en */ = {isa = PBXFileReference; lastKnownFileType = text.plist.strings; name = en; path = en.lproj/Consent.strings; sourceTree = "<group>"; };
@@ -717,8 +715,7 @@
 				58F2E143276A13F300A79513 /* StartTunnelOperation.swift */,
 				58F2E145276A2C9900A79513 /* StopTunnelOperation.swift */,
 				58F2E147276A307400A79513 /* MapConnectionStatusOperation.swift */,
-				58F2E149276A43AA00A79513 /* RegenerateTunnelPrivateKeyOperation.swift */,
-				58F2E14B276A61C000A79513 /* RotatePrivateKeyOperation.swift */,
+				58F2E14B276A61C000A79513 /* ReplaceKeyOperation.swift */,
 				588527B5276B58B300BAA373 /* SetTunnelSettingsOperation.swift */,
 			);
 			path = TunnelManager;
@@ -1403,7 +1400,7 @@
 				5806766C27048E3E00C858CB /* AnyOptional.swift in Sources */,
 				58ACF64D26567A5000ACE4B7 /* CustomSwitch.swift in Sources */,
 				5850367F25A481D800A43E93 /* IPAddressRange+Codable.swift in Sources */,
-				58F2E14C276A61C000A79513 /* RotatePrivateKeyOperation.swift in Sources */,
+				58F2E14C276A61C000A79513 /* ReplaceKeyOperation.swift in Sources */,
 				5871FB96254ADE4E0051A0A4 /* ConsolidatedApplicationLog.swift in Sources */,
 				58FEEB58260B662E00A621A8 /* AutomaticKeyboardResponder.swift in Sources */,
 				5846227326E22A160035F7C2 /* AppStorePaymentObserver.swift in Sources */,
@@ -1533,7 +1530,6 @@
 				58561C99239A5D1500BD6B5E /* IPEndpoint.swift in Sources */,
 				58FD5BF22424F7D700112C88 /* UserInterfaceInteractionRestriction.swift in Sources */,
 				5811DE50239014550011EB53 /* NEVPNStatus+Debug.swift in Sources */,
-				58F2E14A276A43AA00A79513 /* RegenerateTunnelPrivateKeyOperation.swift in Sources */,
 				58C3A4B222456F1B00340BDB /* AccountInputGroupView.swift in Sources */,
 				58F840B22464491D0044E708 /* ChainedError.swift in Sources */,
 				58FAEDFF24533A7000CB0F5B /* KeychainReturn.swift in Sources */,
diff --git a/ios/MullvadVPN/Operations/OperationCompletion.swift b/ios/MullvadVPN/Operations/OperationCompletion.swift
index 23d206d32d..21cb92e169 100644
--- a/ios/MullvadVPN/Operations/OperationCompletion.swift
+++ b/ios/MullvadVPN/Operations/OperationCompletion.swift
@@ -31,6 +31,17 @@ enum OperationCompletion<Success, Failure: Error> {
         }
     }
 
+    func map<NewSuccess>(_ block: (Success) -> NewSuccess) -> OperationCompletion<NewSuccess, Failure> {
+        switch self {
+        case .success(let value):
+            return .success(block(value))
+        case .failure(let error):
+            return .failure(error)
+        case .cancelled:
+            return .cancelled
+        }
+    }
+
     func mapError<NewFailure: Error>(_ block: (Failure) -> NewFailure) -> OperationCompletion<Success, NewFailure> {
         switch self {
         case .success(let value):
diff --git a/ios/MullvadVPN/TunnelManager/ReplaceKeyOperation.swift b/ios/MullvadVPN/TunnelManager/ReplaceKeyOperation.swift
new file mode 100644
index 0000000000..a3569a400e
--- /dev/null
+++ b/ios/MullvadVPN/TunnelManager/ReplaceKeyOperation.swift
@@ -0,0 +1,214 @@
+//
+//  ReplaceKeyOperation.swift
+//  MullvadVPN
+//
+//  Created by pronebird on 15/12/2021.
+//  Copyright © 2021 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+import Logging
+
+class ReplaceKeyOperation: AsyncOperation {
+    typealias CompletionHandler = (OperationCompletion<TunnelManager.KeyRotationResult, TunnelManager.Error>) -> Void
+
+    private let queue: DispatchQueue
+    private let state: TunnelManager.State
+
+    private let restClient: REST.Client
+    private var restRequest: Cancellable?
+
+    private let rotationInterval: TimeInterval?
+    private var completionHandler: CompletionHandler?
+
+    private let logger = Logger(label: "TunnelManager.ReplaceKeyOperation")
+
+    class func operationForKeyRotation(
+        queue: DispatchQueue,
+        state: TunnelManager.State,
+        restClient: REST.Client,
+        rotationInterval: TimeInterval,
+        completionHandler: @escaping CompletionHandler
+    ) -> ReplaceKeyOperation {
+        return ReplaceKeyOperation(
+            queue: queue,
+            state: state,
+            restClient: restClient,
+            rotationInterval: rotationInterval,
+            completionHandler: completionHandler
+        )
+    }
+
+    class func operationForKeyRegeneration(
+        queue: DispatchQueue,
+        state: TunnelManager.State,
+        restClient: REST.Client,
+        completionHandler: @escaping (OperationCompletion<(), TunnelManager.Error>) -> Void
+    ) -> ReplaceKeyOperation {
+        return ReplaceKeyOperation(
+            queue: queue,
+            state: state,
+            restClient: restClient,
+            rotationInterval: nil
+        ) { completion in
+            let mappedCompletion = completion.map { keyRotationResult -> () in
+                switch keyRotationResult {
+                case .finished:
+                    return ()
+                case .throttled:
+                    fatalError("ReplaceKeyOperation.operationForKeyRegeneration() must never recieve throttled!")
+                }
+            }
+
+            completionHandler(mappedCompletion)
+        }
+    }
+
+    private init(
+        queue: DispatchQueue,
+        state: TunnelManager.State,
+        restClient: REST.Client,
+        rotationInterval: TimeInterval?,
+        completionHandler: @escaping CompletionHandler
+    ) {
+        self.queue = queue
+        self.state = state
+
+        self.restClient = restClient
+        self.rotationInterval = rotationInterval
+
+        self.completionHandler = completionHandler
+    }
+
+    override func main() {
+        queue.async {
+            self.execute { completion in
+                self.completionHandler?(completion)
+                self.completionHandler = nil
+
+                self.finish()
+            }
+        }
+    }
+
+    override func cancel() {
+        super.cancel()
+
+        queue.async {
+            self.restRequest?.cancel()
+        }
+    }
+
+    private func execute(completionHandler: @escaping CompletionHandler) {
+        guard !isCancelled else {
+            completionHandler(.cancelled)
+            return
+        }
+
+        guard let tunnelInfo = state.tunnelInfo else {
+            completionHandler(.failure(.missingAccount))
+            return
+        }
+
+        if let rotationInterval = rotationInterval {
+            let creationDate = tunnelInfo.tunnelSettings.interface.privateKey.creationDate
+            let timeElapsed = Date().timeIntervalSince(creationDate)
+
+            if timeElapsed < rotationInterval {
+                logger.debug("Throttle private key rotation.")
+
+                completionHandler(.success(.throttled(creationDate)))
+                return
+            } else {
+                logger.debug("Private key is old enough, rotate right away.")
+            }
+        } else {
+            logger.debug("Rotate private key right away.")
+        }
+
+        let newPrivateKey: PrivateKeyWithMetadata
+        let oldPublicKey = tunnelInfo.tunnelSettings.interface.publicKey
+
+        if let nextPrivateKey = tunnelInfo.tunnelSettings.interface.nextPrivateKey {
+            newPrivateKey = nextPrivateKey
+
+            logger.debug("Next private key is already created.")
+        } else {
+            newPrivateKey = PrivateKeyWithMetadata()
+
+            logger.debug("Create next private key.")
+
+            let saveResult = TunnelSettingsManager.update(searchTerm: .accountToken(tunnelInfo.token)) { newTunnelSettings in
+                newTunnelSettings.interface.nextPrivateKey = newPrivateKey
+            }
+
+            switch saveResult {
+            case .success(let newTunnelSettings):
+                logger.debug("Saved next private key.")
+
+                state.tunnelInfo?.tunnelSettings = newTunnelSettings
+
+            case .failure(let error):
+                logger.error(chainedError: error, message: "Failed to save next private key.")
+
+                completionHandler(.failure(.updateTunnelSettings(error)))
+                return
+            }
+        }
+
+        logger.debug("Replacing old key with new key on server...")
+
+        let requestAdapter = self.restClient.replaceWireguardKey(
+            token: tunnelInfo.token,
+            oldPublicKey: oldPublicKey,
+            newPublicKey: newPrivateKey.publicKey
+        )
+
+        restRequest = requestAdapter.execute(retryStrategy: .default) { result in
+            self.queue.async {
+                self.didReceiveResponse(
+                    result: result,
+                    accountToken: tunnelInfo.token,
+                    newPrivateKey: newPrivateKey,
+                    completionHandler: completionHandler
+                )
+            }
+        }
+    }
+
+    private func didReceiveResponse(result: Result<REST.WireguardAddressesResponse, REST.Error>, accountToken: String, newPrivateKey: PrivateKeyWithMetadata, completionHandler: @escaping CompletionHandler) {
+        switch result {
+        case .success(let associatedAddresses):
+            logger.debug("Replaced old key with new key on server.")
+
+            let saveResult = TunnelSettingsManager.update(searchTerm: .accountToken(accountToken)) { newTunnelSettings in
+                newTunnelSettings.interface.privateKey = newPrivateKey
+                newTunnelSettings.interface.nextPrivateKey = nil
+
+                newTunnelSettings.interface.addresses = [
+                    associatedAddresses.ipv4Address,
+                    associatedAddresses.ipv6Address
+                ]
+            }
+
+            switch saveResult {
+            case .success(let newTunnelSettings):
+                logger.debug("Saved associated addresses.")
+
+                state.tunnelInfo?.tunnelSettings = newTunnelSettings
+
+                completionHandler(.success(.finished))
+
+            case .failure(let error):
+                logger.error(chainedError: error, message: "Failed to save associated addresses.")
+
+                completionHandler(.failure(.updateTunnelSettings(error)))
+            }
+
+        case .failure(let restError):
+            logger.error(chainedError: restError, message: "Failed to replace old key with new key on server.")
+
+            completionHandler(.failure(.replaceWireguardKey(restError)))
+        }
+    }
+}
diff --git a/ios/MullvadVPN/TunnelManager/RotatePrivateKeyOperation.swift b/ios/MullvadVPN/TunnelManager/RotatePrivateKeyOperation.swift
deleted file mode 100644
index 1ec6d5cc2a..0000000000
--- a/ios/MullvadVPN/TunnelManager/RotatePrivateKeyOperation.swift
+++ /dev/null
@@ -1,123 +0,0 @@
-//
-//  RotatePrivateKeyOperation.swift
-//  MullvadVPN
-//
-//  Created by pronebird on 15/12/2021.
-//  Copyright © 2021 Mullvad VPN AB. All rights reserved.
-//
-
-import Foundation
-
-class RotatePrivateKeyOperation: AsyncOperation {
-    typealias CompletionHandler = (OperationCompletion<TunnelManager.KeyRotationResult, TunnelManager.Error>) -> Void
-
-    private let queue: DispatchQueue
-    private let state: TunnelManager.State
-    private let restClient: REST.Client
-    private let rotationInterval: TimeInterval
-    private var completionHandler: CompletionHandler?
-    private var restRequest: Cancellable?
-
-    init(queue: DispatchQueue,
-         state: TunnelManager.State,
-         restClient: REST.Client,
-         rotationInterval: TimeInterval,
-         completionHandler: @escaping CompletionHandler)
-    {
-        self.queue = queue
-        self.state = state
-        self.restClient = restClient
-        self.rotationInterval = rotationInterval
-        self.completionHandler = completionHandler
-    }
-
-    override func main() {
-        queue.async {
-            self.execute { completion in
-                self.completionHandler?(completion)
-                self.completionHandler = nil
-
-                self.finish()
-            }
-        }
-    }
-
-    override func cancel() {
-        super.cancel()
-
-        queue.async {
-            self.restRequest?.cancel()
-        }
-    }
-
-    private func execute(completionHandler: @escaping CompletionHandler) {
-        guard !isCancelled else {
-            completionHandler(.cancelled)
-            return
-        }
-
-        guard let tunnelInfo = state.tunnelInfo else {
-            completionHandler(.failure(.missingAccount))
-            return
-        }
-
-        let creationDate = tunnelInfo.tunnelSettings.interface.privateKey.creationDate
-        let timeInterval = Date().timeIntervalSince(creationDate)
-
-        guard timeInterval >= rotationInterval else {
-            completionHandler(.success(.throttled(creationDate)))
-            return
-        }
-
-        let newPrivateKey = PrivateKeyWithMetadata()
-        let oldPublicKey = tunnelInfo.tunnelSettings.interface.publicKey
-
-        let requestAdapter = self.restClient.replaceWireguardKey(
-            token: tunnelInfo.token,
-            oldPublicKey: oldPublicKey,
-            newPublicKey: newPrivateKey.publicKey
-        )
-
-        restRequest = requestAdapter.execute(retryStrategy: .default) { result in
-            self.queue.async {
-                self.didRotatePrivateKey(
-                    result: result,
-                    accountToken: tunnelInfo.token,
-                    newPrivateKey: newPrivateKey,
-                    completionHandler: completionHandler
-                )
-            }
-        }
-    }
-
-    private func didRotatePrivateKey(result: Result<REST.WireguardAddressesResponse, REST.Error>, accountToken: String, newPrivateKey: PrivateKeyWithMetadata, completionHandler: @escaping CompletionHandler) {
-        let saveResult = Self.handleResponse(accountToken: accountToken, newPrivateKey: newPrivateKey, result: result)
-
-        switch saveResult {
-        case .success(let tunnelSettings):
-            state.tunnelInfo?.tunnelSettings = tunnelSettings
-
-            completionHandler(.success(.finished))
-
-        case .failure(let error):
-            completionHandler(.failure(error))
-        }
-    }
-
-    private class func handleResponse(accountToken: String, newPrivateKey: PrivateKeyWithMetadata, result: Result<REST.WireguardAddressesResponse, REST.Error>) -> Result<TunnelSettings, TunnelManager.Error> {
-        return result.flatMapError { restError in
-            return .failure(.replaceWireguardKey(restError))
-        }
-        .flatMap { associatedAddresses in
-            return TunnelSettingsManager.update(searchTerm: .accountToken(accountToken)) { newTunnelSettings in
-                newTunnelSettings.interface.privateKey = newPrivateKey
-                newTunnelSettings.interface.addresses = [
-                    associatedAddresses.ipv4Address,
-                    associatedAddresses.ipv6Address
-                ]
-            }.mapError { error -> TunnelManager.Error in
-                return .updateTunnelSettings(error)
-            }
-        }
-    }
-}
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index d0eecaa4df..6b62d9d39b 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -317,7 +317,7 @@ class TunnelManager: TunnelManagerStateDelegate
     }
 
     func regeneratePrivateKey(completionHandler: ((TunnelManager.Error?) -> Void)? = nil) {
-        let operation = RegeneratePrivateKeyOperation(queue: stateQueue, state: state, restClient: restClient) { [weak self] completion in
+        let operation = ReplaceKeyOperation.operationForKeyRegeneration(queue: stateQueue, state: state, restClient: restClient) { [weak self] completion in
             guard let self = self else { return }
 
             dispatchPrecondition(condition: .onQueue(self.stateQueue))
@@ -353,7 +353,7 @@ class TunnelManager: TunnelManagerStateDelegate
     }
 
     func rotatePrivateKey(completionHandler: @escaping (KeyRotationResult?, TunnelManager.Error?) -> Void) {
-        let operation = RotatePrivateKeyOperation(
+        let operation = ReplaceKeyOperation.operationForKeyRotation(
             queue: stateQueue,
             state: state,
             restClient: restClient,
diff --git a/ios/MullvadVPN/TunnelSettings.swift b/ios/MullvadVPN/TunnelSettings.swift
index fcbdebbe81..c9a96fad80 100644
--- a/ios/MullvadVPN/TunnelSettings.swift
+++ b/ios/MullvadVPN/TunnelSettings.swift
@@ -13,6 +13,8 @@ import struct WireGuardKit.IPAddressRange
 /// A struct that holds a tun interface configuration.
 struct InterfaceSettings: Codable, Equatable {
     var privateKey: PrivateKeyWithMetadata
+    var nextPrivateKey: PrivateKeyWithMetadata?
+
     var addresses: [IPAddressRange]
     var dnsSettings: DNSSettings
 
@@ -21,11 +23,12 @@ struct InterfaceSettings: Codable, Equatable {
     }
 
     private enum CodingKeys: String, CodingKey {
-        case privateKey, addresses, dnsSettings
+        case privateKey, nextPrivateKey, addresses, dnsSettings
     }
 
-    init(privateKey: PrivateKeyWithMetadata = PrivateKeyWithMetadata(), addresses: [IPAddressRange] = [], dnsSettings: DNSSettings = DNSSettings()) {
+    init(privateKey: PrivateKeyWithMetadata = PrivateKeyWithMetadata(), nextPrivateKey: PrivateKeyWithMetadata? = nil, addresses: [IPAddressRange] = [], dnsSettings: DNSSettings = DNSSettings()) {
         self.privateKey = privateKey
+        self.nextPrivateKey = nextPrivateKey
         self.addresses = addresses
         self.dnsSettings = dnsSettings
     }
@@ -36,6 +39,9 @@ struct InterfaceSettings: Codable, Equatable {
         privateKey = try container.decode(PrivateKeyWithMetadata.self, forKey: .privateKey)
         addresses = try container.decode([IPAddressRange].self, forKey: .addresses)
 
+        // Added in 2022.1
+        nextPrivateKey = try container.decodeIfPresent(PrivateKeyWithMetadata.self, forKey: .nextPrivateKey)
+
         // Provide default value, since `dnsSettings` key does not exist in <= 2021.2
         dnsSettings = try container.decodeIfPresent(DNSSettings.self, forKey: .dnsSettings)
             ?? DNSSettings()
@@ -45,6 +51,7 @@ struct InterfaceSettings: Codable, Equatable {
         var container = encoder.container(keyedBy: CodingKeys.self)
 
         try container.encode(privateKey, forKey: .privateKey)
+        try container.encode(nextPrivateKey, forKey: .nextPrivateKey)
         try container.encode(addresses, forKey: .addresses)
         try container.encode(dnsSettings, forKey: .dnsSettings)
     }