Prevent endless rekey and tunnel update cycle

Sources/Common/Concurrency/TaskExtension.swift

@@ -42,6 +42,29 @@ public extension Task where Success == Never, Failure == Error {
             }
         }
     }
+
+    static func periodic(delay: TimeInterval? = nil,
+                         interval: TimeInterval,
+                         operation: @escaping @Sendable () async throws -> Void,
+                         cancellationHandler: (@Sendable () async -> Void)? = nil) -> Task {
+
+        Task {
+            do {
+                if let delay {
+                    try await Task<Never, Never>.sleep(interval: delay)
+                }
+
+                repeat {
+                    try await operation()
+
+                    try await Task<Never, Never>.sleep(interval: interval)
+                } while true
+            } catch {
+                await cancellationHandler?()
+                throw error
+            }
+        }
+    }
 }
 
 public extension Task where Success == Never, Failure == Never {

Sources/NetworkProtection/Diagnostics/NetworkProtectionEntitlementMonitor.swift

@@ -52,7 +52,7 @@ public actor NetworkProtectionEntitlementMonitor {
 
     // MARK: - Start/Stop monitoring
 
-    public func start(entitlementCheck: @escaping () async -> Swift.Result<Bool, Error>, callback: @escaping (Result) -> Void) {
+    public func start(entitlementCheck: @escaping () async -> Swift.Result<Bool, Error>, callback: @escaping (Result) async -> Void) {
         os_log("⚫️ Starting entitlement monitor", log: .networkProtectionEntitlementMonitorLog)
 
         task = Task.periodic(interval: Self.monitoringInterval) {
@@ -61,14 +61,14 @@ public actor NetworkProtectionEntitlementMonitor {
             case .success(let hasEntitlement):
                 if hasEntitlement {
                     os_log("⚫️ Valid entitlement", log: .networkProtectionEntitlementMonitorLog)
-                    callback(.validEntitlement)
+                    await callback(.validEntitlement)
                 } else {
                     os_log("⚫️ Invalid entitlement", log: .networkProtectionEntitlementMonitorLog)
-                    callback(.invalidEntitlement)
+                    await callback(.invalidEntitlement)
                 }
             case .failure(let error):
                 os_log("⚫️ Error retrieving entitlement: %{public}@", log: .networkProtectionEntitlementMonitorLog, error.localizedDescription)
-                callback(.error(error))
+                await callback(.error(error))
             }
         }
     }

Sources/NetworkProtection/PacketTunnelProvider.swift

@@ -271,10 +271,20 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             providerEvents.fire(.rekeyAttempt(.success))
         } catch {
             providerEvents.fire(.rekeyAttempt(.failure(error)))
+            await subscriptionAccessErrorHandler(error)
             throw error
         }
     }
 
+    private func subscriptionAccessErrorHandler(_ error: Error) async {
+        switch error {
+        case TunnelError.vpnAccessRevoked:
+            await handleAccessRevoked(attemptsShutdown: true)
+        default:
+            break
+        }
+    }
+
     private func setKeyValidity(_ interval: TimeInterval?) {
         if let interval {
             let firstExpirationDate = Date().addingTimeInterval(interval)
@@ -595,7 +605,18 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         settings.changePublisher
             .receive(on: DispatchQueue.main)
             .sink { [weak self] change in
-                self?.handleSettingsChange(change)
+                guard let self else { return }
+                let handleSettingsChange = handleSettingsChange
+                let subscriptionAccessErrorHandler = subscriptionAccessErrorHandler
+
+                Task { @MainActor in
+                    do {
+                        try await handleSettingsChange(change)
+                    } catch {
+                        await subscriptionAccessErrorHandler(error)
+                        throw error
+                    }
+                }
             }.store(in: &cancellables)
     }
 
@@ -656,7 +677,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             }
 
             os_log("🔴 Stopping VPN due to no auth token: %{public}s", log: .networkProtection)
-            await attemptShutdown()
+            await attemptShutdownDueToRevokedAccess()
 
             throw error
         }
@@ -943,7 +964,6 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             )
         } catch {
             if isSubscriptionEnabled, let error = error as? NetworkProtectionError, case .vpnAccessRevoked = error {
-                await handleInvalidEntitlement(attemptsShutdown: true)
                 throw TunnelError.vpnAccessRevoked
             }
 
@@ -1044,16 +1064,14 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         settings.apply(change: change)
     }
 
