Delay tunnel reconnection after private key rotation

9 files changed, 105 insertions, 9 deletions

diff --git a/ios/CHANGELOG.md b/ios/CHANGELOG.md
index 5a12de6830..faa22d3ac1 100644
--- a/ios/CHANGELOG.md
+++ b/ios/CHANGELOG.md
@@ -25,6 +25,8 @@ Line wrap the file at 100 chars.                                              Th
 ## [Unreleased]
 ### Changed
 - Changed key rotation interval from 4 to 7 days.
+- Delay tunnel reconnection after a WireGuard private key rotates. Accounts for latency in key
+  propagation to relays.
 
 
 ## [2023.1] - 2023-03-21
diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 72be02bc76..00ea634a9d 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -1207,17 +1207,17 @@
 				58F2E143276A13F300A79513 /* StartTunnelOperation.swift */,
 				58F2E145276A2C9900A79513 /* StopTunnelOperation.swift */,
 				58E0A98727C8F46300FE6BDD /* Tunnel.swift */,
-				5803B4B12940A48700C23744 /* TunnelStore.swift */,
-				5803B4AF2940A47300C23744 /* TunnelConfiguration.swift */,
 				5875960926F371FC00BF6711 /* Tunnel+Messaging.swift */,
+				5878A27229091D6D0096FC88 /* TunnelBlockObserver.swift */,
+				5803B4AF2940A47300C23744 /* TunnelConfiguration.swift */,
 				58968FAD28743E2000B799DC /* TunnelInteractor.swift */,
 				5835B7CB233B76CB0096D79F /* TunnelManager.swift */,
 				5820676326E771DB00655B05 /* TunnelManagerErrors.swift */,
 				5823FA5326CE49F600283BF8 /* TunnelObserver.swift */,
 				58B93A1226C3F13600A55733 /* TunnelState.swift */,
+				5803B4B12940A48700C23744 /* TunnelStore.swift */,
 				5842102F282D8A3C00F24E46 /* UpdateAccountDataOperation.swift */,
 				58421031282E42B000F24E46 /* UpdateDeviceDataOperation.swift */,
-				5878A27229091D6D0096FC88 /* TunnelBlockObserver.swift */,
 			);
 			path = TunnelManager;
 			sourceTree = "<group>";
diff --git a/ios/MullvadVPN/SimulatorTunnelProvider/SimulatorTunnelProviderHost.swift b/ios/MullvadVPN/SimulatorTunnelProvider/SimulatorTunnelProviderHost.swift
index 3841a3bd8b..3c2c8f85d0 100644
--- a/ios/MullvadVPN/SimulatorTunnelProvider/SimulatorTunnelProviderHost.swift
+++ b/ios/MullvadVPN/SimulatorTunnelProvider/SimulatorTunnelProviderHost.swift
@@ -143,6 +143,9 @@ final class SimulatorTunnelProviderHost: SimulatorTunnelProviderDelegate {
             urlRequestProxy.cancelRequest(identifier: id)
 
             completionHandler?(nil)
+
+        case .privateKeyRotation:
+            completionHandler?(nil)
         }
     }
 
