Attempt shutdown

duckduckgo · Mar 6, 2024 · 1ab1e29 · 1ab1e29
1 parent 5092a36
commit 1ab1e29
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 69 deletions.
diff --git a/Sources/NetworkProtection/ExtensionMessage/ExtensionRequest.swift b/Sources/NetworkProtection/ExtensionMessage/ExtensionRequest.swift
@@ -23,7 +23,6 @@ public enum DebugCommand: Codable {
     case removeSystemExtension
     case removeVPNConfiguration
     case sendTestNotification
-    case blockAllTraffic
     case disableConnectOnDemandAndShutDown
 }
 

diff --git a/Sources/NetworkProtection/PacketTunnelProvider.swift b/Sources/NetworkProtection/PacketTunnelProvider.swift
@@ -806,52 +806,6 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         return configurationResult.0
     }
 
-    /// Placeholder configuration to switch to when the entitlement expires
-    /// This will block all traffic
-    @MainActor
-    private func updatePlaceholderTunnelConfiguration() async throws {
-        let interface = InterfaceConfiguration(
-            privateKey: PrivateKey(),
-            addresses: [IPAddressRange(from: "0.0.0.0/0")!],
-            includedRoutes: [],
-            excludedRoutes: [],
-            listenPort: 0,
-            dns: [DNSServer(address: IPv4Address.loopback)]
-        )
-
-        var peerConfiguration = PeerConfiguration(publicKey: PrivateKey().publicKey)
-        peerConfiguration.endpoint = Endpoint(host: "127.0.0.1", port: 9090)
-
-        let tunnelConfiguration = TunnelConfiguration(name: "Placeholder", interface: interface, peers: [peerConfiguration])
-
-        try await withCheckedThrowingContinuation { [weak self] (continuation: CheckedContinuation<Void, Error>) in
-            guard let self = self else {
-                continuation.resume()
-                return
-            }
-
-            self.adapter.update(tunnelConfiguration: tunnelConfiguration, reassert: true) { [weak self] error in
-                if let error = error {
-                    os_log("🔵 Failed to update the placeholder configuration: %{public}@", type: .error, error.localizedDescription)
-                    self?.debugEvents?.fire(error.networkProtectionError)
-                    continuation.resume(throwing: error)
-                    return
-                }
-
-                Task { [weak self] in
-                    do {
-                        try await self?.handleAdapterStarted(startReason: .reconnected)
-                    } catch {
-                        continuation.resume(throwing: error)
-                        return
-                    }
-
-                    continuation.resume()
-                }
-            }
-        }
-    }
-
     // MARK: - App Messages
 
     // swiftlint:disable:next cyclomatic_complexity
@@ -982,11 +936,14 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             handleExpireRegistrationKey(completionHandler: completionHandler)
         case .sendTestNotification:
             handleSendTestNotification(completionHandler: completionHandler)
-        case .blockAllTraffic:
-            handleBlockAllTraffic(completionHandler: completionHandler)
         case .disableConnectOnDemandAndShutDown:
             if #available(iOS 17, *) {
                 handleShutDown(completionHandler: completionHandler)
+            } else {
+                Task {
+                    await rekey()
+                    completionHandler?(nil)
+                }
             }
         case .removeVPNConfiguration:
             // Since the VPN configuration is being removed we may as well reset all state
@@ -1082,14 +1039,6 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         completionHandler?(nil)
     }
 
-    public func handleBlockAllTraffic(completionHandler: ((Data?) -> Void)? = nil) {
-        Task { @MainActor [weak self] in
-            await self?.stopMonitors()
-            try? await self?.updatePlaceholderTunnelConfiguration()
-            completionHandler?(nil)
-        }
-    }
-
     @available(iOS 17, *)
     public func handleShutDown(completionHandler: ((Data?) -> Void)? = nil) {
         Task { @MainActor in
@@ -1243,23 +1192,31 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         guard isSubscriptionEnabled, let entitlementCheck else { return }
 
         await entitlementMonitor.start(entitlementCheck: entitlementCheck) { [weak self] result in
+            /// Attempt tunnel shutdown & show messaging iff the entitlement is verified to be invalid
+            /// Ignore otherwise
             switch result {
-            case .validEntitlement:
-                self?.defaults.resetEntitlementMessaging()
             case .invalidEntitlement:
-                self?.defaults.enableEntitlementMessaging()
-                self?.notificationsPresenter.showEntitlementNotification()
+                self?.handleInvalidEntitlement()
+            case .validEntitlement, .error:
+                break
+            }
+        }
+    }
 
-                Task { @MainActor [weak self] in
-                    await self?.stopMonitors()
+    private func handleInvalidEntitlement() {
+        defaults.enableEntitlementMessaging()
+        notificationsPresenter.showEntitlementNotification()
 
-                    // We add a delay here so the notification has a chance to show up
-                    try? await Task.sleep(interval: .seconds(5))
+        Task { @MainActor [weak self] in
+            await self?.stopMonitors()
 
-                    try? await self?.updatePlaceholderTunnelConfiguration()
-                }
-            case .error:
-                break
+            // We add a delay here so the notification has a chance to show up
+            try? await Task.sleep(interval: .seconds(5))
+
+            if #available(iOS 17, *) {
+                self?.handleShutDown()
+            } else {
+                await self?.rekey()
             }
         }
     }