Ensure atomicity between (re)connection attempts

4 files changed, 50 insertions, 32 deletions

diff --git a/ios/PacketTunnelCore/Actor/Actor+ConnectionMonitoring.swift b/ios/PacketTunnelCore/Actor/Actor+ConnectionMonitoring.swift
index 2980d4d11d..c2bc921b05 100644
--- a/ios/PacketTunnelCore/Actor/Actor+ConnectionMonitoring.swift
+++ b/ios/PacketTunnelCore/Actor/Actor+ConnectionMonitoring.swift
@@ -47,28 +47,14 @@ extension PacketTunnelActor {
         }
     }
 
-    /// Increment connection attempt counter and reconnect the tunnel.
+    /// Tell the tunnel to reconnect providing the correct reason to ensure that the attempt counter is incremented before reconnect.
     private func onHandleConnectionRecovery() async {
         switch state {
-        case var .connecting(connState):
-            connState.incrementAttemptCount()
-            state = .connecting(connState)
-
-        case var .reconnecting(connState):
-            connState.incrementAttemptCount()
-            state = .reconnecting(connState)
-
-        case var .connected(connState):
-            connState.incrementAttemptCount()
-            state = .connected(connState)
+        case .connecting, .reconnecting, .connected:
+            commandChannel.send(.reconnect(.random, reason: .connectionLoss))
 
         case .initial, .disconnected, .disconnecting, .error:
-            // Explicit return to prevent reconnecting the tunnel.
-            return
+            break
         }
-
-        // Tunnel monitor should already be paused at this point so don't stop it to avoid a reset of its internal
-        // counters.
-        commandChannel.send(.reconnect(.random, stopTunnelMonitor: false))
     }
 }
diff --git a/ios/PacketTunnelCore/Actor/Actor.swift b/ios/PacketTunnelCore/Actor/Actor.swift
index fafeba2403..546ae3f59e 100644
--- a/ios/PacketTunnelCore/Actor/Actor.swift
+++ b/ios/PacketTunnelCore/Actor/Actor.swift
@@ -87,8 +87,8 @@ public actor PacketTunnelActor {
                 case .stop:
                     await stop()
 
-                case let .reconnect(nextRelay, stopTunnelMonitor):
-                    await reconnect(to: nextRelay, shouldStopTunnelMonitor: stopTunnelMonitor)
+                case let .reconnect(nextRelay, reason):
+                    await reconnect(to: nextRelay, reason: reason)
 
                 case let .error(reason):
                     await setErrorStateInternal(with: reason)
@@ -173,16 +173,22 @@ extension PacketTunnelActor {
 
      - Parameters:
          - nextRelay: next relay to connect to
-         - shouldStopTunnelMonitor: whether tunnel monitor should be stopped
+         - reason: reason for reconnect
      */
-    private func reconnect(to nextRelay: NextRelay, shouldStopTunnelMonitor: Bool) async {
+    private func reconnect(to nextRelay: NextRelay, reason: ReconnectReason) async {
         do {
             switch state {
             case .connecting, .connected, .reconnecting, .error:
-                if shouldStopTunnelMonitor {
+                switch reason {
+                case .connectionLoss:
+                    // Tunnel monitor is already paused at this point. Avoid calling stop() to prevent the reset of
+                    // internal state
+                    break
+                case .userInitiated:
                     tunnelMonitor.stop()
                 }
-                try await tryStart(nextRelay: nextRelay)
+
+                try await tryStart(nextRelay: nextRelay, reason: reason)
 
             case .disconnected, .disconnecting, .initial:
                 break
@@ -205,12 +211,14 @@ extension PacketTunnelActor {
      - Start tunnel monitor.
      - Reactivate default path observation (disabled when configuring tunnel adapter)
 
-     - Parameter nextRelay: which relay should be selected next.
+     - Parameters:
+         - nextRelay: which relay should be selected next.
+         - reason: reason for reconnect
      */
-    private func tryStart(nextRelay: NextRelay = .random) async throws {
+    private func tryStart(nextRelay: NextRelay = .random, reason: ReconnectReason = .userInitiated) async throws {
         let settings: Settings = try settingsReader.read()
 
-        guard let connectionState = try makeConnectionState(nextRelay: nextRelay, settings: settings),
+        guard let connectionState = try makeConnectionState(nextRelay: nextRelay, settings: settings, reason: reason),
               let targetState = state.targetStateForReconnect else { return }
 
         let activeKey: PrivateKey
@@ -261,10 +269,15 @@ extension PacketTunnelActor {
      - Parameters:
          - nextRelay: relay preference that should be used when selecting next relay.
          - settings: current settings
+         - reason: reason for reconnect
 
      - Returns: New connection state or `nil` if current state is at or past `.disconnecting` phase.
      */
-    private func makeConnectionState(nextRelay: NextRelay, settings: Settings) throws -> ConnectionState? {
+    private func makeConnectionState(
+        nextRelay: NextRelay,
+        settings: Settings,
+        reason: ReconnectReason
+    ) throws -> ConnectionState? {
         let relayConstraints = settings.relayConstraints
         let privateKey = settings.privateKey
 
@@ -284,7 +297,18 @@ extension PacketTunnelActor {
                 connectionAttemptCount: 0
             )
 
-        case var .connecting(connState), var .connected(connState), var .reconnecting(connState):
+        case var .connecting(connState), var .reconnecting(connState):
+            switch reason {
+            case .connectionLoss:
+                // Increment attempt counter when reconnection is requested due to connectivity loss.
+                connState.incrementAttemptCount()
+            case .userInitiated:
+                break
+            }
+            // Explicit fallthrough
+            fallthrough
+
+        case var .connected(connState):
             connState.selectedRelay = try selectRelay(
                 nextRelay: nextRelay,
                 relayConstraints: relayConstraints,
diff --git a/ios/PacketTunnelCore/Actor/Command.swift b/ios/PacketTunnelCore/Actor/Command.swift
index d7cf121684..668f444d49 100644
--- a/ios/PacketTunnelCore/Actor/Command.swift
+++ b/ios/PacketTunnelCore/Actor/Command.swift
@@ -17,9 +17,7 @@ enum Command {
     case stop
 
     /// Reconnect tunnel.
-    /// `stopTunnelMonitor = false` is only used when tunnel monitor is paused in response to connectivity loss and shouldn't be stopped explicitly,
-    /// as this would reset its internal counters.
-    case reconnect(NextRelay, stopTunnelMonitor: Bool = true)
+    case reconnect(NextRelay, reason: ReconnectReason = .userInitiated)
 
     /// Enter blocked state.
     case error(BlockedStateReason)
diff --git a/ios/PacketTunnelCore/Actor/State.swift b/ios/PacketTunnelCore/Actor/State.swift
index 79cfc19ce6..36a427e496 100644
--- a/ios/PacketTunnelCore/Actor/State.swift
+++ b/ios/PacketTunnelCore/Actor/State.swift
@@ -212,3 +212,13 @@ public enum NextRelay: Equatable, Codable {
     /// Use pre-selected relay.
     case preSelected(SelectedRelay)
 }
+
+/// Describes the reason for reconnection request.
+public enum ReconnectReason {
+    /// 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
+}