Execute print user script in page content world

The print user script needs to execute in the page content world so that it mutates the globals used by the page which the user is interacting with. The existing setup runs scripts under the default content world (which is what you generally want), but this script is a special case as we want to overwrite a global function on the window object.
duckduckgo · Oct 15, 2021 · 0a68d1c · 0a68d1c
1 parent 2a430d9
commit 0a68d1c
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 13 deletions.
diff --git a/DuckDuckGo/BrowserTab/Model/PrintingUserScript.swift b/DuckDuckGo/BrowserTab/Model/PrintingUserScript.swift
@@ -32,18 +32,24 @@ public class PrintingUserScript: NSObject, UserScript {
 
     public var source: String = """
 (function() {
-    document.addEventListener("click", function(event) {
-        event = event || window.event;
-
-        let onClickAttribute = event.target.getAttribute('onclick');
-
-        if (onClickAttribute.includes('window.print()')) {
-            webkit.messageHandlers.printHandler.postMessage({});
-        }
-    });
+    window.print = function() {
+        webkit.messageHandlers.printHandler.postMessage({});
+    };
 }) ();
 """
 
+    private func makeWKUserScriptInPage(source: String, injectionTime: WKUserScriptInjectionTime, forMainFrameOnly: Bool) -> WKUserScript {
+        if #available(macOS 11.0, iOS 14.0, *) {
+            return WKUserScript(source: source, injectionTime: injectionTime, forMainFrameOnly: forMainFrameOnly, in: .page)
+        } else {
+            return WKUserScript(source: source, injectionTime: injectionTime, forMainFrameOnly: forMainFrameOnly)
+        }
+    }
+
+    public func makeWKUserScript() -> WKUserScript {
+        return makeWKUserScriptInPage(source: source, injectionTime: injectionTime, forMainFrameOnly: forMainFrameOnly)
+    }
+
     public var injectionTime: WKUserScriptInjectionTime = .atDocumentStart
     public var forMainFrameOnly: Bool = false
     public var messageNames: [String] = ["printHandler"]

diff --git a/DuckDuckGo/BrowserTab/Model/UserScripts.swift b/DuckDuckGo/BrowserTab/Model/UserScripts.swift
@@ -41,6 +41,10 @@ final class UserScripts {
         scripts.append(autofillScript.makeWKUserScript())
     }
 
+    lazy var pageMutationScripts: [UserScript] = [
+        self.printingUserScript
+    ]
+
     lazy var userScripts: [UserScript] = [
         self.debugScript,
         self.faviconScript,
@@ -49,24 +53,25 @@ final class UserScripts {
         self.contentBlockerScript,
         self.contentBlockerRulesScript,
         self.pageObserverScript,
-        self.printingUserScript,
         self.hoverUserScript,
         self.autofillScript
     ]
 
-    lazy var scripts = userScripts.map { $0.makeWKUserScript() }
+    lazy var scripts = (pageMutationScripts + userScripts).map { $0.makeWKUserScript() }
 
 }
 
 extension UserScripts {
 
     func install(into controller: WKUserContentController) {
         scripts.forEach(controller.addUserScript)
+        pageMutationScripts.forEach(controller.addHandlerPageContentWorld)
         userScripts.forEach(controller.addHandler)
     }
 
     func remove(from controller: WKUserContentController) {
         controller.removeAllUserScripts()
+        pageMutationScripts.forEach(controller.removeHandlerPageContentWorld)
         userScripts.forEach(controller.removeHandler)
     }
 

diff --git a/DuckDuckGo/Common/Extensions/WKUserContentControllerExtension.swift b/DuckDuckGo/Common/Extensions/WKUserContentControllerExtension.swift
@@ -27,6 +27,26 @@ extension WKUserContentController {
             add(userScript, name: messageName)
         }
     }
+
+    func addHandlerPageContentWorld(_ userScript: UserScript) {
+        for messageName in userScript.messageNames {
+            if #available(macOS 11.0, *) {
+                add(userScript, contentWorld: .page, name: messageName)
+            } else {
+                add(userScript, name: messageName)
+            }
+        }
+    }
+
+    func removeHandlerPageContentWorld(_ userScript: UserScript) {
+        userScript.messageNames.forEach {
+            if #available(macOS 11.0, *) {
+                removeScriptMessageHandler(forName: $0, contentWorld: .page)
+            } else {
+                removeScriptMessageHandler(forName: $0)
+            }
+        }
+    }
 
     func addHandler(_ userScript: UserScript) {
         for messageName in userScript.messageNames {

diff --git a/Unit Tests/UserScripts/UserScriptsTest.swift b/Unit Tests/UserScripts/UserScriptsTest.swift
@@ -29,7 +29,7 @@ final class UserScriptsTests: XCTestCase {
 
         // swiftlint:disable force_cast
         let expected = Mirror(reflecting: userScripts).children.compactMap { $0.value as? UserScript } as! [NSObject]
-        let actual = Set(userScripts.userScripts as! [NSObject])
+        let actual = Set((userScripts.userScripts + userScripts.pageMutationScripts) as! [NSObject])
         // swiftlint:enable force_cast
 
         XCTAssertEqual(actual, Set(expected))
@@ -83,6 +83,7 @@ final class UserScriptsTests: XCTestCase {
 
         let scripts = Set(userScripts.scripts)
         let messageHandlersCount = userScripts.userScripts.reduce(0) { $0 + $1.messageNames.count }
+        let mutationScriptsCount = userScripts.pageMutationScripts.reduce(0) { $0 + $1.messageNames.count }
 
         let e = expectation(description: "Should receive message")
         let instrumentation = TabInstrumentationMock {
@@ -103,9 +104,10 @@ final class UserScriptsTests: XCTestCase {
                                                                           didReceive: message)
 
         let installedScripts = Set(userContentController.userScripts)
+        let count = messageHandlersCount + mutationScriptsCount
 
         XCTAssertEqual(scripts, installedScripts)
-        XCTAssertEqual(messageHandlersCount, userContentController.handlers.count + userContentController.handlersWithReply.count)
+        XCTAssertEqual(count, userContentController.handlers.count + userContentController.handlersWithReply.count)
         waitForExpectations(timeout: 0.3)
     }