Merge branch 'packettunnel-state-refactor-port'

mullvad · Apr 3, 2024 · 3686aae · 3686aae
2 parents 895646f + 8e8910d
commit 3686aae
Show file tree

Hide file tree

Showing 8 changed files with 205 additions and 216 deletions.
diff --git 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
@@ -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
@@ -23,65 +23,19 @@ 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
-
-                    // 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
+        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.
+            stateData.keyPolicy = .usePrior(currentKey, startKeySwitchTask())
+            stateData.currentKey = nil
         }
     }
 
@@ -123,36 +77,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 +90,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
@@ -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
@@ -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,