DBP: Fix memory leak on WebViewHandler (#2483)

duckduckgo · Mar 22, 2024 · 66fdc78 · 66fdc78
1 parent 6f4150c
commit 66fdc78
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 19 deletions.
diff --git a/LocalPackages/DataBrokerProtection/Sources/DataBrokerProtection/CCF/WebViewHandler.swift b/LocalPackages/DataBrokerProtection/Sources/DataBrokerProtection/CCF/WebViewHandler.swift
@@ -27,7 +27,7 @@ protocol WebViewHandler: NSObject {
     func load(url: URL) async throws
     func takeSnaphost(path: String, fileName: String) async throws
     func saveHTML(path: String, fileName: String) async throws
-    func waitForWebViewLoad(timeoutInSeconds: Int) async throws
+    func waitForWebViewLoad() async throws
     func finish() async
     func execute(action: Action, data: CCFRequestData) async
     func evaluateJavaScript(_ javaScript: String) async throws
@@ -38,7 +38,7 @@ final class DataBrokerProtectionWebViewHandler: NSObject, WebViewHandler {
     private var activeContinuation: CheckedContinuation<Void, Error>?
 
     private let isFakeBroker: Bool
-    private let webViewConfiguration: WKWebViewConfiguration
+    private var webViewConfiguration: WKWebViewConfiguration?
     private var userContentController: DataBrokerUserContentController?
 
     private var webView: WebView?
@@ -59,7 +59,11 @@ final class DataBrokerProtectionWebViewHandler: NSObject, WebViewHandler {
     }
 
     func initializeWebView(showWebView: Bool) async {
-        webView = WebView(frame: CGRect(origin: .zero, size: CGSize(width: 1024, height: 1024)), configuration: webViewConfiguration)
+        guard let configuration = self.webViewConfiguration else {
+            return
+        }
+
+        webView = WebView(frame: CGRect(origin: .zero, size: CGSize(width: 1024, height: 1024)), configuration: configuration)
         webView?.navigationDelegate = self
 
         if showWebView {
@@ -78,40 +82,30 @@ final class DataBrokerProtectionWebViewHandler: NSObject, WebViewHandler {
     func load(url: URL) async throws {
         webView?.load(url)
         os_log("Loading URL: %@", log: .action, String(describing: url.absoluteString))
-        try await waitForWebViewLoad(timeoutInSeconds: 120)
+        try await waitForWebViewLoad()
     }
 
     func finish() {
         os_log("WebViewHandler finished", log: .action)
-
         webView?.stopLoading()
         userContentController?.cleanUpBeforeClosing()
         WKWebsiteDataStore.default().removeData(ofTypes: [WKWebsiteDataTypeDiskCache, WKWebsiteDataTypeMemoryCache], modifiedSince: Date(timeIntervalSince1970: 0)) {
             os_log("WKWebView data store deleted correctly", log: .action)
         }
 
+        webViewConfiguration = nil
         userContentController = nil
         webView?.navigationDelegate = nil
         webView = nil
     }
 
     deinit {
-        print("WebViewHandler Deinit")
+        os_log("WebViewHandler Deinit", log: .action)
     }
 
-    func waitForWebViewLoad(timeoutInSeconds: Int = 0) async throws {
+    func waitForWebViewLoad() async throws {
         try await withCheckedThrowingContinuation { continuation in
             self.activeContinuation = continuation
-
-            if timeoutInSeconds > 0 {
-                Task {
-                    try await Task.sleep(nanoseconds: UInt64(timeoutInSeconds) * NSEC_PER_SEC)
-                    if self.activeContinuation != nil {
-                        self.activeContinuation?.resume()
-                        self.activeContinuation = nil
-                    }
-                }
-            }
         }
     }
 

diff --git a/...ckages/DataBrokerProtection/Sources/DataBrokerProtection/DebugUI/DebugScanOperation.swift b/...ckages/DataBrokerProtection/Sources/DataBrokerProtection/DebugUI/DebugScanOperation.swift
@@ -166,6 +166,8 @@ final class DebugScanOperation: DataBrokerOperation {
         } else {
             os_log("Releasing the web view", log: .action)
             await webViewHandler?.finish() // If we executed all steps we release the web view
+            continuation = nil
+            webViewHandler = nil
         }
     }
 
@@ -178,4 +180,8 @@ final class DebugScanOperation: DataBrokerOperation {
             await completeWith(error: error)
         }
     }
+
+    deinit {
+        os_log("DebugScanOperation Deinit", log: .action)
+    }
 }
diff --git a/...es/DataBrokerProtection/Sources/DataBrokerProtection/Operations/DataBrokerOperation.swift b/...es/DataBrokerProtection/Sources/DataBrokerProtection/Operations/DataBrokerOperation.swift
@@ -186,7 +186,9 @@ extension DataBrokerOperation {
         switch actionType {
         case .click:
             stageCalculator?.fireOptOutFillForm()
-            try? await webViewHandler?.waitForWebViewLoad(timeoutInSeconds: 30)
+            try? await webViewHandler?.waitForWebViewLoad()
+            // We wait 10 seconds before tapping
+            try? await Task.sleep(nanoseconds: UInt64(10) * 1_000_000_000)
             await executeNextStep()
         case .fillForm:
             stageCalculator?.fireOptOutFillForm()

diff --git a/LocalPackages/DataBrokerProtection/Tests/DataBrokerProtectionTests/Mocks.swift b/LocalPackages/DataBrokerProtection/Tests/DataBrokerProtectionTests/Mocks.swift
@@ -191,7 +191,7 @@ final class WebViewHandlerMock: NSObject, WebViewHandler {
         wasLoadCalledWithURL = url
     }
 
-    func waitForWebViewLoad(timeoutInSeconds: Int) async throws {
+    func waitForWebViewLoad() async throws {
         wasWaitForWebViewLoadCalled = true
     }