Fix #5859: Fix Security Tokens issue

Client/Frontend/Browser/Handlers/BraveSearchScriptHandler.swift

@@ -32,7 +32,9 @@ class BraveSearchScriptHandler: TabContentScript {
 
   static func name() -> String { "BraveSearchHelper" }
 
-  func scriptMessageHandlerName() -> String? { BraveSearchScriptHandler.name() }
+  func scriptMessageHandlerName() -> String? {
+    "BraveSearchHelper_\(UserScriptManager.messageHandlerTokenString)"
+  }
 
   private enum Method: Int {
     case canSetBraveSearchAsDefault = 1
@@ -52,46 +54,46 @@ class BraveSearchScriptHandler: TabContentScript {
     didReceiveScriptMessage message: WKScriptMessage,
     replyHandler: (Any?, String?) -> Void
   ) {
-    defer { replyHandler(nil, nil) }
     let allowedHosts = DomainUserScript.braveSearchHelper.associatedDomains
 
     guard let requestHost = message.frameInfo.request.url?.host,
       allowedHosts.contains(requestHost),
       message.frameInfo.isMainFrame
     else {
       log.error("Backup search request called from disallowed host")
+      replyHandler(nil, nil)
       return
     }
 
     guard let data = try? JSONSerialization.data(withJSONObject: message.body, options: []),
       let method = try? JSONDecoder().decode(MethodModel.self, from: data).methodId
     else {
       log.error("Failed to retrieve method id")
+      replyHandler(nil, nil)
       return
     }
 
     switch method {
     case Method.canSetBraveSearchAsDefault.rawValue:
-      handleCanSetBraveSearchAsDefault(methodId: method)
+      handleCanSetBraveSearchAsDefault(replyHandler: replyHandler)
     case Method.setBraveSearchDefault.rawValue:
-      handleSetBraveSearchDefault(methodId: method)
+      handleSetBraveSearchDefault(replyHandler: replyHandler)
     default:
       break
     }
   }
 
-  private func handleCanSetBraveSearchAsDefault(methodId: Int) {
-
+  private func handleCanSetBraveSearchAsDefault(replyHandler: (Any?, String?) -> Void) {
     if PrivateBrowsingManager.shared.isPrivateBrowsing {
       log.debug("Private mode detected, skipping setting Brave Search as a default")
-      callback(methodId: methodId, result: false)
+      replyHandler(false, nil)
       return
     }
 
     let maximumPromptCount = Preferences.Search.braveSearchDefaultBrowserPromptCount
     if Self.canSetAsDefaultCounter >= maxCountOfDefaultBrowserPromptsPerSession || maximumPromptCount.value >= maxCountOfDefaultBrowserPromptsTotal {
       log.debug("Maximum number of tries of Brave Search website prompts reached")
-      callback(methodId: methodId, result: false)
+      replyHandler(false, nil)
       return
     }
 
@@ -100,27 +102,11 @@ class BraveSearchScriptHandler: TabContentScript {
 
     let defaultEngine = profile.searchEngines.defaultEngine(forType: .standard).shortName
     let canSetAsDefault = defaultEngine != OpenSearchEngine.EngineNames.brave
-
-    callback(methodId: methodId, result: canSetAsDefault)
+    replyHandler(canSetAsDefault, nil)
   }
 
-  private func handleSetBraveSearchDefault(methodId: Int) {
+  private func handleSetBraveSearchDefault(replyHandler: (Any?, String?) -> Void) {
     profile.searchEngines.updateDefaultEngine(OpenSearchEngine.EngineNames.brave, forType: .standard)
-    callback(methodId: methodId, result: nil)
-  }
-
-  private func callback(methodId: Int, result: Bool?) {
-    let functionName =
-      "window.__firefox__.BSH\(UserScriptManager.messageHandlerTokenString).resolve"
-
-    var args: [Any] = [methodId]
-    if let result = result {
-      args.append(result)
-    }
-
-    self.tab?.webView?.evaluateSafeJavaScript(
-      functionName: functionName,
-      args: args,
-      contentWorld: .page)
+    replyHandler(nil, nil)
   }
 }

Client/Frontend/Browser/Handlers/BraveTalkScriptHandler.swift

@@ -26,7 +26,7 @@ class BraveTalkScriptHandler: TabContentScript {
 
   static func name() -> String { "BraveTalkHelper" }
 
-  func scriptMessageHandlerName() -> String? { BraveTalkScriptHandler.name() }
+  func scriptMessageHandlerName() -> String? { "BraveTalkHelper_\(UserScriptManager.messageHandlerTokenString)" }
 
   func userContentController(
     _ userContentController: WKUserContentController,

Client/Frontend/Browser/User Scripts/ScriptFactory.swift

@@ -70,7 +70,7 @@ class ScriptFactory {
     case .farblingProtection(let etld):
       let randomConfiguration = RandomConfiguration(etld: etld)
       let fakeParams = try FarblingProtectionHelper.makeFarblingParams(from: randomConfiguration)
-      source = "\(source)\nwindow.braveFarble(\(fakeParams))\ndelete window.braveFarble"
+      source = source.replacingOccurrences(of: "$<farbling_protection_args>", with: fakeParams)
 
     case .nacl:
       // No modifications needed
@@ -95,15 +95,15 @@ class ScriptFactory {
 
       case .braveSearchHelper:
         let securityToken = UserScriptManager.securityTokenString
-        let messageToken = "BSH\(UserScriptManager.messageHandlerTokenString)"
+        let messageToken = "BraveSearchHelper_\(UserScriptManager.messageHandlerTokenString)"
 
         source = source
           .replacingOccurrences(of: "$<brave-search-helper>", with: messageToken, options: .literal)
           .replacingOccurrences(of: "$<security_token>", with: securityToken)
 
       case .braveTalkHelper:
         let securityToken = UserScriptManager.securityTokenString
-        let messageToken = "BT\(UserScriptManager.messageHandlerTokenString)"
+        let messageToken = "BraveTalkHelper_\(UserScriptManager.messageHandlerTokenString)"
 
         source = source
           .replacingOccurrences(of: "$<brave-talk-helper>", with: messageToken, options: .literal)

Client/Frontend/Browser/UserScriptManager.swift

@@ -190,11 +190,7 @@ class UserScriptManager {
         return nil
       }
 
-      source = [
-        source,
-        "window.braveBlockRequests(\(fakeParams))",
-        "delete window.braveBlockRequests"
-      ].joined(separator: "\n")
+      source = source.replacingOccurrences(of: "$<request_blocking_args>", with: fakeParams)
 
       return WKUserScript.create(
         source: source,
@@ -225,12 +221,8 @@ class UserScriptManager {
         assertionFailure("A nil here is impossible")
         return nil
       }
-
-      source = [
-        source,
-        "window.braveDeAmp(\(arguments))",
-        "delete window.braveDeAmp"
-      ].joined(separator: "\n")
+
+      source = source.replacingOccurrences(of: "$<deamp_args>", with: arguments, options: .literal)
 
       return WKUserScript.create(
         source: source,

Client/Frontend/UserContent/UserScripts/AllFrames/AtDocumentEnd/PrintHandler.js → Client/Frontend/UserContent/UserScripts/AllFrames/AtDocumentStart/PrintHandler.js

@@ -8,7 +8,11 @@
 // Ensure this module only gets included once. This is
 // required for user scripts injected into all frames.
 window.__firefox__.includeOnce("PrintHandler", function() {
-  window.print = function() {
+  function postMessage() {
     webkit.messageHandlers.printHandler.postMessage({"securitytoken": SECURITY_TOKEN});
+  }
+
+  window.print = function() {
+    postMessage();
   };
 });

Client/Frontend/UserContent/UserScripts/BraveSearchHelper.js

@@ -5,45 +5,22 @@
 
 'use strict';
 
-Object.defineProperty(window.__firefox__, '$<brave-search-helper>', {
+window.__firefox__.includeOnce("BraveSearchHelper", function() {
+  function sendMessage(method_id) {
+    return webkit.messageHandlers.$<brave-search-helper>.postMessage({ 'securitytoken': '$<security_token>' ,'method_id': method_id});
+  }
+
+  Object.defineProperty(window, 'brave', {
     enumerable: false,
     configurable: true,
     writable: false,
     value: {
-        id: 1,
-        resolution_handlers: {},
-        resolve(id, data, error) {
-            if (error && window.__firefox__.$<brave-search-helper>.resolution_handlers[id].reject) {
-                window.__firefox__.$<brave-search-helper>.resolution_handlers[id].reject(error);
-            } else if (window.__firefox__.$<brave-search-helper>.resolution_handlers[id].resolve) {
-                window.__firefox__.$<brave-search-helper>.resolution_handlers[id].resolve(data);
-            } else if (window.__firefox__.$<brave-search-helper>.resolution_handlers[id].reject) {
-                window.__firefox__.$<brave-search-helper>.resolution_handlers[id].reject(new Error("Invalid Data!"));
-            } else {
-                console.log("Invalid Promise ID: ", id);
-            }
-
-            delete window.__firefox__.$<brave-search-helper>.resolution_handlers[id];
-        },
-        sendMessage(method_id) {
-            return new Promise((resolve, reject) => {
-               window.__firefox__.$<brave-search-helper>.resolution_handlers[method_id] = { resolve, reject };
-               webkit.messageHandlers.BraveSearchHelper.postMessage({ 'securitytoken': '$<security_token>' ,'method_id': method_id});
-           });
-        }
-    }
-});
-
-Object.defineProperty(window, 'brave', {
-    enumerable: false,
-    configurable: true,
-    writable: false,
-        value: {
-        getCanSetDefaultSearchProvider() {
-            return window.__firefox__.$<brave-search-helper>.sendMessage(1);
-        },
-        setIsDefaultSearchProvider() {
-            return window.__firefox__.$<brave-search-helper>.sendMessage(2);
-        }
+      getCanSetDefaultSearchProvider() {
+        return sendMessage(1);
+      },
+      setIsDefaultSearchProvider() {
+        return sendMessage(2);
+      }
     }
+  });
 });

Client/Frontend/UserContent/UserScripts/BraveTalkHelper.js

@@ -5,42 +5,19 @@
 
 'use strict';
 
-Object.defineProperty(window.__firefox__, '$<brave-talk-helper>', {
+window.__firefox__.includeOnce("BraveTalkHelper", function() {
+  function sendMessage() {
+    return webkit.messageHandlers.$<brave-talk-helper>.postMessage({ 'securitytoken': '$<security_token>' });
+  }
+
+  Object.defineProperty(window, 'chrome', {
     enumerable: false,
     configurable: true,
     writable: false,
     value: {
-        id: 1,
-        resolution_handlers: {},
-        resolve(id, data, error) {
-            if (error && window.__firefox__.$<brave-talk-helper>.resolution_handlers[id].reject) {
-                window.__firefox__.$<brave-talk-helper>.resolution_handlers[id].reject(error);
-            } else if (window.__firefox__.$<brave-talk-helper>.resolution_handlers[id].resolve) {
-                window.__firefox__.$<brave-talk-helper>.resolution_handlers[id].resolve(data);
-            } else if (window.__firefox__.$<brave-talk-helper>.resolution_handlers[id].reject) {
-                window.__firefox__.$<brave-talk-helper>.resolution_handlers[id].reject(new Error("Invalid Data!"));
-            } else {
-                console.log("Invalid Promise ID: ", id);
-            }
-
-            delete window.__firefox__.$<brave-talk-helper>.resolution_handlers[id];
-        },
-        sendMessage() {
-            return new Promise((resolve, reject) => {
-               window.__firefox__.$<brave-talk-helper>.resolution_handlers[1] = { resolve, reject };
-               webkit.messageHandlers.BraveTalkHelper.postMessage({ 'securitytoken': '$<security_token>' });
-           });
-        }
+      braveRequestAdsEnabled() {
+        return sendMessage();
+      }
     }
-});
-
-Object.defineProperty(window, 'chrome', {
-enumerable: false,
-configurable: true,
-writable: false,
-    value: {
-    braveRequestAdsEnabled() {
-        return window.__firefox__.$<brave-talk-helper>.sendMessage();
-    }
-}
+  });
 });