Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

allow to register early scripts #930

Merged
merged 3 commits into from
Aug 7, 2024
Merged

allow to register early scripts #930

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
02b16fe
allow to register early scripts
SabrinaTardio Aug 7, 2024
79c90b7
fix lint
SabrinaTardio Aug 7, 2024
b44d314
refactor name and add comments
SabrinaTardio Aug 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion Sources/BrowserServicesKit/ContentScopeScript/UserContentController.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,12 +113,14 @@ final public class UserContentController: WKUserContentController {
private let scriptMessageHandler = PermanentScriptMessageHandler()

@MainActor
public init<Pub, Content>(assetsPublisher: Pub, privacyConfigurationManager: PrivacyConfigurationManaging)
public init<Pub, Content>(assetsPublisher: Pub, privacyConfigurationManager: PrivacyConfigurationManaging, earlyAccessScripts: [UserScript] = [])
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the naming here may be misleading since we‘re only installing the handlers and the scripts are ignored. A previous naming (earlyAccessHandlers) and some comment would do a better job IMO

where Pub: Publisher, Content: UserContentControllerNewContent, Pub.Output == Content, Pub.Failure == Never {

self.privacyConfigurationManager = privacyConfigurationManager
super.init()

installUserScripts([], handlers: earlyAccessScripts)

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
Expand Down
53 changes: 53 additions & 0 deletions Tests/BrowserServicesKitTests/ContentBlocker/UserContentControllerTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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(), earlyAccessScripts: [script1, script2])
ucc.delegate = self

XCTAssertTrue(ucc.registeredScriptHandlerNames.contains("message1"))
XCTAssertTrue(ucc.registeredScriptHandlerNames.contains("message2"))
}

@MainActor
func testWhenContentBlockingAssetsPublished_contentRuleListsAreInstalled() async throws {
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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) {
}
}
Loading