1 files changed, 120 insertions, 0 deletions
diff --git a/WireGuard/WireGuard/Tunnel/ActivateOnDemandOption.swift b/WireGuard/WireGuard/Tunnel/ActivateOnDemandOption.swift
new file mode 100644
index 0000000..8b39f8c
--- /dev/null
+++ b/WireGuard/WireGuard/Tunnel/ActivateOnDemandOption.swift
@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: MIT
+// Copyright © 2018-2019 WireGuard LLC. All Rights Reserved.
+
+import NetworkExtension
+
+enum ActivateOnDemandOption: Equatable {
+    case off
+    case wiFiInterfaceOnly(ActivateOnDemandSSIDOption)
+    case nonWiFiInterfaceOnly
+    case anyInterface(ActivateOnDemandSSIDOption)
+}
+
+#if os(iOS)
+private let nonWiFiInterfaceType: NEOnDemandRuleInterfaceType = .cellular
+#elseif os(macOS)
+private let nonWiFiInterfaceType: NEOnDemandRuleInterfaceType = .ethernet
+#else
+#error("Unimplemented")
+#endif
+
+enum ActivateOnDemandSSIDOption: Equatable {
+    case anySSID
+    case onlySpecificSSIDs([String])
+    case exceptSpecificSSIDs([String])
+}
+
+extension ActivateOnDemandOption {
+    func apply(on tunnelProviderManager: NETunnelProviderManager) {
+        let rules: [NEOnDemandRule]?
+        switch self {
+        case .off:
+            rules = nil
+        case .wiFiInterfaceOnly(let ssidOption):
+            rules = ssidOnDemandRules(option: ssidOption) + [NEOnDemandRuleDisconnect(interfaceType: nonWiFiInterfaceType)]
+        case .nonWiFiInterfaceOnly:
+            rules = [NEOnDemandRuleConnect(interfaceType: nonWiFiInterfaceType), NEOnDemandRuleDisconnect(interfaceType: .wiFi)]
+        case .anyInterface(let ssidOption):
+            if case .anySSID = ssidOption {
+                rules = [NEOnDemandRuleConnect(interfaceType: .any)]
+            } else {
+                rules = ssidOnDemandRules(option: ssidOption) + [NEOnDemandRuleConnect(interfaceType: nonWiFiInterfaceType)]
+            }
+        }
+        tunnelProviderManager.onDemandRules = rules
+        tunnelProviderManager.isOnDemandEnabled = self != .off
+    }
+
+    init(from tunnelProviderManager: NETunnelProviderManager) {
+        let rules = tunnelProviderManager.onDemandRules ?? []
+        let activateOnDemandOption: ActivateOnDemandOption
+        switch rules.count {
+        case 0:
+            activateOnDemandOption = .off
+        case 1:
+            let rule = rules[0]
+            precondition(rule.action == .connect)
+            activateOnDemandOption = .anyInterface(.anySSID)
+        case 2:
+            let connectRule = rules.first(where: { $0.action == .connect })!
+            let disconnectRule = rules.first(where: { $0.action == .disconnect })!
+            if connectRule.interfaceTypeMatch == .wiFi && disconnectRule.interfaceTypeMatch == nonWiFiInterfaceType {
+                activateOnDemandOption = .wiFiInterfaceOnly(.anySSID)
+            } else if connectRule.interfaceTypeMatch == nonWiFiInterfaceType && disconnectRule.interfaceTypeMatch == .wiFi {
+                activateOnDemandOption = .nonWiFiInterfaceOnly
+            } else {
+                fatalError("Unexpected onDemandRules set on tunnel provider manager")
+            }
+        case 3:
+            let ssidRule = rules.first(where: { $0.interfaceTypeMatch == .wiFi && $0.ssidMatch != nil })!
+            let nonWiFiRule = rules.first(where: { $0.interfaceTypeMatch == nonWiFiInterfaceType })!
+            let ssids = ssidRule.ssidMatch!
+            switch (ssidRule.action, nonWiFiRule.action) {
+            case (.connect, .connect):
+                activateOnDemandOption = .anyInterface(.onlySpecificSSIDs(ssids))
+            case (.connect, .disconnect):
+                activateOnDemandOption = .wiFiInterfaceOnly(.onlySpecificSSIDs(ssids))
+            case (.disconnect, .connect):
+                activateOnDemandOption = .anyInterface(.exceptSpecificSSIDs(ssids))
+            case (.disconnect, .disconnect):
+                activateOnDemandOption = .wiFiInterfaceOnly(.exceptSpecificSSIDs(ssids))
+            default:
+                fatalError("Unexpected SSID onDemandRules set on tunnel provider manager")
+            }
+        default:
+            fatalError("Unexpected number of onDemandRules set on tunnel provider manager")
+        }
+
+        self = activateOnDemandOption
+    }
+}
+
+private extension NEOnDemandRuleConnect {
+    convenience init(interfaceType: NEOnDemandRuleInterfaceType, ssids: [String]? = nil) {
+        self.init()
+        interfaceTypeMatch = interfaceType
+        ssidMatch = ssids
+    }
+}
+
+private extension NEOnDemandRuleDisconnect {
+    convenience init(interfaceType: NEOnDemandRuleInterfaceType, ssids: [String]? = nil) {
+        self.init()
+        interfaceTypeMatch = interfaceType
+        ssidMatch = ssids
+    }
+}
+
+private func ssidOnDemandRules(option: ActivateOnDemandSSIDOption) -> [NEOnDemandRule] {
+    switch option {
+    case .anySSID:
+        return [NEOnDemandRuleConnect(interfaceType: .wiFi)]
+    case .onlySpecificSSIDs(let ssids):
+        assert(!ssids.isEmpty)
+        return [NEOnDemandRuleConnect(interfaceType: .wiFi, ssids: ssids),
+                NEOnDemandRuleDisconnect(interfaceType: .wiFi)]
+    case .exceptSpecificSSIDs(let ssids):
+        return [NEOnDemandRuleDisconnect(interfaceType: .wiFi, ssids: ssids),
+                NEOnDemandRuleConnect(interfaceType: .wiFi)]
+    }
+}