Bat.js fix (fix for installing content blocking rules multiple times)

Sources/BrowserServicesKit/ContentBlocking/ContentBlockerRulesIdentifier.swift

@@ -30,7 +30,7 @@ public class ContentBlockerRulesIdentifier: Equatable, Codable {
         return name + tdsEtag + tempListId + allowListId + unprotectedSitesHash
     }
 
-    public struct Difference: OptionSet {
+    public struct Difference: OptionSet, CustomDebugStringConvertible {
         public let rawValue: Int
 
         public init(rawValue: Int) {
@@ -43,6 +43,29 @@ public class ContentBlockerRulesIdentifier: Equatable, Codable {
         public static let unprotectedSites = Difference(rawValue: 1 << 3)
 
         public static let all: Difference = [.tdsEtag, .tempListId, .allowListId, .unprotectedSites]
+
+        public var debugDescription: String {
+            if self == .all {
+                return "all"
+            }
+            var result = "["
+            for i in 0...Int(log2(Double(max(self.rawValue, Self.all.rawValue)))) where self.contains(Self(rawValue: 1 << i)) {
+                if result.count > 1 {
+                    result += ", "
+                }
+                result += {
+                    switch Self(rawValue: 1 << i) {
+                    case .tdsEtag: ".tdsEtag"
+                    case .tempListId: ".tempListId"
+                    case .allowListId: ".allowListId"
+                    case .unprotectedSites: ".unprotectedSites"
+                    default: "1<<\(i)"
+                    }
+                }()
+            }
+            result += "]"
+            return result
+        }
     }
 
     private class func normalize(identifier: String?) -> String {

Sources/BrowserServicesKit/ContentBlocking/ContentBlockerRulesManager.swift

@@ -96,7 +96,7 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
     private let cache: ContentBlockerRulesCaching?
     public let exceptionsSource: ContentBlockerRulesExceptionsSource
 
-    public struct UpdateEvent {
+    public struct UpdateEvent: CustomDebugStringConvertible {
         public let rules: [ContentBlockerRulesManager.Rules]
         public let changes: [String: ContentBlockerRulesIdentifier.Difference]
         public let completionTokens: [ContentBlockerRulesManager.CompletionToken]
@@ -108,6 +108,14 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
             self.changes = changes
             self.completionTokens = completionTokens
         }
+
+        public var debugDescription: String {
+            """
+              rules: \(rules.map { "\($0.name):\($0.identifier) – \($0.rulesList) (\($0.etag))" }.joined(separator: ", "))
+              changes: \(changes)
+              completionTokens: \(completionTokens)
+            """
+        }
     }
     private let updatesSubject = PassthroughSubject<UpdateEvent, Never>()
     public var updatesPublisher: AnyPublisher<UpdateEvent, Never> {
@@ -193,6 +201,7 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
     @discardableResult
     public func scheduleCompilation() -> CompletionToken {
         let token = UUID().uuidString
+        os_log("Scheduling compilation with %{public}s", log: log, type: .default, token)
         workQueue.async {
             let shouldStartCompilation = self.updateCompilationState(token: token)
             if shouldStartCompilation {
@@ -228,19 +237,25 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
      Returns true if rules were found, false otherwise.
      */
     private func lookupCompiledRules() -> Bool {
+        os_log("Lookup compiled rules", log: log, type: .debug)
         prepareSourceManagers()
         let initialCompilationTask = LookupRulesTask(sourceManagers: Array(sourceManagers.values))
         let mutex = DispatchSemaphore(value: 0)
 
-        Task {
-            try? await initialCompilationTask.lookupCachedRulesLists()
+        Task { [log] in
+            do {
+                try await initialCompilationTask.lookupCachedRulesLists()
+            } catch {
+                os_log("❌ Lookup failed: %{public}s", log: log, type: .debug, error.localizedDescription)
+            }
             mutex.signal()
         }
         // We want to confine Compilation work to WorkQueue, so we wait to come back from async Task
         mutex.wait()
 
         if let result = initialCompilationTask.result {
             let rules = result.map(Rules.init(compilationResult:))
+            os_log("🟩 Found %{public}d rules", log: log, type: .debug, rules.count)
             applyRules(rules)
             return true
         }
@@ -252,6 +267,8 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
      Returns true if rules were found, false otherwise.
      */
     private func fetchLastCompiledRules(with lastCompiledRules: [LastCompiledRules]) {
+        os_log("Fetch last compiled rules: %{public}d", log: log, type: .debug, lastCompiledRules.count)
+
         let initialCompilationTask = LastCompiledRulesLookupTask(sourceRules: rulesSource.contentBlockerRulesLists,
                                                                  lastCompiledRules: lastCompiledRules)
         let mutex = DispatchSemaphore(value: 0)
@@ -294,6 +311,7 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
     }
 
     private func startCompilationProcess() {
+        os_log("Starting compilataion process", log: log, type: .debug)
         prepareSourceManagers()
 
         // Prepare compilation tasks based on the sources

Sources/BrowserServicesKit/ContentBlocking/ContentBlockingRulesCompilationTask.swift

@@ -53,12 +53,14 @@ extension ContentBlockerRulesManager {
         func start(ignoreCache: Bool = false, completionHandler: @escaping Completion) {
             self.workQueue.async {
                 guard let model = self.sourceManager.makeModel() else {
+                    os_log("❌ compilation impossible", log: self.log, type: .default)
                     self.compilationImpossible = true
                     completionHandler(self, false)
                     return
                 }
 
                 guard !ignoreCache else {
+                    os_log("❗️ ignoring cache", log: self.log, type: .default)
                     self.workQueue.async {
                         self.compile(model: model, completionHandler: completionHandler)
                     }
@@ -68,8 +70,10 @@ extension ContentBlockerRulesManager {
                 // Delegate querying to main thread - crashes were observed in background.
                 DispatchQueue.main.async {
                     let identifier = model.rulesIdentifier.stringValue
+                    os_log("Lookup CBR with %{public}s", log: self.log, type: .default, identifier)
                     WKContentRuleListStore.default()?.lookUpContentRuleList(forIdentifier: identifier) { ruleList, _ in
                         if let ruleList = ruleList {
+                            os_log("🟢 CBR loaded from cache: %{public}s", log: self.log, type: .default, self.rulesList.name)
                             self.compilationSucceeded(with: ruleList, model: model, completionHandler: completionHandler)
                         } else {
                             self.workQueue.async {
@@ -94,7 +98,7 @@ extension ContentBlockerRulesManager {
                                        with error: Error,
                                        completionHandler: @escaping Completion) {
             workQueue.async {
-                os_log("Failed to compile %{public}s rules %{public}s",
+                os_log("❌ Failed to compile %{public}s rules %{public}s",
                        log: self.log,
                        type: .error,
                        self.rulesList.name,
@@ -125,7 +129,7 @@ extension ContentBlockerRulesManager {
             do {
                 data = try JSONEncoder().encode(rules)
             } catch {
-                os_log("Failed to encode content blocking rules %{public}s", log: log, type: .error, rulesList.name)
+                os_log("❌ Failed to encode content blocking rules %{public}s", log: log, type: .error, rulesList.name)
                 compilationFailed(for: model, with: error, completionHandler: completionHandler)
                 return
             }
@@ -136,6 +140,7 @@ extension ContentBlockerRulesManager {
                                                                         encodedContentRuleList: ruleList) { ruleList, error in
 
                     if let ruleList = ruleList {
+                        os_log("🟢 CBR compilation for %{public}s succeeded", log: self.log, type: .default, self.rulesList.name)
                         self.compilationSucceeded(with: ruleList, model: model, completionHandler: completionHandler)
                     } else if let error = error {
                         self.compilationFailed(for: model, with: error, completionHandler: completionHandler)

Sources/BrowserServicesKit/ContentScopeScript/UserContentController.swift

@@ -16,9 +16,11 @@
 //  limitations under the License.
 //
 
-import WebKit
 import Combine
+import Common
 import UserScript
+import WebKit
+import QuartzCore
 
 public protocol UserContentControllerDelegate: AnyObject {
     @MainActor
@@ -37,12 +39,13 @@ public protocol UserContentControllerNewContent {
     var makeUserScripts: @MainActor (SourceProvider) -> UserScripts { get }
 }
 
+@objc(UserContentController)
 final public class UserContentController: WKUserContentController {
     public let privacyConfigurationManager: PrivacyConfigurationManaging
     @MainActor
     public weak var delegate: UserContentControllerDelegate?
 
-    public struct ContentBlockingAssets {
+    public struct ContentBlockingAssets: CustomDebugStringConvertible {
         public let globalRuleLists: [String: WKContentRuleList]
         public let userScripts: UserScriptsProvider
         public let wkUserScripts: [WKUserScript]
@@ -58,32 +61,52 @@ final public class UserContentController: WKUserContentController {
 
             self.wkUserScripts = await userScripts.loadWKUserScripts()
         }
+
+        public var debugDescription: String {
+            """
+            <ContentBlockingAssets
+            globalRuleLists: \(globalRuleLists)
+            wkUserScripts: \(wkUserScripts)
+            updateEvent: (
+            \(updateEvent.debugDescription)
+            )>
+            """
+        }
     }
 
     @Published @MainActor public private(set) var contentBlockingAssets: ContentBlockingAssets? {
         willSet {
             self.removeAllContentRuleLists()
             self.removeAllUserScripts()
+
+            if let contentBlockingAssets = newValue {
+                os_log(.debug, log: .contentBlocking, "\(self): 📚 installing \(contentBlockingAssets)")
+                self.installGlobalContentRuleLists(contentBlockingAssets.globalRuleLists)
+                os_log(.debug, log: .userScripts, "\(self): 📜 installing user scripts")
+                self.installUserScripts(contentBlockingAssets.wkUserScripts, handlers: contentBlockingAssets.userScripts.userScripts)
+                os_log(.debug, log: .contentBlocking, "\(self): ✅ installing content blocking assets done")
+            }
         }
     }
     @MainActor
     private func installContentBlockingAssets(_ contentBlockingAssets: ContentBlockingAssets) {
         // don‘t install ContentBlockingAssets (especially Message Handlers retaining `self`) after cleanUpBeforeClosing was called
         guard assetsPublisherCancellable != nil else { return }
-
+        // installation should happen in `contentBlockingAssets.willSet`
+        // so the $contentBlockingAssets subscribers receive an update only after everything is set
         self.contentBlockingAssets = contentBlockingAssets
 
-        self.installGlobalContentRuleLists(contentBlockingAssets.globalRuleLists)
-        self.installUserScripts(contentBlockingAssets.wkUserScripts, handlers: contentBlockingAssets.userScripts.userScripts)
-
         delegate?.userContentController(self,
                                         didInstallContentRuleLists: contentBlockingAssets.globalRuleLists,
                                         userScripts: contentBlockingAssets.userScripts,
                                         updateEvent: contentBlockingAssets.updateEvent)
     }
 
+    enum ContentRuleListIdentifier: Hashable {
+        case global(String), local(String)
+    }
     @MainActor
-    private var localRuleLists = [String: WKContentRuleList]()
+    private var contentRuleLists = [ContentRuleListIdentifier: WKContentRuleList]()
     @MainActor
     private var assetsPublisherCancellable: AnyCancellable?
     @MainActor
@@ -96,7 +119,8 @@ final public class UserContentController: WKUserContentController {
         self.privacyConfigurationManager = privacyConfigurationManager
         super.init()
 
-        assetsPublisherCancellable = assetsPublisher.sink { [weak self] content in
+        assetsPublisherCancellable = assetsPublisher.sink { [weak self, selfDescr=self.debugDescription] content in
+            os_log(.debug, log: .contentBlocking, "\(selfDescr): 📚 received content blocking assets")
             Task.detached { [weak self] in
                 let contentBlockingAssets = await ContentBlockingAssets(content: content)
                 await self?.installContentBlockingAssets(contentBlockingAssets)
@@ -116,50 +140,73 @@ final public class UserContentController: WKUserContentController {
     }
 
     @MainActor
-    private func installGlobalContentRuleLists(_ contentRuleLists: [String: WKContentRuleList]) {
+    private func installGlobalContentRuleLists(_ globalContentRuleLists: [String: WKContentRuleList]) {
+        assert(contentRuleLists.isEmpty, "installGlobalContentRuleLists should be called after removing all Content Rule Lists")
         guard self.privacyConfigurationManager.privacyConfig.isEnabled(featureKey: .contentBlocking) else {
+            os_log(.debug, log: .contentBlocking, "\(self): ❗️ content blocking disabled, removing all content rule lists")
             removeAllContentRuleLists()
             return
         }
 
-        contentRuleLists.values.forEach(self.add)
+        os_log(.debug, log: .contentBlocking, "\(self): ❇️ installing global rule lists: \(globalContentRuleLists))")
+        contentRuleLists = globalContentRuleLists.reduce(into: [:]) {
+            $0[.global($1.key)] = $1.value
+        }
+        globalContentRuleLists.values.forEach(self.add)
     }
 
     public struct ContentRulesNotFoundError: Error {}
     @MainActor
     public func enableGlobalContentRuleList(withIdentifier identifier: String) throws {
-        guard let ruleList = self.contentBlockingAssets?.globalRuleLists[identifier] else {
+        guard let ruleList = contentBlockingAssets?.globalRuleLists[identifier]
+                // when enabling from a $contentBlockingAssets subscription, the ruleList gets
+                // to contentRuleLists before contentBlockingAssets value is set
+                ?? contentRuleLists[.global(identifier)] else {
+            os_log(.debug, log: .contentBlocking, "\(self): ❗️ can‘t enable rule list `\(identifier)` as it‘s not available")
             throw ContentRulesNotFoundError()
         }
-        self.add(ruleList)
+        guard contentRuleLists[.global(identifier)] == nil else { return /* already enabled */ }
+
+        os_log(.debug, log: .contentBlocking, "\(self): 🟩 enabling rule list `\(identifier)`")
+        contentRuleLists[.global(identifier)] = ruleList
+        add(ruleList)
     }
 
     public struct ContentRulesNotEnabledError: Error {}
     @MainActor
     public func disableGlobalContentRuleList(withIdentifier identifier: String) throws {
-        guard let ruleList = self.contentBlockingAssets?.globalRuleLists[identifier] else {
+        guard let ruleList = contentRuleLists[.global(identifier)] else {
+            os_log(.debug, log: .contentBlocking, "\(self): ❗️ can‘t disable rule list `\(identifier)` as it‘s not enabled")
             throw ContentRulesNotEnabledError()
         }
-        self.remove(ruleList)
+
+        os_log(.debug, log: .contentBlocking, "\(self): 🔻 disabling rule list `\(identifier)`")
+        contentRuleLists[.global(identifier)] = nil
+        remove(ruleList)
     }
 
     @MainActor
     public func installLocalContentRuleList(_ ruleList: WKContentRuleList, identifier: String) {
-        localRuleLists[identifier] = ruleList
-        self.add(ruleList)
+        // replace if already installed
+        removeLocalContentRuleList(withIdentifier: identifier)
+
+        os_log(.debug, log: .contentBlocking, "\(self): 🔸 installing local rule list `\(identifier)`")
+        contentRuleLists[.local(identifier)] = ruleList
+        add(ruleList)
     }
 
     @MainActor
     public func removeLocalContentRuleList(withIdentifier identifier: String) {
-        guard let ruleList = localRuleLists.removeValue(forKey: identifier) else {
-            return
-        }
-        self.remove(ruleList)
+        guard let ruleList = contentRuleLists.removeValue(forKey: .local(identifier)) else { return }
+
+        os_log(.debug, log: .contentBlocking, "\(self): 🔻 removing local rule list `\(identifier)`")
+        remove(ruleList)
     }
 
     @MainActor
     public override func removeAllContentRuleLists() {
-        localRuleLists = [:]
+        os_log(.debug, log: .contentBlocking, "\(self): 🧹 removing all content rule lists")
+        contentRuleLists.removeAll(keepingCapacity: true)
         super.removeAllContentRuleLists()
     }
 
@@ -171,6 +218,8 @@ final public class UserContentController: WKUserContentController {
 
     @MainActor
     public func cleanUpBeforeClosing() {
+        os_log(.debug, log: .contentBlocking, "\(self): 💀 cleanUpBeforeClosing")
+
         self.removeAllUserScripts()
 
         if #available(macOS 11.0, *) {
@@ -222,7 +271,9 @@ public extension UserContentController {
     @MainActor
     var awaitContentBlockingAssetsInstalled: () async -> Void {
         guard !contentBlockingAssetsInstalled else { return {} }
-        return { [weak self] in
+        os_log(.debug, log: .contentBlocking, "\(self): 🛑 will wait for content blocking assets installed")
+        let startTime = CACurrentMediaTime()
+        return { [weak self, selfDescr=self.description] in
             // merge $contentBlockingAssets with Task cancellation completion event publisher
             let taskCancellationSubject = PassthroughSubject<ContentBlockingAssets?, Error>()
             guard let assetsPublisher = self?.$contentBlockingAssets else { return }
@@ -237,14 +288,21 @@ public extension UserContentController {
             try? await withTaskCancellationHandler {
                 try await withCheckedThrowingContinuation { c in
                     var cancellable: AnyCancellable!
+                    var elapsedTime: String {
+                        String(format: "%.2fs.", CACurrentMediaTime() - startTime)
+                    }
                     cancellable = throwingPublisher.sink /* completion: */ { _ in
                         withExtendedLifetime(cancellable) {
+                            os_log(.debug, log: .contentBlocking, "\(selfDescr): ❌ wait cancelled after \(elapsedTime)")
+
                             c.resume(with: .failure(CancellationError()))
                             cancellable.cancel()
                         }
                     } receiveValue: { assets in
                         guard assets != nil else { return }
                         withExtendedLifetime(cancellable) {
+                            os_log(.debug, log: .contentBlocking, "\(selfDescr): 🏁 content blocking assets installed (\(elapsedTime))")
+
                             c.resume(with: .success( () ))
                             cancellable.cancel()
                         }