Check for entitlement in DBP agent

DuckDuckGo/Application/AppDelegate.swift

@@ -94,7 +94,13 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
 
     private var networkProtectionSubscriptionEventHandler: NetworkProtectionSubscriptionEventHandler?
 #if DBP
-    private var dataBrokerProtectionSubscriptionEventHandler: DataBrokerProtectionSubscriptionEventHandler?
+    private lazy var dataBrokerProtectionSubscriptionEventHandler: DataBrokerProtectionSubscriptionEventHandler = {
+        let authManager = DataBrokerAuthenticationManagerBuilder.buildAuthenticationManager(subscriptionManager: subscriptionManager)
+        return DataBrokerProtectionSubscriptionEventHandler(featureDisabler: DataBrokerProtectionFeatureDisabler(),
+                                                            authenticationManager: authManager,
+                                                            pixelHandler: DataBrokerProtectionPixelsHandler())
+    }()
+
 #endif
 
     private var didFinishLaunching = false
@@ -216,9 +222,6 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
         networkProtectionSubscriptionEventHandler = NetworkProtectionSubscriptionEventHandler(subscriptionManager: subscriptionManager,
                                                                                               tunnelController: tunnelController,
                                                                                               vpnUninstaller: vpnUninstaller)
-#if DBP
-        dataBrokerProtectionSubscriptionEventHandler = DataBrokerProtectionSubscriptionEventHandler(subscriptionManager: subscriptionManager)
-#endif
     }
 
     // swiftlint:disable:next function_body_length
@@ -310,7 +313,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
         UNUserNotificationCenter.current().delegate = self
 
 #if DBP
-        dataBrokerProtectionSubscriptionEventHandler?.registerForSubscriptionAccountManagerEvents()
+        dataBrokerProtectionSubscriptionEventHandler.registerForSubscriptionAccountManagerEvents()
 #endif
 
 #if DBP

DuckDuckGo/DBP/DataBrokerProtectionPixelsHandler.swift

@@ -98,7 +98,11 @@ public class DataBrokerProtectionPixelsHandler: EventMapping<DataBrokerProtectio
                     .homeViewShowWebUI,
                     .homeViewShowBadPathError,
                     .homeViewCTAMoveApplicationClicked,
-                    .homeViewCTAGrantPermissionClicked:
+                    .homeViewCTAGrantPermissionClicked,
+
+                    .entitlementCheckValid,
+                    .entitlementCheckInvalid,
+                    .entitlementCheckError:
                 PixelKit.fire(event, frequency: .dailyAndCount)
             }
         }

DuckDuckGo/DBP/DataBrokerProtectionSubscriptionEventHandler.swift

@@ -21,28 +21,55 @@ import Foundation
 import Subscription
 import DataBrokerProtection
 import PixelKit
