diff --git a/Sources/BrowserServicesKit/ContentScopeScript/UserContentController.swift b/Sources/BrowserServicesKit/ContentScopeScript/UserContentController.swift index 762f036c1..5a8942efe 100644 --- a/Sources/BrowserServicesKit/ContentScopeScript/UserContentController.swift +++ b/Sources/BrowserServicesKit/ContentScopeScript/UserContentController.swift @@ -112,13 +112,17 @@ final public class UserContentController: WKUserContentController { @MainActor private let scriptMessageHandler = PermanentScriptMessageHandler() + /// if earlyAccessHandlers (WKScriptMessageHandlers) are provided they are installed without waiting for contentBlockingAssets to be loaded if. @MainActor - public init(assetsPublisher: Pub, privacyConfigurationManager: PrivacyConfigurationManaging) + public init(assetsPublisher: Pub, privacyConfigurationManager: PrivacyConfigurationManaging, earlyAccessHandlers: [UserScript] = []) where Pub: Publisher, Content: UserContentControllerNewContent, Pub.Output == Content, Pub.Failure == Never { self.privacyConfigurationManager = privacyConfigurationManager super.init() + // Install initial WKScriptMessageHandlers if any. Currently, no WKUserScript are provided at initialization. + installUserScripts([], handlers: earlyAccessHandlers) + 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 diff --git a/Tests/BrowserServicesKitTests/ContentBlocker/UserContentControllerTests.swift b/Tests/BrowserServicesKitTests/ContentBlocker/UserContentControllerTests.swift index 470aae6ae..c14459a60 100644 --- a/Tests/BrowserServicesKitTests/ContentBlocker/UserContentControllerTests.swift +++ b/Tests/BrowserServicesKitTests/ContentBlocker/UserContentControllerTests.swift @@ -55,6 +55,7 @@ final class UserContentControllerTests: XCTestCase { @MainActor override func setUp() async throws { _=WKUserContentController.swizzleContentRuleListsMethodsOnce + _=WKUserContentController.swizzleScriptMessageHandlerMethodsOnce ucc = UserContentController(assetsPublisher: assetsSubject, privacyConfigurationManager: PrivacyConfigurationManagerMock()) ucc.delegate = self } @@ -69,6 +70,16 @@ final class UserContentControllerTests: XCTestCase { } // MARK: - Tests + @MainActor + func testWhenUserContentControllerInitialisedWithEarlyAccessScriptsThenHandlersAreRegistered() async throws { + let script1 = MockUserScript(messageNames: ["message1"]) + let script2 = MockUserScript(messageNames: ["message2"]) + ucc = UserContentController(assetsPublisher: assetsSubject, privacyConfigurationManager: PrivacyConfigurationManagerMock(), earlyAccessHandlers: [script1, script2]) + ucc.delegate = self + + XCTAssertTrue(ucc.registeredScriptHandlerNames.contains("message1")) + XCTAssertTrue(ucc.registeredScriptHandlerNames.contains("message2")) + } @MainActor func testWhenContentBlockingAssetsPublished_contentRuleListsAreInstalled() async throws { @@ -250,6 +261,34 @@ extension WKUserContentController { } } +extension WKUserContentController { + private static let scriptHandlersKey = UnsafeRawPointer(bitPattern: "scriptHandlersKey".hashValue)! + + private static var installedScriptHandlers: [(WKScriptMessageHandler, WKContentWorld, String)] { + get { + objc_getAssociatedObject(self, scriptHandlersKey) as? [(WKScriptMessageHandler, WKContentWorld, String)] ?? [] + } + set { + objc_setAssociatedObject(self, scriptHandlersKey, newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC) + } + } + + static let swizzleScriptMessageHandlerMethodsOnce: Void = { + let originalAddMethod = class_getInstanceMethod(WKUserContentController.self, #selector(WKUserContentController.add(_:contentWorld:name:)))! + let swizzledAddMethod = class_getInstanceMethod(WKUserContentController.self, #selector(swizzled_add(_:contentWorld:name:)))! + method_exchangeImplementations(originalAddMethod, swizzledAddMethod) + }() + + @objc dynamic private func swizzled_add(_ scriptMessageHandler: WKScriptMessageHandler, contentWorld: WKContentWorld, name: String) { + Self.installedScriptHandlers.append((scriptMessageHandler, contentWorld, name)) + swizzled_add(scriptMessageHandler, contentWorld: contentWorld, name: name) // calling the original method + } + + var registeredScriptHandlerNames: [String] { + return Self.installedScriptHandlers.map { $0.2 } + } +} + extension UserContentControllerTests: UserContentControllerDelegate { func userContentController(_ userContentController: UserContentController, didInstallContentRuleLists contentRuleLists: [String: WKContentRuleList], userScripts: any UserScriptsProvider, updateEvent: ContentBlockerRulesManager.UpdateEvent) { onAssetsInstalled?((contentRuleLists, userScripts, updateEvent)) @@ -305,3 +344,17 @@ class PrivacyConfigurationMock: PrivacyConfiguration { func userEnabledProtection(forDomain: String) {} func userDisabledProtection(forDomain: String) {} } + +class MockUserScript: NSObject, UserScript { + var source: String = "MockUserScript" + var injectionTime: WKUserScriptInjectionTime = .atDocumentEnd + var forMainFrameOnly: Bool = false + var messageNames: [String] + + init(messageNames: [String]) { + self.messageNames = messageNames + } + + func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) { + } +}