webxdc vc leak fix

deltachat · Nov 2, 2024 · f43074d · f43074d
1 parent 87dcb1d
commit f43074d
Showing 1 changed file with 51 additions and 13 deletions.
diff --git a/deltachat-ios/Controller/WebxdcViewController.swift b/deltachat-ios/Controller/WebxdcViewController.swift
@@ -233,12 +233,12 @@ class WebxdcViewController: WebViewViewController {
         let preferences = WKPreferences()
         let contentController = WKUserContentController()
 
-        contentController.add(self, name: WebxdcHandler.sendStatusUpdate.rawValue)
-        contentController.add(self, name: WebxdcHandler.log.rawValue)
-        contentController.add(self, name: WebxdcHandler.sendToChat.rawValue)
-        contentController.add(self, name: WebxdcHandler.sendRealtimeAdvertisement.rawValue)
-        contentController.add(self, name: WebxdcHandler.sendRealtimeData.rawValue)
-        contentController.add(self, name: WebxdcHandler.leaveRealtime.rawValue)
+        contentController.add(weak: self, name: WebxdcHandler.sendStatusUpdate.rawValue)
+        contentController.add(weak: self, name: WebxdcHandler.log.rawValue)
+        contentController.add(weak: self, name: WebxdcHandler.sendToChat.rawValue)
+        contentController.add(weak: self, name: WebxdcHandler.sendRealtimeAdvertisement.rawValue)
+        contentController.add(weak: self, name: WebxdcHandler.sendRealtimeData.rawValue)
+        contentController.add(weak: self, name: WebxdcHandler.leaveRealtime.rawValue)
 
         let scriptSource = """
             window.RTCPeerConnection = ()=>{};
@@ -252,8 +252,8 @@ class WebxdcViewController: WebViewViewController {
         contentController.addUserScript(script)
 
         config.userContentController = contentController
-        config.setURLSchemeHandler(self, forURLScheme: INTERNALSCHEMA)
-        
+        config.setURLSchemeHandler(weak: self, forURLScheme: INTERNALSCHEMA)
+
         config.mediaTypesRequiringUserActionForPlayback = []
         config.allowsInlineMediaPlayback = true
 
@@ -406,11 +406,7 @@ class WebxdcViewController: WebViewViewController {
         WKContentRuleListStore.default().compileContentRuleList(
             forIdentifier: "WebxdcContentBlockingRules",
             encodedContentRuleList: blockRules) { (contentRuleList, error) in
-
-                guard let contentRuleList = contentRuleList, error == nil else {
-                    return
-                }
-
+                guard let contentRuleList, error == nil else { return }
                 let configuration = self.webView.configuration
                 configuration.userContentController.add(contentRuleList)
                 self.loadHtml()
@@ -591,3 +587,45 @@ extension WebxdcViewController: WKURLSchemeHandler {
     func webView(_ webView: WKWebView, stop urlSchemeTask: WKURLSchemeTask) {
     }
 }
+
+// MARK: Memory Leak Prevention
+
+extension WKUserContentController {
+    func add(weak scriptMessageHandler: any WKScriptMessageHandler, name: String) {
+        add(WeakScriptMessageHandler(scriptMessageHandler), name: name)
+    }
+
+    private class WeakScriptMessageHandler: NSObject, WKScriptMessageHandler {
+        weak var scriptMessageHandler: WKScriptMessageHandler?
+
+        init(_ scriptMessageHandler: WKScriptMessageHandler) {
+            self.scriptMessageHandler = scriptMessageHandler
+        }
+
+        func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
+            scriptMessageHandler?.userContentController(userContentController, didReceive: message)
+        }
+    }
+}
+
+extension WKWebViewConfiguration {
+    func setURLSchemeHandler(weak urlSchemeHandler: (any WKURLSchemeHandler), forURLScheme urlScheme: String) {
+        setURLSchemeHandler(WeakURLSchemeHandler(urlSchemeHandler), forURLScheme: urlScheme)
+    }
+
+    private class WeakURLSchemeHandler: NSObject, WKURLSchemeHandler {
+        weak var urlSchemeHandler: WKURLSchemeHandler?
+
+        init(_ urlSchemeHandler: WKURLSchemeHandler) {
+            self.urlSchemeHandler = urlSchemeHandler
+        }
+
+        func webView(_ webView: WKWebView, start urlSchemeTask: WKURLSchemeTask) {
+            urlSchemeHandler?.webView(webView, start: urlSchemeTask)
+        }
+
+        func webView(_ webView: WKWebView, stop urlSchemeTask: WKURLSchemeTask) {
+            urlSchemeHandler?.webView(webView, stop: urlSchemeTask)
+        }
+    }
+}