on-demand: Infrastructure for supporting SSID-based rules

Signed-off-by: Roopesh Chander <roop@roopc.net>

4 files changed, 71 insertions, 21 deletions

diff --git a/WireGuard/WireGuard/Tunnel/ActivateOnDemandSetting.swift b/WireGuard/WireGuard/Tunnel/ActivateOnDemandSetting.swift
index 0cb3f98..28612c7 100644
--- a/WireGuard/WireGuard/Tunnel/ActivateOnDemandSetting.swift
+++ b/WireGuard/WireGuard/Tunnel/ActivateOnDemandSetting.swift
@@ -8,11 +8,11 @@ struct ActivateOnDemandSetting {
     var activateOnDemandOption: ActivateOnDemandOption
 }
 
-enum ActivateOnDemandOption {
+enum ActivateOnDemandOption: Equatable {
     case none // Valid only when isActivateOnDemandEnabled is false
-    case wiFiInterfaceOnly
+    case wiFiInterfaceOnly(ActivateOnDemandSSIDOption)
     case nonWiFiInterfaceOnly
-    case anyInterface
+    case anyInterface(ActivateOnDemandSSIDOption)
 }
 
 #if os(iOS)
@@ -23,25 +23,29 @@ private let nonWiFiInterfaceType: NEOnDemandRuleInterfaceType = .ethernet
 #error("Unimplemented")
 #endif
 
+enum ActivateOnDemandSSIDOption: Equatable {
+    case anySSID
+    case onlySpecificSSIDs([String])
+    case exceptSpecificSSIDs([String])
+}
+
 extension ActivateOnDemandSetting {
     func apply(on tunnelProviderManager: NETunnelProviderManager) {
         tunnelProviderManager.isOnDemandEnabled = isActivateOnDemandEnabled
         let rules: [NEOnDemandRule]?
-        let connectRule = NEOnDemandRuleConnect()
-        let disconnectRule = NEOnDemandRuleDisconnect()
         switch activateOnDemandOption {
         case .none:
             rules = nil
-        case .wiFiInterfaceOnly:
-            connectRule.interfaceTypeMatch = .wiFi
-            disconnectRule.interfaceTypeMatch = nonWiFiInterfaceType
-            rules = [connectRule, disconnectRule]
+        case .wiFiInterfaceOnly(let ssidOption):
+            rules = ssidOnDemandRules(option: ssidOption) + [NEOnDemandRuleDisconnect(interfaceType: nonWiFiInterfaceType)]
         case .nonWiFiInterfaceOnly:
-            connectRule.interfaceTypeMatch = nonWiFiInterfaceType
-            disconnectRule.interfaceTypeMatch = .wiFi
-            rules = [connectRule, disconnectRule]
-        case .anyInterface:
-            rules = [connectRule]
+            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
     }
@@ -55,17 +59,33 @@ extension ActivateOnDemandSetting {
         case 1:
             let rule = rules[0]
             precondition(rule.action == .connect)
-            activateOnDemandOption = .anyInterface
+            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
+                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")
         }
@@ -82,3 +102,33 @@ extension ActivateOnDemandSetting {
 extension ActivateOnDemandSetting {
     static var defaultSetting = ActivateOnDemandSetting(isActivateOnDemandEnabled: false, activateOnDemandOption: .none)
 }
+
+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)]
+    }
+}
diff --git a/WireGuard/WireGuard/UI/TunnelViewModel.swift b/WireGuard/WireGuard/UI/TunnelViewModel.swift
index 8a9b007..956bfd1 100644
--- a/WireGuard/WireGuard/UI/TunnelViewModel.swift
+++ b/WireGuard/WireGuard/UI/TunnelViewModel.swift
@@ -661,7 +661,7 @@ extension TunnelViewModel {
     }
 
     static func defaultActivateOnDemandOption() -> ActivateOnDemandOption {
-        return .anyInterface
+        return .anyInterface(.anySSID)
     }
 }
 
diff --git a/WireGuard/WireGuard/UI/iOS/ViewController/TunnelEditTableViewController.swift b/WireGuard/WireGuard/UI/iOS/ViewController/TunnelEditTableViewController.swift
index 1a20ffe..22c3ec4 100644
--- a/WireGuard/WireGuard/UI/iOS/ViewController/TunnelEditTableViewController.swift
+++ b/WireGuard/WireGuard/UI/iOS/ViewController/TunnelEditTableViewController.swift
@@ -44,8 +44,8 @@ class TunnelEditTableViewController: UITableViewController {
     ]
 
     let activateOnDemandOptions: [ActivateOnDemandOption] = [
-        .anyInterface,
-        .wiFiInterfaceOnly,
+        .anyInterface(.anySSID),
+        .wiFiInterfaceOnly(.anySSID),
         .nonWiFiInterfaceOnly
     ]
 
diff --git a/WireGuard/WireGuard/UI/macOS/ViewController/TunnelEditViewController.swift b/WireGuard/WireGuard/UI/macOS/ViewController/TunnelEditViewController.swift
index 2c1c538..1c1c054 100644
--- a/WireGuard/WireGuard/UI/macOS/ViewController/TunnelEditViewController.swift
+++ b/WireGuard/WireGuard/UI/macOS/ViewController/TunnelEditViewController.swift
@@ -82,8 +82,8 @@ class TunnelEditViewController: NSViewController {
 
     let activateOnDemandOptions: [ActivateOnDemandOption] = [
         .none,
-        .anyInterface,
-        .wiFiInterfaceOnly,
+        .anyInterface(.anySSID),
+        .wiFiInterfaceOnly(.anySSID),
         .nonWiFiInterfaceOnly
     ]