From c2b64f6d6cbacd77b814f0f933b29a49bd8bb3bf Mon Sep 17 00:00:00 2001
From: Diego Rey Mendez <dreymendez@duckduckgo.com>
Date: Wed, 21 Feb 2024 14:45:31 +0100
Subject: [PATCH] Fix on demand so that it's not enabled too soon

---
 .../NetworkProtectionTunnelController.swift   | 38 +++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/DuckDuckGo/NetworkProtection/AppTargets/BothAppTargets/NetworkProtectionTunnelController.swift b/DuckDuckGo/NetworkProtection/AppTargets/BothAppTargets/NetworkProtectionTunnelController.swift
index d5eb04eecf..2bea9ba3a4 100644
--- a/DuckDuckGo/NetworkProtection/AppTargets/BothAppTargets/NetworkProtectionTunnelController.swift
+++ b/DuckDuckGo/NetworkProtection/AppTargets/BothAppTargets/NetworkProtectionTunnelController.swift
@@ -70,6 +70,10 @@ final class NetworkProtectionTunnelController: NetworkProtection.TunnelControlle
     private let networkExtensionBundleID: String
     private let networkExtensionController: NetworkExtensionController
 
+    // MARK: - Notification Center
+
+    private let notificationCenter: NotificationCenter
+
     // MARK: - User Defaults
 
     /* Temporarily disabled - https://app.asana.com/0/0/1205766100762904/f
@@ -129,13 +133,41 @@ final class NetworkProtectionTunnelController: NetworkProtection.TunnelControlle
         self.logger = logger
         self.networkExtensionBundleID = networkExtensionBundleID
         self.networkExtensionController = networkExtensionController
+        self.notificationCenter = notificationCenter
         self.settings = settings
         self.tokenStore = tokenStore
 
         subscribeToSettingsChanges()
+        subscribeToStatusChanges()
+    }
+
+    // MARK: - Observing Status Changes
+
+    private func subscribeToStatusChanges() {
+        notificationCenter.publisher(for: .NEVPNStatusDidChange)
+            .sink(receiveValue: handleStatusChange(_:))
+            .store(in: &cancellables)
+    }
+
+    private func handleStatusChange(_ notification: Notification) {
+        guard let session = (notification.object as? NETunnelProviderSession),
+            let manager = session.manager as? NETunnelProviderManager else {
+
+            return
+        }
+
+        Task { @MainActor in
+            switch session.status {
+            case .connected:
+                try await enableOnDemand(tunnelManager: manager)
+            default:
+                break
+            }
+
+        }
     }
 
-    // MARK: - Tunnel Settings
+    // MARK: - Subscriptions
 
     private func subscribeToSettingsChanges() {
         settings.changePublisher
@@ -154,6 +186,8 @@ final class NetworkProtectionTunnelController: NetworkProtection.TunnelControlle
             .store(in: &cancellables)
     }
 
+    // MARK: - Handling Settings Changes
+
     /// This is where the tunnel owner has a chance to handle the settings change locally.
     ///
     /// The extension can also handle these changes so not everything needs to be handled here.
@@ -443,8 +477,6 @@ final class NetworkProtectionTunnelController: NetworkProtection.TunnelControlle
                 guard let self, error == nil, fired else { return }
                 self.settings.vpnFirstEnabled = PixelKit.pixelLastFireDate(event: NetworkProtectionPixelEvent.networkProtectionNewUser)
             }
-
-        try await enableOnDemand(tunnelManager: tunnelManager)
     }
 
     /// Stops the VPN connection used for Network Protection