+import Common
 
 final class DataBrokerProtectionSubscriptionEventHandler {
 
-    private let subscriptionManager: SubscriptionManaging
-    private let authRepository: AuthenticationRepository
+    private let authenticationManager: DataBrokerProtectionAuthenticationManaging
     private let featureDisabler: DataBrokerProtectionFeatureDisabling
+    private let pixelHandler: EventMapping<DataBrokerProtectionPixels>
 
-    init(subscriptionManager: SubscriptionManaging,
-         authRepository: AuthenticationRepository = KeychainAuthenticationData(),
-         featureDisabler: DataBrokerProtectionFeatureDisabling = DataBrokerProtectionFeatureDisabler()) {
-        self.subscriptionManager = subscriptionManager
-        self.authRepository = authRepository
+    init(featureDisabler: DataBrokerProtectionFeatureDisabling,
+         authenticationManager: DataBrokerProtectionAuthenticationManaging,
+         pixelHandler: EventMapping<DataBrokerProtectionPixels>) {
         self.featureDisabler = featureDisabler
+        self.authenticationManager = authenticationManager
+        self.pixelHandler = pixelHandler
     }
 
     func registerForSubscriptionAccountManagerEvents() {
-        NotificationCenter.default.addObserver(self, selector: #selector(handleAccountDidSignOut), name: .accountDidSignOut, object: nil)
+        NotificationCenter.default.addObserver(self,
+                                               selector: #selector(handleAccountDidSignOut),
+                                               name: .accountDidSignOut,
+                                               object: nil)
+
+        NotificationCenter.default.addObserver(self,
+                                               selector: #selector(entitlementsDidChange),
+                                               name: .entitlementsDidChange,
+                                               object: nil)
     }
 
     @objc private func handleAccountDidSignOut() {
         featureDisabler.disableAndDelete()
     }
+
+    @objc private func entitlementsDidChange() {
+        Task { @MainActor in
+            do {
+                if try await authenticationManager.hasValidEntitlement() {
+                    pixelHandler.fire(.entitlementCheckValid)
+                } else {
+                    pixelHandler.fire(.entitlementCheckInvalid)
+                    featureDisabler.disableAndDelete()
+                }
+            } catch {
+                /// We don't want to disable the agent in case of an error while checking for entitlements.
+                /// Since this is a destructive action, the only situation that should cause the data to be deleted and the agent to be removed is .success(false)
+                pixelHandler.fire(.entitlementCheckError)
+                assertionFailure("Error validating entitlement \(error)")
+            }
+        }
+    }
 }
 
 #endif

LocalPackages/DataBrokerProtection/Sources/DataBrokerProtection/Authentication/DataBrokerProtectionEntitlementMonitoring.swift

@@ -0,0 +1,56 @@
+//
+//  DataBrokerProtectionEntitlementMonitoring.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
+
+protocol DataBrokerProtectionEntitlementMonitoring {
+    func start(checkEntitlementFunction: @escaping () async throws -> Bool, interval: TimeInterval, callback: @escaping (DataBrokerProtectionEntitlementMonitorResult) -> Void)
+    func stop()
+}
+
+public enum DataBrokerProtectionEntitlementMonitorResult {
+    case enabled
+    case disabled
+    case error
+}
+
+final class DataBrokerProtectionEntitlementMonitor: DataBrokerProtectionEntitlementMonitoring {
+    private var timer: Timer?
+
+    func start(checkEntitlementFunction: @escaping () async throws -> Bool, interval: TimeInterval, callback: @escaping (DataBrokerProtectionEntitlementMonitorResult) -> Void) {
+        timer = Timer.scheduledTimer(withTimeInterval: interval, repeats: true) { _ in
+            Task {
+                do {
+                    switch try await checkEntitlementFunction() {
+                    case true:
+                        callback(.enabled)
+                    case false:
+                        callback(.disabled)
+                    }
+                } catch {
+                    callback(.error)
+                }
+            }
+        }
+    }
+
+    func stop() {
+        timer?.invalidate()
+        timer = nil
+    }
+}

LocalPackages/DataBrokerProtection/Sources/DataBrokerProtection/Pixels/DataBrokerProtectionPixels.swift

@@ -164,6 +164,11 @@ public enum DataBrokerProtectionPixels {
     case initialScanSiteLoadDuration(duration: Double, hasError: Bool, brokerURL: String, sleepDuration: Double)
     case initialScanPostLoadingDuration(duration: Double, hasError: Bool, brokerURL: String, sleepDuration: Double)
     case initialScanPreStartDuration(duration: Double)
+
+    // Entitlements
+    case entitlementCheckValid
+    case entitlementCheckInvalid
+    case entitlementCheckError
 }
 
 extension DataBrokerProtectionPixels: PixelKitEvent {
@@ -267,6 +272,11 @@ extension DataBrokerProtectionPixels: PixelKitEvent {
         case .initialScanSiteLoadDuration: return "m_mac_dbp_scan_broker_site_loaded"
         case .initialScanPostLoadingDuration: return "m_mac_dbp_initial_scan_broker_post_loading"
         case .initialScanPreStartDuration: return "m_mac_dbp_initial_scan_pre_start_duration"
+
+            // Entitlements
+        case .entitlementCheckValid: return "m_mac_dbp_macos_entitlement_valid"
+        case .entitlementCheckInvalid: return "m_mac_dbp_macos_entitlement_invalid"
+        case .entitlementCheckError: return "m_mac_dbp_macos_entitlement_error"
         }
     }
 
@@ -361,7 +371,9 @@ extension DataBrokerProtectionPixels: PixelKitEvent {
                 .homeViewShowBadPathError,
                 .homeViewCTAMoveApplicationClicked,
                 .homeViewCTAGrantPermissionClicked,
-
+                .entitlementCheckValid,
+                .entitlementCheckInvalid,
+                .entitlementCheckError,
                 .secureVaultInitError,
                 .secureVaultError:
             return [:]
@@ -486,7 +498,10 @@ public class DataBrokerProtectionPixelsHandler: EventMapping<DataBrokerProtectio
                     .homeViewShowWebUI,
                     .homeViewShowBadPathError,
                     .homeViewCTAMoveApplicationClicked,
-                    .homeViewCTAGrantPermissionClicked:
+                    .homeViewCTAGrantPermissionClicked,
+                    .entitlementCheckValid,
+                    .entitlementCheckInvalid,
+                    .entitlementCheckError:
                 PixelKit.fire(event, frequency: .dailyAndCount)
 
             }

LocalPackages/DataBrokerProtection/Sources/DataBrokerProtection/Scheduler/DataBrokerProtectionAgentManager.swift

@@ -71,7 +71,13 @@ public class DataBrokerProtectionAgentManagerProvider {
                                                          contentScopeProperties: contentScopeProperties,
                                                          emailService: emailService,
                                                          captchaService: captchaService)
-         let operationDependencies = DefaultDataBrokerOperationDependencies(
+
+        let agentstopper = DefaultDataBrokerProtectionAgentStopper(dataManager: dataManager,
+                                                                   entitlementMonitor: DataBrokerProtectionEntitlementMonitor(),
+                                                                   authenticationManager: authenticationManager,
+                                                                   pixelHandler: pixelHandler)
+
+        let operationDependencies = DefaultDataBrokerOperationDependencies(
             database: dataManager.database,
             config: executionConfig,
             runnerProvider: runnerProvider,
@@ -86,7 +92,8 @@ public class DataBrokerProtectionAgentManagerProvider {
             queueManager: queueManager,
             dataManager: dataManager,
             operationDependencies: operationDependencies,
-            pixelHandler: pixelHandler)
+            pixelHandler: pixelHandler,
+            agentStopper: agentstopper)
     }
 }
 
@@ -99,6 +106,7 @@ public final class DataBrokerProtectionAgentManager {
     private let dataManager: DataBrokerProtectionDataManaging
     private let operationDependencies: DataBrokerOperationDependencies
     private let pixelHandler: EventMapping<DataBrokerProtectionPixels>
+    private let agentStopper: DataBrokerProtectionAgentStopper
 
     // Used for debug functions only, so not injected
     private lazy var browserWindowManager = BrowserWindowManager()
@@ -111,14 +119,17 @@ public final class DataBrokerProtectionAgentManager {
          queueManager: DataBrokerProtectionQueueManager,
          dataManager: DataBrokerProtectionDataManaging,
          operationDependencies: DataBrokerOperationDependencies,
-         pixelHandler: EventMapping<DataBrokerProtectionPixels>) {
+         pixelHandler: EventMapping<DataBrokerProtectionPixels>,
+         agentStopper: DataBrokerProtectionAgentStopper
+    ) {
         self.userNotificationService = userNotificationService
         self.activityScheduler = activityScheduler
         self.ipcServer = ipcServer
         self.queueManager = queueManager
         self.dataManager = dataManager
         self.operationDependencies = operationDependencies
         self.pixelHandler = pixelHandler
+        self.agentStopper = agentStopper
 
         self.activityScheduler.delegate = self
         self.ipcServer.serverDelegate = self
@@ -127,20 +138,20 @@ public final class DataBrokerProtectionAgentManager {
 
     public func agentFinishedLaunching() {
 
-        do {
-            // If there's no saved profile we don't need to start the scheduler
-            // Theoretically this should never happen, if there's no data, the agent shouldn't be running
-            guard (try dataManager.fetchProfile()) != nil else {
-                return
-            }
-        } catch {
-            os_log("Error during AgentManager.agentFinishedLaunching when trying to fetchProfile, error: %{public}@", log: .dataBrokerProtection, error.localizedDescription)
-            return
+        Task { @MainActor in
+            // The browser shouldn't start the agent if these prerequisites aren't met.
+            // However, since the agent can auto-start after a reboot without the browser, we need to validate it again.
+            // If the agent needs to be stopped, this function will stop it, so the subsequent calls after it will not be made.
+            await agentStopper.validateRunPrerequisitesAndStopAgentIfNecessary()
+
+            activityScheduler.startScheduler()
+            didStartActivityScheduler = true
+            queueManager.startScheduledOperationsIfPermitted(showWebView: false, operationDependencies: operationDependencies, completion: nil)
+
+            /// Monitors entitlement changes every 60 minutes to optimize system performance and resource utilization by avoiding unnecessary operations when entitlement is invalid.
+            /// While keeping the agent active with invalid entitlement has no significant risk, setting the monitoring interval at 60 minutes is a good balance to minimize backend checks.
+            agentStopper.monitorEntitlementAndStopAgentIfEntitlementIsInvalid(interval: .minutes(60))
         }
-
-        activityScheduler.startScheduler()
-        didStartActivityScheduler = true
-        queueManager.startScheduledOperationsIfPermitted(showWebView: false, operationDependencies: operationDependencies, completion: nil)
     }
 }