Add selected environment preference

Sources/NetworkProtection/ExtensionMessage/ExtensionRequest.swift

@@ -20,6 +20,8 @@ import Foundation
 
 public enum DebugCommand: Codable {
     case expireRegistrationKey
+    case removeSystemExtension
+    case removeVPNConfiguration
     case sendTestNotification
 }
 

Sources/NetworkProtection/NetworkProtectionOptionKey.swift

@@ -20,6 +20,7 @@ import Foundation
 
 public enum NetworkProtectionOptionKey {
     public static let keyValidity = "keyValidity"
+    public static let selectedEnvironment = "selectedEnvironment"
     public static let selectedServer = "selectedServer"
     public static let authToken = "authToken"
     public static let isOnDemand = "is-on-demand"

Sources/NetworkProtection/Networking/NetworkProtectionClient.swift

@@ -74,25 +74,20 @@ struct RedeemResponse: Decodable {
 
 public final class NetworkProtectionBackendClient: NetworkProtectionClient {
 
-    enum Constants {
-        static let productionEndpoint = URL(string: "https://controller.netp.duckduckgo.com")!
-        static let stagingEndpoint = URL(string: "https://staging.netp.duckduckgo.com")!
-    }
-
     private enum DecoderError: Error {
         case failedToDecode(key: String)
     }
 
     var serversURL: URL {
-        Constants.productionEndpoint.appending("/servers")
+        endpointURL.appending("/servers")
     }
 
     var registerKeyURL: URL {
-        Constants.productionEndpoint.appending("/register")
+        endpointURL.appending("/register")
     }
 
     var redeemURL: URL {
-        Constants.productionEndpoint.appending("/redeem")
+        endpointURL.appending("/redeem")
     }
 
     private let decoder: JSONDecoder = {
@@ -114,7 +109,11 @@ public final class NetworkProtectionBackendClient: NetworkProtectionClient {
         return decoder
     }()
 
-    public init() {}
+    private let endpointURL: URL
+
+    public init(environment: TunnelSettings.SelectedEnvironment = .default) {
+        endpointURL = environment.endpointURL
+    }
 
     public func getServers(authToken: String) async -> Result<[NetworkProtectionServer], NetworkProtectionClientError> {
         var request = URLRequest(url: serversURL)

Sources/NetworkProtection/PacketTunnelProvider.swift

@@ -348,6 +348,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
 
     private func load(options: StartupOptions) throws {
         loadKeyValidity(from: options)
+        loadSelectedEnvironment(from: options)
         loadSelectedServer(from: options)
         loadTesterEnabled(from: options)
         try loadAuthToken(from: options)
@@ -370,6 +371,17 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         }
     }
 
+    private func loadSelectedEnvironment(from options: StartupOptions) {
+        switch options.selectedEnvironment {
+        case .set(let selectedEnvironment):
+            settings.selectedEnvironment = selectedEnvironment
+        case .useExisting:
+            break
+        case .reset:
+            settings.selectedEnvironment = .default
+        }
+    }
+
     private func loadSelectedServer(from options: StartupOptions) {
         switch options.selectedServer {
         case .set(let selectedServer):
@@ -477,10 +489,13 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         let onDemand = options.startupMethod == .automaticOnDemand
 
         os_log("Starting tunnel %{public}@", log: .networkProtection, options.startupMethod.debugDescription)
-        startTunnel(selectedServer: settings.selectedServer, onDemand: onDemand, completionHandler: completionHandler)
+        startTunnel(environment: settings.selectedEnvironment, 
+                    selectedServer: settings.selectedServer,
+                    onDemand: onDemand,
+                    completionHandler: completionHandler)
     }
 
-    private func startTunnel(selectedServer: TunnelSettings.SelectedServer, onDemand: Bool, completionHandler: @escaping (Error?) -> Void) {
+    private func startTunnel(environment: TunnelSettings.SelectedEnvironment, selectedServer: TunnelSettings.SelectedServer, onDemand: Bool, completionHandler: @escaping (Error?) -> Void) {
 
         Task {
             let serverSelectionMethod: NetworkProtectionServerSelectionMethod
@@ -494,7 +509,8 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
 
             do {
                 os_log("🔵 Generating tunnel config", log: .networkProtection, type: .info)
-                let tunnelConfiguration = try await generateTunnelConfiguration(serverSelectionMethod: serverSelectionMethod,
+                let tunnelConfiguration = try await generateTunnelConfiguration(environment: environment,
+                                                                                serverSelectionMethod: serverSelectionMethod,
                                                                                 includedRoutes: includedRoutes ?? [],
                                                                                 excludedRoutes: excludedRoutes ?? [])
                 startTunnel(with: tunnelConfiguration, onDemand: onDemand, completionHandler: completionHandler)
@@ -619,7 +635,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
 
     // MARK: - Tunnel Configuration
 
-    public func updateTunnelConfiguration(selectedServer: TunnelSettings.SelectedServer, reassert: Bool = true) async throws {
+    public func updateTunnelConfiguration(environment: TunnelSettings.SelectedEnvironment = .default, selectedServer: TunnelSettings.SelectedServer, reassert: Bool = true) async throws {
         let serverSelectionMethod: NetworkProtectionServerSelectionMethod
 
         switch settings.selectedServer {
@@ -629,12 +645,13 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             serverSelectionMethod = .preferredServer(serverName: serverName)
         }
 
-        try await updateTunnelConfiguration(serverSelectionMethod: serverSelectionMethod, reassert: reassert)
+        try await updateTunnelConfiguration(environment: environment, serverSelectionMethod: serverSelectionMethod, reassert: reassert)
     }
 
-    public func updateTunnelConfiguration(serverSelectionMethod: NetworkProtectionServerSelectionMethod, reassert: Bool = true) async throws {
+    public func updateTunnelConfiguration(environment: TunnelSettings.SelectedEnvironment = .default, serverSelectionMethod: NetworkProtectionServerSelectionMethod, reassert: Bool = true) async throws {
 
-        let tunnelConfiguration = try await generateTunnelConfiguration(serverSelectionMethod: serverSelectionMethod,
+        let tunnelConfiguration = try await generateTunnelConfiguration(environment: environment,
+                                                                        serverSelectionMethod: serverSelectionMethod,
                                                                         includedRoutes: includedRoutes ?? [],
                                                                         excludedRoutes: excludedRoutes ?? [])
 
@@ -666,12 +683,14 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
         }
     }
 
-    private func generateTunnelConfiguration(serverSelectionMethod: NetworkProtectionServerSelectionMethod, includedRoutes: [IPAddressRange], excludedRoutes: [IPAddressRange]) async throws -> TunnelConfiguration {
+    private func generateTunnelConfiguration(environment: TunnelSettings.SelectedEnvironment = .default, serverSelectionMethod: NetworkProtectionServerSelectionMethod, includedRoutes: [IPAddressRange], excludedRoutes: [IPAddressRange]) async throws -> TunnelConfiguration {
 
         let configurationResult: (TunnelConfiguration, NetworkProtectionServerInfo)
 
         do {
-            let deviceManager = NetworkProtectionDeviceManager(tokenStore: tokenStore,
+            let networkClient = NetworkProtectionBackendClient(environment: environment)
+            let deviceManager = NetworkProtectionDeviceManager(networkClient: networkClient,
+                                                               tokenStore: tokenStore,
                                                                keyStore: keyStore,
                                                                errorEvents: debugEvents)
 
@@ -767,24 +786,31 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             }
 
             Task {
-                try? await updateTunnelConfiguration(serverSelectionMethod: serverSelectionMethod)
+                try? await updateTunnelConfiguration(environment: settings.selectedEnvironment, serverSelectionMethod: serverSelectionMethod)
                 completionHandler?(nil)
             }
         case .setIncludeAllNetworks,
                 .setEnforceRoutes,
                 .setExcludeLocalNetworks,
-                .setRegistrationKeyValidity:
+                .setRegistrationKeyValidity,
+                .setSelectedEnvironment:
             // Intentional no-op, as some setting changes don't require any further operation
             break
         }
     }
 
     private func handleDebugCommand(_ command: DebugCommand, completionHandler: ((Data?) -> Void)? = nil) {
         switch command {
+        case .removeSystemExtension:
+            // Since the system extension is being removed we may as well reset all state
+            handleResetAllState(completionHandler: completionHandler)
         case .expireRegistrationKey:
             handleExpireRegistrationKey(completionHandler: completionHandler)
         case .sendTestNotification:
             handleSendTestNotification(completionHandler: completionHandler)
+        case .removeVPNConfiguration:
+            // Since the VPN configuration is being removed we may as well reset all state
+            handleResetAllState(completionHandler: completionHandler)
         }
     }
 

Sources/NetworkProtection/Settings/Extensions/UserDefaults+selectedEnvironment.swift

@@ -0,0 +1,57 @@
+//
+//  UserDefaults+selectedEnvironment.swift
+//
+//  Copyright © 2023 DuckDuckGo. All rights reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+import Combine
+import Foundation
+
+extension UserDefaults {
+    private var selectedEnvironmentKey: String {
+        "networkProtectionSettingSelectedEnvironmentRawValue"
+    }
+
+    @objc
+    dynamic var networkProtectionSettingSelectedEnvironmentRawValue: String {
+        get {
+            value(forKey: selectedEnvironmentKey) as? String ?? TunnelSettings.SelectedEnvironment.default.rawValue
+        }
+
+        set {
+            set(newValue, forKey: selectedEnvironmentKey)
+        }
+    }
+
+    var networkProtectionSettingSelectedEnvironment: TunnelSettings.SelectedEnvironment {
+        get {
+            TunnelSettings.SelectedEnvironment(rawValue: networkProtectionSettingSelectedEnvironmentRawValue) ?? .default
+        }
+
+        set {
+             networkProtectionSettingSelectedEnvironmentRawValue = newValue.rawValue
+        }
+    }
+
+    var networkProtectionSettingSelectedEnvironmentPublisher: AnyPublisher<TunnelSettings.SelectedEnvironment, Never> {
+        publisher(for: \.networkProtectionSettingSelectedEnvironmentRawValue).map { value in
+            TunnelSettings.SelectedEnvironment(rawValue: value) ?? .default
+        }.eraseToAnyPublisher()
+    }
+
+    func resetNetworkProtectionSettingSelectedEnvironment() {
+        networkProtectionSettingSelectedEnvironment = .default
+    }
+}

Sources/NetworkProtection/Settings/TunnelSettings.swift

@@ -33,6 +33,7 @@ public final class TunnelSettings {
         case setExcludeLocalNetworks(_ excludeLocalNetworks: Bool)
         case setRegistrationKeyValidity(_ validity: RegistrationKeyValidity)
         case setSelectedServer(_ selectedServer: SelectedServer)
+        case setSelectedEnvironment(_ selectedEnvironment: SelectedEnvironment)
     }
 
     public enum RegistrationKeyValidity: Codable {
@@ -52,6 +53,22 @@ public final class TunnelSettings {
         }
     }
 
+    public enum SelectedEnvironment: String, Codable {
+        case production
+        case staging
+
+        public static var `default`: SelectedEnvironment = .production
+
+        public var endpointURL: URL {
+            switch self {
+            case .production:
+                return URL(string: "https://controller.netp.duckduckgo.com")!
+            case .staging:
+                return URL(string: "https://staging1.netp.duckduckgo.com")!
+            }
+        }
+    }
+
     private let defaults: UserDefaults
 
     private(set) public lazy var changePublisher: AnyPublisher<Change, Never> = {
@@ -76,11 +93,16 @@ public final class TunnelSettings {
             Change.setSelectedServer(server)
         }.eraseToAnyPublisher()
 
+        let environmentChangePublisher = selectedEnvironmentPublisher.map { environment in
+            Change.setSelectedEnvironment(environment)
+        }.eraseToAnyPublisher()
+
         return Publishers.MergeMany(
             includeAllNetworksPublisher,
             enforceRoutesPublisher,
             excludeLocalNetworksPublisher,
-            serverChangePublisher).eraseToAnyPublisher()
+            serverChangePublisher,
+            environmentChangePublisher).eraseToAnyPublisher()
     }()
 
     public init(defaults: UserDefaults) {
@@ -95,6 +117,7 @@ public final class TunnelSettings {
         defaults.resetNetworkProtectionSettingIncludeAllNetworks()
         defaults.resetNetworkProtectionSettingRegistrationKeyValidity()
         defaults.resetNetworkProtectionSettingSelectedServer()
+        defaults.resetNetworkProtectionSettingSelectedEnvironment()
     }
 
     // MARK: - Applying Changes
@@ -111,6 +134,8 @@ public final class TunnelSettings {
             self.registrationKeyValidity = registrationKeyValidity
         case .setSelectedServer(let selectedServer):
             self.selectedServer = selectedServer
+        case .setSelectedEnvironment(let selectedEnvironment):
+            self.selectedEnvironment = selectedEnvironment
         }
     }
 
@@ -198,6 +223,22 @@ public final class TunnelSettings {
         }
     }
 
+    // MARK: - Environment
+
+    public var selectedEnvironmentPublisher: AnyPublisher<SelectedEnvironment, Never> {
+        defaults.networkProtectionSettingSelectedEnvironmentPublisher
+    }
+
+    public var selectedEnvironment: SelectedEnvironment {
+        get {
+            defaults.networkProtectionSettingSelectedEnvironment
+        }
+
+        set {
+            defaults.networkProtectionSettingSelectedEnvironment = newValue
+        }
+    }
+
     // MARK: - Routes
 
     public enum ExclusionListItem {

Sources/NetworkProtection/StartupOptions.swift

@@ -94,6 +94,7 @@ struct StartupOptions {
     let simulateCrash: Bool
     let simulateMemoryCrash: Bool
     let keyValidity: StoredOption<TimeInterval>
+    let selectedEnvironment: StoredOption<TunnelSettings.SelectedEnvironment>
     let selectedServer: StoredOption<TunnelSettings.SelectedServer>
     let authToken: StoredOption<String>
     let enableTester: StoredOption<Bool>
@@ -121,6 +122,7 @@ struct StartupOptions {
         authToken = Self.readAuthToken(from: options, resetIfNil: resetStoredOptionsIfNil)
         enableTester = Self.readEnableTester(from: options, resetIfNil: resetStoredOptionsIfNil)
         keyValidity = Self.readKeyValidity(from: options, resetIfNil: resetStoredOptionsIfNil)
+        selectedEnvironment = Self.readSelectedEnvironment(from: options, resetIfNil: resetStoredOptionsIfNil)
         selectedServer = Self.readSelectedServer(from: options, resetIfNil: resetStoredOptionsIfNil)
     }
 
@@ -150,6 +152,17 @@ struct StartupOptions {
         }
     }
 
+    private static func readSelectedEnvironment(from options: [String: Any], resetIfNil: Bool) -> StoredOption<TunnelSettings.SelectedEnvironment> {
+
+        StoredOption(resetIfNil: resetIfNil) {
+            guard let environment = options[NetworkProtectionOptionKey.selectedEnvironment] as? String else {
+                return nil
+            }
+
+            return TunnelSettings.SelectedEnvironment(rawValue: environment) ?? .default
+        }
+    }
+
     private static func readSelectedServer(from options: [String: Any], resetIfNil: Bool) -> StoredOption<TunnelSettings.SelectedServer> {
 
         StoredOption(resetIfNil: resetIfNil) {