Reuse the connection attempts count logic for PQ PSK negotiation

Author: buggmagnet

ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift

@@ -307,6 +307,6 @@ extension PacketTunnelProvider: PostQuantumKeyReceiving {
         postQuantumActor.endCurrentNegotiation()
         // Do not try reconnecting to the `.current` relay, else the actor's `State` equality check will fail
         // and it will not try to reconnect
-        actor.reconnect(to: .random)
+        actor.reconnect(to: .random, reconnectReason: .connectionLoss)
     }
 }

ios/PacketTunnelCore/Actor/PacketTunnelActor+PostQuantum.swift

@@ -17,10 +17,10 @@ extension PacketTunnelActor {
     internal func tryStartPostQuantumNegotiation(
         withSettings settings: Settings,
         nextRelay: NextRelay,
-        reason: ReconnectReason
+        reason: ActorReconnectReason
     ) async throws {
-        if let connectionState = try makeConnectionState(nextRelay: nextRelay, settings: settings, reason: reason) {
-            let selectedEndpoint = connectionState.selectedRelay.endpoint
+        if let connectionState = try obfuscateConnection(nextRelay: nextRelay, settings: settings, reason: reason) {
+            let selectedEndpoint = connectionState.connectedEndpoint
             let activeKey = activeKey(from: connectionState, in: settings)
 
             let configurationBuilder = ConfigurationBuilder(

ios/PacketTunnelCore/Actor/PacketTunnelActor+Public.swift

@@ -38,8 +38,8 @@ extension PacketTunnelActor {
 
      - Parameter nextRelay: next relay to connect to.
      */
-    public nonisolated func reconnect(to nextRelay: NextRelay) {
-        eventChannel.send(.reconnect(nextRelay))
+    public nonisolated func reconnect(to nextRelay: NextRelay, reconnectReason: ActorReconnectReason = .userInitiated) {
+        eventChannel.send(.reconnect(nextRelay, reason: reconnectReason))
     }
 
     /**

ios/PacketTunnelCore/Actor/PacketTunnelActor.swift

@@ -148,16 +148,6 @@ public actor PacketTunnelActor {
 // MARK: -
 
 extension PacketTunnelActor {
-    /// Describes the reason for reconnection request.
-    enum ReconnectReason: Equatable {
-        /// Initiated by user.
-        case userInitiated
-
-        /// Initiated by tunnel monitor due to loss of connectivity.
-        /// Actor will increment the connection attempt counter before picking next relay.
-        case connectionLoss
-    }
-
     /**
      Start the tunnel.
 
@@ -221,7 +211,7 @@ extension PacketTunnelActor {
          - nextRelay: next relay to connect to
          - reason: reason for reconnect
      */
-    private func reconnect(to nextRelay: NextRelay, reason: ReconnectReason) async {
+    private func reconnect(to nextRelay: NextRelay, reason: ActorReconnectReason) async {
         do {
             switch state {
             // There is no connection monitoring going on when exchanging keys.
@@ -256,7 +246,7 @@ extension PacketTunnelActor {
      */
     private func tryStart(
         nextRelay: NextRelay,
-        reason: ReconnectReason = .userInitiated
+        reason: ActorReconnectReason = .userInitiated
     ) async throws {
         let settings: Settings = try settingsReader.read()
 
@@ -284,7 +274,7 @@ extension PacketTunnelActor {
     private func tryStartConnection(
         withSettings settings: Settings,
         nextRelay: NextRelay,
-        reason: ReconnectReason
+        reason: ActorReconnectReason
     ) async throws {
         guard let connectionState = try obfuscateConnection(nextRelay: nextRelay, settings: settings, reason: reason),
               let targetState = state.targetStateForReconnect else { return }
@@ -341,7 +331,7 @@ extension PacketTunnelActor {
     internal func makeConnectionState(
         nextRelay: NextRelay,
         settings: Settings,
-        reason: ReconnectReason
+        reason: ActorReconnectReason
     ) throws -> State.ConnectionData? {
         var keyPolicy: State.KeyPolicy = .useCurrent
         var networkReachability = defaultPathObserver.defaultPath?.networkReachability ?? .undetermined
@@ -359,19 +349,23 @@ extension PacketTunnelActor {
         switch state {
         case .initial:
             break
-        case var .connecting(connectionState), var .reconnecting(connectionState):
+        // Handle PQ PSK separately as it doesn't interfere with either the `.connecting` or `.reconnecting` states.
+        case var .negotiatingPostQuantumKey(connectionState, _):
             if reason == .connectionLoss {
                 connectionState.incrementAttemptCount()
             }
-            fallthrough
-        case var .negotiatingPostQuantumKey(connectionState, _):
             let selectedRelay = try callRelaySelector(
                 connectionState.selectedRelay,
                 connectionState.connectionAttemptCount
             )
             connectionState.selectedRelay = selectedRelay
             connectionState.relayConstraints = settings.relayConstraints
             return connectionState
+        case var .connecting(connectionState), var .reconnecting(connectionState):
+            if reason == .connectionLoss {
+                connectionState.incrementAttemptCount()
+            }
+            fallthrough
         case var .connected(connectionState):
             let selectedRelay = try callRelaySelector(
                 connectionState.selectedRelay,
@@ -416,7 +410,7 @@ extension PacketTunnelActor {
     internal func obfuscateConnection(
         nextRelay: NextRelay,
         settings: Settings,
-        reason: ReconnectReason
+        reason: ActorReconnectReason
     ) throws -> State.ConnectionData? {
         guard let connectionState = try makeConnectionState(nextRelay: nextRelay, settings: settings, reason: reason)
         else { return nil }

ios/PacketTunnelCore/Actor/PacketTunnelActorCommand.swift

@@ -19,7 +19,7 @@ extension PacketTunnelActor {
         case stop
 
         /// Reconnect tunnel.
-        case reconnect(NextRelay, reason: ReconnectReason = .userInitiated)
+        case reconnect(NextRelay, reason: ActorReconnectReason = .userInitiated)
 
         /// Enter blocked state.
         case error(BlockedStateReason)

ios/PacketTunnelCore/Actor/PacketTunnelActorProtocol.swift

@@ -11,6 +11,6 @@ import Foundation
 public protocol PacketTunnelActorProtocol {
     var observedState: ObservedState { get async }
 
-    func reconnect(to nextRelay: NextRelay)
+    func reconnect(to nextRelay: NextRelay, reconnectReason: ActorReconnectReason)
     func notifyKeyRotation(date: Date?)
 }

ios/PacketTunnelCore/Actor/PacketTunnelActorReducer.swift

@@ -18,7 +18,7 @@ extension PacketTunnelActor {
         case stopTunnelMonitor
         case updateTunnelMonitorPath(NetworkPath)
         case startConnection(NextRelay)
-        case restartConnection(NextRelay, ReconnectReason)
+        case restartConnection(NextRelay, ActorReconnectReason)
         // trigger a reconnect, which becomes several effects depending on the state
         case reconnect(NextRelay)
         case stopTunnelAdapter
@@ -123,7 +123,7 @@ extension PacketTunnelActor {
 
         fileprivate static func subreducerForReconnect(
             _ state: State,
-            _ reason: PacketTunnelActor.ReconnectReason,
+            _ reason: ActorReconnectReason,
             _ nextRelay: NextRelay
         ) -> [PacketTunnelActor.Effect] {
             switch state {

ios/PacketTunnelCore/Actor/State.swift

@@ -238,3 +238,13 @@ public enum NextRelay: Equatable, Codable {
     /// Use pre-selected relay.
     case preSelected(SelectedRelay)
 }
+
+/// Describes the reason for reconnection request.
+public enum ActorReconnectReason: Equatable {
+    /// Initiated by user.
+    case userInitiated
+
+    /// Initiated by tunnel monitor due to loss of connectivity, or if post quantum key negotiation times out.
+    /// Actor will increment the connection attempt counter before picking next relay.
+    case connectionLoss
+}

ios/PacketTunnelCore/IPC/AppMessageHandler.swift

@@ -54,7 +54,7 @@ public struct AppMessageHandler {
             return nil
 
         case let .reconnectTunnel(nextRelay):
-            packetTunnelActor.reconnect(to: nextRelay)
+            packetTunnelActor.reconnect(to: nextRelay, reconnectReason: ActorReconnectReason.userInitiated)
             return nil
         }
     }

ios/PacketTunnelCoreTests/Mocks/PacketTunnelActorStub.swift

@@ -23,7 +23,7 @@ struct PacketTunnelActorStub: PacketTunnelActorProtocol {
         }
     }
 
-    func reconnect(to nextRelay: NextRelay) {
+    func reconnect(to nextRelay: PacketTunnelCore.NextRelay, reconnectReason: ActorReconnectReason) {
         reconnectExpectation?.fulfill()
     }
 

ios/PacketTunnelCoreTests/Mocks/SettingsReaderStub.swift

@@ -37,4 +37,18 @@ extension SettingsReaderStub {
             return staticSettings
         }
     }
+
+    static func postQuantumConfiguration() -> SettingsReaderStub {
+        let staticSettings = Settings(
+            privateKey: PrivateKey(),
+            interfaceAddresses: [IPAddressRange(from: "127.0.0.1/32")!],
+            relayConstraints: RelayConstraints(),
+            dnsServers: .gateway,
+            obfuscation: WireGuardObfuscationSettings(state: .off, port: .automatic),
+            quantumResistance: .on
+        )
+        return SettingsReaderStub {
+            return staticSettings
+        }
+    }
 }

ios/PacketTunnelCoreTests/PacketTunnelActorTests.swift

@@ -123,6 +123,37 @@ final class PacketTunnelActorTests: XCTestCase {
         await fulfillment(of: [connectingStateExpectation], timeout: 1)
     }
 
+    func testPostQuantumReconnectionTransition() async throws {
+        let tunnelMonitor = TunnelMonitorStub { _, _ in }
+        let actor = PacketTunnelActor.mock(
+            tunnelMonitor: tunnelMonitor,
+            settingsReader: SettingsReaderStub.postQuantumConfiguration()
+        )
+        let negotiatingPostQuantumKeyStateExpectation = expectation(description: "Expect post quantum state")
+        negotiatingPostQuantumKeyStateExpectation.expectedFulfillmentCount = 5
+        var nextAttemptCount: UInt = 0
+        stateSink = await actor.$observedState
+            .receive(on: DispatchQueue.main)
+            .sink { newState in
+                switch newState {
+                case .initial:
+                    break
+                case let .negotiatingPostQuantumKey(connState, _):
+                    XCTAssertEqual(connState.connectionAttemptCount, nextAttemptCount)
+                    nextAttemptCount += 1
+                    negotiatingPostQuantumKeyStateExpectation.fulfill()
+                    if nextAttemptCount < negotiatingPostQuantumKeyStateExpectation.expectedFulfillmentCount {
+                        actor.reconnect(to: .random, reconnectReason: .connectionLoss)
+                    }
+                default:
+                    XCTFail("Received invalid state: \(newState.name).")
+                }
+            }
+
+        actor.start(options: StartOptions(launchSource: .app))
+        await fulfillment(of: [negotiatingPostQuantumKeyStateExpectation], timeout: 1)
+    }
+
     /**
      Each subsequent re-connection attempt should produce a single change to `state` containing the incremented attempt counter and new relay.
      .reconnecting (attempt: 0) → .reconnecting (attempt: 1) → .reconnecting (attempt: 2) → ...