Fix #5461: Add more robust previousUrl check for de-amping

Client/Frontend/Browser/BrowserViewController.swift

@@ -2225,6 +2225,7 @@ extension BrowserViewController: TabDelegate {
     tab.addContentScript(RewardsReporting(rewards: rewards, tab: tab), name: RewardsReporting.name(), contentWorld: .page)
     tab.addContentScript(AdsMediaReporting(rewards: rewards, tab: tab), name: AdsMediaReporting.name(), contentWorld: .defaultClient)
     tab.addContentScript(ReadyStateScriptHelper(tab: tab), name: ReadyStateScriptHelper.name(), contentWorld: .page)
+    tab.addContentScript(DeAmpHelper(tab: tab), name: DeAmpHelper.name(), contentWorld: .defaultClient)
   }
 
   func tab(_ tab: Tab, willDeleteWebView webView: WKWebView) {

Client/Frontend/Browser/BrowserViewController/BrowserViewController+WKNavigationDelegate.swift

@@ -455,8 +455,9 @@ extension BrowserViewController: WKNavigationDelegate {
 
   public func webView(_ webView: WKWebView, didCommit navigation: WKNavigation!) {
     guard let tab = tabManager[webView] else { return }
-
-    tab.url = webView.url
+    // Set the committed url which will also set tab.url
+    tab.committedURL = webView.url
+
     // Need to evaluate Night mode script injection after url is set inside the Tab
     tab.nightMode = Preferences.General.nightModeEnabled.value
 

Client/Frontend/Browser/Helpers/DeAmpHelper.swift

@@ -0,0 +1,70 @@
+// Copyright 2022 The Brave Authors. All rights reserved.
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+import Foundation
+import Shared
+import WebKit
+
+public class DeAmpHelper: TabContentScript {
+  private struct DeAmpDTO: Decodable {
+    enum CodingKeys: String, CodingKey {
+      case destURL, securityToken
+    }
+
+    let securityToken: String
+    let destURL: URL
+
+    init(from decoder: Decoder) throws {
+      let container = try decoder.container(keyedBy: CodingKeys.self)
+      self.securityToken = try container.decode(String.self, forKey: .securityToken)
+      let destURLString = try container.decode(String.self, forKey: .destURL)
+
+      guard let destURL = URL(string: destURLString) else {
+        throw DecodingError.dataCorrupted(DecodingError.Context(codingPath: decoder.codingPath, debugDescription: "`resourceURL` is not a valid URL. Fix the `RequestBlocking.js` script"))
+      }
+
+      self.destURL = destURL
+    }
+  }
+
+  private weak var tab: Tab?
+
+  init(tab: Tab) {
+    self.tab = tab
+  }
+
+  static func name() -> String {
+    return "DeAmpHelper"
+  }
+
+  static func scriptMessageHandlerName() -> String {
+    return ["deAmpHelper", UserScriptManager.messageHandlerTokenString].joined(separator: "_")
+  }
+
+  func scriptMessageHandlerName() -> String? {
+    return Self.scriptMessageHandlerName()
+  }
+
+  func userContentController(_ userContentController: WKUserContentController, didReceiveScriptMessage message: WKScriptMessage, replyHandler: @escaping (Any?, String?) -> Void) {
+    do {
+      let data = try JSONSerialization.data(withJSONObject: message.body)
+      let dto = try JSONDecoder().decode(DeAmpDTO.self, from: data)
+
+      guard dto.securityToken == UserScriptManager.securityTokenString else {
+        assertionFailure("Invalid security token. Fix the `RequestBlocking.js` script")
+        replyHandler(false, nil)
+        return
+      }
+
+      // Check that the destination is not the same as the previousURL
+      // or that previousURL is nil
+      let shouldRedirect = dto.destURL != tab?.previousComittedURL
+      replyHandler(shouldRedirect, nil)
+    } catch {
+      assertionFailure("Invalid type of message. Fix the `RequestBlocking.js` script")
+      replyHandler(false, nil)
+    }
+  }
+}

Client/Frontend/Browser/Tab.swift

@@ -117,6 +117,21 @@ class Tab: NSObject {
   fileprivate var lastRequest: URLRequest?
   var restoring: Bool = false
   var pendingScreenshot = false
+
+  /// The url set after a successful navigation. This will also set the `url` property.
+  ///
+  /// - Note: Unlike the `url` property, which may be set during pre-navigation,
+  /// the `committedURL` is only assigned when navigation was committed..
+  var committedURL: URL? {
+    willSet {
+      url = newValue
+      previousComittedURL = committedURL
+    }
+  }
+
+  /// The previous url that was set before `comittedURL` was set again
+  private(set) var previousComittedURL: URL?
+
   var url: URL? {
     didSet {
       if let _url = url, let internalUrl = InternalURL(_url), internalUrl.isAuthorized {

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

@@ -72,4 +72,15 @@ class UserScriptHelper {
 
     return userScriptTypes
   }
+
+  static func makeDeAmpScriptParamters() throws -> String {
+    let arguments: [String: String] = [
+      "handlerName": DeAmpHelper.scriptMessageHandlerName(),
+      "securityToken": UserScriptManager.securityTokenString
+    ]
+
+    let encoder = JSONEncoder()
+    let data = try encoder.encode(arguments)
+    return String(data: data, encoding: .utf8)!
+  }
 }

Client/Frontend/Browser/UserScriptManager.swift

@@ -174,16 +174,29 @@ class UserScriptManager {
   /// The first part is handled by an ad-block rule and enabled via a `deAmpEnabled` boolean in `AdBlockStats`
   /// The third part is handled by debouncing amp links and handled by debouncing logic in `DebouncingResourceDownloader`
   private let deAMPUserScript: WKUserScript? = {
-    guard let path = Bundle.current.path(forResource: "DeAMP", ofType: "js"), let source: String = try? String(contentsOfFile: path) else {
+    guard let path = Bundle.current.path(forResource: "DeAMP", ofType: "js"), var source: String = try? String(contentsOfFile: path) else {
       log.error("Failed to load cookie control user script")
       return nil
     }
-
-    return WKUserScript.create(
-      source: source,
-      injectionTime: .atDocumentStart,
-      forMainFrameOnly: true,
-      in: .defaultClient)
+
+    do {
+      let arguments = try UserScriptHelper.makeDeAmpScriptParamters()
+
+      source = [
+        source,
+        "window.braveDeAmp(\(arguments))",
+        "delete window.braveDeAmp"
+      ].joined(separator: "\n")
+
+      return WKUserScript.create(
+        source: source,
+        injectionTime: .atDocumentStart,
+        forMainFrameOnly: true,
+        in: .defaultClient)
+    } catch {
+      assertionFailure(error.localizedDescription)
+      return nil
+    }
   }()
 
   // PaymentRequestUserScript is injected at document start to handle

Client/Frontend/UserContent/UserScripts/DeAMP.js

@@ -5,12 +5,20 @@
 
 "use strict";
 
-(_ => {
+window.braveDeAmp = (args) => {
   const W = window
   const D = W.document
+  const securityToken = args.securityToken
 
   let timesToCheck = 20
   let intervalId
+
+  const sendMessage = (destURL) => {
+    return webkit.messageHandlers[args.handlerName].postMessage({
+      securityToken: securityToken,
+      destURL: destURL.href
+    })
+  }
 
   const checkIfShouldStopChecking = _ => {
     timesToCheck -= 1
@@ -45,14 +53,19 @@
     const targetHref = canonicalLinkElm.getAttribute('href')
     try {
       const destUrl = new URL(targetHref)
-      const locationUrl = W.location
       W.clearInterval(intervalId)
-      if (locationUrl.href == destUrl.href || destUrl.href == D.referrer || !(destUrl.protocol === 'http:' || destUrl.protocol === 'https:')) {
+
+      if (W.location.href == destUrl.href || !(destUrl.protocol === 'http:' || destUrl.protocol === 'https:')) {
         // Only handle http/https and only if the canoncial url is different than the current url
         // Also add a check the referrer to prevent an infinite load loop in some cases
         return
       }
-      W.location.replace(destUrl.href)
+
+      sendMessage(destUrl).then(deAmp => {
+        if (deAmp) {
+          W.location.replace(destUrl.href)
+        }
+      })
     } catch (_) {
       // Invalid canonical URL detected
       W.clearInterval(intervalId)
@@ -62,4 +75,6 @@
 
   intervalId = W.setInterval(checkForAmp, 250)
   checkForAmp()
-})()
+}
+
+// Invoke the method and delete the function