add experiment logic

Sources/BrowserServicesKit/PrivacyConfig/AppPrivacyConfiguration.swift

@@ -33,19 +33,26 @@ public struct AppPrivacyConfiguration: PrivacyConfiguration {
     private let locallyUnprotected: DomainsProtectionStore
     private let internalUserDecider: InternalUserDecider
     private let userDefaults: UserDefaults
+    private let locale: Locale
+    private let experimentManager: ExperimentCohortsManaging
     private let installDate: Date?
+    static let experimentManagerQueue = DispatchQueue(label: "com.experimentManager.queue")
 
     public init(data: PrivacyConfigurationData,
                 identifier: String,
                 localProtection: DomainsProtectionStore,
                 internalUserDecider: InternalUserDecider,
                 userDefaults: UserDefaults = UserDefaults(),
+                locale: Locale = Locale.current,
+                experimentManager: ExperimentCohortsManaging = ExperimentCohortsManager(),
                 installDate: Date? = nil) {
         self.data = data
         self.identifier = identifier
         self.locallyUnprotected = localProtection
         self.internalUserDecider = internalUserDecider
         self.userDefaults = userDefaults
+        self.locale = locale
+        self.experimentManager = experimentManager
         self.installDate = installDate
     }
 
@@ -137,13 +144,14 @@ public struct AppPrivacyConfiguration: PrivacyConfiguration {
         }
     }
 
