From 2700cb9d6e0a55cc55c20c38fd2711aec1c087c0 Mon Sep 17 00:00:00 2001
From: Diego Rey Mendez <diegoreymendez@users.noreply.github.com>
Date: Fri, 23 Feb 2024 21:36:45 +0100
Subject: [PATCH] Fix on demand so that it's not enabled too soon (#2499)

Task/Issue URL: https://app.asana.com/0/1203137811378537/1206543593200557/f

## Description:

Enables on-demand only after the VPN has connected.
---
 .../NetworkProtectionTunnelController.swift   | 42 ++++++++++++++++---
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/DuckDuckGo/NetworkProtectionTunnelController.swift b/DuckDuckGo/NetworkProtectionTunnelController.swift
index 595b0afa13..77a09a12e1 100644
--- a/DuckDuckGo/NetworkProtectionTunnelController.swift
+++ b/DuckDuckGo/NetworkProtectionTunnelController.swift
@@ -31,6 +31,9 @@ final class NetworkProtectionTunnelController: TunnelController {
     private let debugFeatures = NetworkProtectionDebugFeatures()
     private let tokenStore = NetworkProtectionKeychainTokenStore()
     private let errorStore = NetworkProtectionTunnelErrorStore()
+    private let notificationCenter: NotificationCenter = .default
+    private var previousStatus: NEVPNStatus = .invalid
+    private var cancellables = Set<AnyCancellable>()
 
     // MARK: - Starting & Stopping the VPN
 
@@ -39,6 +42,10 @@ final class NetworkProtectionTunnelController: TunnelController {
         case simulateControllerFailureError
     }
 
+    init() {
+        subscribeToStatusChanges()
+    }
+
     /// Starts the VPN connection used for Network Protection
     ///
     func start() async {
@@ -142,12 +149,6 @@ final class NetworkProtectionTunnelController: TunnelController {
             Pixel.fire(pixel: .networkProtectionActivationRequestFailed, error: error)
             throw error
         }
-
-        if !debugFeatures.alwaysOnDisabled {
-            Task {
-                try await enableOnDemand(tunnelManager: tunnelManager)
-            }
-        }
     }
 
     /// The actual storage for our tunnel manager.
@@ -231,6 +232,35 @@ final class NetworkProtectionTunnelController: TunnelController {
         tunnelManager.onDemandRules = [NEOnDemandRuleConnect()]
     }
 
+    // MARK: - Observing Status Changes
+
+    private func subscribeToStatusChanges() {
+        notificationCenter.publisher(for: .NEVPNStatusDidChange)
+            .sink(receiveValue: handleStatusChange(_:))
+            .store(in: &cancellables)
+    }
+
+    private func handleStatusChange(_ notification: Notification) {
+        guard !debugFeatures.alwaysOnDisabled,
+              let session = (notification.object as? NETunnelProviderSession),
+              session.status != previousStatus,
+              let manager = session.manager as? NETunnelProviderManager else {
+            return
+        }
+
+        Task { @MainActor in
+            previousStatus = session.status
+
+            switch session.status {
+            case .connected:
+                try await enableOnDemand(tunnelManager: manager)
+            default:
+                break
+            }
+
+        }
+    }
+
     // MARK: - On Demand
 
     @MainActor