Merge pull request #57 from klee0kai/fix_multi_subscribe_crash

protect multi subscribe on chart
tradingview · Aug 29, 2023 · 5394d40 · 5394d40
2 parents 73e5800 + aabd6dc
commit 5394d40
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 5 deletions.
diff --git a/Sources/LightweightCharts/Implementations/API/Chart.swift b/Sources/LightweightCharts/Implementations/API/Chart.swift
@@ -11,6 +11,11 @@ public protocol ChartDelegate: AnyObject {
 // MARK: -
 class Chart: JavaScriptObject {
 
+    enum SubscribeState: CaseIterable {
+        case declared
+        case active
+    }
+
     typealias Context = JavaScriptEvaluator & JavaScriptMessageProducer
 
     let jsName = "chart" + .uniqueString
@@ -20,6 +25,7 @@ class Chart: JavaScriptObject {
     private unowned var context: Context
     private let messageHandler: MessageHandler
     private weak var closureStore: ClosuresStore?
+    private var activeSubscriptions: Dictionary<Subscription,SubscribeState> = [:]
 
     init(context: Context, closureStore: ClosuresStore?) {
         self.context = context
@@ -54,17 +60,30 @@ class Chart: JavaScriptObject {
     }
 
     private func subscribe(subscription: Subscription) {
+        if (activeSubscriptions[subscription] == .active) {
+            NSLog("LWChart: double subscribe detected \(subscription)")
+            return
+        }
         let name = subscriberName(for: subscription)
-        let subscriberScript = subsriberScript(forName: name, subscription: subscription)
+        var subscriberScript = ""
+        if (activeSubscriptions[subscription] != .declared) {
+            subscriberScript = subsriberScript(forName: name, subscription: subscription)
+            context.addMessageHandler(messageHandler, name: name)
+        }
         let script = subscriberScript + "\n\(jsName).subscribe\(subscription.jsRepresentation)(\(name));"
-        context.addMessageHandler(messageHandler, name: name)
         context.evaluateScript(script, completion: nil)
+        activeSubscriptions[subscription] = .active
     }
 
     private func unsubscribe(subsription: Subscription) {
+        if (activeSubscriptions[subsription] != .active) {
+            NSLog("LWChart: double unsubscribe detected \(subsription)")
+            return
+        }
         let name = subscriberName(for: subsription)
         let script = "\(jsName).unsubscribe\(subsription.jsRepresentation)(\(name));"
         context.evaluateScript(script, completion: nil)
+        activeSubscriptions[subsription] = .declared
     }
 
     private func unsubscribeAll() {

diff --git a/Sources/LightweightCharts/Implementations/API/TimeScale.swift b/Sources/LightweightCharts/Implementations/API/TimeScale.swift
@@ -11,6 +11,11 @@ public protocol TimeScaleDelegate: AnyObject {
 // MARK: -
 class TimeScale: JavaScriptObject {
 
+    enum SubscribeState: CaseIterable {
+        case declared
+        case active
+    }
+
     typealias Context = JavaScriptEvaluator & JavaScriptMessageProducer
 
     let jsName = "timeScale" + .uniqueString
@@ -20,7 +25,8 @@ class TimeScale: JavaScriptObject {
     private weak var context: Context?
     weak var closureStore: ClosuresStore?
     private let messageHandler: MessageHandler
-
+    private var activeSubscriptions: Dictionary<Subscription,SubscribeState> = [:]
+
     init(context: Context, closureStore: ClosuresStore?) {
         self.context = context
         self.closureStore = closureStore
@@ -42,17 +48,30 @@ class TimeScale: JavaScriptObject {
     }
 
     private func subscribe(subscription: Subscription) {
+        if (activeSubscriptions[subscription] == .active) {
+            NSLog("LWChart: double subscribe detected \(subscription)")
+            return
+        }
         let name = subscriberName(for: subscription)
-        let subscriberScript = subsriberScript(forName: name, subscription: subscription)
+        var subscriberScript = ""
+        if (activeSubscriptions[subscription] != .declared) {
+            subscriberScript = subsriberScript(forName: name, subscription: subscription)
+            context?.addMessageHandler(messageHandler, name: name)
+        }
         let script = subscriberScript + "\n\(jsName).subscribe\(subscription.jsRepresentation)(\(name));"
-        context?.addMessageHandler(messageHandler, name: name)
         context?.evaluateScript(script, completion: nil)
+        activeSubscriptions[subscription] = .active
     }
 
     private func unsubscribe(subsription: Subscription) {
+        if (activeSubscriptions[subsription] != .active) {
+            NSLog("LWChart: double unsubscribe detected \(subsription)")
+            return
+        }
         let name = subscriberName(for: subsription)
         let script = "\(jsName).unsubscribe\(subsription.jsRepresentation)(\(name));"
         context?.evaluateScript(script, completion: nil)
+        activeSubscriptions[subsription] = .declared
     }
 
 }