-    private func isRolloutEnabled(subfeature: any PrivacySubfeature,
+    private func isRolloutEnabled(subfeatureID: SubfeatureID,
+                                  parentID: ParentFeatureID,
                                   rolloutSteps: [PrivacyConfigurationData.PrivacyFeature.Feature.RolloutStep],
                                   randomizer: (Range<Double>) -> Double) -> Bool {
         // Empty rollouts should be default enabled
         guard !rolloutSteps.isEmpty else { return true }
 
-        let defsPrefix = "config.\(subfeature.parent.rawValue).\(subfeature.rawValue)"
+        let defsPrefix = "config.\(parentID).\(subfeatureID)"
         if userDefaults.bool(forKey: "\(defsPrefix).\(Constants.enabledKey)") {
             return true
         }
@@ -180,28 +188,76 @@ public struct AppPrivacyConfiguration: PrivacyConfiguration {
     }
 
     public func isSubfeatureEnabled(_ subfeature: any PrivacySubfeature,
+                                    cohortID: CohortID?,
                                     versionProvider: AppVersionProvider,
                                     randomizer: (Range<Double>) -> Double) -> Bool {
-
-        switch stateFor(subfeature, versionProvider: versionProvider, randomizer: randomizer) {
+        switch stateFor(subfeature, cohortID: cohortID, versionProvider: versionProvider, randomizer: randomizer) {
         case .enabled:
             return true
         case .disabled:
             return false
         }
     }
 
-    public func stateFor(_ subfeature: any PrivacySubfeature, versionProvider: AppVersionProvider, randomizer: (Range<Double>) -> Double) -> PrivacyConfigurationFeatureState {
+    public func getAllActiveExperiments(versionProvider: AppVersionProvider,
+                                        randomizer: (Range<Double>) -> Double) -> Experiments {
+        Self.experimentManagerQueue.sync {
+               guard let assignedExperiments = experimentManager.experiments else { return [:] }
+               return assignedExperiments.filter { key, value in
+                   stateFor(subfeatureID: key, experimentData: value, versionProvider: versionProvider, randomizer: randomizer) == .enabled
+               }
+           }
+    }
 
-        let parentState = stateFor(featureKey: subfeature.parent, versionProvider: versionProvider)
-        guard case .enabled = parentState else { return parentState }
+    private func stateFor(subfeatureID: SubfeatureID, experimentData: ExperimentData, versionProvider: AppVersionProvider,
+                          randomizer: (Range<Double>) -> Double) -> PrivacyConfigurationFeatureState {
+        guard let parentFeature = PrivacyFeature(rawValue: experimentData.parentID),
+              let subfeatureData = subfeatures(for: parentFeature)[subfeatureID] else {
+            return .disabled(.featureMissing)
+        }
+        return stateFor(parentFeature: parentFeature,
+                        subfeatureData: subfeatureData,
+                        subfeatureID: subfeatureID,
+                        cohortID: experimentData.cohort,
+                        assignCohortEnabled: false,
+                        versionProvider: versionProvider,
+                        randomizer: randomizer)
+    }
+
+    public func stateFor(_ subfeature: any PrivacySubfeature,
+                         cohortID: CohortID?,
+                         versionProvider: AppVersionProvider,
+                         randomizer: (Range<Double>) -> Double) -> PrivacyConfigurationFeatureState {
+        Self.experimentManagerQueue.sync {
+            guard let subfeatureData = subfeatures(for: subfeature.parent)[subfeature.rawValue] else {
+                return .disabled(.featureMissing)
+            }
+
+            return stateFor(parentFeature: subfeature.parent,
+                            subfeatureData: subfeatureData,
+                            subfeatureID: subfeature.rawValue,
+                            cohortID: cohortID,
+                            versionProvider: versionProvider,
+                            randomizer: randomizer)
+        }
+    }
 
-        let subfeatures = subfeatures(for: subfeature.parent)
-        let subfeatureData = subfeatures[subfeature.rawValue]
+    private func stateFor(parentFeature: PrivacyFeature,
+                          subfeatureData: PrivacyConfigurationData.PrivacyFeature.Feature,
+                          subfeatureID: SubfeatureID,
+                          cohortID: CohortID?,
+                          assignCohortEnabled: Bool = true,
+                          versionProvider: AppVersionProvider,
+                          randomizer: (Range<Double>) -> Double) -> PrivacyConfigurationFeatureState {
+        // Step 1: Check parent feature state
+        let parentState = stateFor(featureKey: parentFeature, versionProvider: versionProvider)
+        guard case .enabled = parentState else { return parentState }
 
-        let satisfiesMinVersion = satisfiesMinVersion(subfeatureData?.minSupportedVersion, versionProvider: versionProvider)
+        // Step 2: Check  version
+        let satisfiesMinVersion = satisfiesMinVersion(subfeatureData.minSupportedVersion, versionProvider: versionProvider)
 
-        switch subfeatureData?.state {
+        // Step 3: Check sub-feature state
+        switch subfeatureData.state {
         case PrivacyConfigurationData.State.enabled:
             guard satisfiesMinVersion else { return .disabled(.appVersionNotSupported) }
         case PrivacyConfigurationData.State.internal:
@@ -210,15 +266,66 @@ public struct AppPrivacyConfiguration: PrivacyConfiguration {
         default: return .disabled(.disabledInConfig)
         }
 
-        // Handle Rollouts
-        if let rollout = subfeatureData?.rollout,
-           !isRolloutEnabled(subfeature: subfeature, rolloutSteps: rollout.steps, randomizer: randomizer) {
+        // Step 4: Handle Rollouts
+        if let rollout = subfeatureData.rollout,
+           !isRolloutEnabled(subfeatureID: subfeatureID, parentID: parentFeature.rawValue, rolloutSteps: rollout.steps, randomizer: randomizer) {
             return .disabled(.stillInRollout)
         }
 
+        // Step 5: Check if a cohort was passed in the func
+        // If no corhort passed check for Target and Rollout
+        guard let passedCohort = cohortID else {
+            return checkTargets(subfeatureData)
+        }
+
+        // Step 5: Cohort handling
+        // Check if cohort assigned and matches passed cohort
+        // If cohort not assigned
+        // Tries to assign if matching target
+        // Check if cohort assigned and matches passed cohort
+        return checkCohortState(subfeatureData,
+                         passedCohort: passedCohort,
+                         assignCohortEnabled: assignCohortEnabled,
+                         subfeatureID: subfeatureID,
+                         parentFeature: parentFeature)
+    }
+
+    // Check if cohort assigned and matches passed cohort
+    // If cohort not assigned
+    // Tries to assign if matching target
+    // Check if cohort assigned and matches passed cohort
+    private func checkCohortState(_ subfeatureData: PrivacyConfigurationData.PrivacyFeature.Feature,
+                                  passedCohort: CohortID?,
+                                  assignCohortEnabled: Bool,
+                                  subfeatureID: SubfeatureID,
+                                  parentFeature: PrivacyFeature) -> PrivacyConfigurationFeatureState {
+        let cohorts = subfeatureData.cohorts ?? []
+        let targetsState = checkTargets(subfeatureData)
+        let assignIfEnabled = assignCohortEnabled && targetsState == .enabled
+        let assignedCohortResponse = experimentManager.cohort(for: ExperimentSubfeature(parentID: parentFeature.rawValue, subfeatureID: subfeatureID, cohorts: cohorts), assignIfEnabled: assignIfEnabled)
+        let possibleDisabledReason: PrivacyConfigurationFeatureDisabledReason = assignedCohortResponse.didAttemptAssignment &&  targetsState != .enabled ? .targetDoesNotMatch : .experimentCohortDoesNotMatch
+        if let assignedCohort = assignedCohortResponse.cohortID {
+            return (assignedCohort == passedCohort) ? .enabled : .disabled(possibleDisabledReason)
+        } else {
+            return .disabled(possibleDisabledReason)
+        }
+    }
+
+    private func checkTargets(_ subfeatureData: PrivacyConfigurationData.PrivacyFeature.Feature?) -> PrivacyConfigurationFeatureState {
+        // Check Targets
+        if let targets = subfeatureData?.targets, !matchTargets(targets: targets){
+            return .disabled(.targetDoesNotMatch)
+        }
         return .enabled
     }
 
+    private func matchTargets(targets: [PrivacyConfigurationData.PrivacyFeature.Feature.Target]) -> Bool {
+        return targets.contains { target in
+            (target.localeCountry == nil || target.localeCountry == locale.regionCode) &&
+            (target.localeLanguage == nil || target.localeLanguage == locale.languageCode)
+        }
+    }
+
     private func subfeatures(for feature: PrivacyFeature) -> PrivacyConfigurationData.PrivacyFeature.Features {
         return data.features[feature.rawValue]?.features ?? [:]
     }

Sources/BrowserServicesKit/PrivacyConfig/ExperimentCohortsManager.swift

@@ -18,63 +18,78 @@
 
 import Foundation
 
-struct ExperimentSubfeature {
+public struct ExperimentSubfeature {
+    let parentID: ParentFeatureID
     let subfeatureID: SubfeatureID
     let cohorts: [PrivacyConfigurationData.Cohort]
 }
 
-typealias CohortID = String
-typealias SubfeatureID = String
+public typealias CohortID = String
+public typealias SubfeatureID = String
+public typealias ParentFeatureID = String
 
-struct ExperimentData: Codable, Equatable {
-    let cohort: String
-    let enrollmentDate: Date
+public struct ExperimentData: Codable, Equatable {
+    public let parentID: String
+    public let cohort: String
+    public let enrollmentDate: Date
 }
 
-typealias Experiments = [String: ExperimentData]
-
-protocol ExperimentCohortsManaging {
-    /// Retrieves the cohort ID associated with the specified subfeature.
-    /// - Parameter subfeatureID: The name of the experiment subfeature for which the cohort ID is needed.
-    /// - Returns: The cohort ID as a `String` if one exists; otherwise, returns `nil`.
-    func cohort(for subfeatureID: SubfeatureID) -> CohortID?
-
-    /// Retrieves the enrollment date for the specified subfeature.
-    /// - Parameter subfeatureID: The name of the experiment subfeature for which the enrollment date is needed.
-    /// - Returns: The `Date` of enrollment if one exists; otherwise, returns `nil`.
-    func enrollmentDate(for subfeatureID: SubfeatureID) -> Date?
-
-    /// Assigns a cohort to the given subfeature based on defined weights and saves it to UserDefaults.
-    /// - Parameter subfeature: The ExperimentSubfeature to which a cohort needs to be assigned to.
-    /// - Returns: The name of the assigned cohort, or `nil` if no cohort could be assigned.
-    func assignCohort(to subfeature: ExperimentSubfeature) -> CohortID?
-
-    /// Removes the assigned cohort data for the specified subfeature.
-    /// - Parameter subfeatureID: The name of the experiment subfeature for which the cohort data should be removed.
-    func removeCohort(from subfeatureID: SubfeatureID)
+public typealias Experiments = [String: ExperimentData]
+
+public protocol ExperimentCohortsManaging {
+    /// Retrieves all the experiments a user is enrolled into
+    var experiments: Experiments? { get }
+
+    /// Retrieves the assigned cohort for a given experiment subfeature, or attempts to assign a new cohort if none is currently assigned
+    /// and `assignIfEnabled` is set to true. If a cohort is already assigned but does not match any valid cohorts for the experiment,
+    /// the cohort will be removed.
+    ///
+    /// - Parameters:
+    ///   - experiment: The `ExperimentSubfeature` for which to retrieve, assign, or remove a cohort. This subfeature includes
+    ///     relevant identifiers and potential cohorts that may be assigned.
+    ///   - assignIfEnabled: A Boolean value that determines whether a new cohort should be assigned if none is currently assigned.
+    ///     If `true`, the function will attempt to assign a cohort from the available options; otherwise, it will only check for existing assignments.
+    ///
+    /// - Returns: A tuple containing:
+    ///   - `cohortID`: The identifier of the assigned cohort if one exists, or `nil` if no cohort was assigned, if assignment failed, or if the cohort was removed.
+    ///   - `didAttemptAssignment`: A Boolean indicating whether an assignment attempt was made. This will be `true` if `assignIfEnabled`
+    ///     is `true` and no cohort was previously assigned, and `false` otherwise.
+    func cohort(for experiment: ExperimentSubfeature, assignIfEnabled: Bool) -> (cohortID: CohortID?, didAttemptAssignment: Bool)
 }
 
-final class ExperimentCohortsManager: ExperimentCohortsManaging {
+public class ExperimentCohortsManager: ExperimentCohortsManaging {
 
     private var store: ExperimentsDataStoring
     private let randomizer: (Range<Double>) -> Double
+    private let queue = DispatchQueue(label: "com.ExperimentCohortsManager.queue")
 
-    init(store: ExperimentsDataStoring = ExperimentsDataStore(), randomizer: @escaping (Range<Double>) -> Double) {
+    public var experiments: Experiments? {
+        get {
+            queue.sync {
+                store.experiments
+            }
+        }
+    }
+
+    public init(store: ExperimentsDataStoring = ExperimentsDataStore(), randomizer: @escaping (Range<Double>) -> Double = Double.random(in:)) {
         self.store = store
         self.randomizer = randomizer
     }
 
-    func cohort(for subfeatureID: SubfeatureID) -> CohortID? {
-        guard let experiments = store.experiments else { return nil }
-        return experiments[subfeatureID]?.cohort
-    }
+    public func cohort(for experiment: ExperimentSubfeature, assignIfEnabled: Bool) -> (cohortID: CohortID?, didAttemptAssignment: Bool) {
+        queue.sync {
+            let assignedCohort = cohort(for: experiment.subfeatureID)
+            if experiment.cohorts.contains(where: { $0.name == assignedCohort }) {
+                return (assignedCohort, false)
+            } else {
+                removeCohort(from: experiment.subfeatureID)
+            }
 
-    func enrollmentDate(for subfeatureID: SubfeatureID) -> Date? {
-        guard let experiments = store.experiments else { return nil }
-        return experiments[subfeatureID]?.enrollmentDate
+            return assignIfEnabled ? (assignCohort(to: experiment), true) : (nil, true)
+        }
     }
 
-    func assignCohort(to subfeature: ExperimentSubfeature) -> CohortID? {
+    private func assignCohort(to subfeature: ExperimentSubfeature) -> CohortID? {
         let cohorts = subfeature.cohorts
         let totalWeight = cohorts.map(\.weight).reduce(0, +)
         guard totalWeight > 0 else { return nil }
@@ -85,23 +100,34 @@ final class ExperimentCohortsManager: ExperimentCohortsManaging {
         for cohort in cohorts {
             cumulativeWeight += Double(cohort.weight)
             if randomValue < cumulativeWeight {
-                saveCohort(cohort.name, in: subfeature.subfeatureID)
+                saveCohort(cohort.name, in: subfeature.subfeatureID, parentID: subfeature.parentID)
                 return cohort.name
             }
         }
         return nil
     }
 
-    func removeCohort(from subfeatureID: SubfeatureID) {
+    func cohort(for subfeatureID: SubfeatureID) -> CohortID? {
+        guard let experiments = store.experiments else { return nil }
+        return experiments[subfeatureID]?.cohort
+    }
+
+    private  func enrollmentDate(for subfeatureID: SubfeatureID) -> Date? {
+        guard let experiments = store.experiments else { return nil }
+        return experiments[subfeatureID]?.enrollmentDate
+    }
+
+    private func removeCohort(from subfeatureID: SubfeatureID) {
         guard var experiments = store.experiments else { return }
         experiments.removeValue(forKey: subfeatureID)
         store.experiments = experiments
     }
 
-    private func saveCohort(_ cohort: CohortID, in experimentID: SubfeatureID) {
+    private func saveCohort(_ cohort: CohortID, in experimentID: SubfeatureID, parentID: ParentFeatureID) {
         var experiments = store.experiments ?? Experiments()
-        let experimentData = ExperimentData(cohort: cohort, enrollmentDate: Date())
+        let experimentData = ExperimentData(parentID: parentID, cohort: cohort, enrollmentDate: Date())
         experiments[experimentID] = experimentData
         store.experiments = experiments
     }
+
 }

Sources/BrowserServicesKit/PrivacyConfig/ExperimentsDataStore.swift

@@ -18,16 +18,16 @@
 
 import Foundation
 
-protocol ExperimentsDataStoring {
+public protocol ExperimentsDataStoring {
     var experiments: Experiments? { get set }
 }
 
-protocol LocalDataStoring {
+public protocol LocalDataStoring {
     func data(forKey defaultName: String) -> Data?
     func set(_ value: Any?, forKey defaultName: String)
 }
 
-struct ExperimentsDataStore: ExperimentsDataStoring {
+public struct ExperimentsDataStore: ExperimentsDataStoring {
 
     private enum Constants {
         static let experimentsDataKey = "ExperimentsData"
@@ -36,13 +36,13 @@ struct ExperimentsDataStore: ExperimentsDataStoring {
     private let decoder = JSONDecoder()
     private let encoder = JSONEncoder()
 
-    init(localDataStoring: LocalDataStoring = UserDefaults.standard) {
+    public init(localDataStoring: LocalDataStoring = UserDefaults.standard) {
         self.localDataStoring = localDataStoring
         encoder.dateEncodingStrategy = .secondsSince1970
         decoder.dateDecodingStrategy = .secondsSince1970
     }
 
-    var experiments: Experiments? {
+    public var experiments: Experiments? {
         get {
             guard let savedData = localDataStoring.data(forKey: Constants.experimentsDataKey) else { return nil }
             return try? decoder.decode(Experiments.self, from: savedData)

Sources/BrowserServicesKit/PrivacyConfig/Features/PrivacyFeature.swift

@@ -19,7 +19,7 @@
 import Foundation
 
 /// Features whose `rawValue` should be the key to access their corresponding `PrivacyConfigurationData.PrivacyFeature` object
-public enum PrivacyFeature: String {
+public enum PrivacyFeature: String, CaseIterable {
     case contentBlocking
     case duckPlayer
     case fingerprintingTemporaryStorage