Implement protocol for common State connection/error data

ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift

@@ -23,51 +23,19 @@ extension PacketTunnelActor {
      - Parameter lastKeyRotation: date when last key rotation took place.
      */
     func cacheActiveKey(lastKeyRotation: Date?) {
-        func connectionStateMutator(_ connState: inout ConnectionState) -> Bool {
-            switch connState.keyPolicy {
-            case .useCurrent:
-                if let currentKey = connState.currentKey {
-                    connState.lastKeyRotation = lastKeyRotation
-
-                    // Move currentKey into keyPolicy.
-                    connState.keyPolicy = .usePrior(currentKey, startKeySwitchTask())
-                    connState.currentKey = nil
-
-                    return true
-                } else {
-                    return false
-                }
-
-            case .usePrior:
-                // It's unlikely that we'll see subsequent key rotations happen frequently.
-                return false
-            }
+        // There should be a way to rationalise these two identical functions into one.
+        func connectionStateMutator(_ connState: inout ConnectionOrBlockedState) {
+            guard connState.keyPolicy == .useCurrent, let currentKey = connState.currentKey else { return }
+            // Key policy is preserved between states and key rotation may still happen while in blocked state.
+            // Therefore perform the key switch as normal with one exception that it shouldn't reconnect the tunnel
+            // automatically.
+            connState.lastKeyRotation = lastKeyRotation
+
+            // Move currentKey into keyPolicy.
+            connState.keyPolicy = .usePrior(currentKey, startKeySwitchTask())
+            connState.currentKey = nil
         }
-
-        func blockedStateMutator(_ blockedState: inout BlockedState) -> Bool {
-            switch blockedState.keyPolicy {
-            case .useCurrent:
-                // Key policy is preserved between states and key rotation may still happen while in blocked state.
-                // Therefore perform the key switch as normal with one exception that it shouldn't reconnect the tunnel
-                // automatically.
-                if let currentKey = blockedState.currentKey {
-                    blockedState.lastKeyRotation = lastKeyRotation
-
-                    // Move currentKey into keyPolicy.
-                    blockedState.keyPolicy = .usePrior(currentKey, startKeySwitchTask())
-                    blockedState.currentKey = nil
-
-                    return true
-                }
-
-            default:
-                return false
-            }
-            return false
-        }
-
-        _ = state.mutateConnectionState(connectionStateMutator) ||
-            state.mutateBlockedState(blockedStateMutator)
+        state.mutateConnectionOrBlockedState(connectionStateMutator)
     }
 
     /**
@@ -108,9 +76,10 @@ extension PacketTunnelActor {
      - Returns: `true` if the tunnel should reconnect, otherwise `false`.
      */
     private func switchToCurrentKeyInner() -> Bool {
-        let changed = state.mutateKeyPolicy(setCurrentKeyPolicy)
+        let oldKeyPolicy = state.keyPolicy
+        state.mutateKeyPolicy(setCurrentKeyPolicy)
         // Prevent tunnel from reconnecting when in blocked state.
-        guard case .error = state else { return changed }
+        guard case .error = state else { return state.keyPolicy != oldKeyPolicy }
         return false
     }
 
@@ -120,14 +89,9 @@ extension PacketTunnelActor {
      - Parameter keyPolicy: a reference to key policy held either in connection state or blocked state struct.
      - Returns: `true` when the policy was modified, otherwise `false`.
      */
-    private func setCurrentKeyPolicy(_ keyPolicy: inout KeyPolicy) -> Bool {
-        switch keyPolicy {
-        case .useCurrent:
-            return false
-
-        case .usePrior:
+    private func setCurrentKeyPolicy(_ keyPolicy: inout KeyPolicy) {
+        if case .usePrior = keyPolicy {
             keyPolicy = .useCurrent
-            return true
         }
     }
 }

ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift

@@ -43,16 +43,6 @@ extension PacketTunnelActor {
 
         let newReachability = networkPath.networkReachability
 
-        func connectionStateMutator(_ connState: inout ConnectionState) -> Bool {
-            connState.networkReachability = newReachability
-            return true
-        }
-
-        func blockedStateMutator(_ blockedState: inout BlockedState) -> Bool {
-            blockedState.networkReachability = newReachability
-            return true
-        }
-
-        _ = state.mutateConnectionState(connectionStateMutator) || state.mutateBlockedState(blockedStateMutator)
+        state.mutateConnectionOrBlockedState { $0.networkReachability = newReachability }
     }
 }

ios/PacketTunnelCore/Actor/State+Extensions.swift

@@ -78,6 +78,32 @@ extension State {
         }
     }
 
+    var connectionState: ConnectionState? {
+        switch self {
+        case
+            let .connecting(connState),
+            let .connected(connState),
+            let .reconnecting(connState),
+            let .disconnecting(connState): connState
+        default: nil
+        }
+    }
+
+    var blockedState: BlockedState? {
+        switch self {
+        case let .error(blockedState): blockedState
+        default: nil
+        }
+    }
+
+    var connectionOrErrorState: ConnectionOrBlockedState? {
+        self.connectionState ?? self.blockedState
+    }
+
+    var keyPolicy: KeyPolicy? {
+        connectionOrErrorState?.keyPolicy
+    }
+
     /// Return a copy of this state with the associated value (if appropriate) replaced with a new value.
     /// If the value does not apply, this just returns the state as is, ignoring it.
 
@@ -91,44 +117,31 @@ extension State {
         }
     }
 
-    /// Apply a mutating function to the connection state if this state has one, and replace its value. If not, this is a no-op.
-    /// - parameter modifier: A function that takes an `inout ConnectionState` and returns `true`if it has been mutated
-    /// - returns: `true` if the state's value has been changed
-    @discardableResult mutating func mutateConnectionState(_ modifier: (inout ConnectionState) -> Bool) -> Bool {
-        var modified = false
-        switch self {
-        case var .connecting(connState), var .connected(connState), var .reconnecting(connState),
-             var .disconnecting(connState):
-            modified = modifier(&connState)
-            self = self.replacingConnectionState(with: connState)
-        default:
-            return false
-        }
-        return modified
-    }
-
-    /// Apply a mutating function to the blocked state if this state has one, and replace its value. If not, this is a no-op.
-    /// - parameter modifier: A function that takes an `inout ConnectionState` and returns `true`if it has been mutated
-    /// - returns: `true` if the state's value has been changed
-
-    @discardableResult mutating func mutateBlockedState(_ modifier: (inout BlockedState) -> Bool) -> Bool {
+    /// Apply a mutating function to the connection/error state's associated data if this state has one,
+    /// and replace its value. If not, this is a no-op.
+    /// - parameter modifier: A function that takes an `inout ConnectionOrBlockedState` and modifies it
+    mutating func mutateConnectionOrBlockedState(_ modifier: (inout ConnectionOrBlockedState) -> Void) {
         switch self {
-        case var .error(blockedState):
-            let modified = modifier(&blockedState)
-            self = .error(blockedState)
-            return modified
+        case let .connecting(connState),
+             let .connected(connState),
+             let .reconnecting(connState),
+             let .disconnecting(connState):
+            var connOrErrorState: ConnectionOrBlockedState = connState
+            modifier(&connOrErrorState)
+            self = self.replacingConnectionState(with: connOrErrorState as! ConnectionState)
+        case let .error(blockedState):
+            var connOrErrorState: ConnectionOrBlockedState = blockedState
+            modifier(&connOrErrorState)
+            self = .error(connOrErrorState as! BlockedState)
         default:
-            return false
+            break
         }
     }
 
     /// Apply a mutating function to the state's key policy
-    /// - parameter modifier: A function that takes an `inout KeyPolicy` and returns `true`if it has been mutated
-    /// - returns: `true` if the state's value has been changed
-    @discardableResult mutating func mutateKeyPolicy(_ modifier: (inout KeyPolicy) -> Bool) -> Bool {
-        self.mutateConnectionState { modifier(&$0.keyPolicy) }
-            || self.mutateBlockedState { modifier(&$0.keyPolicy) }
-            || false
+    /// - parameter modifier: A function that takes an `inout KeyPolicy` and modifies it
+    mutating func mutateKeyPolicy(_ modifier: (inout KeyPolicy) -> Void) {
+        self.mutateConnectionOrBlockedState { modifier(&$0.keyPolicy) }
     }
 }
 

ios/PacketTunnelCore/Actor/State.swift

@@ -96,8 +96,16 @@ public enum NetworkReachability: Equatable, Codable {
     case undetermined, reachable, unreachable
 }
 
+// perhaps this should have a better name?
+protocol ConnectionOrBlockedState {
+    var currentKey: PrivateKey? { get set }
+    var keyPolicy: KeyPolicy { get set }
+    var networkReachability: NetworkReachability { get set }
+    var lastKeyRotation: Date? { get set }
+}
+
 /// Data associated with states that hold connection data.
-struct ConnectionState: Equatable {
+struct ConnectionState: Equatable, ConnectionOrBlockedState {
     /// Current selected relay.
     public var selectedRelay: SelectedRelay
 
@@ -138,7 +146,7 @@ struct ConnectionState: Equatable {
 }
 
 /// Data associated with error state.
-struct BlockedState {
+struct BlockedState: ConnectionOrBlockedState {
     /// Reason why block state was entered.
     public var reason: BlockedStateReason