diff --git a/Sources/NetworkProtection/Diagnostics/KeyExpirationTester.swift b/Sources/NetworkProtection/Diagnostics/KeyExpirationTester.swift new file mode 100644 index 000000000..56dd1e10e --- /dev/null +++ b/Sources/NetworkProtection/Diagnostics/KeyExpirationTester.swift @@ -0,0 +1,164 @@ +// +// KeyExpirationTester.swift +// +// Copyright © 2024 DuckDuckGo. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +import Foundation +import Network +import NetworkExtension +import Common +import os.log + +/// Rekey timer for the VPN +/// +final actor KeyExpirationTester { + + private let canRekey: @MainActor () async -> Bool + + /// The interval of time between the start of each TCP connection test. + /// + private let intervalBetweenTests: TimeInterval = .seconds(15) + + /// Provides a simple mechanism to synchronize an `isRunning` flag for the tester to know if it needs to interrupt its operation. + /// The reason why this is necessary is that the tester may be stopped while the connection tests are already executing, in a bit + /// of a race condition which could result in the tester returning results when it's already stopped. + /// + private(set) var isRunning = false + private var isTestingExpiration = false + private let keyStore: NetworkProtectionKeyStore + private let rekey: @MainActor () async throws -> Void + private let settings: VPNSettings + private var task: Task? + + // MARK: - Init & deinit + + init(keyStore: NetworkProtectionKeyStore, + settings: VPNSettings, + canRekey: @escaping @MainActor () async -> Bool, + rekey: @escaping @MainActor () async throws -> Void) { + + self.keyStore = keyStore + self.rekey = rekey + self.canRekey = canRekey + self.settings = settings + + Logger.networkProtectionMemory.debug("[+] \(String(describing: self), privacy: .public)") + } + + deinit { + Logger.networkProtectionMemory.debug("[-] \(String(describing: self), privacy: .public)") + task?.cancel() + } + + // MARK: - Starting & Stopping the tester + + func start(testImmediately: Bool) async { + guard !isRunning else { + Logger.networkProtectionKeyManagement.log("Will not start the key expiration tester as it's already running") + return + } + + isRunning = true + + Logger.networkProtectionKeyManagement.log("🟢 Starting rekey timer") + await scheduleTimer(testImmediately: testImmediately) + } + + func stop() { + Logger.networkProtectionKeyManagement.log("🔴 Stopping rekey timer") + stopScheduledTimer() + isRunning = false + } + + // MARK: - Timer scheduling + + private func scheduleTimer(testImmediately: Bool) async { + stopScheduledTimer() + + if testImmediately { + await rekeyIfExpired() + } + + task = Task.periodic(interval: intervalBetweenTests) { [weak self] in + await self?.rekeyIfExpired() + } + } + + private func stopScheduledTimer() { + task?.cancel() + task = nil + } + + // MARK: - Testing the connection + + private var isKeyExpired: Bool { + guard let currentExpirationDate = keyStore.currentExpirationDate else { + return true + } + + return currentExpirationDate <= Date() + } + + // MARK: - Expiration check + + func rekeyIfExpired() async { + + guard !isTestingExpiration else { + return + } + + isTestingExpiration = true + + defer { + isTestingExpiration = false + } + + guard await canRekey() else { + Logger.networkProtectionKeyManagement.log("Can't rekey right now as some preconditions aren't met.") + return + } + + Logger.networkProtectionKeyManagement.log("Checking if rekey is necessary...") + + guard isKeyExpired else { + Logger.networkProtectionKeyManagement.log("The key is not expired") + return + } + + Logger.networkProtectionKeyManagement.log("Rekeying now.") + do { + try await rekey() + Logger.networkProtectionKeyManagement.log("Rekeying completed.") + } catch { + Logger.networkProtectionKeyManagement.error("Rekeying failed with error: \(error, privacy: .public).") + } + } + + // MARK: - Key Validity + + func setKeyValidity(_ interval: TimeInterval?) { + if let interval { + let firstExpirationDate = Date().addingTimeInterval(interval) + Logger.networkProtectionKeyManagement.log("Setting key validity interval to \(String(describing: interval), privacy: .public) seconds (next expiration date \(String(describing: firstExpirationDate), privacy: .public))") + settings.registrationKeyValidity = .custom(interval) + } else { + Logger.networkProtectionKeyManagement.log("Resetting key validity interval") + settings.registrationKeyValidity = .automatic + } + + keyStore.setValidityInterval(interval) + } +} diff --git a/Sources/NetworkProtection/Diagnostics/NetworkProtectionConnectionTester.swift b/Sources/NetworkProtection/Diagnostics/NetworkProtectionConnectionTester.swift index 382532893..ff2a137d8 100644 --- a/Sources/NetworkProtection/Diagnostics/NetworkProtectionConnectionTester.swift +++ b/Sources/NetworkProtection/Diagnostics/NetworkProtectionConnectionTester.swift @@ -216,7 +216,7 @@ final class NetworkProtectionConnectionTester { Logger.networkProtectionConnectionTester.log("👎 VPN is DOWN") handleDisconnected() } else { - Logger.networkProtectionConnectionTester.log("👍 VPN: \(vpnIsConnected ? "UP" : "DOWN") local: \(localIsConnected ? "UP" : "DOWN")") + Logger.networkProtectionConnectionTester.log("👍 VPN: \(vpnIsConnected ? "UP" : "DOWN", privacy: .public) local: \(localIsConnected ? "UP" : "DOWN", privacy: .public)") handleConnected() } } diff --git a/Sources/NetworkProtection/KeyManagement/NetworkProtectionKeyStore.swift b/Sources/NetworkProtection/KeyManagement/NetworkProtectionKeyStore.swift index 0df61e950..d8350e917 100644 --- a/Sources/NetworkProtection/KeyManagement/NetworkProtectionKeyStore.swift +++ b/Sources/NetworkProtection/KeyManagement/NetworkProtectionKeyStore.swift @@ -22,6 +22,10 @@ import os.log public protocol NetworkProtectionKeyStore { + /// Obtain the current expiration date + /// + var currentExpirationDate: Date? { get } + /// Obtain the current `KeyPair`. /// func currentKeyPair() -> KeyPair? @@ -145,7 +149,7 @@ public final class NetworkProtectionKeychainKeyStore: NetworkProtectionKeyStore // MARK: - UserDefaults - var currentExpirationDate: Date? { + public var currentExpirationDate: Date? { get { return userDefaults.object(forKey: UserDefaultKeys.expirationDate) as? Date } diff --git a/Sources/NetworkProtection/PacketTunnelProvider.swift b/Sources/NetworkProtection/PacketTunnelProvider.swift index 7e82291d7..28351b1ae 100644 --- a/Sources/NetworkProtection/PacketTunnelProvider.swift +++ b/Sources/NetworkProtection/PacketTunnelProvider.swift @@ -268,25 +268,6 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { keyStore.resetCurrentKeyPair() } - private var isKeyExpired: Bool { - guard let currentExpirationDate = keyStore.currentExpirationDate else { - return true - } - - return currentExpirationDate <= Date() - } - - private func rekeyIfExpired() async { - Logger.networkProtectionKeyManagement.log("Checking if rekey is necessary...") - - guard isKeyExpired else { - Logger.networkProtectionKeyManagement.log("The key is not expired") - return - } - - try? await rekey() - } - private func rekey() async throws { providerEvents.fire(.userBecameActive) @@ -320,36 +301,8 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { } } - private func setKeyValidity(_ interval: TimeInterval?) { - if let interval { - let firstExpirationDate = Date().addingTimeInterval(interval) - Logger.networkProtectionKeyManagement.log("Setting key validity interval to \(String(describing: interval), privacy: .public) seconds (next expiration date \(String(describing: firstExpirationDate), privacy: .public))") - settings.registrationKeyValidity = .custom(interval) - } else { - Logger.networkProtectionKeyManagement.log("Resetting key validity interval") - settings.registrationKeyValidity = .automatic - } - - keyStore.setValidityInterval(interval) - } - // MARK: - Bandwidth Analyzer - private func updateBandwidthAnalyzerAndRekeyIfExpired() { - Task { - await updateBandwidthAnalyzer() - - // This provides a more frequent active user pixel check - providerEvents.fire(.userBecameActive) - - guard self.bandwidthAnalyzer.isConnectionIdle() else { - return - } - - await rekeyIfExpired() - } - } - /// Updates the bandwidth analyzer with the latest data from the WireGuard Adapter /// public func updateBandwidthAnalyzer() async { @@ -366,6 +319,21 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { private static let connectionTesterExtendedFailuresCount = 8 private var isConnectionTesterEnabled: Bool = true + @MainActor + private lazy var keyExpirationTester: KeyExpirationTester = { + KeyExpirationTester(keyStore: keyStore, settings: settings) { @MainActor [weak self] in + guard let self else { return false } + + // This provides a more frequent active user pixel check + providerEvents.fire(.userBecameActive) + + await updateBandwidthAnalyzer() + return bandwidthAnalyzer.isConnectionIdle() + } rekey: { @MainActor [weak self] in + try await self?.rekey() + } + }() + @MainActor private lazy var connectionTester: NetworkProtectionConnectionTester = { NetworkProtectionConnectionTester(timerQueue: timerQueue) { @MainActor [weak self] result in @@ -376,7 +344,6 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { switch result { case .connected: self.tunnelHealth.isHavingConnectivityIssues = false - self.updateBandwidthAnalyzerAndRekeyIfExpired() case .reconnected(let failureCount): providerEvents.fire( @@ -392,7 +359,6 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { } self.tunnelHealth.isHavingConnectivityIssues = false - self.updateBandwidthAnalyzerAndRekeyIfExpired() case .disconnected(let failureCount): if failureCount == 1 { @@ -534,11 +500,15 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { private func loadKeyValidity(from options: StartupOptions) { switch options.keyValidity { case .set(let validity): - setKeyValidity(validity) + Task { @MainActor in + await keyExpirationTester.setKeyValidity(validity) + } case .useExisting: break case .reset: - setKeyValidity(nil) + Task { @MainActor in + await keyExpirationTester.setKeyValidity(nil) + } } } @@ -821,7 +791,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { } #endif } catch { - self.cancelTunnelWithError(error) + Logger.networkProtection.log("Connection Tester failed to start... will run without it: \(error, privacy: .public)") return } } @@ -1198,7 +1168,8 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { private func handleExpireRegistrationKey(completionHandler: ((Data?) -> Void)? = nil) { Task { - try? await rekey() + keyStore.currentExpirationDate = Date() + await keyExpirationTester.rekeyIfExpired() completionHandler?(nil) } } @@ -1311,7 +1282,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { private func handleSetKeyValidity(_ keyValidity: TimeInterval?, completionHandler: ((Data?) -> Void)? = nil) { Task { - setKeyValidity(keyValidity) + await keyExpirationTester.setKeyValidity(keyValidity) completionHandler?(nil) } } @@ -1593,6 +1564,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { await startLatencyMonitor() await startEntitlementMonitor() await startServerStatusMonitor() + await keyExpirationTester.start(testImmediately: testImmediately) do { try await startConnectionTester(testImmediately: testImmediately) @@ -1604,7 +1576,8 @@ open class PacketTunnelProvider: NEPacketTunnelProvider { @MainActor public func stopMonitors() async { - self.connectionTester.stop() + connectionTester.stop() + await keyExpirationTester.stop() await self.tunnelFailureMonitor.stop() await self.latencyMonitor.stop() await self.entitlementMonitor.stop() diff --git a/Tests/NetworkProtectionTests/Mocks/NetworkProtectionKeyStoreMocks.swift b/Tests/NetworkProtectionTests/Mocks/NetworkProtectionKeyStoreMocks.swift index a1ba18466..2f74f184d 100644 --- a/Tests/NetworkProtectionTests/Mocks/NetworkProtectionKeyStoreMocks.swift +++ b/Tests/NetworkProtectionTests/Mocks/NetworkProtectionKeyStoreMocks.swift @@ -26,6 +26,10 @@ final class NetworkProtectionKeyStoreMock: NetworkProtectionKeyStore { // MARK: - NetworkProtectionKeyStore + var currentExpirationDate: Date? { + Date() + } + func currentKeyPair() -> NetworkProtection.KeyPair? { keyPair }