Fix a race condition bug involving timers and async tasks

2 files changed, 33 insertions, 0 deletions

diff --git a/ios/PacketTunnelCore/Actor/Task+Duration.swift b/ios/PacketTunnelCore/Actor/Task+Duration.swift
index 32f7a3d866..fdbc6be038 100644
--- a/ios/PacketTunnelCore/Actor/Task+Duration.swift
+++ b/ios/PacketTunnelCore/Actor/Task+Duration.swift
@@ -40,10 +40,21 @@ extension Task where Success == Never, Failure == Never {
 
         try await withTaskCancellationHandler {
             try await withCheckedThrowingContinuation { continuation in
+                // The `continuation` handler should never be `resume`'d more than once.
+                // Setting the eventHandler on the timer after it has been cancelled will be ignored.
+                // https://github.com/apple/swift-corelibs-libdispatch/blob/77766345740dfe3075f2f60bead854b29b0cfa24/src/source.c#L338
+                // Therefore, set a flag indicating `resume` has already been called to avoid `resume`ing more than once.
+                // Cancelling the timer does not cancel an  event handler that is already running however,
+                // the cancel handler will be scheduled after the event handler has finished.
+                // https://developer.apple.com/documentation/dispatch/1385604-dispatch_source_cancel
+                // Therefore, there it is safe to do this here.
+                var hasResumed = false
                 timer.setEventHandler {
+                    hasResumed = true
                     continuation.resume()
                 }
                 timer.setCancelHandler {
+                    guard hasResumed == false else { return }
                     continuation.resume(throwing: TaskCancellationError())
                 }
                 timer.schedule(wallDeadline: .now() + DispatchTimeInterval.milliseconds(duration.milliseconds))
diff --git a/ios/PacketTunnelCoreTests/TaskSleepTests.swift b/ios/PacketTunnelCoreTests/TaskSleepTests.swift
index ad2cb7a361..b71dc14f10 100644
--- a/ios/PacketTunnelCoreTests/TaskSleepTests.swift
+++ b/ios/PacketTunnelCoreTests/TaskSleepTests.swift
@@ -25,4 +25,26 @@ final class TaskSleepTests: XCTestCase {
             XCTAssert(error is CancellationError)
         }
     }
+
+    /// This test triggers a race condition in `sleepUsingContinuousClock` where an `AutoCancellingTask` will
+    /// cancel a `DispatchSourceTimer` in a `Task` trying to call `resume` on its continuation handler more than once
+    func testSuccessfulEventHandlerRemovesCancellation() async throws {
+        for _ in 0 ... 20 {
+            var task = recoveryTask()
+            try await Task.sleep(duration: .milliseconds(10))
+            task.doAnythingToSilenceAWarning()
+        }
+    }
+
+    private func recoveryTask() -> AutoCancellingTask {
+        AutoCancellingTask(Task.detached {
+            while Task.isCancelled == false {
+                try await Task.sleepUsingContinuousClock(for: .milliseconds(10))
+            }
+        })
+    }
+}
+
+private extension AutoCancellingTask {
+    func doAnythingToSilenceAWarning() {}
 }