Fix #6254: Get cosmetic filters on serial queue (#7148)

Co-authored-by: Michał Buczek <[email protected]>

Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController+WKNavigationDelegate.swift

@@ -279,7 +279,7 @@ extension BrowserViewController: WKNavigationDelegate {
       // Check if custom user scripts must be added to or removed from the web view.
       if let targetFrame = navigationAction.targetFrame {
         tab?.currentPageData?.addSubframeURL(forRequestURL: url, isForMainFrame: targetFrame.isMainFrame)
-        let scriptTypes = tab?.currentPageData?.makeUserScriptTypes(domain: domainForMainFrame) ?? []
+        let scriptTypes = await tab?.currentPageData?.makeUserScriptTypes(domain: domainForMainFrame) ?? []
         tab?.setCustomUserScript(scripts: scriptTypes)
       }
     }
@@ -401,7 +401,7 @@ extension BrowserViewController: WKNavigationDelegate {
     if let responseURL = responseURL,
        let domain = tab?.currentPageData?.domain(persistent: !isPrivateBrowsing),
        tab?.currentPageData?.upgradeFrameURL(forResponseURL: responseURL, isForMainFrame: navigationResponse.isForMainFrame) == true {
-      let scriptTypes = tab?.currentPageData?.makeUserScriptTypes(domain: domain) ?? []
+      let scriptTypes = await tab?.currentPageData?.makeUserScriptTypes(domain: domain) ?? []
       tab?.setCustomUserScript(scripts: scriptTypes)
     }
 

Sources/Brave/Frontend/Browser/PageData.swift

@@ -70,7 +70,7 @@ struct PageData {
   }
 
   /// Return all the user script types for this page. The number of script types grows as more frames are loaded.
-  @MainActor func makeUserScriptTypes(domain: Domain) -> Set<UserScriptType> {
+  @MainActor func makeUserScriptTypes(domain: Domain) async -> Set<UserScriptType> {
     var userScriptTypes: Set<UserScriptType> = [.siteStateListener]
 
     // Handle dynamic domain level scripts on the main document.
@@ -98,18 +98,26 @@ struct PageData {
       }
     }
 
+    let allEngineScriptTypes = await makeAllEngineScripts(for: domain)
+    return userScriptTypes.union(allEngineScriptTypes)
+  }
+
+  func makeMainFrameEngineScriptTypes(domain: Domain) async -> Set<UserScriptType> {
+    return await adBlockStats.makeEngineScriptTypes(frameURL: mainFrameURL, isMainFrame: true, domain: domain)
+  }
+
+  func makeAllEngineScripts(for domain: Domain) async -> Set<UserScriptType> {
     // Add engine scripts for the main frame
-    userScriptTypes = userScriptTypes.union(
-      adBlockStats.makeEngineScriptTypes(frameURL: mainFrameURL, isMainFrame: true, domain: domain)
-    )
+    async let engineScripts = adBlockStats.makeEngineScriptTypes(frameURL: mainFrameURL, isMainFrame: true, domain: domain)
 
     // Add engine scripts for all of the known sub-frames
-    let additionalScriptTypes = allSubframeURLs.map({ frameURL in
-      return self.adBlockStats.makeEngineScriptTypes(frameURL: frameURL, isMainFrame: false, domain: domain)
+    async let additionalScriptTypes = allSubframeURLs.asyncConcurrentCompactMap({ frameURL in
+      return await self.adBlockStats.makeEngineScriptTypes(frameURL: frameURL, isMainFrame: false, domain: domain)
     }).reduce(Set<UserScriptType>(), { partialResult, scriptTypes in
       return partialResult.union(scriptTypes)
     })
 
-    return userScriptTypes.union(additionalScriptTypes)
+    let allEngineScripts = await (mainFrame: engineScripts, subFrames: additionalScriptTypes)
+    return allEngineScripts.mainFrame.union(allEngineScripts.subFrames)
   }
 }

Sources/Brave/Frontend/Browser/Playlist/Managers & Cache/PlaylistCacheLoader.swift

@@ -558,7 +558,7 @@ extension PlaylistWebLoader: WKNavigationDelegate {
       if let requestURL = navigationAction.request.url,
          let targetFrame = navigationAction.targetFrame {
         tab.currentPageData?.addSubframeURL(forRequestURL: requestURL, isForMainFrame: targetFrame.isMainFrame)
-        let scriptTypes = tab.currentPageData?.makeUserScriptTypes(domain: domainForMainFrame) ?? []
+        let scriptTypes = await tab.currentPageData?.makeUserScriptTypes(domain: domainForMainFrame) ?? []
         tab.setCustomUserScript(scripts: scriptTypes)
       }
 
@@ -615,7 +615,7 @@ extension PlaylistWebLoader: WKNavigationDelegate {
     if let responseURL = responseURL,
        tab.currentPageData?.upgradeFrameURL(forResponseURL: responseURL, isForMainFrame: navigationResponse.isForMainFrame) == true,
        let domain = tab.currentPageData?.domain(persistent: false) {
-      let scriptTypes = tab.currentPageData?.makeUserScriptTypes(domain: domain) ?? []
+      let scriptTypes = await tab.currentPageData?.makeUserScriptTypes(domain: domain) ?? []
       tab.setCustomUserScript(scripts: scriptTypes)
     }
 

Sources/Brave/Frontend/UserContent/UserScripts/Scripts_Dynamic/ScriptHandlers/Paged/CosmeticFiltersScriptHandler.swift

@@ -58,7 +58,7 @@ class CosmeticFiltersScriptHandler: TabContentScript {
 
         let selectorArrays = await cachedEngines.asyncConcurrentMap { cachedEngine -> [String] in
           do {
-            return try cachedEngine.selectorsForCosmeticRules(
+            return try await cachedEngine.selectorsForCosmeticRules(
               frameURL: frameURL,
               ids: dto.data.ids,
               classes: dto.data.classes

Sources/Brave/WebFilters/ShieldStats/Adblock/AdBlockStats.swift

@@ -43,13 +43,13 @@ public class AdBlockStats {
   }
 
   /// This returns all the user script types for the given frame
-  @MainActor func makeEngineScriptTypes(frameURL: URL, isMainFrame: Bool, domain: Domain) -> Set<UserScriptType> {
+  func makeEngineScriptTypes(frameURL: URL, isMainFrame: Bool, domain: Domain) async -> Set<UserScriptType> {
     // Add any engine scripts for this frame
-    return cachedEngines.enumerated().map({ (index, cachedEngine) -> Set<UserScriptType> in
-      guard cachedEngine.isEnabled(for: domain) else { return [] }
+    return await cachedEngines.enumerated().asyncMap({ (index, cachedEngine) -> Set<UserScriptType> in
+      guard await cachedEngine.isEnabled(for: domain) else { return [] }
 
       do {
-        return try cachedEngine.makeEngineScriptTypes(
+        return try await cachedEngine.makeEngineScriptTypes(
           frameURL: frameURL, isMainFrame: isMainFrame, domain: domain, index: index
         )
       } catch {

Sources/Brave/WebFilters/ShieldStats/Adblock/CachedAdBlockEngine.swift

@@ -40,31 +40,48 @@ public class CachedAdBlockEngine {
 
   /// Returns all the models for this frame URL
   /// The results are cached per url, so you may call this method as many times for the same url without any performance implications.
-  @MainActor func cosmeticFilterModel(forFrameURL frameURL: URL) throws -> CosmeticFilterModel? {
-    if let model = self.cachedCosmeticFilterModels.getElement(frameURL) {
-      return model
+  func cosmeticFilterModel(forFrameURL frameURL: URL) async throws -> CosmeticFilterModel? {
+    return try await withCheckedThrowingContinuation { continuation in
+      serialQueue.async { [weak self] in
+        if let model = self?.cachedCosmeticFilterModels.getElement(frameURL) {
+          continuation.resume(returning: model)
+          return
+        }
+
+        do {
+          let model = try self?.engine.cosmeticFilterModel(forFrameURL: frameURL)
+          self?.cachedCosmeticFilterModels.addElement(model, forKey: frameURL)
+          continuation.resume(returning: model)
+        } catch {
+          continuation.resume(throwing: error)
+        }
+      }
     }
-
-    let model = try self.engine.cosmeticFilterModel(forFrameURL: frameURL)
-    self.cachedCosmeticFilterModels.addElement(model, forKey: frameURL)
-    return model
   }
 
   /// Return the selectors that need to be hidden given the frameURL, ids and classes
-  @MainActor func selectorsForCosmeticRules(frameURL: URL, ids: [String], classes: [String]) throws -> [String] {
-    let model = try cosmeticFilterModel(forFrameURL: frameURL)
-    let selectorsJSON = self.engine.stylesheetForCosmeticRulesIncluding(
-      classes: classes,
-      ids: ids,
-      exceptions: model?.exceptions ?? []
-    )
-
-    guard let data = selectorsJSON.data(using: .utf8) else {
-      return []
+  func selectorsForCosmeticRules(frameURL: URL, ids: [String], classes: [String]) async throws -> [String] {
+    let exceptions = try await cosmeticFilterModel(forFrameURL: frameURL)?.exceptions ?? []
+
+    return try await withCheckedThrowingContinuation { continuation in
+      serialQueue.async { [weak self] in
+        let selectorsJSON = self?.engine.stylesheetForCosmeticRulesIncluding(classes: classes, ids: ids, exceptions: exceptions)
+
+        guard let data = selectorsJSON?.data(using: .utf8) else {
+          continuation.resume(returning: [])
+          return
+        }
+
+        let decoder = JSONDecoder()
+
+        do {
+          let result = try decoder.decode([String].self, from: data)
+          continuation.resume(returning: result)
+        } catch {
+          continuation.resume(throwing: error)
+        }
+      }
     }
-
-    let decoder = JSONDecoder()
-    return try decoder.decode([String].self, from: data)
   }
 
   /// Checks the general and regional engines to see if the request should be blocked
@@ -86,15 +103,15 @@ public class CachedAdBlockEngine {
   }
 
   /// This returns all the user script types for the given frame
-  @MainActor func makeEngineScriptTypes(frameURL: URL, isMainFrame: Bool, domain: Domain, index: Int) throws -> Set<UserScriptType> {
+  @MainActor func makeEngineScriptTypes(frameURL: URL, isMainFrame: Bool, domain: Domain, index: Int) async throws -> Set<UserScriptType> {
     if let userScriptTypes = cachedFrameScriptTypes.getElement(frameURL) {
       return userScriptTypes
     }
 
     // Add the selectors poller scripts for this frame
     var userScriptTypes: Set<UserScriptType> = []
 
-    if let source = try cosmeticFilterModel(forFrameURL: frameURL)?.injectedScript, !source.isEmpty {
+    if let source = try await cosmeticFilterModel(forFrameURL: frameURL)?.injectedScript, !source.isEmpty {
       let configuration = UserScriptType.EngineScriptConfiguration(
         frameURL: frameURL, isMainFrame: isMainFrame, source: source, order: index,
         isDeAMPEnabled: Preferences.Shields.autoRedirectAMPPages.value

Tests/ClientTests/PageDataTests.swift

@@ -22,7 +22,7 @@ final class PageDataTests: XCTestCase {
       // When
       // We get the script types for the main frame
       let domain = pageData.domain(persistent: false)
-      let mainFrameRequestTypes = pageData.makeUserScriptTypes(domain: domain)
+      let mainFrameRequestTypes = await pageData.makeUserScriptTypes(domain: domain)
 
       // Then
       // We get only entries of the main frame
@@ -34,7 +34,7 @@ final class PageDataTests: XCTestCase {
 
       // When
       // Nothing has changed
-      let unchangedScriptTypes = pageData.makeUserScriptTypes(domain: domain)
+      let unchangedScriptTypes = await pageData.makeUserScriptTypes(domain: domain)
 
       // Then
       // We get the same result as before
@@ -48,7 +48,7 @@ final class PageDataTests: XCTestCase {
       // We get no aditional scripts for sub-frame
       // NOTE: This is because we have no engines on AdBlockStats.
       // If we were to add some engines we might see additional types
-      let addedSubFrameFrameRequestTypes = pageData.makeUserScriptTypes(domain: domain)
+      let addedSubFrameFrameRequestTypes = await pageData.makeUserScriptTypes(domain: domain)
       let expectedMainAndSubFrameTypes: Set<UserScriptType> = [
         .siteStateListener, .nacl, .farblingProtection(etld: "example.com")
       ]
@@ -61,7 +61,7 @@ final class PageDataTests: XCTestCase {
       // Then
       // We get the same result as before
       XCTAssertTrue(isUpgradedMainFrame)
-      let upgradedMainFrameRequestTypes = pageData.makeUserScriptTypes(domain: domain)
+      let upgradedMainFrameRequestTypes = await pageData.makeUserScriptTypes(domain: domain)
       XCTAssertEqual(upgradedMainFrameRequestTypes, unchangedScriptTypes)
 
       // When
@@ -71,7 +71,7 @@ final class PageDataTests: XCTestCase {
       // Then
       // We get the same result as before
       XCTAssertTrue(isUpgradedSubFrame)
-      let upgradedSubFrameFrameRequestTypes = pageData.makeUserScriptTypes(domain: domain)
+      let upgradedSubFrameFrameRequestTypes = await pageData.makeUserScriptTypes(domain: domain)
       XCTAssertEqual(upgradedSubFrameFrameRequestTypes, unchangedScriptTypes)
 
       expectation.fulfill()

Tests/ClientTests/Web Filters/AdBlockEngineManagerTests.swift

@@ -89,13 +89,13 @@ class AdBlockEngineManagerTests: XCTestCase {
         let cachedEngines = stats.cachedEngines(for: domain)
         XCTAssertEqual(cachedEngines.count, 3)
 
-        let types = stats.makeEngineScriptTypes(frameURL: url, isMainFrame: true, domain: domain)
+        let types = await stats.makeEngineScriptTypes(frameURL: url, isMainFrame: true, domain: domain)
         // We should have no scripts injected
         XCTAssertEqual(types.count, 0)
 
         let url2 = URL(string:  "https://stackoverflow.com")!
         let domain2 = Domain.getOrCreate(forUrl: url2, persistent: false)
-        let types2 = stats.makeEngineScriptTypes(frameURL: URL(string:  "https://reddit.com")!, isMainFrame: true, domain: domain2)
+        let types2 = await stats.makeEngineScriptTypes(frameURL: URL(string:  "https://reddit.com")!, isMainFrame: true, domain: domain2)
         // We should have 1 engine script injected
         XCTAssertEqual(types2.count, 1)
         XCTAssertTrue(types2.contains(where: { scriptType in