Refactor PacketTunnelActor State, decoupling state from ancilliary data

7 files changed, 203 insertions, 214 deletions

diff --git a/ios/PacketTunnelCore/Actor/ObservedState.swift b/ios/PacketTunnelCore/Actor/ObservedState.swift
index 94a1aa19c8..3f92b300e5 100644
--- a/ios/PacketTunnelCore/Actor/ObservedState.swift
+++ b/ios/PacketTunnelCore/Actor/ObservedState.swift
@@ -86,8 +86,8 @@ extension State {
     }
 }
 
-extension ConnectionState {
-    /// Map `ConnectionState` to `ObservedConnectionState`.
+extension State.ConnectionData {
+    /// Map `State.ConnectionData` to `ObservedConnectionState`.
     var observedConnectionState: ObservedConnectionState {
         ObservedConnectionState(
             selectedRelay: selectedRelay,
@@ -101,8 +101,8 @@ extension ConnectionState {
     }
 }
 
-extension BlockedState {
-    /// Map `BlockedState` to `ObservedBlockedState`
+extension State.BlockingData {
+    /// Map `State.BlockingData` to `ObservedBlockedState`
     var observedBlockedState: ObservedBlockedState {
         return ObservedBlockedState(reason: reason, relayConstraints: relayConstraints)
     }
diff --git a/ios/PacketTunnelCore/Actor/PacketTunnelActor+ErrorState.swift b/ios/PacketTunnelCore/Actor/PacketTunnelActor+ErrorState.swift
index 94888c8115..40d0431d1a 100644
--- a/ios/PacketTunnelCore/Actor/PacketTunnelActor+ErrorState.swift
+++ b/ios/PacketTunnelCore/Actor/PacketTunnelActor+ErrorState.swift
@@ -54,10 +54,10 @@ extension PacketTunnelActor {
      - Returns: New blocked state that should be assigned to error state, otherwise `nil` when actor is past or at `disconnecting` phase or
                 when actor is already in the error state and no changes need to be made.
      */
-    private func makeBlockedState(reason: BlockedStateReason) -> BlockedState? {
+    private func makeBlockedState(reason: BlockedStateReason) -> State.BlockingData? {
         switch state {
         case .initial:
-            return BlockedState(
+            return State.BlockingData(
                 reason: reason,
                 relayConstraints: nil,
                 currentKey: nil,
@@ -93,11 +93,11 @@ extension PacketTunnelActor {
      Map connection state to blocked state.
      */
     private func mapConnectionState(
-        _ connState: ConnectionState,
+        _ connState: State.ConnectionData,
         reason: BlockedStateReason,
         priorState: StatePriorToBlockedState
-    ) -> BlockedState {
-        BlockedState(
+    ) -> State.BlockingData {
+        State.BlockingData(
             reason: reason,
             relayConstraints: connState.relayConstraints,
             currentKey: connState.currentKey,
diff --git a/ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift b/ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift
index 5ab13353e9..3084b21676 100644
--- a/ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift
+++ b/ios/PacketTunnelCore/Actor/PacketTunnelActor+KeyPolicy.swift
@@ -23,65 +23,20 @@ extension PacketTunnelActor {
      - Parameter lastKeyRotation: date when last key rotation took place.
      */
     func cacheActiveKey(lastKeyRotation: Date?) {
-        func mutateConnectionState(_ connState: inout ConnectionState) -> Bool {
-            switch connState.keyPolicy {
-            case .useCurrent:
-                if let currentKey = connState.currentKey {
-                    connState.lastKeyRotation = lastKeyRotation
+        // There should be a way to rationalise these two identical functions into one.
+        state.mutateAssociatedData { stateData in
+            guard
+                stateData.keyPolicy == .useCurrent,
+                let currentKey = stateData.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.
+            stateData.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
-            }
-        }
-
-        switch state {
-        case var .connecting(connState):
-            if mutateConnectionState(&connState) {
-                state = .connecting(connState)
-            }
-
-        case var .connected(connState):
-            if mutateConnectionState(&connState) {
-                state = .connected(connState)
-            }
-
-        case var .reconnecting(connState):
-            if mutateConnectionState(&connState) {
-                state = .reconnecting(connState)
-            }
-
-        case var .error(blockedState):
-            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
-
-                    state = .error(blockedState)
-                }
-
-            case .usePrior:
-                break
-            }
-
-        case .initial, .disconnected, .disconnecting:
-            break
+            // Move currentKey into keyPolicy.
+            stateData.keyPolicy = .usePrior(currentKey, startKeySwitchTask())
+            stateData.currentKey = nil
         }
     }
 
@@ -123,36 +78,10 @@ extension PacketTunnelActor {
      - Returns: `true` if the tunnel should reconnect, otherwise `false`.
      */
     private func switchToCurrentKeyInner() -> Bool {
-        switch state {
-        case var .connecting(connState):
-            if setCurrentKeyPolicy(&connState.keyPolicy) {
-                state = .connecting(connState)
-                return true
-            }
-
-        case var .connected(connState):
-            if setCurrentKeyPolicy(&connState.keyPolicy) {
-                state = .connected(connState)
-                return true
-            }
-
-        case var .reconnecting(connState):
-            if setCurrentKeyPolicy(&connState.keyPolicy) {
-                state = .reconnecting(connState)
-                return true
-            }
-
-        case var .error(blockedState):
-            if setCurrentKeyPolicy(&blockedState.keyPolicy) {
-                state = .error(blockedState)
-
-                // Prevent tunnel from reconnecting when in blocked state.
-                return false
-            }
-
-        case .disconnected, .disconnecting, .initial:
-            break
-        }
+        let oldKeyPolicy = state.keyPolicy
+        state.mutateKeyPolicy(setCurrentKeyPolicy)
+        // Prevent tunnel from reconnecting when in blocked state.
+        guard case .error = state else { return state.keyPolicy != oldKeyPolicy }
         return false
     }
 
@@ -162,14 +91,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
         }
     }
 }
diff --git a/ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift b/ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift
index 2d3b600239..e60a0081ff 100644
--- a/ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift
+++ b/ios/PacketTunnelCore/Actor/PacketTunnelActor+NetworkReachability.swift
@@ -43,43 +43,6 @@ extension PacketTunnelActor {
 
         let newReachability = networkPath.networkReachability
 
-        func mutateConnectionState(_ connState: inout ConnectionState) -> Bool {
-            if connState.networkReachability != newReachability {
-                connState.networkReachability = newReachability
-                return true
-            }
-            return false
-        }
-
-        switch state {
-        case var .connecting(connState):
-            if mutateConnectionState(&connState) {
-                state = .connecting(connState)
-            }
-
-        case var .connected(connState):
-            if mutateConnectionState(&connState) {
-                state = .connected(connState)
-            }
-
-        case var .reconnecting(connState):
-            if mutateConnectionState(&connState) {
-                state = .reconnecting(connState)
-            }
-
-        case var .disconnecting(connState):
-            if mutateConnectionState(&connState) {
-                state = .disconnecting(connState)
-            }
-
-        case var .error(blockedState):
-            if blockedState.networkReachability != newReachability {
-                blockedState.networkReachability = newReachability
-                state = .error(blockedState)
-            }
-
-        case .initial, .disconnected:
-            break
-        }
+        state.mutateAssociatedData { $0.networkReachability = newReachability }
     }
 }
diff --git a/ios/PacketTunnelCore/Actor/PacketTunnelActor.swift b/ios/PacketTunnelCore/Actor/PacketTunnelActor.swift
index e7a384d8b6..360bfb8f55 100644
--- a/ios/PacketTunnelCore/Actor/PacketTunnelActor.swift
+++ b/ios/PacketTunnelCore/Actor/PacketTunnelActor.swift
@@ -27,7 +27,8 @@ import WireGuardKitTypes
  */
 public actor PacketTunnelActor {
     var state: State = .initial {
-        didSet {
+        didSet(oldValue) {
+            guard state != oldValue else { return }
             logger.debug("\(state.logFormat())")
             observedState = state.observedState
         }
@@ -293,7 +294,7 @@ extension PacketTunnelActor {
         nextRelay: NextRelay,
         settings: Settings,
         reason: ReconnectReason
-    ) throws -> ConnectionState? {
+    ) throws -> State.ConnectionData? {
         var keyPolicy: KeyPolicy = .useCurrent
         var networkReachability = defaultPathObserver.defaultPath?.networkReachability ?? .undetermined
         var lastKeyRotation: Date?
@@ -332,7 +333,7 @@ extension PacketTunnelActor {
             return nil
         }
         let selectedRelay = try callRelaySelector(nil, 0)
-        return ConnectionState(
+        return State.ConnectionData(
             selectedRelay: selectedRelay,
             relayConstraints: settings.relayConstraints,
             currentKey: settings.privateKey,
@@ -350,7 +351,7 @@ extension PacketTunnelActor {
         nextRelay: NextRelay,
         settings: Settings,
         reason: ReconnectReason
-    ) throws -> ConnectionState? {
+    ) throws -> State.ConnectionData? {
         guard let connectionState = try makeConnectionState(nextRelay: nextRelay, settings: settings, reason: reason)
         else { return nil }
 
@@ -361,7 +362,7 @@ extension PacketTunnelActor {
         )
 
         let transportLayer = protocolObfuscator.transportLayer.map { $0 } ?? .udp
-        return ConnectionState(
+        return State.ConnectionData(
             selectedRelay: connectionState.selectedRelay,
             relayConstraints: connectionState.relayConstraints,
             currentKey: settings.privateKey,
diff --git a/ios/PacketTunnelCore/Actor/State+Extensions.swift b/ios/PacketTunnelCore/Actor/State+Extensions.swift
index dc7580fdf1..97d2deb388 100644
--- a/ios/PacketTunnelCore/Actor/State+Extensions.swift
+++ b/ios/PacketTunnelCore/Actor/State+Extensions.swift
@@ -62,21 +62,91 @@ extension State {
     var name: String {
         switch self {
         case .connected:
-            return "Connected"
+            "Connected"
         case .connecting:
-            return "Connecting"
+            "Connecting"
         case .reconnecting:
-            return "Reconnecting"
+            "Reconnecting"
         case .disconnecting:
-            return "Disconnecting"
+            "Disconnecting"
         case .disconnected:
-            return "Disconnected"
+            "Disconnected"
         case .initial:
-            return "Initial"
+            "Initial"
         case .error:
-            return "Error"
+            "Error"
         }
     }
+
+    var connectionData: State.ConnectionData? {
+        switch self {
+        case
+            let .connecting(connState),
+            let .connected(connState),
+            let .reconnecting(connState),
+            let .disconnecting(connState): connState
+        default: nil
+        }
+    }
+
+    var blockedData: State.BlockingData? {
+        switch self {
+        case let .error(blockedState): blockedState
+        default: nil
+        }
+    }
+
+    var associatedData: StateAssociatedData? {
+        self.connectionData ?? self.blockedData
+    }
+
+    var keyPolicy: KeyPolicy? {
+        associatedData?.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.
+
+    internal func replacingConnectionData(with newValue: State.ConnectionData) -> State {
+        switch self {
+        case .connecting: .connecting(newValue)
+        case .connected: .connected(newValue)
+        case .reconnecting: .reconnecting(newValue)
+        case .disconnecting: .disconnecting(newValue)
+        default: self
+        }
+    }
+
+    /// 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 mutateAssociatedData(_ modifier: (inout StateAssociatedData) -> Void) {
+        switch self {
+        case let .connecting(connState),
+             let .connected(connState),
+             let .reconnecting(connState),
+             let .disconnecting(connState):
+            var associatedData: StateAssociatedData = connState
+            modifier(&associatedData)
+            // swiftlint:disable:next force_cast
+            self = self.replacingConnectionData(with: associatedData as! ConnectionData)
+
+        case let .error(blockedState):
+            var associatedData: StateAssociatedData = blockedState
+            modifier(&associatedData)
+            // swiftlint:disable:next force_cast
+            self = .error(associatedData as! BlockingData)
+
+        default:
+            break
+        }
+    }
+
+    /// Apply a mutating function to the state's key policy
+    /// - parameter modifier: A function that takes an `inout KeyPolicy` and modifies it
+    mutating func mutateKeyPolicy(_ modifier: (inout KeyPolicy) -> Void) {
+        self.mutateAssociatedData { modifier(&$0.keyPolicy) }
+    }
 }
 
 extension KeyPolicy {
@@ -90,6 +160,16 @@ extension KeyPolicy {
     }
 }
 
+extension KeyPolicy: Equatable {
+    static func == (lhs: KeyPolicy, rhs: KeyPolicy) -> Bool {
+        switch (lhs, rhs) {
+        case (.useCurrent, .useCurrent): true
+        case let (.usePrior(priorA, _), .usePrior(priorB, _)): priorA == priorB
+        default: false
+        }
+    }
+}
+
 extension BlockedStateReason {
     /**
      Returns true if the tunnel should attempt to restart periodically to recover from error that does not require explicit restart to be initiated by user.
@@ -111,3 +191,15 @@ extension BlockedStateReason {
         }
     }
 }
+
+extension State.BlockingData: Equatable {
+    static func == (lhs: State.BlockingData, rhs: State.BlockingData) -> Bool {
+        lhs.reason == rhs.reason
+            && lhs.relayConstraints == rhs.relayConstraints
+            && lhs.currentKey == rhs.currentKey
+            && lhs.keyPolicy == rhs.keyPolicy
+            && lhs.networkReachability == rhs.networkReachability
+            && lhs.lastKeyRotation == rhs.lastKeyRotation
+            && lhs.priorState == rhs.priorState
+    }
+}
diff --git a/ios/PacketTunnelCore/Actor/State.swift b/ios/PacketTunnelCore/Actor/State.swift
index 3cca82d865..198ac45c96 100644
--- a/ios/PacketTunnelCore/Actor/State.swift
+++ b/ios/PacketTunnelCore/Actor/State.swift
@@ -54,22 +54,22 @@ import WireGuardKitTypes
  `.connecting`, `.reconnecting`, `.error` can be interrupted if the tunnel is requested to stop, which should segue actor towards `.disconnected` state.
 
  */
-enum State {
+enum State: Equatable {
     /// Initial state at the time when actor is initialized but before the first connection attempt.
     case initial
 
     /// Tunnel is attempting to connect.
     /// The actor should remain in this state until the very first connection is established, i.e determined by tunnel monitor.
-    case connecting(ConnectionState)
+    case connecting(ConnectionData)
 
     /// Tunnel is connected.
-    case connected(ConnectionState)
+    case connected(ConnectionData)
 
     /// Tunnel is attempting to reconnect.
-    case reconnecting(ConnectionState)
+    case reconnecting(ConnectionData)
 
     /// Tunnel is disconnecting.
-    case disconnecting(ConnectionState)
+    case disconnecting(ConnectionData)
 
     /// Tunnel is disconnected.
     /// Normally the process is shutdown after entering this state.
@@ -79,7 +79,7 @@ enum State {
     /// This state is normally entered when the tunnel is unable to start or reconnect.
     /// In this state the tunnel blocks all nework connectivity by setting up a peerless WireGuard tunnel, and either awaits user action or, in certain
     /// circumstances, attempts to recover automatically using a repeating timer.
-    case error(BlockedState)
+    case error(BlockingData)
 }
 
 /// Policy describing what WG key to use for tunnel communication.
@@ -96,76 +96,85 @@ public enum NetworkReachability: Equatable, Codable {
     case undetermined, reachable, unreachable
 }
 
-/// Data associated with states that hold connection data.
-struct ConnectionState {
-    /// Current selected relay.
-    public var selectedRelay: SelectedRelay
+protocol StateAssociatedData {
+    var currentKey: PrivateKey? { get set }
+    var keyPolicy: KeyPolicy { get set }
+    var networkReachability: NetworkReachability { get set }
+    var lastKeyRotation: Date? { get set }
+}
 
-    /// Last relay constraints read from settings.
-    /// This is primarily used by packet tunnel for updating constraints in tunnel provider.
-    public var relayConstraints: RelayConstraints
+extension State {
+    /// Data associated with states that hold connection data.
+    struct ConnectionData: Equatable, StateAssociatedData {
+        /// Current selected relay.
+        public var selectedRelay: SelectedRelay
 
-    /// Last WG key read from settings.
-    /// Can be `nil` if moved to `keyPolicy`.
-    public var currentKey: PrivateKey?
+        /// Last relay constraints read from settings.
+        /// This is primarily used by packet tunnel for updating constraints in tunnel provider.
+        public var relayConstraints: RelayConstraints
 
-    /// Policy describing the current key that should be used by the tunnel.
-    public var keyPolicy: KeyPolicy
+        /// Last WG key read from settings.
+        /// Can be `nil` if moved to `keyPolicy`.
+        public var currentKey: PrivateKey?
 
-    /// Whether network connectivity outside of tunnel is available.
-    public var networkReachability: NetworkReachability
+        /// Policy describing the current key that should be used by the tunnel.
+        public var keyPolicy: KeyPolicy
 
-    /// Connection attempt counter.
-    /// Reset to zero once connection is established.
-    public var connectionAttemptCount: UInt
+        /// Whether network connectivity outside of tunnel is available.
+        public var networkReachability: NetworkReachability
 
-    /// Last time packet tunnel rotated the key.
-    public var lastKeyRotation: Date?
+        /// Connection attempt counter.
+        /// Reset to zero once connection is established.
+        public var connectionAttemptCount: UInt
 
-    /// Increment connection attempt counter by one, wrapping to zero on overflow.
-    public mutating func incrementAttemptCount() {
-        let (value, isOverflow) = connectionAttemptCount.addingReportingOverflow(1)
-        connectionAttemptCount = isOverflow ? 0 : value
-    }
+        /// Last time packet tunnel rotated the key.
+        public var lastKeyRotation: Date?
 
-    /// The actual endpoint fed to WireGuard, can be a local endpoint if obfuscation is used.
-    public let connectedEndpoint: MullvadEndpoint
-    /// Via which transport protocol was the connection made to the relay
-    public let transportLayer: TransportLayer
+        /// Increment connection attempt counter by one, wrapping to zero on overflow.
+        public mutating func incrementAttemptCount() {
+            let (value, isOverflow) = connectionAttemptCount.addingReportingOverflow(1)
+            connectionAttemptCount = isOverflow ? 0 : value
+        }
 
-    /// The remote port that was chosen to connect to `connectedEndpoint`
-    public let remotePort: UInt16
-}
+        /// The actual endpoint fed to WireGuard, can be a local endpoint if obfuscation is used.
+        public let connectedEndpoint: MullvadEndpoint
+        /// Via which transport protocol was the connection made to the relay
+        public let transportLayer: TransportLayer
 
-/// Data associated with error state.
-struct BlockedState {
-    /// Reason why block state was entered.
-    public var reason: BlockedStateReason
+        /// The remote port that was chosen to connect to `connectedEndpoint`
+        public let remotePort: UInt16
+    }
+
+    /// Data associated with error state.
+    struct BlockingData: StateAssociatedData {
+        /// Reason why block state was entered.
+        public var reason: BlockedStateReason
 
-    /// Last relay constraints read from settings.
-    /// This is primarily used by packet tunnel for updating constraints in tunnel provider.
-    public var relayConstraints: RelayConstraints?
+        /// Last relay constraints read from settings.
+        /// This is primarily used by packet tunnel for updating constraints in tunnel provider.
+        public var relayConstraints: RelayConstraints?
 
-    /// Last WG key read from setings.
-    /// Can be `nil` if moved to `keyPolicy` or when it's uknown.
-    public var currentKey: PrivateKey?
+        /// Last WG key read from setings.
+        /// Can be `nil` if moved to `keyPolicy` or when it's uknown.
+        public var currentKey: PrivateKey?
 
-    /// Policy describing the current key that should be used by the tunnel.
-    public var keyPolicy: KeyPolicy
+        /// Policy describing the current key that should be used by the tunnel.
+        public var keyPolicy: KeyPolicy
 
-    /// Whether network connectivity outside of tunnel is available.
-    public var networkReachability: NetworkReachability
+        /// Whether network connectivity outside of tunnel is available.
+        public var networkReachability: NetworkReachability
 
-    /// Last time packet tunnel rotated or attempted to rotate the key.
-    /// This is used by `TunnelManager` to detect when it needs to refresh device state from Keychain.
-    public var lastKeyRotation: Date?
+        /// Last time packet tunnel rotated or attempted to rotate the key.
+        /// This is used by `TunnelManager` to detect when it needs to refresh device state from Keychain.
+        public var lastKeyRotation: Date?
 
-    /// Task responsible for periodically calling actor to restart the tunnel.
-    /// Initiated based on the error that led to blocked state.
-    public var recoveryTask: AutoCancellingTask?
+        /// Task responsible for periodically calling actor to restart the tunnel.
+        /// Initiated based on the error that led to blocked state.
+        public var recoveryTask: AutoCancellingTask?
 
-    /// Prior state of the actor before entering blocked state
-    public var priorState: StatePriorToBlockedState
+        /// Prior state of the actor before entering blocked state
+        public var priorState: StatePriorToBlockedState
+    }
 }
 
 /// Reason why packet tunnel entered error state.
@@ -206,7 +215,7 @@ public enum BlockedStateReason: String, Codable, Equatable {
 }
 
 /// Legal states that can precede error state.
-enum StatePriorToBlockedState {
+enum StatePriorToBlockedState: Equatable {
     case initial, connecting, connected, reconnecting
 }