BSK support for FB CTL update

Package.resolved

@@ -86,8 +86,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/duckduckgo/TrackerRadarKit",
       "state" : {
-        "revision" : "c01e6a59d000356b58ec77053e0a99d538be56a5",
-        "version" : "2.1.1"
+        "revision" : "1403e17eeeb8493b92fb9d11eb8c846bb9776581",
+        "version" : "2.1.2"
       }
     },
     {

Package.swift

@@ -40,7 +40,7 @@ let package = Package(
     dependencies: [
         .package(url: "https://github.com/duckduckgo/duckduckgo-autofill.git", exact: "11.0.2"),
         .package(url: "https://github.com/duckduckgo/GRDB.swift.git", exact: "2.3.0"),
-        .package(url: "https://github.com/duckduckgo/TrackerRadarKit", exact: "2.1.1"),
+        .package(url: "https://github.com/duckduckgo/TrackerRadarKit", exact: "2.1.2"),
         .package(url: "https://github.com/duckduckgo/sync_crypto", exact: "0.2.0"),
         .package(url: "https://github.com/gumob/PunycodeSwift.git", exact: "2.1.0"),
         .package(url: "https://github.com/duckduckgo/content-scope-scripts", exact: "5.15.0"),

Sources/BrowserServicesKit/ContentBlocking/ClickToLoadRulesSplitter.swift

@@ -0,0 +1,158 @@
+//
+//  ClickToLoadRulesSplitter.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 TrackerRadarKit
+
+public struct ClickToLoadRulesSplitter {
+
+    public enum Constants {
+
+        public static let clickToLoadRuleListPrefix = "CTL_"
+        public static let tdsRuleListPrefix = "TDS_"
+
+    }
+
+    private let rulesList: ContentBlockerRulesList
+
+    public init(rulesList: ContentBlockerRulesList) {
+        self.rulesList = rulesList
+    }
+
+    public func split() -> (withoutBlockCTL: ContentBlockerRulesList, withBlockCTL: ContentBlockerRulesList)? {
+        var withoutBlockCTL: (tds: TrackerData, etag: String)?
+        var withBlockCTL: (tds: TrackerData, etag: String)?
+
+        if let trackerData = rulesList.trackerData {
+            let splitTDS = split(trackerData: trackerData)
+
+            withoutBlockCTL = splitTDS?.withoutBlockCTL
+            withBlockCTL = splitTDS?.withBlockCTL
+        }
+
+        return (
+            ContentBlockerRulesList(name: rulesList.name,
+                                    trackerData: withoutBlockCTL,
+                                    fallbackTrackerData: split(trackerData: rulesList.fallbackTrackerData)!.withoutBlockCTL),
+            ContentBlockerRulesList(name: DefaultContentBlockerRulesListsSource.Constants.clickToLoadRulesListName,
+                                    trackerData: withBlockCTL,
+                                    fallbackTrackerData: split(trackerData: rulesList.fallbackTrackerData)!.withBlockCTL)
+        )
+    }
+
+    private func split(trackerData: TrackerDataManager.DataSet) -> (withoutBlockCTL: TrackerDataManager.DataSet, withBlockCTL: TrackerDataManager.DataSet)? {
+        let (mainTrackers, ctlTrackers) = processCTLActions(trackerData.tds.trackers)
+        guard !ctlTrackers.isEmpty else { return nil }
+
+        let trackerDataWithoutBlockCTL = makeTrackerData(using: mainTrackers, originalTDS: trackerData.tds)
+        let trackerDataWithBlockCTL = makeTrackerData(using: ctlTrackers, originalTDS: trackerData.tds)
+
+        return (
+           (tds: trackerDataWithoutBlockCTL, etag: Constants.tdsRuleListPrefix + trackerData.etag),
+           (tds: trackerDataWithBlockCTL, etag: Constants.clickToLoadRuleListPrefix + trackerData.etag)
+        )
+    }
+
+    private func makeTrackerData(using trackers: [String: KnownTracker], originalTDS: TrackerData) -> TrackerData {
+        let entities = originalTDS.extractEntities(for: trackers)
+        let domains = extractDomains(from: entities)
+        return TrackerData(trackers: trackers,
+                           entities: entities,
+                           domains: domains,
+                           cnames: originalTDS.cnames)
+    }
+
+    private func processCTLActions(_ trackers: [String: KnownTracker]) -> (mainTrackers: [String: KnownTracker], ctlTrackers: [String: KnownTracker]) {
+        var mainTDSTrackers: [String: KnownTracker] = [:]
+        var ctlTrackers: [String: KnownTracker] = [:]
+
+        for (key, tracker) in trackers {
+            guard tracker.containsCTLActions, let rules = tracker.rules else {
+                mainTDSTrackers[key] = tracker
+                continue
+            }
+
+            // if we found some CTL rules, split out into its own list
+            var mainRules: [KnownTracker.Rule] = []
+            var ctlRules: [KnownTracker.Rule] = []
+
+            for rule in rules.reversed() {
+                if let action = rule.action, action == .blockCTLFB {
+                    ctlRules.insert(rule, at: 0)
+                } else {
+                    ctlRules.insert(rule, at: 0)
+                    mainRules.insert(rule, at: 0)
+                }
+            }
+
+            let mainTracker = KnownTracker(domain: tracker.domain,
+                                           defaultAction: tracker.defaultAction,
+                                           owner: tracker.owner,
+                                           prevalence: tracker.prevalence,
+                                           subdomains: tracker.subdomains,
+                                           categories: tracker.categories,
+                                           rules: mainRules)
+            let ctlTracker = KnownTracker(domain: tracker.domain,
+                                          defaultAction: tracker.defaultAction,
+                                          owner: tracker.owner,
+                                          prevalence: tracker.prevalence,
+                                          subdomains: tracker.subdomains,
+                                          categories: tracker.categories,
+                                          rules: ctlRules)
+            mainTDSTrackers[key] = mainTracker
+            ctlTrackers[key] = ctlTracker
+        }
+
+        return (mainTDSTrackers, ctlTrackers)
+    }
+
+    private func extractDomains(from entities: [String: Entity]) -> [String: String] {
+        var domains = [String: String]()
+        for (key, entity) in entities {
+            for domain in entity.domains ?? [] {
+                domains[domain] = key
+            }
+        }
+        return domains
+    }
+
+}
+
+private extension TrackerData {
+
+    func extractEntities(for trackers: [String: KnownTracker]) -> [String: Entity] {
+        let trackerOwners = Set(trackers.values.compactMap { $0.owner?.name })
+        let entities = entities.filter { trackerOwners.contains($0.key) }
+        return entities
+    }
+
+}
+
+private extension KnownTracker {
+
+    var containsCTLActions: Bool {
+        if let rules = rules {
+            for rule in rules {
+                if let action = rule.action, action == .blockCTLFB {
+                    return true
+                }
+            }
+        }
+        return false
+    }
+
+}

Sources/BrowserServicesKit/ContentBlocking/ContentBlockerRulesManager.swift

@@ -335,8 +335,7 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
         }
     }
 
-    static func extractSurrogates(from tds: TrackerData) -> TrackerData {
-
+    public static func extractSurrogates(from tds: TrackerData) -> TrackerData {
         let trackers = tds.trackers.filter { pair in
             return pair.value.rules?.first(where: { rule in
                 rule.surrogate != nil
@@ -363,7 +362,6 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
     }
 
     private func compilationCompleted() {
-
         var changes = [String: ContentBlockerRulesIdentifier.Difference]()
 
         lock.lock()

Sources/BrowserServicesKit/ContentBlocking/ContentBlockerRulesSource.swift

@@ -69,6 +69,7 @@ open class DefaultContentBlockerRulesListsSource: ContentBlockerRulesListsSource
 
     public struct Constants {
         public static let trackerDataSetRulesListName = "TrackerDataSet"
+        public static let clickToLoadRulesListName = "ClickToLoad"
     }
 
     private let trackerDataManager: TrackerDataManager

Sources/BrowserServicesKit/ContentBlocking/UserScripts/SurrogatesUserScript.swift

@@ -25,6 +25,7 @@ import Common
 public protocol SurrogatesUserScriptDelegate: NSObjectProtocol {
 
     func surrogatesUserScriptShouldProcessTrackers(_ script: SurrogatesUserScript) -> Bool
+    func surrogatesUserScriptShouldProcessCTLTrackers(_ script: SurrogatesUserScript) -> Bool
     func surrogatesUserScript(_ script: SurrogatesUserScript,
                               detectedTracker tracker: DetectedRequest,
                               withSurrogate host: String)
@@ -83,7 +84,7 @@ public class DefaultSurrogatesUserScriptConfig: SurrogatesUserScriptConfig {
     }
 }
 
-open class SurrogatesUserScript: NSObject, UserScript {
+open class SurrogatesUserScript: NSObject, UserScript, WKScriptMessageHandlerWithReply {
     struct TrackerDetectedKey {
         static let protectionId = "protectionId"
         static let blocked = "blocked"
@@ -111,25 +112,47 @@ open class SurrogatesUserScript: NSObject, UserScript {
 
     public var requiresRunInPageContentWorld: Bool = true
 
-    public var messageNames: [String] = [ "trackerDetectedMessage" ]
+    public var messageNames: [String] = [
+        "trackerDetectedMessage",
+        "isCTLEnabled"
+    ]
 
     public weak var delegate: SurrogatesUserScriptDelegate?
 
-    public func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
+    public func userContentController(_ userContentController: WKUserContentController,
+                                      didReceive message: WKScriptMessage,
+                                      replyHandler: @escaping (Any?, String?) -> Void) {
+
         guard let delegate = delegate else { return }
-        guard delegate.surrogatesUserScriptShouldProcessTrackers(self) else { return }
 
-        guard let dict = message.body as? [String: Any] else { return }
-        guard let blocked = dict[TrackerDetectedKey.blocked] as? Bool else { return }
-        guard let urlString = dict[TrackerDetectedKey.url] as? String else { return }
-        guard let pageUrlStr = dict[TrackerDetectedKey.pageUrl] as? String else { return }
+        if message.name == "isCTLEnabled" {
+            let ctlEnabled = delegate.surrogatesUserScriptShouldProcessCTLTrackers(self)
+            replyHandler(ctlEnabled, nil)
+            return
+        } else if message.name == "trackerDetectedMessage"	{
+            guard delegate.surrogatesUserScriptShouldProcessTrackers(self) else { return }
+
+            guard let dict = message.body as? [String: Any] else { return }
+            guard let blocked = dict[TrackerDetectedKey.blocked] as? Bool else { return }
+            guard let urlString = dict[TrackerDetectedKey.url] as? String else { return }
+            guard let pageUrlStr = dict[TrackerDetectedKey.pageUrl] as? String else { return }
 
-        let tracker = trackerFromUrl(urlString.trimmingWhitespace(), pageUrlString: pageUrlStr, blocked)
+            let tracker = trackerFromUrl(urlString.trimmingWhitespace(), pageUrlString: pageUrlStr, blocked)
 
-        if let isSurrogate = dict[TrackerDetectedKey.isSurrogate] as? Bool, isSurrogate, let host = URL(string: urlString)?.host {
-            delegate.surrogatesUserScript(self, detectedTracker: tracker, withSurrogate: host)
+            if let isSurrogate = dict[TrackerDetectedKey.isSurrogate] as? Bool, isSurrogate, let host = URL(string: urlString)?.host {
+                delegate.surrogatesUserScript(self, detectedTracker: tracker, withSurrogate: host)
+            }
+            replyHandler(nil, nil)
+            return
         }
+
+        replyHandler(nil, "Unknown message")
     }
+
+    public func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
+        assertionFailure("Should never be here!")
+    }
+
     private func trackerFromUrl(_ urlString: String, pageUrlString: String, _ blocked: Bool) -> DetectedRequest {
         let currentTrackerData = configuration.trackerData
         let knownTracker = currentTrackerData?.findTracker(forUrl: urlString)

Sources/BrowserServicesKit/ContentBlocking/UserScripts/TrackerResolver.swift

@@ -206,7 +206,7 @@ fileprivate extension KnownTracker.Rule {
     func action(type: String, host: String) -> TrackerResolver.RuleAction? {
         // If there is a rule its default action is always block
         var resultAction: KnownTracker.ActionType? = action ?? .block
-        if resultAction == .block {
+        if resultAction == .block || resultAction == .blockCTLFB {
             if let options = options, !TrackerResolver.isMatching(options, host: host, resourceType: type) {
                 resultAction = nil
             } else if let exceptions = exceptions, TrackerResolver.isMatching(exceptions, host: host, resourceType: type) {
@@ -221,7 +221,7 @@ fileprivate extension KnownTracker.Rule {
 private extension KnownTracker.ActionType {
 
     func toTrackerResolverRuleAction() -> TrackerResolver.RuleAction {
-        self == .block ? .blockRequest : .allowRequest
+        self == .block || self == .blockCTLFB ? .blockRequest : .allowRequest
     }
 
 }