Move tab helper to be owned by tabs.

brave · Jan 6, 2025 · ad04152 · ad04152
1 parent eeea5b6
commit ad04152
Show file tree

Hide file tree

Showing 7 changed files with 82 additions and 45 deletions.
diff --git a/ios/brave-ios/Sources/Brave/Frontend/Brave Translate/BraveTranslateTabHelper.swift b/ios/brave-ios/Sources/Brave/Frontend/Brave Translate/BraveTranslateTabHelper.swift
@@ -20,7 +20,7 @@ enum BraveTranslateError: String, Error {
   case otherError
 }
 
-class BraveTranslateTabHelper: NSObject {
+class BraveTranslateTabHelper: NSObject, TabHelper {
   private weak var tab: Tab?
   private weak var delegate: BraveTranslateScriptHandlerDelegate?
   private let recognizer = NLLanguageRecognizer()
@@ -35,6 +35,8 @@ class BraveTranslateTabHelper: NSObject {
 
   var currentLanguageInfo = BraveTranslateLanguageInfo()
 
+  static let tabHelperName = String(describing: BraveTranslateTabHelper.self)
+
   // All TabHelpers in Chromium have a `WebState* web_state` parameter in their constructor
   // WebState in Brave, is the same as `Tab`.
   init(tab: Tab, delegate: BraveTranslateScriptHandlerDelegate) {

diff --git a/ios/brave-ios/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+ShareActivity.swift b/ios/brave-ios/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+ShareActivity.swift
@@ -106,10 +106,11 @@ extension BrowserViewController {
           callback: { [weak self] in
             guard let self = self, let tab = tab else { return }
 
-            let scriptHandler =
-              tab.getContentScript(name: BraveTranslateScriptHandler.scriptName)
-              as? BraveTranslateScriptHandler
-            if let tabHelper = scriptHandler?.tabHelper {
+            let tabHelper =
+              tab.getTabHelper(named: BraveTranslateTabHelper.tabHelperName)
+              as? BraveTranslateTabHelper
+
+            if let tabHelper {
               tabHelper.presentUI(on: self)
 
               if tab.translationState == .active {

diff --git a/ios/brave-ios/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+ToolbarDelegate.swift b/ios/brave-ios/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+ToolbarDelegate.swift
@@ -263,11 +263,10 @@ extension BrowserViewController: TopToolbarDelegate {
 
   func topToolbarDidPressTranslateButton(_ urlBar: TopToolbarView) {
     guard let tab = tabManager.selectedTab else { return }
-    let scriptHandler =
-      tab.getContentScript(name: BraveTranslateScriptHandler.scriptName)
-      as? BraveTranslateScriptHandler
+    let tabHelper =
+      tab.getTabHelper(named: BraveTranslateTabHelper.tabHelperName) as? BraveTranslateTabHelper
 
-    if let tabHelper = scriptHandler?.tabHelper {
+    if let tabHelper {
       tabHelper.presentUI(on: self)
 
       if tab.translationState == .active {

diff --git a/ios/brave-ios/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+Translate.swift b/ios/brave-ios/Sources/Brave/Frontend/Browser/BrowserViewController/BVC+Translate.swift
@@ -21,10 +21,10 @@ extension BrowserViewController: BraveTranslateScriptHandlerDelegate {
       topToolbar.updateTranslateButtonState(state)
 
       showTranslateOnboarding(tab: tab) { [weak tab] translateEnabled in
-        let scriptHandler =
-          tab?.getContentScript(name: BraveTranslateScriptHandler.scriptName)
-          as? BraveTranslateScriptHandler
-        if let tabHelper = scriptHandler?.tabHelper {
+        let tabHelper =
+          tab?.getTabHelper(named: BraveTranslateTabHelper.tabHelperName)
+          as? BraveTranslateTabHelper
+        if let tabHelper {
           tabHelper.startTranslation(canShowToast: true)
         }
       }
@@ -63,12 +63,11 @@ extension BrowserViewController: BraveTranslateScriptHandlerDelegate {
               self.topToolbar.locationView.translateButton.setOnboardingState(enabled: false)
               Preferences.Translate.translateEnabled.value = true
 
-              let scriptHandler =
-                tab.getContentScript(
-                  name: BraveTranslateScriptHandler.scriptName
-                ) as? BraveTranslateScriptHandler
+              let tabHelper =
+                tab.getTabHelper(named: BraveTranslateTabHelper.tabHelperName)
+                as? BraveTranslateTabHelper
 
-              scriptHandler?.tabHelper.presentUI(on: self)
+              tabHelper?.presentUI(on: self)
             },
             onDisableFeature: { [weak self, weak tab] in
               guard let tab = tab, let self = self else { return }
@@ -102,13 +101,11 @@ extension BrowserViewController: BraveTranslateScriptHandlerDelegate {
           self.topToolbar.locationView.translateButton.setOnboardingState(enabled: false)
 
           if Preferences.Translate.translateEnabled.value {
-            let scriptHandler =
-              tab.getContentScript(
-                name: BraveTranslateScriptHandler.scriptName
-              )
-              as? BraveTranslateScriptHandler
+            let tabHelper =
+              tab.getTabHelper(named: BraveTranslateTabHelper.tabHelperName)
+              as? BraveTranslateTabHelper
 
-            scriptHandler?.tabHelper.presentUI(on: self)
+            tabHelper?.presentUI(on: self)
 
             completion(true)
             return
@@ -127,13 +124,11 @@ extension BrowserViewController: BraveTranslateScriptHandlerDelegate {
     let popover = PopoverController(
       content: TranslateToast(languageInfo: languageInfo) { [weak tab] languageInfo in
 
-        let scriptHandler =
-          tab?.getContentScript(
-            name: BraveTranslateScriptHandler.scriptName
-          )
-          as? BraveTranslateScriptHandler
+        let tabHelper =
+          tab?.getTabHelper(named: BraveTranslateTabHelper.tabHelperName)
+          as? BraveTranslateTabHelper
 
-        scriptHandler?.tabHelper.startTranslation(canShowToast: false)
+        tabHelper?.startTranslation(canShowToast: false)
       },
       autoLayoutConfiguration: nil
     )

diff --git a/...rave-ios/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController.swift b/...rave-ios/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController.swift
@@ -2659,12 +2659,8 @@ extension BrowserViewController: TabDelegate {
       tab.requestBlockingContentHelper,
     ]
 
-    let braveTranslateTabHelper = BraveTranslateTabHelper(tab: tab, delegate: self)
-
-    injectedScripts.append(
-      BraveTranslateScriptLanguageDetectionHandler(tabHelper: braveTranslateTabHelper)
-    )
-    injectedScripts.append(BraveTranslateScriptHandler(tabHelper: braveTranslateTabHelper))
+    injectedScripts.append(BraveTranslateScriptLanguageDetectionHandler(tab: tab))
+    injectedScripts.append(BraveTranslateScriptHandler(tab: tab))
 
     #if canImport(BraveTalk)
     injectedScripts.append(
@@ -2714,6 +2710,9 @@ extension BrowserViewController: TabDelegate {
       as? PlaylistFolderSharingScriptHandler)?.delegate = self
     (tab.getContentScript(name: Web3NameServiceScriptHandler.scriptName)
       as? Web3NameServiceScriptHandler)?.delegate = self
+
+    // Tab Helpers
+    tab.addTabHelper(BraveTranslateTabHelper(tab: tab, delegate: self))
   }
 
   func tab(_ tab: Tab, willDeleteWebView webView: WKWebView) {

diff --git a/ios/brave-ios/Sources/Brave/Frontend/Browser/Tab.swift b/ios/brave-ios/Sources/Brave/Frontend/Browser/Tab.swift
@@ -45,6 +45,10 @@ protocol TabContentScript: TabContentScriptLoader {
   )
 }
 
+protocol TabHelper {
+  static var tabHelperName: String { get }
+}
+
 protocol TabDelegate {
   func tab(_ tab: Tab, didAddSnackbar bar: SnackBar)
   func tab(_ tab: Tab, didRemoveSnackbar bar: SnackBar)
@@ -172,6 +176,8 @@ class Tab: NSObject {
 
   var sslPinningError: Error?
 
+  private var tabHelpers = [String: TabHelper]()
+
   private let _syncTab: BraveSyncTab?
   private let _faviconDriver: FaviconDriver?
   private var _walletEthProvider: BraveWalletEthereumProvider?
@@ -628,6 +634,7 @@ class Tab: NSObject {
 
   func deleteWebView() {
     contentScriptManager.uninstall(from: self)
+    tabHelpers.removeAll()
 
     if let webView = webView {
       webView.removeObserver(self, forKeyPath: KVOConstants.url.keyPath)
@@ -879,6 +886,23 @@ class Tab: NSObject {
     webView.customUserAgent = desktopMode ? UserAgent.desktop : UserAgent.mobile
   }
 
+  func addTabHelper(_ helper: TabHelper) {
+    let helperName = Swift.type(of: helper).tabHelperName
+    if tabHelpers[helperName] != nil {
+      assertionFailure("Tab Helper: \(helperName) already attached to this tab.")
+    }
+
+    self.tabHelpers[helperName] = helper
+  }
+
+  func getTabHelper(named name: String) -> TabHelper? {
+    self.tabHelpers[name]
+  }
+
+  func removeTabHelper(named name: String) {
+    self.tabHelpers.removeValue(forKey: name)
+  }
+
   func addContentScript(_ helper: TabContentScript, name: String, contentWorld: WKContentWorld) {
     contentScriptManager.addContentScript(
       helper,

diff --git a/...nt/UserScripts/Scripts_Dynamic/ScriptHandlers/Sandboxed/BraveTranslateScriptHandler.swift b/...nt/UserScripts/Scripts_Dynamic/ScriptHandlers/Sandboxed/BraveTranslateScriptHandler.swift
@@ -18,11 +18,11 @@ protocol BraveTranslateScriptHandlerDelegate: NSObject {
 }
 
 class BraveTranslateScriptHandler: NSObject, TabContentScript {
-  let tabHelper: BraveTranslateTabHelper
+  private weak var tab: Tab?
   private static var elementScriptTask: Task<String, Error> = downloadElementScript()
 
-  init(tabHelper: BraveTranslateTabHelper) {
-    self.tabHelper = tabHelper
+  init(tab: Tab) {
+    self.tab = tab
     super.init()
   }
 
@@ -93,22 +93,32 @@ class BraveTranslateScriptHandler: NSObject, TabContentScript {
     }
 
     if command == "ready" {
-      Task { @MainActor [weak tabHelper] in
-        try await tabHelper?.setupTranslate()
+      Task { @MainActor [weak tab] in
+        try await
+          (tab?.getTabHelper(named: BraveTranslateTabHelper.tabHelperName)
+          as? BraveTranslateTabHelper)?.setupTranslate()
         replyHandler(nil, nil)
       }
 
       return
     }
 
     if command == "request" {
-      Task { @MainActor in
+      Task { @MainActor [weak tab] in
         do {
           let message = try JSONDecoder().decode(
             BraveTranslateSession.RequestMessage.self,
             from: JSONSerialization.data(withJSONObject: body, options: .fragmentsAllowed)
           )
 
+          guard let tab = tab,
+            let tabHelper = tab.getTabHelper(named: BraveTranslateTabHelper.tabHelperName)
+              as? BraveTranslateTabHelper
+          else {
+            replyHandler(nil, BraveTranslateError.otherError.rawValue)
+            return
+          }
+
           let (data, response) = try await tabHelper.processTranslationRequest(message)
 
           replyHandler(
@@ -163,11 +173,11 @@ class BraveTranslateScriptHandler: NSObject, TabContentScript {
 }
 
 class BraveTranslateScriptLanguageDetectionHandler: NSObject, TabContentScript {
-  private let tabHelper: BraveTranslateTabHelper
+  private weak var tab: Tab?
   private static let namespace = "translate_\(uniqueID)"
 
-  init(tabHelper: BraveTranslateTabHelper) {
-    self.tabHelper = tabHelper
+  init(tab: Tab) {
+    self.tab = tab
     super.init()
   }
 
@@ -190,6 +200,13 @@ class BraveTranslateScriptLanguageDetectionHandler: NSObject, TabContentScript {
       return
     }
 
+    guard
+      let tabHelper = tab?.getTabHelper(named: BraveTranslateTabHelper.tabHelperName)
+        as? BraveTranslateTabHelper
+    else {
+      return
+    }
+
     do {
       let message = try JSONDecoder().decode(
         Message.self,