-    private func handleSettingsChange(_ change: VPNSettings.Change, completionHandler: ((Data?) -> Void)? = nil) {
+    @MainActor
+    private func handleSettingsChange(_ change: VPNSettings.Change) async throws {
         switch change {
         case .setExcludeLocalNetworks:
-            Task { @MainActor in
-                if case .connected = connectionStatus {
-                    try? await updateTunnelConfiguration(
-                        updateMethod: .selectServer(currentServerSelectionMethod),
-                        reassert: false)
-                }
-                completionHandler?(nil)
+            if case .connected = connectionStatus {
+                try await updateTunnelConfiguration(
+                    updateMethod: .selectServer(currentServerSelectionMethod),
+                    reassert: false)
             }
         case .setSelectedServer(let selectedServer):
             let serverSelectionMethod: NetworkProtectionServerSelectionMethod
@@ -1065,13 +1083,10 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
                 serverSelectionMethod = .preferredServer(serverName: serverName)
             }
 
-            Task { @MainActor in
-                if case .connected = connectionStatus {
-                    try? await updateTunnelConfiguration(
-                        updateMethod: .selectServer(serverSelectionMethod),
-                        reassert: true)
-                }
-                completionHandler?(nil)
+            if case .connected = connectionStatus {
+                try? await updateTunnelConfiguration(
+                    updateMethod: .selectServer(serverSelectionMethod),
+                    reassert: true)
             }
         case .setSelectedLocation(let selectedLocation):
             let serverSelectionMethod: NetworkProtectionServerSelectionMethod
@@ -1083,22 +1098,16 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
                 serverSelectionMethod = .preferredLocation(location)
             }
 
-            Task { @MainActor in
-                if case .connected = connectionStatus {
-                    try? await updateTunnelConfiguration(
-                        updateMethod: .selectServer(serverSelectionMethod),
-                        reassert: true)
-                }
-                completionHandler?(nil)
+            if case .connected = connectionStatus {
+                try? await updateTunnelConfiguration(
+                    updateMethod: .selectServer(serverSelectionMethod),
+                    reassert: true)
             }
         case .setDNSSettings:
-            Task { @MainActor in
-                if case .connected = connectionStatus {
-                    try? await updateTunnelConfiguration(
-                        updateMethod: .selectServer(currentServerSelectionMethod),
-                        reassert: true)
-                }
-                completionHandler?(nil)
+            if case .connected = connectionStatus {
+                try? await updateTunnelConfiguration(
+                    updateMethod: .selectServer(currentServerSelectionMethod),
+                    reassert: true)
             }
         case .setConnectOnLogin,
                 .setIncludeAllNetworks,
@@ -1109,7 +1118,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
                 .setShowInMenuBar,
                 .setDisableRekeying:
             // Intentional no-op, as some setting changes don't require any further operation
-            completionHandler?(nil)
+            break
         }
     }
 
@@ -1124,7 +1133,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             handleSendTestNotification(completionHandler: completionHandler)
         case .disableConnectOnDemandAndShutDown:
             Task { [weak self] in
-                await self?.attemptShutdown()
+                await self?.attemptShutdownDueToRevokedAccess()
                 completionHandler?(nil)
             }
         case .removeVPNConfiguration:
@@ -1245,26 +1254,23 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
     }
 
     @available(iOS 17, *)
-    public func handleShutDown(completionHandler: ((Data?) -> Void)? = nil) {
-        Task { @MainActor in
-            let managers = try await NETunnelProviderManager.loadAllFromPreferences()
-
-            guard let manager = managers.first else {
-                completionHandler?(nil)
-                return
-            }
-
-            os_log("⚪️ Disabling Connect On Demand and shutting down the tunnel", log: .networkProtection)
+    @MainActor
+    public func handleShutDown() async throws {
+        os_log("🔴 Disabling Connect On Demand and shutting down the tunnel", log: .networkProtection)
+        let managers = try await NETunnelProviderManager.loadAllFromPreferences()
 
-            manager.isOnDemandEnabled = false
-            try await manager.saveToPreferences()
-            try await manager.loadFromPreferences()
+        guard let manager = managers.first else {
+            os_log("Could not find a viable manager, bailing out of shutdown", log: .networkProtection)
+            // Doesn't matter a lot what error we throw here, since we'll try cancelling the
+            // tunnel.
+            throw TunnelError.vpnAccessRevoked
+        }
 
-            let error = NSError(domain: "com.duckduckgo.vpn", code: 0)
-            await cancelTunnel(with: error)
+        manager.isOnDemandEnabled = false
+        try await manager.loadFromPreferences()
+        try await manager.saveToPreferences()
 
-            completionHandler?(nil)
-        }
+        await cancelTunnel(with: TunnelError.vpnAccessRevoked)
     }
 
     private func setIncludedRoutes(_ includedRoutes: [IPAddressRange], completionHandler: ((Data?) -> Void)? = nil) {
@@ -1451,9 +1457,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             /// Ignore otherwise
             switch result {
             case .invalidEntitlement:
-                Task { [weak self] in
-                    await self?.handleInvalidEntitlement(attemptsShutdown: true)
-                }
+                await self?.handleAccessRevoked(attemptsShutdown: true)
             case .validEntitlement, .error:
                 break
             }
@@ -1492,7 +1496,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
     }
 
     @MainActor
-    private func handleInvalidEntitlement(attemptsShutdown: Bool) async {
+    private func handleAccessRevoked(attemptsShutdown: Bool) async {
         defaults.enableEntitlementMessaging()
         notificationsPresenter.showEntitlementNotification()
 
@@ -1502,18 +1506,32 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         try? await Task.sleep(interval: .seconds(5))
 
         if attemptsShutdown {
-            await attemptShutdown()
+            await attemptShutdownDueToRevokedAccess()
         }
     }
 
-    // Attempt to shut down the tunnel
-    // On iOS 16 and below, as a workaround, we rekey to force a 403 error so that the tunnel fails to restart
+    /// Tries to shut down the tunnel after access has been revoked.
+    ///
+    /// iOS 17+ supports disabling on-demand, but macOS does not... so we resort to removing the subscription token
+    /// which should prevent the VPN from even trying to start.
+    ///
     @MainActor
-    private func attemptShutdown() async {
+    private func attemptShutdownDueToRevokedAccess() async {
+        let cancelTunnel = {
+#if os(macOS)
+            try? self.tokenStore.deleteToken()
+#endif
+            self.cancelTunnelWithError(TunnelError.vpnAccessRevoked)
+        }
+
         if #available(iOS 17, *) {
-            handleShutDown()
+            do {
+                try await handleShutDown()
+            } catch {
+                cancelTunnel()
+            }
         } else {
-            try? await rekey()
+            cancelTunnel()
         }
     }