From ff9059c4cf999c5fedd9562d61aa3230b8003fab Mon Sep 17 00:00:00 2001
From: Domenico Corvasce <dcorvasce@outlook.com>
Date: Fri, 7 Jul 2023 20:35:09 +0200
Subject: [PATCH] Move STWebpageController inside Tab

---
 Sources/Brave/Frontend/Browser/Tab.swift | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/Sources/Brave/Frontend/Browser/Tab.swift b/Sources/Brave/Frontend/Browser/Tab.swift
index 747a5903b6b..f20bc3bbff0 100644
--- a/Sources/Brave/Frontend/Browser/Tab.swift
+++ b/Sources/Brave/Frontend/Browser/Tab.swift
@@ -13,6 +13,7 @@ import Data
 import os.log
 import BraveWallet
 import Favicon
+import ScreenTime
 
 protocol TabContentScriptLoader {
   static func loadUserScript(named: String) -> String?
@@ -65,6 +66,7 @@ enum TabSecureContentState {
 class Tab: NSObject {
   let id: UUID
   let rewardsId: UInt32
+  let screenTimeViewController = STWebpageController()
 
   var onScreenshotUpdated: (() -> Void)?
   var rewardsEnabledCallback: ((Bool) -> Void)?
@@ -74,7 +76,6 @@ class Tab: NSObject {
   
   private(set) var type: TabType = .regular
   
-  
   var redirectURLs = [URL]()
 
   var isPrivate: Bool {
@@ -176,9 +177,10 @@ class Tab: NSObject {
     willSet {
       url = newValue
       previousComittedURL = committedURL
+      screenTimeViewController.url = isPrivate ? nil : url
     }
   }
-  
+
   /// The previous url that was set before `comittedURL` was set again
   private(set) var previousComittedURL: URL?
   
@@ -307,6 +309,7 @@ class Tab: NSObject {
 
     super.init()
     self.type = type
+    self.screenTimeViewController.suppressUsageRecording = type.isPrivate
   }
 
   weak var navigationDelegate: WKNavigationDelegate? {
@@ -343,6 +346,11 @@ class Tab: NSObject {
       let webView = TabWebView(frame: .zero, tab: self, configuration: configuration!, isPrivate: isPrivate)
       webView.delegate = self
 
+      webView.addSubview(screenTimeViewController.view)
+      screenTimeViewController.view.snp.makeConstraints {
+        $0.edges.equalToSuperview()
+      }
+
       webView.accessibilityLabel = Strings.webContentAccessibilityLabel
       webView.allowsBackForwardNavigationGestures = true
       webView.allowsLinkPreview = true