Check for entitlement in DBP agent

DuckDuckGo/Application/AppDelegate.swift

@@ -99,7 +99,12 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
     }()
 
 #if DBP
-    private let dataBrokerProtectionSubscriptionEventHandler = DataBrokerProtectionSubscriptionEventHandler()
+    private lazy var dataBrokerProtectionSubscriptionEventHandler: DataBrokerProtectionSubscriptionEventHandler = {
+        let authManager = DataBrokerAuthenticationManagerBuilder.buildAuthenticationManager()
+        return DataBrokerProtectionSubscriptionEventHandler(featureDisabler: DataBrokerProtectionFeatureDisabler(),
+                                                            authenticationManager: authManager,
+                                                            pixelHandler: DataBrokerProtectionPixelsHandler())
+    }()
 #endif
 
     private var didFinishLaunching = false

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,21 +21,54 @@ import Foundation
 import Subscription
 import DataBrokerProtection
 import PixelKit
+import Common
 
 final class DataBrokerProtectionSubscriptionEventHandler {
     private let featureDisabler: DataBrokerProtectionFeatureDisabling
+    private let authenticationManager: DataBrokerProtectionAuthenticationManaging
+    private let pixelHandler: EventMapping<DataBrokerProtectionPixels>
 
-    init(featureDisabler: DataBrokerProtectionFeatureDisabling = DataBrokerProtectionFeatureDisabler()) {
+    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,36 +119,40 @@ 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
         self.ipcServer.activate()
     }
 
-    public func agentFinishedLaunching() {
+    public func agentFinishedLaunching(completion: (() -> Void)? = nil) {
 
-        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))
+            completion?()
         }
-
-        activityScheduler.startScheduler()
-        didStartActivityScheduler = true
-        queueManager.startScheduledOperationsIfPermitted(showWebView: false, operationDependencies: operationDependencies, completion: nil)
     }
 }
 

LocalPackages/DataBrokerProtection/Sources/DataBrokerProtection/Utils/DataBrokerProtectionAgentStopper.swift

@@ -0,0 +1,110 @@
+//
+//  DataBrokerProtectionAgentStopper.swift
+//
+//  Copyright © 2023 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 Common
+
+protocol DataBrokerProtectionAgentStopper {
+    /// Validates if the user has profile data, is authenticated, and has valid entitlement. If any of these conditions are not met, the agent will be stopped.
+    func validateRunPrerequisitesAndStopAgentIfNecessary() async
+
+    /// Monitors the entitlement package. If the entitlement check returns false, the agent will be stopped.
+    /// This function ensures that the agent is stopped if the user's subscription has expired, even if the browser is not active. Regularly checking for entitlement is required since notifications are not posted to agents.
+    func monitorEntitlementAndStopAgentIfEntitlementIsInvalid(interval: TimeInterval)
+}
+
+struct DefaultDataBrokerProtectionAgentStopper: DataBrokerProtectionAgentStopper {
+    private let dataManager: DataBrokerProtectionDataManaging
+    private let entitlementMonitor: DataBrokerProtectionEntitlementMonitoring
+    private let authenticationManager: DataBrokerProtectionAuthenticationManaging
+    private let pixelHandler: EventMapping<DataBrokerProtectionPixels>
+    private let stopAction: DataProtectionStopAction
+
+    init(dataManager: DataBrokerProtectionDataManaging,
+         entitlementMonitor: DataBrokerProtectionEntitlementMonitoring,
+         authenticationManager: DataBrokerProtectionAuthenticationManaging,
+         pixelHandler: EventMapping<DataBrokerProtectionPixels>,
+         stopAction: DataProtectionStopAction = DefaultDataProtectionStopAction()) {
+        self.dataManager = dataManager
+        self.entitlementMonitor = entitlementMonitor
+        self.authenticationManager = authenticationManager
+        self.pixelHandler = pixelHandler
+        self.stopAction = stopAction
+    }
+
+    public func validateRunPrerequisitesAndStopAgentIfNecessary() async {
+        do {
+            guard try dataManager.fetchProfile() != nil,
+                  authenticationManager.isUserAuthenticated else {
+                os_log("Prerequisites are invalid", log: .dataBrokerProtection)
+                stopAgent()
+                return
+            }
+            os_log("Prerequisites are valid", log: .dataBrokerProtection)
+        } catch {
+            os_log("Error validating prerequisites, error: %{public}@", log: .dataBrokerProtection, error.localizedDescription)
+            stopAgent()
+        }
+
+        do {
+            let result = try await authenticationManager.hasValidEntitlement()
+            stopAgentBasedOnEntitlementCheckResult(result ? .enabled : .disabled)
+        } catch {
+            stopAgentBasedOnEntitlementCheckResult(.error)
+        }
+    }
+
+    public func monitorEntitlementAndStopAgentIfEntitlementIsInvalid(interval: TimeInterval) {
+        entitlementMonitor.start(checkEntitlementFunction: authenticationManager.hasValidEntitlement,
+                                 interval: interval) { result in
+            stopAgentBasedOnEntitlementCheckResult(result)
+        }
+    }
+
+    private func stopAgent() {
+        stopAction.stopAgent()
+    }
+
+    private func stopAgentBasedOnEntitlementCheckResult(_ result: DataBrokerProtectionEntitlementMonitorResult) {
+        switch result {
+        case .enabled:
+            os_log("Valid entitlement", log: .dataBrokerProtection)
+            pixelHandler.fire(.entitlementCheckValid)
+        case .disabled:
+            os_log("Invalid entitlement", log: .dataBrokerProtection)
+            pixelHandler.fire(.entitlementCheckInvalid)
+            stopAgent()
+        case .error:
+            os_log("Error when checking entitlement", log: .dataBrokerProtection)
+            /// 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)
+        }
+    }
+}
+
+protocol DataProtectionStopAction {
+    func stopAgent()
+}
+
+struct DefaultDataProtectionStopAction: DataProtectionStopAction {
+    func stopAgent() {
+        os_log("Stopping DataBrokerProtection Agent", log: .dataBrokerProtection)
+        exit(EXIT_SUCCESS)
+    }
+}