2 files changed, 36 insertions, 33 deletions

diff --git a/WireGuard/WireGuard/UI/iOS/View/SwitchCell.swift b/WireGuard/WireGuard/UI/iOS/View/SwitchCell.swift
index 758083f..6169593 100644
--- a/WireGuard/WireGuard/UI/iOS/View/SwitchCell.swift
+++ b/WireGuard/WireGuard/UI/iOS/View/SwitchCell.swift
@@ -22,6 +22,8 @@ class SwitchCell: UITableViewCell {
 
     var onSwitchToggled: ((Bool) -> Void)?
 
+    var observationToken: AnyObject?
+
     let switchView = UISwitch()
 
     override init(style: UITableViewCell.CellStyle, reuseIdentifier: String?) {
@@ -45,5 +47,6 @@ class SwitchCell: UITableViewCell {
         isEnabled = true
         message = ""
         isOn = false
+        observationToken = nil
     }
 }
diff --git a/WireGuard/WireGuard/UI/iOS/ViewController/TunnelDetailTableViewController.swift b/WireGuard/WireGuard/UI/iOS/ViewController/TunnelDetailTableViewController.swift
index 3e80d94..31ad111 100644
--- a/WireGuard/WireGuard/UI/iOS/ViewController/TunnelDetailTableViewController.swift
+++ b/WireGuard/WireGuard/UI/iOS/ViewController/TunnelDetailTableViewController.swift
@@ -32,7 +32,6 @@ class TunnelDetailTableViewController: UITableViewController {
     private var interfaceFieldIsVisible = [Bool]()
     private var peerFieldIsVisible = [[Bool]]()
 
-    private weak var statusCell: SwitchCell?
     private var statusObservationToken: AnyObject?
     private var reloadRuntimeConfigurationTimer: Timer?
 
@@ -45,9 +44,6 @@ class TunnelDetailTableViewController: UITableViewController {
         loadVisibleFields()
         statusObservationToken = tunnel.observe(\.status) { [weak self] _, _ in
             guard let self = self else { return }
-            if let cell = self.statusCell {
-                self.updateStatus(statusCell: cell)
-            }
             if tunnel.status == .active {
                 self.startUpdatingRuntimeConfiguration()
             } else if tunnel.status == .inactive {
@@ -129,33 +125,6 @@ class TunnelDetailTableViewController: UITableViewController {
         present(alert, animated: true, completion: nil)
     }
 
-    func updateStatus(statusCell cell: SwitchCell) {
-        let status = tunnel.status
-        let text: String
-        switch status {
-        case .inactive:
-            text = tr("tunnelStatusInactive")
-        case .activating:
-            text = tr("tunnelStatusActivating")
-        case .active:
-            text = tr("tunnelStatusActive")
-        case .deactivating:
-            text = tr("tunnelStatusDeactivating")
-        case .reasserting:
-            text = tr("tunnelStatusReasserting")
-        case .restarting:
-            text = tr("tunnelStatusRestarting")
-        case .waiting:
-            text = tr("tunnelStatusWaiting")
-        }
-        cell.textLabel?.text = text
-        DispatchQueue.main.asyncAfter(deadline: .now() + .milliseconds(200)) { [weak cell] in
-            cell?.switchView.isOn = !(status == .deactivating || status == .inactive)
-            cell?.switchView.isUserInteractionEnabled = (status == .inactive || status == .active)
-        }
-        cell.isEnabled = status == .active || status == .inactive
-    }
-
     func startUpdatingRuntimeConfiguration() {
         reloadRuntimeConfiguration()
         reloadRuntimeConfigurationTimer?.invalidate()
@@ -312,7 +281,39 @@ extension TunnelDetailTableViewController {
 
     private func statusCell(for tableView: UITableView, at indexPath: IndexPath) -> UITableViewCell {
         let cell: SwitchCell = tableView.dequeueReusableCell(for: indexPath)
-        updateStatus(statusCell: cell)
+
+        let statusUpdate: (SwitchCell, TunnelStatus) -> Void = { cell, status in
+            let text: String
+            switch status {
+            case .inactive:
+                text = tr("tunnelStatusInactive")
+            case .activating:
+                text = tr("tunnelStatusActivating")
+            case .active:
+                text = tr("tunnelStatusActive")
+            case .deactivating:
+                text = tr("tunnelStatusDeactivating")
+            case .reasserting:
+                text = tr("tunnelStatusReasserting")
+            case .restarting:
+                text = tr("tunnelStatusRestarting")
+            case .waiting:
+                text = tr("tunnelStatusWaiting")
+            }
+            cell.textLabel?.text = text
+            DispatchQueue.main.asyncAfter(deadline: .now() + .milliseconds(200)) { [weak cell] in
+                cell?.switchView.isOn = !(status == .deactivating || status == .inactive)
+                cell?.switchView.isUserInteractionEnabled = (status == .inactive || status == .active)
+            }
+            cell.isEnabled = status == .active || status == .inactive
+        }
+
+        statusUpdate(cell, tunnel.status)
+        cell.observationToken = tunnel.observe(\.status) { [weak cell] tunnel, _ in
+            guard let cell = cell else { return }
+            statusUpdate(cell, tunnel.status)
+        }
+
         cell.onSwitchToggled = { [weak self] isOn in
             guard let self = self else { return }
             if isOn {
@@ -321,7 +322,6 @@ extension TunnelDetailTableViewController {
                 self.tunnelsManager.startDeactivation(of: self.tunnel)
             }
         }
-        self.statusCell = cell
         return cell
     }