diff --git a/ios/MullvadVPN/TunnelManager/ReconnectTunnelOperation.swift b/ios/MullvadVPN/TunnelManager/ReconnectTunnelOperation.swift
index bbaa83850e..2c38566fc1 100644
--- a/ios/MullvadVPN/TunnelManager/ReconnectTunnelOperation.swift
+++ b/ios/MullvadVPN/TunnelManager/ReconnectTunnelOperation.swift
@@ -39,7 +39,9 @@ class ReconnectTunnelOperation: ResultOperation<Void> {
             let selectorResult = selectNewRelay ? try interactor.selectRelay() : nil
 
             task = tunnel
-                .reconnectTunnel(relaySelectorResult: selectorResult) { [weak self] result in
+                .reconnectTunnel(
+                    relaySelectorResult: selectorResult
+                ) { [weak self] result in
                     self?.finish(result: result)
                 }
         } catch {
diff --git a/ios/MullvadVPN/TunnelManager/RotateKeyOperation.swift b/ios/MullvadVPN/TunnelManager/RotateKeyOperation.swift
index 5f74e5ee2a..1793453107 100644
--- a/ios/MullvadVPN/TunnelManager/RotateKeyOperation.swift
+++ b/ios/MullvadVPN/TunnelManager/RotateKeyOperation.swift
@@ -100,7 +100,13 @@ class RotateKeyOperation: ResultOperation<Bool> {
 
                 interactor.setDeviceState(.loggedIn(accountData, deviceData), persist: true)
 
-                finish(result: .success(true))
+                if let tunnel = interactor.tunnel {
+                    _ = tunnel.notifyKeyRotation { [weak self] _ in
+                        self?.finish(result: .success(true))
+                    }
+                } else {
+                    finish(result: .success(true))
+                }
             default:
                 finish(result: .failure(InvalidDeviceStateError()))
             }
diff --git a/ios/MullvadVPN/TunnelManager/Tunnel+Messaging.swift b/ios/MullvadVPN/TunnelManager/Tunnel+Messaging.swift
index 9859da99b8..3b69dd34eb 100644
--- a/ios/MullvadVPN/TunnelManager/Tunnel+Messaging.swift
+++ b/ios/MullvadVPN/TunnelManager/Tunnel+Messaging.swift
@@ -93,4 +93,21 @@ extension Tunnel {
 
         return operation
     }
+
+    /// Notify tunnel about private key rotation.
+    func notifyKeyRotation(
+        completionHandler: @escaping (Result<Void, Error>) -> Void
+    ) -> Cancellable {
+        let operation = SendTunnelProviderMessageOperation(
+            dispatchQueue: dispatchQueue,
+            application: .shared,
+            tunnel: self,
+            message: .privateKeyRotation,
+            completionHandler: completionHandler
+        )
+
+        operationQueue.addOperation(operation)
+
+        return operation
+    }
 }
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index 1607cd75df..247715f879 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -504,9 +504,7 @@ final class TunnelManager: StorePaymentObserver {
 
             switch result {
             case .success:
-                self.reconnectTunnel(selectNewRelay: true) { _ in
-                    completionHandler(result)
-                }
+                completionHandler(result)
 
             case let .failure(error):
                 if !error.isOperationCancellationError {
diff --git a/ios/PacketTunnel/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider.swift
index b33b013fe6..380293695b 100644
--- a/ios/PacketTunnel/PacketTunnelProvider.swift
+++ b/ios/PacketTunnel/PacketTunnelProvider.swift
@@ -21,6 +21,9 @@ import WireGuardKit
 /// Restart interval (in seconds) for the tunnel that failed to start early on.
 private let tunnelStartupFailureRestartInterval: TimeInterval = 2
 
+/// Delay before trying to reconnect tunnel after private key rotation.
+private let keyRotationTunnelReconnectionDelay = 60 * 2
+
 class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
     /// Tunnel provider logger.
     private let providerLogger: Logger
@@ -93,6 +96,16 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
     /// Internal operation queue.
     private let operationQueue = AsyncOperationQueue()
 
+    /// Timer for tunnel reconnection. Used to delay reconnection when a private key has just been
+    /// rotated, to account for latency in key propagation to relays.
+    private var tunnelReconnectionTimer: DispatchSourceTimer?
+
+    /// Current device state for the tunnel.
+    private var cachedDeviceState: DeviceState?
+
+    /// Whether to use the cached device state.
+    private var useCachedDeviceState = false
+
     /// Returns `PacketTunnelStatus` used for sharing with main bundle process.
     private var packetTunnelStatus: PacketTunnelStatus {
         let errors: [PacketTunnelErrorWrapper?] = [
@@ -263,6 +276,7 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
             self.providerLogger.debug("Stop the tunnel: \(reason)")
 
             self.isStopping = true
+            self.cancelTunnelReconnectionTimer()
             self.cancelTunnelStartupFailureRecovery()
             self.startTunnelCompletionHandler = nil
 
@@ -348,6 +362,12 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
             case let .cancelURLRequest(id):
                 self.urlRequestProxy.cancelRequest(identifier: id)
                 completionHandler?(nil)
+
+            case .privateKeyRotation:
+                self.startTunnelReconnectionTimer(
+                    reconnectionDelay: keyRotationTunnelReconnectionDelay
+                )
+                completionHandler?(nil)
             }
         }
     }
@@ -415,6 +435,42 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
 
     // MARK: - Private
 
+    private func startTunnelReconnectionTimer(reconnectionDelay: Int) {
+        dispatchPrecondition(condition: .onQueue(dispatchQueue))
+
+        providerLogger.debug("Delaying tunnel reconnection by \(reconnectionDelay) seconds...")
+        useCachedDeviceState = true
+
+        let timer = DispatchSource.makeTimerSource(queue: dispatchQueue)
+
+        timer.setEventHandler { [weak self] in
+            self?.providerLogger.debug("Reconnecting the tunnel...")
+
+            let nextRelay: NextRelay = self?.selectorResult
+                .map { .set($0) } ?? .automatic
+
+            self?.useCachedDeviceState = false
+            self?.reconnectTunnel(to: nextRelay, shouldStopTunnelMonitor: true)
+        }
+
+        timer.setCancelHandler { [weak self] in
+            self?.useCachedDeviceState = false
+        }
+
+        timer.schedule(deadline: .now() + .seconds(reconnectionDelay))
+        timer.activate()
+
+        tunnelReconnectionTimer?.cancel()
+        tunnelReconnectionTimer = timer
+    }
+
+    private func cancelTunnelReconnectionTimer() {
+        dispatchPrecondition(condition: .onQueue(dispatchQueue))
+
+        tunnelReconnectionTimer?.cancel()
+        tunnelReconnectionTimer = nil
+    }
+
     private func beginTunnelStartupFailureRecovery() {
         dispatchPrecondition(condition: .onQueue(dispatchQueue))
 
@@ -493,9 +549,16 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
         throws -> PacketTunnelConfiguration
     {
         let tunnelSettings = try SettingsManager.readSettings()
-        let deviceState = try SettingsManager.readDeviceState()
         let selectorResult: RelaySelectorResult
 
+        var deviceState: DeviceState
+        if let cachedDeviceState = cachedDeviceState, useCachedDeviceState {
+            deviceState = cachedDeviceState
+        } else {
+            deviceState = try SettingsManager.readDeviceState()
+            cachedDeviceState = deviceState
+        }
+
         switch nextRelay {
         case .automatic:
             selectorResult = try selectRelayEndpoint(
diff --git a/ios/TunnelProviderMessaging/TunnelProviderMessage.swift b/ios/TunnelProviderMessaging/TunnelProviderMessage.swift
index 03e55be01f..877d7c5047 100644
--- a/ios/TunnelProviderMessaging/TunnelProviderMessage.swift
+++ b/ios/TunnelProviderMessaging/TunnelProviderMessage.swift
@@ -24,6 +24,9 @@ public enum TunnelProviderMessage: Codable, CustomStringConvertible {
     /// Cancel HTTP request sent outside of VPN tunnel.
     case cancelURLRequest(UUID)
 
+    /// Notify tunnel about private key rotation.
+    case privateKeyRotation
+
     public var description: String {
         switch self {
         case .reconnectTunnel:
@@ -34,6 +37,8 @@ public enum TunnelProviderMessage: Codable, CustomStringConvertible {
             return "send-http-request"
         case .cancelURLRequest:
             return "cancel-http-request"
+        case .privateKeyRotation:
+            return "private-key-rotation"
         }
     }